Add files using upload-large-folder tool

flax_model/alphafold3/_tools/hmmer/share/man/man1/esl-alimask.1

@@ -0,0 +1,618 @@
+.TH "esl\-alimask" 1 "Aug 2023" "Easel 0.49" "Easel Manual"
+
+.SH NAME
+esl\-alimask \- remove columns from a multiple sequence alignment
+
+.SH SYNOPSIS
+
+.nf
+\fBesl\-alimask \fR[\fIoptions\fR] \fImsafile maskfile\fR
+  (remove columns based on a mask in an input file)
+
+\fBesl\-alimask \-t \fR[\fIoptions\fR] \fImsafile coords\fR
+  (remove a contiguous set of columns at the start and end of an alignment)
+
+\fBesl\-alimask \-g \fR[\fIoptions\fR] \fImsafile\fR
+  (remove columns based on their frequency of gaps)
+
+\fBesl\-alimask \-p \fR[\fIoptions\fR] \fImsafile\fR
+  (remove columns based on their posterior probability annotation)
+
+\fBesl\-alimask \-\-rf\-is\-mask \fR[\fIoptions\fR] \fImsafile\fR
+  (only remove columns that are gaps in the RF annotation)
+
+The \fB\-g\fR and \fB\-p\fR options may be used in combination. 
+.fi
+
+
+.SH DESCRIPTION
+
+.B esl\-alimask
+reads a single input alignment, removes some columns from it
+(i.e. masks it), and outputs the masked alignment.
+
+.PP
+.B esl\-alimask 
+can be run in several different modes.
+
+.PP
+.B esl\-alimask 
+runs in "mask file mode" by default when two
+command-line arguments (\fImsafile\fR and \fImaskfile\fR)
+are supplied. In this mode, a bit-vector mask in the 
+.I maskfile
+defines which columns to keep/remove.  The mask is a string that may
+only contain the characters '0' and '1'. A '0' at position x of the
+mask indicates that column x is excluded by the mask and should be
+removed during masking.  A '1' at position x of the mask indicates
+that column x is included by the mask and should not be removed during
+masking.  All lines in the
+.I maskfile
+that begin with '#' are considered comment lines and are ignored.  All
+non-whitespace characters in non-comment lines are considered to be
+part of the mask. The length of the mask must equal either the total
+number of columns in the (first) alignment in
+.I msafile,
+or the number of columns that are not gaps in the RF annotation of that
+alignment. The latter case is only valid if
+.I msafile
+is in Stockholm format and contains '#=GC RF' annotation. 
+If the mask length is equal to the non-gap RF length, all gap
+RF columns will automatically be removed.
+
+.PP
+.B esl\-alimask 
+runs in "truncation mode" if the 
+.B \-t 
+option is used along with two command line arguments
+(\fImsafile\fR and \fIcoords\fR). In this mode,
+the alignment will be truncated by removing a contiguous set of
+columns from the beginning and end of the alignment. The second
+command line argument is the 
+.I coords
+string, that specifies what range of columns to keep in the
+alignment, all columns outside of this range will be removed.
+The
+.I coords
+string consists of start and end coordinates separated
+by any nonnumeric, nonwhitespace character or characters you like; for
+example,
+.BR 23..100 ,
+.BR 23/100 ,
+or
+.B 23\-100
+all work. To keep all alignment columns beginning at 23 until the
+end of the alignment, you 
+can omit the end; for example,
+.B 23:
+would work.
+If the 
+.B \-\-t\-rf 
+option is used in combination with 
+.B \-t,
+the coordinates in 
+.I coords
+are interpreted as non-gap RF column coordinates. For example,
+with 
+.B \-\-t\-rf, 
+a 
+.I coords 
+string of
+.B 23\-100 
+would remove all columns before the 23rd non-gap residue in
+the "#=GC RF" annotation and after the 100th non-gap RF residue.
+
+.PP
+.B esl\-alimask 
+runs in "RF mask" mode if the
+.B \-\-rf\-is\-mask
+option is enabled. In this mode, the alignment must be in Stockholm
+format and contain '#=GC RF' annotation. 
+.B esl\-alimask
+will simply remove all columns that are gaps in the RF annotation.
+
+.PP
+.B esl\-alimask
+runs in "gap frequency mode" if 
+.B \-g 
+is enabled. In this mode columns for which greater than 
+.I <f>
+fraction of the aligned sequences have gap residues will be removed. 
+By default, 
+.I <f>
+is 0.5, but this value can be changed to 
+.I <f>
+with the 
+.BI \-\-gapthresh " <f>" 
+option. In this mode, if the alignment is in Stockholm format and
+has RF annotation, then all columns that are gaps in the RF annotation
+will automatically be removed, unless
+.B \-\-saveins
+is enabled.
+
+.PP
+.B esl\-alimask
+runs in "posterior probability mode" if 
+.B \-p 
+is enabled. In this mode,  masking is based on posterior probability annotation,
+and the input alignment must be in Stockholm format and contain '#=GR
+PP' (posterior probability) annotation for all sequences. As a special
+case, if 
+.B \-p 
+is used in combination with 
+.B \-\-ppcons,
+then the input alignment need not have '#=GR PP' annotation, but must
+contain '#=GC PP_cons' (posterior probability consensus) annotation.
+
+.PP
+Characters in Stockholm alignment posterior probability annotation
+(both '#=GR PP' and '#=GC PP_cons') can have 12 possible values: the
+ten digits '0-9', '*', and '.'. If '.', the position must correspond to
+a gap in the sequence (for '#=GR PP') or in the RF annotation (for '#=GC
+PP_cons').  A value of '0' indicates a posterior probability of
+between 0.0 and 0.05, '1' indicates between 0.05 and 0.15, '2'
+indicates between 0.15 and 0.25 and so on up to '9' which indicates
+between 0.85 and 0.95. A value of '*' indicates a posterior
+probability of between 0.95 and 1.0. Higher posterior probabilities
+correspond to greater confidence that the aligned residue belongs
+where it appears in the alignment.
+
+.PP
+When
+.B \-p 
+is enabled with 
+.BI \-\-ppcons " <x>",
+columns which have a consensus posterior probability of less than
+.I <x>
+will be removed during masking, and all other columns will not be removed.
+
+.PP
+When
+.B \-p 
+is enabled without
+.B \-\-ppcons,
+the number of each possible PP value in each column is counted. 
+If 
+.I <x>
+fraction of the sequences that contain aligned residues (i.e. do not
+contain gaps) in a column have a posterior probability 
+greater than or equal to 
+.I <y>,
+then that column will not be removed during masking. All columns that
+do not meet this criterion will be removed. By default, the values of both
+.I <x>
+and 
+.I <y>
+are 0.95, but they can be changed with the 
+.BI \-\-pfract " <x>"
+and 
+.BI \-\-pthresh " <y>" 
+options, respectively.
+
+.PP
+In posterior probability mode, all columns that have 0 residues
+(i.e. that are 100% gaps) will be automatically removed, unless the 
+.B \-\-pallgapok
+option is enabled, in which case such columns will not be removed.
+
+.PP
+Importantly, during posterior probability masking, unless
+.B \-\-pavg 
+is used, PP annotation
+values are always considered to be the minimum numerical value in
+their corresponding range. For example, a PP '9' character is converted
+to a numerical posterior probability of 0.85. If
+.B \-\-pavg 
+is used, PP annotation values are considered to be the average
+numerical value in their range. For example, a PP '9' character is
+converted to a numerical posterior probability of 0.90.
+
+.PP
+In posterior probability mode, if the alignment is in Stockholm format and
+has RF annotation, then all columns that are gaps in the RF annotation
+will automatically be removed, unless
+.B \-\-saveins
+is enabled.
+
+.PP
+A single run of
+.B esl\-alimask
+can perform both gap frequency-based masking and posterior
+probability-based masking if both the 
+.B \-g
+and
+.B \-p
+options are enabled. In this case, a gap frequency-based mask and a
+posterior probability-based mask are independently computed.  These
+two masks are combined to create the final mask using a logical 'and'
+operation. Any column that is to be removed by either the gap or PP
+mask will be removed by the final mask.
+
+.PP
+With the
+.B \-\-small
+option, 
+.B esl\-alimask
+will operate in memory saving mode and the required RAM for the masking
+will be minimal (usually less than a Mb) and independent of the
+alignment size. To use 
+.BR \-\-small ,
+the alignment alphabet must be specified with either
+.BR \-\-amino ,
+.BR \-\-dna , 
+or 
+.BR \-\-rna ,
+and the alignment must be in Pfam format (non-interleaved, 1
+line/sequence Stockholm format). Pfam format is the default output
+format of INFERNAL's
+.B cmalign 
+program. Without 
+.B \-\-small
+the required RAM will be equal to roughly the size of the first input
+alignment (the size of the alignment file itself if it only contains
+one alignment).
+
+
+.SH OUTPUT
+
+By default, 
+.B esl\-alimask
+will print only the masked alignment to stdout and then exit.
+If the
+.BI \-o " <f>"
+option is used, the alignment will be saved to file 
+.I <f>
+, and information on the number of columns kept and removed will be
+printed to stdout. If 
+.B \-q
+is used in combination with 
+.BR \-o ,
+nothing is printed to stdout.
+
+.PP
+The mask(s) computed by 
+.B esl\-alimask
+when the 
+.BR \-t ,
+.BR \-p ,
+.BR \-g ,
+or
+.B \-\-rf\-is\-mask
+options are used can be saved to output files using the options
+\fB\-\-fmask\-rf\fR\fI <f>\fR,
+\fB\-\-fmask\-all\fR\fI <f>\fR,
+\fB\-\-gmask\-rf\fR\fI <f>\fR,
+\fB\-\-gmask\-all\fR\fI <f>\fR,
+\fB\-\-pmask\-rf\fR\fI <f>\fR, and 
+\fB\-\-pmask\-all\fR\fI <f>\fR.
+In all cases, 
+.I <f> 
+will contain a single line, a bit vector of length
+.I <n>,
+where 
+.I <n> 
+is the either the total number of columns in the alignment (for the
+options suffixed with 'all') or the number of non-gap columns in the
+RF annotation (for the options suffixed with 'rf'). The mask will be a
+string of '0' and '1' characters: a '0' at position x in the mask
+indicates column x was removed (excluded) by the mask, and a '1' at
+position x indicates column x was kept (included) by the mask. For
+the 'rf' suffixed options, the mask only applies to non-gap RF
+columns.  The options beginning with 'f' will save the 'final' mask
+used to keep/remove columns from the alignment. The options beginning
+with 'g' save the masks based on gap frequency and require
+.BR \-g .
+The options beginning with 'p' save the masks based on posterior
+probabilities and require 
+.BR \-p .
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Print brief help; includes version number and summary of
+all options, including expert options.
+
+.TP
+.BI \-o " <f>"
+Output the final, masked alignment to file 
+.I <f>
+instead of to stdout.
+When this option is used, information about the number of columns
+kept/removed is printed to stdout.
+
+.TP
+.B \-q
+Be quiet; do not print anything to stdout. 
+This option can only be used in combination with the
+.B \-o 
+option.
+
+.TP
+.B \-\-small
+Operate in memory saving mode. Required RAM will be independent of the
+size of the input alignment to mask, instead of roughly the size of the
+input alignment. When enabled, the alignment must be in
+Pfam Stockholm (non-interleaved 1 line/seq) format (see
+.BR esl\-reformat )
+and the output alignment will be in Pfam format.
+
+.TP 
+.BI \-\-informat " <s>"
+Assert that input
+.I msafile
+is in alignment format
+.IR <s> .
+Common choices for 
+.I <s> 
+include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+The string
+.I <s>
+is case-insensitive (\fBa2m\fR or \fBA2M\fR both work).
+Default is 
+.B stockholm
+format, unless
+.B \-\-small
+is used, in which case
+.B pfam
+format (non-interleaved Stockholm) is assumed.
+
+.TP 
+.BI \-\-outformat " <s>"
+Write the output
+.I msafile
+in alignment format
+.IR <s> .
+Common choices for 
+.I <s> 
+include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+The string
+.I <s>
+is case-insensitive (\fBa2m\fR or \fBA2M\fR both work).
+Default is
+.BR stockholm ,
+unless
+.B \-\-small
+is enabled, in which case
+.B pfam
+(noninterleaved Stockholm) is the default output format.
+
+
+.TP 
+.BI \-\-fmask\-rf " <f>"
+Save the non-gap RF-length final mask used to mask the alignment
+to file
+.IR <f> .
+The input alignment must be in Stockholm format and contain '#=GC RF'
+annotation for this option to be valid. See the OUTPUT section above for
+more details on output mask files.
+
+.TP 
+.BI \-\-fmask\-all " <f>"
+Save the full alignment-length final mask used to mask the alignment
+to file
+.IR <f> .
+See the OUTPUT section above for more details on output mask files.
+
+.TP 
+.B \-\-amino
+Specify that the input alignment is a protein alignment.
+By default,
+.B esl\-alimask
+will try to autodetect the alphabet, but if the alignment is
+sufficiently small it may be ambiguous. This option defines the
+alphabet as protein. Importantly, if 
+.B \-\-small
+is enabled, the alphabet must be specified with either
+.BR \-\-amino ,
+.BR \-\-dna ,
+or 
+.BR \-\-rna .
+
+.TP 
+.B \-\-dna
+Specify that the input alignment is a DNA alignment.
+
+.TP 
+.B \-\-rna
+Specify that the input alignment is an RNA alignment. 
+
+.TP 
+.B \-\-t\-rf
+With
+.BR \-t ,
+specify that the start and end coordinates defined in
+the second command line argument 
+.I coords
+correspond to non-gap RF coordinates. To use this option, the
+alignment must be in Stockholm format and have "#=GC RF"
+annotation. See the DESCRIPTION section for an example of using the
+.B \-\-t\-rf
+option.
+
+.TP 
+.B \-\-t\-rmins
+With
+.BR \-t ,
+specify that all columns that are gaps in the reference (RF)
+annotation in between the specified start and end coordinates be
+removed. By default, these columns will be kept.
+To use this option, the alignment must be in  Stockholm format and
+have "#=GC RF" annotation. 
+
+.TP 
+.BI \-\-gapthresh " <x>"
+With
+.BR \-g ,
+specify that a column is kept (included by mask) if no more
+than 
+.I <f>
+fraction of sequences in the alignment have a gap ('.', '\-', or '_')
+at that position. All other columns are removed (excluded by mask).
+By default, 
+.I <x>
+is 0.5.
+
+.TP 
+.BI \-\-gmask\-rf " <f>"
+Save the non-gap RF-length gap frequency-based mask used to mask the alignment
+to file
+.IR <f> .
+The input alignment must be in Stockholm format and contain '#=GC RF'
+annotation for this option to be valid. See the OUTPUT section above for
+more details on output mask files.
+
+.TP 
+.BI \-\-gmask\-all " <f>"
+Save the full alignment-length gap frequency-based mask used to mask the alignment
+to file
+.IR <f> .
+See the OUTPUT section above for more details on output mask files.
+
+
+.TP 
+.BI \-\-pfract " <x>"
+With
+.BR \-p ,
+specify that a column is kept (included by mask) if the
+fraction of sequences with a non-gap residue in that column with a 
+posterior probability of at least 
+.I <y>
+(from \fB\-\-pthresh\fR\fI <y>\fR) is
+.I <x>
+or greater. All other columns are removed (excluded by mask).
+By default 
+.I <x> 
+is 0.95. 
+
+.TP 
+.BI \-\-pthresh " <y>"
+With
+.BR \-p ,
+specify that a column is kept (included by mask) if 
+.I <x>
+(from \fB\-\-pfract \fR\fI<x>\fR)
+fraction of sequences with a non-gap residue in that column have a 
+posterior probability of at least 
+.IR <y> . 
+All other columns are removed (excluded by mask).
+By default 
+.I <y> 
+is 0.95. See the DESCRIPTION section for more on
+posterior probability (PP) masking. 
+Due to the granularity of the PP annotation, different 
+.I <y>
+values within a range covered by a single PP character will be
+have the same effect on masking. For example, using 
+.B \-\-pthresh 0.86 
+will have the same effect as using
+\fB\-\-pthresh 0.94\fR.
+
+.TP 
+.BI \-\-pavg " <x>"
+With
+.BR \-p ,
+specify that a column is kept (included by mask) if 
+the average posterior probability of non-gap residues in that column
+is at least
+.IR <x> .
+See the DESCRIPTION section for more on
+posterior probability (PP) masking. 
+
+.TP 
+.BI \-\-ppcons " <x>"
+With
+.BR \-p ,
+use the '#=GC PP_cons' annotation to define which columns to
+keep/remove. A column is kept (included by mask) if the PP_cons value
+for that column is 
+.I <x>
+or greater. Otherwise it is removed.
+
+.TP 
+.B \-\-pallgapok
+With
+.BR \-p ,
+do not automatically remove any columns that are 100% gaps
+(i.e. contain 0 aligned residues). By default, such columns will be
+removed.
+
+.TP 
+.BI \-\-pmask\-rf " <f>"
+Save the non-gap RF-length posterior probability-based mask used to mask the alignment
+to file
+.IR <f> .
+The input alignment must be in Stockholm format and contain '#=GC RF'
+annotation for this option to be valid. See the OUTPUT section above for
+more details on output mask files.
+
+.TP 
+.BI \-\-pmask\-all " <f>"
+Save the full alignment-length posterior probability-based mask used to mask the alignment
+to file
+.IR <f> .
+See the OUTPUT section above for more details on output mask files.
+
+
+.TP
+.B \-\-keepins 
+If 
+.B \-p 
+and/or
+.B \-g
+is enabled and the alignment is in Stockholm or Pfam format and has '#=GC RF'
+annotation, then allow columns that are gaps in the RF annotation to
+possibly be kept. By default, all gap RF columns would be removed
+automatically, but with this option enabled gap and non-gap RF columns
+are treated identically. 
+To automatically remove all gap RF columns when using a 
+.I maskfile 
+, then define the mask in 
+.I maskfile
+as having length equal to the non-gap RF length in the alignment.
+To automatically remove all gap RF columns when using 
+.B \-t,
+use the
+.B \-\-t\-rmins
+option.
+
+
+
+
+
+
+
+
+.SH SEE ALSO
+
+.nf
+http://bioeasel.org/
+.fi
+
+.SH COPYRIGHT
+
+.nf 
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi 
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi

flax_model/alphafold3/_tools/hmmer/share/man/man1/hmmalign.1

@@ -0,0 +1,185 @@
+.TH "hmmalign" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+hmmalign \- align sequences to a profile 
+
+.SH SYNOPSIS
+.B hmmalign
+[\fIoptions\fR]
+.I hmmfile
+.I seqfile
+
+.SH DESCRIPTION
+
+.PP
+Perform a multiple sequence alignment of all the sequences in
+.I seqfile
+by aligning them individually to the profile HMM in
+.I hmmfile.
+The new alignment is output to stdout.
+
+.PP 
+The
+.I hmmfile
+should contain only a single profile. If it contains more, only the
+first profile in the file will be used. 
+
+.PP
+Either 
+.I hmmfile
+or 
+.I seqfile 
+(but not both) may be '\-' (dash), which
+means reading this input from stdin rather than a file.  
+
+.PP
+The sequences in 
+.I seqfile
+are aligned in unihit local alignment mode.  Therefore they should
+already be known to contain only a single domain (or a fragment of one).
+The optimal alignment
+may assign some residues as nonhomologous (N and C states), in which
+case these residues are still included in the resulting alignment, but
+shoved to the outer edges. To trim these unaligned nonhomologous residues from
+the result, see the
+.B \-\-trim
+option.
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+.TP
+.BI \-o " <f>"
+Direct the output alignment to file
+.I <f>,
+rather than to stdout.
+
+.TP
+.BI \-\-mapali " <f>"
+Merge the existing alignment in file 
+.I <f>
+into the result, where 
+.I <f> 
+is exactly the same alignment that was used to build the model in  
+.I hmmfile.
+This is done using a map of alignment columns to consensus 
+profile positions that is stored in the
+.I hmmfile.
+The multiple alignment in 
+.I <f>
+will be exactly reproduced in its consensus columns (as defined by the
+profile), but the displayed alignment in insert columns may be
+altered, because insertions relative to a profile are considered by
+convention to be unaligned data.
+
+
+.TP
+.B \-\-trim
+Trim nonhomologous residues (assigned to N and C states in the optimal
+alignments) from the resulting multiple alignment output. 
+
+.TP
+.B \-\-amino
+Assert that sequences in 
+.I seqfile
+are protein, bypassing alphabet autodetection.
+
+.TP
+.B \-\-dna
+Assert that sequences in
+.I seqfile
+are DNA, bypassing alphabet autodetection.
+
+.TP
+.B \-\-rna
+Assert that sequences in 
+.I seqfile
+are RNA, bypassing alphabet autodetection.
+
+
+.TP 
+.BI \-\-informat " <s>"
+Assert that input
+.I seqfile
+is in format
+.IR <s> ,
+bypassing format autodetection.
+Common choices for 
+.I <s> 
+include:
+.BR fasta ,
+.BR embl ,
+.BR genbank.
+Alignment formats also work;
+common choices include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+The string
+.I <s>
+is case-insensitive (\fBfasta\fR or \fBFASTA\fR both work).
+
+.TP 
+.BI \-\-outformat " <s>"
+Write the output alignment
+in format
+.IR <s> .
+Common choices for 
+.I <s> 
+include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+The string
+.I <s>
+is case-insensitive (\fBa2m\fR or \fBA2M\fR both work).
+Default is
+.BR stockholm .
+
+
+
+.SH SEE ALSO 
+
+See 
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi

flax_model/alphafold3/_tools/hmmer/share/man/man1/hmmpgmd_shard.1

@@ -0,0 +1,162 @@
+.TH "hmmpgmd_shard" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+hmmpgmd_shard \- sharded daemon for database search web services 
+
+
+.SH SYNOPSIS
+.B hmmpgmd_shard
+[\fIoptions\fR]
+
+
+.SH DESCRIPTION
+
+.PP
+The 
+.B hmmpgmd_shard 
+program provides a sharded version of the 
+.B hmmpgmd 
+program that we use internally to implement high-performance HMMER services that can be accessed via the internet.  See the 
+.B hmmpgmd
+man page for a discussion of how the base 
+.B hmmpgmd
+program is used.  This man page discusses differences between 
+.B hmmpgmd_shard
+and 
+.B hmmpgmd.   
+The base 
+.B hmmpgmd
+program loads the entirety of its database file into RAM on every worker node, in spite of the fact that each worker node searches a predictable fraction of the database(s) contained in that file when performing searches.  This wastes RAM, particularly when many worker nodes are used to accelerate searches of large databases.
+
+.PP
+.B Hmmpgmd_shard 
+addresses this by dividing protein sequence database files into shards.  Each worker node loads only 1/Nth of the database file, where N is the number of worker nodes attached to the master.  HMM database files are not sharded, meaning that every worker node will load the entire database file into RAM.  Current HMM databases are much smaller than current protein sequence databases, and easily fit into the RAM of modern servers even without sharding.
+
+.PP
+.B Hmmpgmd_shard 
+is used in the same manner as 
+.B hmmpgmd
+, except that it takes one additional argument: 
+.BI \-\-num_shards " <n>"
+, which specifies the number of shards that protein databases will be divided into, and defaults to 1 if unspecified.  This argument is only valid for the master node of a 
+.B hmmpgmd
+system (i.e., when 
+.BI \-\-master
+is passed to the 
+.B hmmpgmd
+program), and must be equal to the number of worker nodes that will connect to the master node.  
+.B Hmmpgmd_shard 
+will signal an error if more than 
+.BI num_shards
+worker nodes attempt to connect to the master node or if a search is started when fewer than 
+.BI num_shards 
+workers are connected to the master.
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+.TP 
+.BI \-\-master
+Run as the master server.
+
+.TP
+.BI \-\-worker " <s>"
+Run as a worker, connecting to the master server that is running on IP
+address
+.IR <s> .
+
+.TP 
+.BI \-\-cport " <n>"
+Port to use for communication between clients and the master server. 
+The default is 51371.
+
+.TP 
+.BI \-\-wport " <n>"
+Port to use for communication between workers and the master server. 
+The default is 51372.
+
+.TP 
+.BI \-\-ccncts " <n>"
+Maximum number of client connections to accept. The default is 16.
+
+.TP 
+.BI \-\-wcncts " <n>"
+Maximum number of worker connections to accept. The default is 32.
+
+.TP 
+.BI \-\-pid " <f>"
+Name of file into which the process id will be written. 
+
+.TP 
+.BI \-\-seqdb " <f>"
+Name of the file (in
+.B hmmpgmd
+format) containing protein sequences.
+The contents of this file will be cached for searches. 
+
+.TP 
+.BI \-\-hmmdb " <f>"
+Name of the file containing protein HMMs. The contents of this file 
+will be cached for searches.
+
+.TP 
+.BI \-\-cpu " <n>"
+Number of parallel threads to use (for 
+.B \-\-worker
+).
+
+.TP 
+.BI \-\-num_shards " <n>"
+Number of shards to divide cached sequence database(s) into.  HMM databases are not sharded, due to their small size.
+This option is only valid when the 
+.B \-\-master 
+option is present, and defaults to 1 if not specified.
+.B Hmmpgmd_shard 
+requires that the number of shards be equal to the number of worker nodes, and will give errors if more than 
+.BI num_shards 
+workers attempt to connect to the master node or if a search is started with fewer than 
+.BI num_shards 
+workers connected to the master.
+
+.SH SEE ALSO 
+
+See 
+.BR hmmmpgmd (1)
+for a description of the base hmmpgmd command and how the daemon should be used.
+
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi
+
+
+

flax_model/alphafold3/_tools/hmmer/share/man/man1/hmmstat.1

@@ -0,0 +1,133 @@
+.TH "hmmstat" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+hmmstat \- summary statistics for a profile file
+
+
+.SH SYNOPSIS
+.B hmmstat
+[\fIoptions\fR]
+.I hmmfile
+
+
+.SH DESCRIPTION
+
+The
+.B hmmstat
+utility prints out a tabular file of summary statistics for each
+profile in
+.IR hmmfile .
+
+
+.PP 
+.I hmmfile 
+may be '\-' (a dash character), in which case profiles
+are read from a
+stdin
+pipe instead of from a file.
+
+.PP
+The columns are:
+
+.TP
+.B idx
+The index of this profile, numbering each profile in the file starting from 1.
+
+.TP
+.B name
+The name of the profile.
+
+.TP
+.B accession
+The optional accession of the profile, or "\-" if there is none.
+
+.TP
+.B nseq
+The number of sequences that the profile was estimated from.
+
+.TP
+.B eff_nseq
+The effective number of sequences that the profile was estimated
+from, after HMMER applied an effective sequence number calculation
+such as the default entropy weighting.
+
+.TP
+.B M
+The length of the model in consensus residues (match states).
+
+.TP
+.B relent
+Mean relative entropy per match state, in bits. This is the expected
+(mean) score per consensus position. This is what the default
+entropy-weighting method for effective sequence number estimation
+focuses on, so for default HMMER3 models, you expect this value to
+reflect the default target for entropy-weighting.
+
+.TP
+.B info
+Mean information content per match state, in bits.
+Probably not useful. Information content is a slightly
+different calculation than relative entropy. 
+
+.TP
+.B "p relE"
+Mean positional relative entropy, in bits.
+This is a fancier version of the per-match-state relative entropy,
+taking into account the transition (insertion/deletion) probabilities;
+it may be a more accurate estimation of the average score contributed
+per model consensus position.
+
+.TP
+.B compKL
+Kullback-Leibler divergence from
+the default background frequency distribution to
+the average composition of the profile's consensus match states, in bits.
+The higher this number, the more biased the residue composition of the
+profile is. Highly biased profiles can slow the HMMER3 acceleration
+pipeline, by causing too many nonhomologous sequences to pass the filters.
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+
+.SH SEE ALSO 
+
+See 
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi
+
+
+
+

flax_model/alphafold3/_tools/hmmer/share/man/man1/makehmmerdb.1

@@ -0,0 +1,133 @@
+.TH "makehmmerdb" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+makehmmerdb \- build nhmmer database from a sequence file
+
+
+.SH SYNOPSIS
+.B makehmmerdb
+[\fIoptions\fR]
+.I seqfile
+.I binaryfile
+
+
+.SH DESCRIPTION
+
+.PP
+.B makehmmerdb 
+is used to create a binary file from a DNA sequence file. This 
+binary file may be used as a target database for the DNA search tool
+.BR nhmmer . 
+Using default settings in 
+.BR nhmmer ,
+this yields a roughly 10-fold acceleration with small loss of 
+sensitivity on benchmarks. 
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+
+
+.SH OTHER OPTIONS
+
+.TP
+.BI \-\-informat " <s>"
+Assert that input
+.I seqfile
+is in format
+.IR <s> ,
+bypassing format autodetection.
+Common choices for 
+.I <s> 
+include:
+.BR fasta ,
+.BR embl ,
+.BR genbank.
+Alignment formats also work;
+common choices include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+The string
+.I <s>
+is case-insensitive (\fBfasta\fR or \fBFASTA\fR both work).
+
+
+.TP 
+.BI \-\-bin_length " <n>"
+Bin length. The binary file depends on a data structure called the 
+FM index, which organizes a permuted copy of the sequence in bins 
+of length
+.IR <n> .
+Longer bin length will lead to smaller files (because data is 
+captured about each bin) and possibly slower query time. The 
+default is 256. Much more than 512 may lead to notable reduction 
+in speed.
+
+
+.TP 
+.BI \-\-sa_freq " <n>"
+Suffix array sample rate. The FM index structure also samples from 
+the underlying suffix array for the sequence database. More frequent 
+sampling (smaller value for 
+.IR <n> )
+will yield larger file size and faster search (until file size becomes
+large enough to cause I/O to be a bottleneck). The default value
+is 8. Must be a power of 2.
+
+
+.TP 
+.BI \-\-block_size " <n>"
+The input sequence is broken into blocks of size
+.I <n>
+million letters. An FM index is built for each block, rather than 
+building an FM index for the entire sequence database. Default is 
+50. Larger blocks do not seem to yield substantial speed increase. 
+
+
+
+.SH SEE ALSO 
+
+See 
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi
+
+
+

flax_model/alphafold3/_tools/hmmer/share/man/man1/nhmmer.1

@@ -0,0 +1,687 @@
+.TH "nhmmer" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+nhmmer \- search DNA queries against a DNA sequence database
+
+
+.SH SYNOPSIS
+.B nhmmer
+[\fIoptions\fR]
+.I queryfile
+.I seqdb
+
+
+.SH DESCRIPTION
+
+.PP
+.B nhmmer 
+is used to search one or more nucleotide queries against a 
+nucleotide sequence database.
+For each query in 
+.IR queryfile ,
+use that query to search the target database of sequences in
+.IR seqdb ,
+and output a ranked list of the hits with the most significant
+matches to the query. A query may be either a profile model 
+built using 
+.BR hmmbuild , 
+a sequence alignment, or a single sequence. Sequence based
+queries can be in a number of formats (see
+.BR \-\-qformat ), 
+and can typically be autodetected. Note that only 
+Stockholm
+format supports queries made up of more than one sequence 
+alignment.
+
+
+
+.PP 
+Either the query
+.I queryfile 
+or the target
+.I seqdb 
+may be '\-' (a dash character), in which case
+the query file or target database input will be read from a <stdin> pipe instead of from a
+file. Only one input source can come through <stdin>, not both.
+If the
+.I queryfile 
+contains more than one query, then
+.I seqdb 
+cannot come from stdin, because we can't rewind the
+streaming target database to search it with another profile. 
+
+.PP
+If the query is sequence-based (unaligned or aligned),
+a new file containing the HMM(s) built from the input(s) in 
+.I queryfile
+may optionally be produced, with the filename set using the 
+.B \-\-hmmout
+flag.
+
+
+.PP
+The output format is designed to be human-readable, but is often so
+voluminous that reading it is impractical, and parsing it is a pain. The
+.B \-\-tblout 
+option saves output in a simple tabular format that is concise and
+easier to parse.
+The 
+.B \-o
+option allows redirecting the main output, including throwing it away
+in /dev/null.
+
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+
+
+.SH OPTIONS FOR CONTROLLING OUTPUT
+
+.TP 
+.BI \-o " <f>"
+Direct the main human-readable output to a file
+.I <f> 
+instead of the default stdout.
+
+.TP
+.BI \-A " <f>"
+Save a multiple alignment of all significant hits (those satisfying
+"inclusion thresholds") to the file 
+.IR <f> .
+
+.TP 
+.BI \-\-tblout " <f>"
+Save a simple tabular (space-delimited) file summarizing the
+per-target output, with one data line per homologous target sequence
+found.
+
+.TP 
+.BI \-\-dfamtblout " <f>"
+Save a tabular (space-delimited) file summarizing the 
+per-hit output, similar to 
+.B \-\-tblout
+but more succinct. 
+
+.TP 
+.BI \-\-aliscoresout " <f>" 
+Save to file a list of per-position scores for each hit.
+This is useful, for example, in identifying regions of high
+score density for use in resolving overlapping hits from 
+different models.
+
+.TP 
+.BI \-\-hmmout " <f>" 
+If
+.I queryfile
+is sequence-based, write the internally-computed HMM(s) to file
+.IR <f> . 
+
+
+.TP 
+.B \-\-acc
+Use accessions instead of names in the main output, where available
+for profiles and/or sequences.
+
+.TP 
+.B \-\-noali
+Omit the alignment section from the main output. This can greatly
+reduce the output volume.
+
+.TP 
+.B \-\-notextw
+Unlimit the length of each line in the main output. The default
+is a limit of 120 characters per line, which helps in displaying
+the output cleanly on terminals and in editors, but can truncate
+target profile description lines.
+
+.TP 
+.BI \-\-textw " <n>"
+Set the main output's line length limit to
+.I <n>
+characters per line. The default is 120.
+
+
+
+.SH OPTIONS CONTROLLING SINGLE SEQUENCE SCORING
+
+By default, if a query is a single sequence from a file in 
+fasta format,
+.B nhmmer 
+uses a search model constructed from that sequence and a standard
+20x20 substitution matrix for residue probabilities, along with two
+additional parameters for position-independent gap open and gap extend
+probabilities. These options allow the default single-sequence scoring
+parameters to be changed, and for single-sequence scoring options to
+be applied to a single sequence coming from an aligned format.
+
+.TP
+.BI \-\-singlemx 
+If a single sequence query comes from a multiple sequence alignment file, 
+such as in Stockholm format, the search model is by default constructed as is typically done 
+for multiple sequence alignments. This option forces 
+.B nhmmer 
+to use the single-sequence method with substitution score matrix.
+
+.TP
+.BI \-\-mxfile "<mxfile"
+Obtain residue alignment probabilities from the substitution matrix
+in file
+.IR mxfile .
+The default score matrix is DNA1 (this matrix is internal to
+HMMER and does not have to be available as a file). 
+The format of a substitution matrix
+.I mxfile
+is the standard format accepted by BLAST, FASTA, and other sequence 
+analysis software.
+See
+.B ftp.ncbi.nlm.nih.gov/blast/matrices/
+for example files. (The only
+exception: we require matrices to be square, so for DNA, use files
+like NCBI's NUC.4.4, not NUC.4.2.)
+
+
+.TP
+.BI \-\-popen " <x>"
+Set the gap open probability for a single sequence query model to 
+.IR <x> .
+The default is 0.02. 
+.I <x> 
+must be >= 0 and < 0.5.
+
+.TP
+.BI \-\-pextend " <x>"
+Set the gap extend probability for a single sequence query model to 
+.IR <x> .
+The default is 0.4. 
+.I <x> 
+must be >= 0 and < 1.0.
+
+
+
+.SH OPTIONS CONTROLLING REPORTING THRESHOLDS
+
+Reporting thresholds control which hits are reported in output files
+(the main output,
+.BR \-\-tblout ,
+and 
+.BR \-\-dfamtblout ).
+Hits are ranked by statistical significance (E-value). 
+
+
+.TP
+.BI \-E " <x>"
+Report target sequences with an E-value of <=
+.IR <x> . 
+The default is 10.0, meaning that on average, about 10 false positives
+will be reported per query, so you can see the top of the noise
+and decide for yourself if it's really noise.
+
+.TP
+.BI \-T " <x>"
+Instead of thresholding output on E-value, instead
+report target sequences with a bit score of >=
+.IR <x> .
+
+
+
+
+.SH OPTIONS FOR INCLUSION THRESHOLDS
+
+Inclusion thresholds are stricter than reporting thresholds.
+Inclusion thresholds control which hits are considered to be reliable
+enough to be included in an output alignment or a subsequent search
+round, or marked as significant ("!") as opposed to questionable ("?")
+in hit output.
+
+.TP
+.BI \-\-incE " <x>"
+Use an E-value of <=
+.I <x>
+as the inclusion threshold.
+The default is 0.01, meaning that on average, about 1 false positive
+would be expected in every 100 searches with different query
+sequences.
+
+.TP
+.BI \-\-incT " <x>"
+Instead of using E-values for setting the inclusion threshold, 
+use a bit score of >= 
+.I <x>
+as the inclusion threshold.
+By default this option is unset.
+
+
+
+.SH OPTIONS FOR MODEL-SPECIFIC SCORE THRESHOLDING
+
+Curated profile databases may define specific bit score thresholds for
+each profile, superseding any thresholding based on statistical
+significance alone.
+
+To use these options, the profile must contain the appropriate (GA,
+TC, and/or NC) optional score threshold annotation; this is picked up
+by 
+.B hmmbuild
+from Stockholm format alignment files. For a nucleotide model, each 
+thresholding option has a single per-hit threshold <x>
+This acts as if
+.BI \-T " <x>"
+.BI \-\-incT " <x>"
+has been applied specifically using each model's curated thresholds.
+
+.TP
+.B \-\-cut_ga
+Use the GA (gathering) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. GA thresholds are generally considered to be the
+reliable curated thresholds defining family membership; for example,
+in Dfam, these thresholds are applied when annotating a genome
+with a model of a family known to be found in that organism. They
+may allow for minimal expected false discovery rate.
+
+.TP
+.B \-\-cut_nc
+Use the NC (noise cutoff) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. NC thresholds are less stringent than GA; in the context
+of Pfam, they are generally used to store the score of the 
+highest-scoring known false positive.
+
+.TP
+.B \-\-cut_tc
+Use the TC (trusted cutoff) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. TC thresholds are more stringent than GA, and are
+generally considered to be the score of the lowest-scoring known 
+true positive that is above all known false positives; for example,
+in Dfam, these thresholds are applied when annotating a genome
+with a model of a family not known to be found in that organism.
+
+
+
+
+.SH OPTIONS CONTROLLING THE ACCELERATION PIPELINE
+
+HMMER3 searches are accelerated in a three-step filter pipeline: the
+scanning-SSV filter, the Viterbi filter, and the Forward filter. The 
+first filter is the fastest and most approximate; the last is the full
+Forward scoring algorithm. There is also a bias filter step between
+SSV and Viterbi. Targets that pass all the steps in the acceleration
+pipeline are then subjected to postprocessing -- domain
+identification and scoring using the Forward/Backward algorithm.
+
+Changing filter thresholds only removes or includes targets from
+consideration; changing filter thresholds does not alter bit scores,
+E-values, or alignments, all of which are determined solely in
+postprocessing.
+
+.TP
+.B \-\-max
+Turn off (nearly) all filters, including the bias filter, and run full
+Forward/Backward postprocessing on most of the target sequence. 
+In contrast to 
+.B phmmer
+and
+.BR hmmsearch ,
+where this flag really does turn off the filters entirely, the
+.B \-\-max
+flag in 
+.B nhmmer
+sets the scanning-SSV filter threshold to 0.4, not 1.0. Use of this
+flag increases sensitivity somewhat, at a large cost in speed.
+
+.TP
+.BI \-\-F1 " <x>"
+Set the P-value threshold for the SSV filter step.  The default is
+0.02, meaning that roughly 2% of the highest scoring nonhomologous
+targets are expected to pass the filter.
+
+.TP
+.BI \-\-F2 " <x>"
+Set the P-value threshold for the Viterbi filter step.
+The default is 0.001. 
+
+.TP
+.BI \-\-F3 " <x>"
+Set the P-value threshold for the Forward filter step.
+The default is 1e-5.
+
+.TP
+.B \-\-nobias
+Turn off the bias filter. This increases sensitivity somewhat, but can
+come at a high cost in speed, especially if the query has biased
+residue composition (such as a repetitive sequence region, or if it is
+a membrane protein with large regions of hydrophobicity). Without the
+bias filter, too many sequences may pass the filter with biased
+queries, leading to slower than expected performance as the
+computationally intensive Forward/Backward algorithms shoulder an
+abnormally heavy load.
+
+
+
+.SH OPTIONS FOR SPECIFYING THE ALPHABET
+
+.TP
+.B \-\-dna
+Assert that sequences in
+.I msafile
+are DNA, bypassing alphabet autodetection.
+
+.TP
+.B \-\-rna
+Assert that sequences in 
+.I msafile
+are RNA, bypassing alphabet autodetection.
+
+
+
+.SH OPTIONS CONTROLLING SEED SEARCH HEURISTIC
+
+When searching with 
+.BR nhmmer ,
+one may optionally precompute a binary version of the target database, using
+.BR makehmmerdb ,
+then search against that database. Using default settings, this yields a
+roughly 10-fold acceleration with small loss of sensitivity on benchmarks.
+This is achieved using a heuristic method that searches for seeds (ungapped 
+alignments) around which full processing is done. This is essentially  
+a replacement to the SSV stage. (This method has been extensively tested, 
+but should still be treated as somewhat experimental.)
+The following options only impact 
+.B nhmmer
+if the value of 
+.B \-\-tformat
+is
+.BR hmmerdb .
+
+Changing parameters for this seed-finding step will impact both speed and 
+sensitivity - typically faster search leads to lower sensitivity. 
+
+.TP
+.BI \-\-seed_max_depth " <n>"
+The seed step requires that a seed reach a specified bit score in length 
+no longer than 
+.IR <n> . 
+By default, this value is 15. Longer seeds allow a greater chance of 
+meeting the bit score threshold, leading to diminished filtering
+(greater sensitivity, slower run time).
+
+.TP
+.BI \-\-seed_sc_thresh " <x>"
+The seed must reach score 
+.I <x>
+(in bits). The default is 15.0 bits. A higher threshold increases 
+filtering stringency, leading to faster run times and lower 
+sensitivity.
+
+.TP
+.BI \-\-seed_sc_density " <x>"
+Either all prefixes or all suffixes of a seed must have 
+bit density (bits per aligned position) of at least 
+.IR <x> . 
+The default is 0.8 bits/position. An increase in the density 
+requirement leads to increased filtering stringency, thus faster 
+run times and lower sensitivity.
+
+.TP
+.BI \-\-seed_drop_max_len " <n>"
+A seed may not have a run of length
+.I <n>
+in which the score drops by 
+.B \-\-seed_drop_lim
+or more. Basically, this prunes seeds that go through long
+slightly-negative seed extensions. The default is 4.  Increasing 
+the limit causes (slightly) diminished filtering efficiency, thus 
+slower run times and higher sensitivity. (minor tuning option)
+
+.TP
+.BI \-\-seed_drop_lim " <x>"
+In a seed, there may be no run of length 
+.B \-\-seed_drop_max_len
+in which the score drops by 
+.BR \-\-seed_drop_lim .
+The default is 0.3 bits. Larger numbers mean less filtering.
+(minor tuning option)
+
+.TP
+.BI \-\-seed_req_pos " <n>"
+A seed must contain a run of at least 
+.I <n>
+positive-scoring matches. The default is 5. Larger values mean
+increased filtering.
+(minor tuning option)
+
+.TP
+.BI \-\-seed_ssv_length " <n>"
+After finding a short seed, an ungapped alignment is extended 
+in both directions in an attempt to meet the 
+.B \-\-F1
+score threshold. The window through which this ungapped alignment
+extends is length 
+.IR <n> .
+The default is 70.  
+Decreasing this value slightly reduces run time, at a small risk of
+reduced sensitivity. (minor tuning option)
+
+
+.SH OTHER OPTIONS
+
+
+.TP
+.BI \-\-qformat " <s>"
+Assert that input
+.I queryfile
+is a sequence file (unaligned or aligned), in format
+.IR <s> ,
+bypassing format autodetection.
+Common choices for 
+.I <s> 
+include:
+.BR fasta ,
+.BR embl ,
+.BR genbank.
+Alignment formats also work, and will serve as the
+basis for automatic creation of a profile HMM used for
+searching;
+common choices include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+
+
+.TP
+.BI \-\-qsingle_seqs
+Force
+.I queryfile
+to be read as individual sequences, even if it is in
+an msa format. For example, if the input is in aligned
+.BR stockholm
+format, the
+.BR \-\-qsingle_seqs
+ flag will cause each sequence in that alignment to be used as a separate query sequence.
+
+.TP
+.BI \-\-tformat " <s>"
+Assert that target sequence database
+.I seqdb
+is in format
+.IR <s> ,
+bypassing format autodetection.
+Common choices for 
+.I <s> 
+include:
+.BR fasta ,
+.BR embl ,
+.BR genbank ,
+.BR fmindex .
+Alignment formats also work;
+common choices include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+The string
+.I <s>
+is case-insensitive (\fBfasta\fR or \fBFASTA\fR both work).
+The format
+.B fmindex
+indicates that the database file is a binary file produced using
+.BR makehmmerdb . 
+
+
+.TP
+.B \-\-nonull2
+Turn off the null2 score corrections for biased composition.
+
+.TP
+.BI \-Z " <x>"
+For the purposes of per-hit E-value calculations,
+Assert that the total size of the target database is
+.I <x>
+million nucleotides, 
+rather than the actual number of targets seen. 
+
+
+.TP
+.BI \-\-seed " <n>"
+Set the random number seed to 
+.IR <n> .
+Some steps in postprocessing require Monte Carlo simulation.  The
+default is to use a fixed seed (42), so that results are exactly
+reproducible. Any other positive integer will give different (but also
+reproducible) results. A choice of 0 uses a randomly chosen seed.
+
+
+.TP 
+.BI \-\-w_beta " <x>"
+Window length tail mass.
+The upper bound,
+.IR W ,
+on the length at which nhmmer expects to find an instance of the 
+model is set such that the fraction of all sequences generated
+by the model with length >= W is less than  
+.IR <x> . 
+The default is 1e-7. 
+This flag may be used to override the value of 
+.I W
+established for the model by 
+.BR hmmbuild ,
+or when the query is sequence-based.
+
+
+
+.TP 
+.BI \-\-w_length " <n>"
+Override the model instance length upper bound, W,
+which is otherwise controlled by
+.BR \-\-w_beta . 
+It should be larger than the model length. The value of W
+is used deep in the acceleration pipeline, and modest changes
+are not expected to impact results (though larger values of W
+do lead to longer run time). 
+This flag may be used to override the value of 
+W established for the model by 
+.BR hmmbuild ,
+or when the query is sequence-based.
+
+
+
+.TP 
+.B \-\-watson 
+Only search the top strand. By default both the query sequence
+and its reverse-complement are searched.
+
+.TP 
+.B \-\-crick 
+Only search the bottom (reverse-complement) strand. By 
+default both the query sequence and its reverse-complement are searched.
+
+
+.TP
+.BI \-\-cpu " <n>"
+Set the number of parallel worker threads to 
+.IR <n> .
+On multicore machines, the default is 2.
+You can also control this number by setting an environment variable, 
+.IR HMMER_NCPU .
+There is also a master thread, so the actual number of threads that
+HMMER spawns is
+.IR <n> +1.
+
+This option is not available if HMMER was compiled with POSIX threads
+support turned off.
+
+
+
+
+.TP
+.BI \-\-stall
+For debugging the MPI master/worker version: pause after start, to
+enable the developer to attach debuggers to the running master and
+worker(s) processes. Send SIGCONT signal to release the pause.
+(Under gdb: 
+.BR "(gdb) signal SIGCONT" )
+(Only available if optional MPI support was enabled at compile-time.)
+
+.TP
+.BI \-\-mpi
+Run under MPI control with master/worker parallelization (using
+.BR mpirun ,
+for example, or equivalent). Only available if optional MPI support
+was enabled at compile-time.
+
+
+
+
+
+
+.SH SEE ALSO 
+
+See 
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi
+
+
+

flax_model/alphafold3/_tools/hmmer/share/man/man1/nhmmscan.1

@@ -0,0 +1,472 @@
+.TH "nhmmscan" 1 "Aug 2023" "HMMER 3.4" "HMMER Manual"
+
+.SH NAME
+nhmmscan \- search DNA sequence(s) against a DNA profile database
+
+
+.SH SYNOPSIS
+.B nhmmscan
+[\fIoptions\fR]
+.I hmmdb
+.I seqfile
+
+
+
+.SH DESCRIPTION
+
+.PP
+.B nhmmscan 
+is used to search nucleotide sequences against collections 
+of nucleotide profiles. For each sequence in 
+.IR seqfile ,
+use that query sequence to search the target database of
+profiles in
+.IR hmmdb ,
+and output ranked lists of the profiles with the
+most significant matches to the sequence.
+
+.PP
+The 
+.I seqfile 
+may contain more than one query sequence. It can be in FASTA format,
+or several other common sequence file formats (genbank, embl, and
+uniprot, among others), or in alignment file formats (stockholm,
+aligned fasta, and others). See the
+.I \-\-qformat 
+option for a complete list.
+
+.PP
+The
+.I hmmdb
+needs to be press'ed using 
+.B hmmpress
+before it can be searched with 
+.BR nhmmscan . 
+This creates four binary files,
+suffixed
+.B .h3{fimp}.
+
+.PP 
+The query
+.I seqfile 
+may be '\-' (a dash character), in which case
+the query sequences are read from a stdin pipe instead of from a
+file.
+The
+.I hmmdb 
+cannot be read from a stdin stream, because it needs to have
+the four auxiliary binary files generated by 
+.BR hmmpress .
+
+.PP
+The output format is designed to be human-readable, but is often so
+voluminous that reading it is impractical, and parsing it is a pain. The
+.B \-\-tblout 
+option saves output in a simple tabular format that is concise and
+easier to parse. 
+The 
+.B \-o
+option allows redirecting the main output, including throwing it away
+in /dev/null.
+
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Help; print a brief reminder of command line usage and all available
+options.
+
+
+
+.SH OPTIONS FOR CONTROLLING OUTPUT
+
+.TP 
+.BI \-o " <f>"
+Direct the main human-readable output to a file
+.I <f> 
+instead of the default stdout.
+
+.TP 
+.BI \-\-tblout " <f>"
+Save a simple tabular (space-delimited) file summarizing the
+per-hit output, with one data line per homologous target model 
+hit found.
+
+.TP 
+.BI \-\-dfamtblout " <f>"
+Save a tabular (space-delimited) file summarizing the 
+per-hit output, similar to 
+.B \-\-tblout
+but more succinct. 
+
+.TP 
+.BI \-\-aliscoresout " <f>" 
+Save to file a list of per-position scores for each hit.
+This is useful, for example, in identifying regions of high
+score density for use in resolving overlapping hits from 
+different models.
+
+
+.TP 
+.B \-\-acc
+Use accessions instead of names in the main output, where available
+for profiles and/or sequences.
+
+.TP 
+.B \-\-noali
+Omit the alignment section from the main output. This can greatly
+reduce the output volume.
+
+.TP 
+.B \-\-notextw
+Unlimit the length of each line in the main output. The default
+is a limit of 120 characters per line, which helps in displaying
+the output cleanly on terminals and in editors, but can truncate
+target profile description lines.
+
+.TP 
+.BI \-\-textw " <n>"
+Set the main output's line length limit to
+.I <n>
+characters per line. The default is 120.
+
+
+
+.SH OPTIONS FOR REPORTING THRESHOLDS
+
+Reporting thresholds control which hits are reported in output files
+(the main output,
+.BR \-\-tblout ,
+and 
+.BR \-\-dfamtblout ).
+Hits are ranked by statistical significance (E-value). 
+
+.TP
+.BI \-E " <x>"
+Report target profiles with an E-value of <=
+.IR <x> . 
+The default is 10.0, meaning that on average, about 10 false positives
+will be reported per query, so you can see the top of the noise
+and decide for yourself if it's really noise.
+
+.TP
+.BI \-T " <x>"
+Instead of thresholding output on E-value, instead
+report target profiles with a bit score of >=
+.IR <x> .
+
+
+
+
+.SH OPTIONS FOR INCLUSION THRESHOLDS
+
+Inclusion thresholds are stricter than reporting thresholds.
+Inclusion thresholds control which hits are considered to be
+reliable enough
+to be included in an output alignment or a subsequent search round.
+In 
+.BR nhmmscan , 
+which does not have any alignment output (like 
+.BR nhmmer ),
+inclusion thresholds have little effect. They only affect what hits
+get marked as significant (!) or questionable (?) in hit
+output. 
+
+.TP
+.BI \-\-incE " <x>"
+Use an E-value of <=
+.I <x>
+as the inclusion threshold.
+The default is 0.01, meaning that on average, about 1 false positive
+would be expected in every 100 searches with different query
+sequences.
+
+.TP
+.BI \-\-incT " <x>"
+Instead of using E-values for setting the inclusion threshold, 
+use a bit score of >= 
+.I <x>
+as the inclusion threshold.
+It would be unusual to use bit score thresholds with
+.BR hmmscan ,
+because you don't expect a single score threshold to work for
+different profiles; different profiles have slightly different
+expected score distributions.
+
+
+
+.SH OPTIONS FOR MODEL-SPECIFIC SCORE THRESHOLDING
+
+Curated profile databases may define specific bit score thresholds for
+each profile, superseding any thresholding based on statistical
+significance alone.
+
+To use these options, the profile must contain the appropriate (GA,
+TC, and/or NC) optional score threshold annotation; this is picked up
+by 
+.B hmmbuild
+from Stockholm format alignment files. For a nucleotide model, each 
+thresholding option has a single per-hit threshold <x>
+This acts as if
+.BI \-T " <x>"
+.BI \-\-incT " <x>"
+has been applied specifically using each model's curated thresholds.
+
+.TP
+.B \-\-cut_ga
+Use the GA (gathering) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. GA thresholds are generally considered to be the
+reliable curated thresholds defining family membership; for example,
+in Dfam, these thresholds are applied when annotating a genome
+with a model of a family known to be found in that organism. They
+may allow for minimal expected false discovery rate.
+
+.TP
+.B \-\-cut_nc
+Use the NC (noise cutoff) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. NC thresholds are less stringent than GA; in the context
+of Pfam, they are generally used to store the score of the 
+highest-scoring known false positive.
+
+.TP
+.B \-\-cut_tc
+Use the TC (trusted cutoff) bit score threshold in the model to set
+per-hit reporting and inclusion
+thresholds. TC thresholds are more stringent than GA, and are
+generally considered to be the score of the lowest-scoring known 
+true positive that is above all known false positives; for example,
+in Dfam, these thresholds are applied when annotating a genome
+with a model of a family not known to be found in that organism.
+
+
+
+.SH CONTROL OF THE ACCELERATION PIPELINE
+
+HMMER3 searches are accelerated in a three-step filter pipeline: the
+scanning-SSV filter, the Viterbi filter, and the Forward filter. The 
+first filter is the fastest and most approximate; the last is the full
+Forward scoring algorithm. There is also a bias filter step between
+SSV and Viterbi. Targets that pass all the steps in the acceleration
+pipeline are then subjected to postprocessing -- domain
+identification and scoring using the Forward/Backward algorithm.
+
+Changing filter thresholds only removes or includes targets from
+consideration; changing filter thresholds does not alter bit scores,
+E-values, or alignments, all of which are determined solely in
+postprocessing.
+
+.TP
+.B \-\-max
+Turn off (nearly) all filters, including the bias filter, and run full
+Forward/Backward postprocessing on most of the target sequence.
+In contrast to  
+.B hmmscan,
+where this flag really does turn off the filters entirely, the 
+.B \-\-max
+flag in 
+.B nhmmscan
+sets the scanning-SSV filter threshold to 0.4, not 1.0. Use of this
+flag increases sensitivity somewhat, at a large cost in speed.
+
+.TP
+.BI \-\-F1 " <x>"
+Set the P-value threshold for the MSV filter step.  The default is
+0.02, meaning that roughly 2% of the highest scoring nonhomologous
+targets are expected to pass the filter.
+
+.TP
+.BI \-\-F2 " <x>"
+Set the P-value threshold for the Viterbi filter step.
+The default is 0.001. 
+
+.TP
+.BI \-\-F3 " <x>"
+Set the P-value threshold for the Forward filter step.
+The default is 1e-5.
+
+.TP
+.B \-\-nobias
+Turn off the bias filter. This increases sensitivity somewhat, but can
+come at a high cost in speed, especially if the query has biased
+residue composition (such as a repetitive sequence region, or if it is
+a membrane protein with large regions of hydrophobicity). Without the
+bias filter, too many sequences may pass the filter with biased
+queries, leading to slower than expected performance as the
+computationally intensive Forward/Backward algorithms shoulder an
+abnormally heavy load.
+
+
+
+.SH OTHER OPTIONS
+
+.TP
+.B \-\-nonull2
+Turn off the null2 score corrections for biased composition.
+
+.TP
+.BI \-Z " <x>"
+Assert that the total number of targets in your searches is
+.IR <x> ,
+for the purposes of per-sequence E-value calculations,
+rather than the actual number of targets seen. 
+
+.TP
+.BI \-\-seed " <n>"
+Set the random number seed to 
+.IR <n> .
+Some steps in postprocessing require Monte Carlo simulation.  The
+default is to use a fixed seed (42), so that results are exactly
+reproducible. Any other positive integer will give different (but also
+reproducible) results. A choice of 0 uses an arbitrarily chosen seed.
+
+.TP
+.BI \-\-qformat " <s>"
+Assert that input query
+.I seqfile
+is in format
+.IR <s> ,
+bypassing format autodetection.
+Common choices for 
+.I <s> 
+include:
+.BR fasta ,
+.BR embl ,
+.BR genbank.
+Alignment formats also work;
+common choices include:
+.BR stockholm , 
+.BR a2m ,
+.BR afa ,
+.BR psiblast ,
+.BR clustal ,
+.BR phylip .
+For more information, and for codes for some less common formats,
+see main documentation.
+The string
+.I <s>
+is case-insensitive (\fBfasta\fR or \fBFASTA\fR both work).
+
+
+.TP 
+.BI \-\-w_beta " <x>"
+Window length tail mass.
+The upper bound, W, 
+on the length at which nhmmer expects to find an instance of the 
+model is set such that the fraction of all sequences generated
+by the model with length >= W is less than  
+.IR <x> . 
+The default is 1e-7. 
+This flag may be used to override the value of W
+established for the model by 
+.BR hmmbuild .
+
+
+
+.TP 
+.BI \-\-w_length " <n>"
+Override the model instance length upper bound, W,
+which is otherwise controlled by
+.BR \-\-w_beta . 
+It should be larger than the model length. The value of  W
+is used deep in the acceleration pipeline, and modest changes
+are not expected to impact results (though larger values of W
+do lead to longer run time). 
+This flag may be used to override the value of W
+established for the model by 
+.BR hmmbuild .
+
+
+.TP 
+.B \-\-watson 
+Only search the top strand. By default both the query sequence
+and its reverse-complement are searched.
+
+.TP 
+.B \-\-crick 
+Only search the bottom (reverse-complement) strand. By 
+default both the query sequence and its reverse-complement are searched.
+
+
+.TP
+.BI \-\-cpu " <n>"
+Set the number of parallel worker threads to 
+.IR <n> .
+The default is 0, meaning off (no thread-level parallelization), because
+.B nhmmscan
+is typically i/o bound and the extra overhead of our current
+multithreaded
+implementation isn't worthwhile.
+You can also control this number by setting an environment variable, 
+.IR HMMER_NCPU .
+There is also a master thread, so the actual number of threads that
+HMMER spawns is at least
+.IR <n> +1.
+
+This option is not available if HMMER was compiled with POSIX threads
+support turned off.
+
+
+
+
+.TP
+.BI \-\-stall
+For debugging the MPI master/worker version: pause after start, to
+enable the developer to attach debuggers to the running master and
+worker(s) processes. Send SIGCONT signal to release the pause.
+(Under gdb: 
+.BR "(gdb) signal SIGCONT" )
+
+(Only available if optional MPI support was enabled at compile-time.)
+
+.TP
+.BI \-\-mpi
+Run under MPI control with master/worker parallelization (using
+.BR mpirun ,
+for example, or equivalent). Only available if optional MPI support
+was enabled at compile-time.
+
+
+
+
+
+
+
+
+
+.SH SEE ALSO 
+
+See 
+.BR hmmer (1)
+for a master man page with a list of all the individual man pages
+for programs in the HMMER package.
+
+.PP
+For complete documentation, see the user guide that came with your
+HMMER distribution (Userguide.pdf); or see the HMMER web page
+(http://hmmer.org/).
+
+
+
+.SH COPYRIGHT
+
+.nf
+Copyright (C) 2023 Howard Hughes Medical Institute.
+Freely distributed under the BSD open source license.
+.fi
+
+For additional information on copyright and licensing, see the file
+called COPYRIGHT in your HMMER source distribution, or see the HMMER
+web page 
+(http://hmmer.org/).
+
+
+.SH AUTHOR
+
+.nf
+http://eddylab.org
+.fi
+

flax_model/alphafold3/data/cpp/msa_profile_pybind.cc

@@ -0,0 +1,79 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <algorithm>
+
+#include "absl/strings/str_cat.h"
+#include "pybind11/cast.h"
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+
+namespace {
+
+namespace py = pybind11;
+
+py::array_t<double> ComputeMsaProfile(
+    const py::array_t<int, py::array::c_style>& msa, int num_residue_types) {
+  if (msa.size() == 0) {
+    throw py::value_error("The MSA must be non-empty.");
+  }
+  if (msa.ndim() != 2) {
+    throw py::value_error(absl::StrCat("The MSA must be rectangular, got ",
+                                       msa.ndim(), "-dimensional MSA array."));
+  }
+  const int msa_depth = msa.shape()[0];
+  const int sequence_length = msa.shape()[1];
+
+  py::array_t<double> profile({sequence_length, num_residue_types});
+  std::fill(profile.mutable_data(), profile.mutable_data() + profile.size(),
+            0.0f);
+  auto profile_unchecked = profile.mutable_unchecked<2>();
+
+  const double normalized_count = 1.0 / msa_depth;
+  const int* msa_it = msa.data();
+  for (int row_index = 0; row_index < msa_depth; ++row_index) {
+    for (int column_index = 0; column_index < sequence_length; ++column_index) {
+      const int residue_code = *(msa_it++);
+      if (residue_code < 0 || residue_code >= num_residue_types) {
+        throw py::value_error(
+            absl::StrCat("All residue codes must be positive and smaller than "
+                         "num_residue_types ",
+                         num_residue_types, ", got ", residue_code));
+      }
+      profile_unchecked(column_index, residue_code) += normalized_count;
+    }
+  }
+  return profile;
+}
+
+constexpr char kComputeMsaProfileDoc[] = R"(
+Computes MSA profile for the given encoded MSA.
+
+Args:
+  msa: A Numpy array of shape (num_msa, num_res) with the integer coded MSA.
+  num_residue_types: Integer that determines the number of unique residue types.
+    This will determine the shape of the output profile.
+
+Returns:
+  A float Numpy array of shape (num_res, num_residue_types) with residue
+  frequency (residue type count normalized by MSA depth) for every column of the
+  MSA.
+)";
+
+}  // namespace
+
+namespace alphafold3 {
+
+void RegisterModuleMsaProfile(pybind11::module m) {
+  m.def("compute_msa_profile", &ComputeMsaProfile, py::arg("msa"),
+        py::arg("num_residue_types"), py::doc(kComputeMsaProfileDoc + 1));
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/data/msa.py

@@ -0,0 +1,348 @@
+
+
+"""Functions for getting MSA and calculating alignment features."""
+
+from collections.abc import MutableMapping, Sequence
+import string
+from typing import Self
+
+from absl import logging
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.data import msa_config
+from flax_model.alphafold3.data import msa_features
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import jackhmmer
+from flax_model.alphafold3.data.tools import msa_tool
+from flax_model.alphafold3.data.tools import nhmmer
+from flax_model.alphafold3.data.tools import mmseqs
+import numpy as np
+
+
+
+class Error(Exception):
+  """Error indicatating a problem with MSA Search."""
+
+
+def _featurize(seq: str, chain_poly_type: str) -> str | list[int]:
+  if mmcif_names.is_standard_polymer_type(chain_poly_type):
+    featurized_seqs, _ = msa_features.extract_msa_features(
+        msa_sequences=[seq], chain_poly_type=chain_poly_type
+    )
+    return featurized_seqs[0].tolist()
+  # For anything else simply require an identical match.
+  return seq
+
+
+def sequences_are_feature_equivalent(
+    sequence1: str,
+    sequence2: str,
+    chain_poly_type: str,
+) -> bool:
+  feat1 = _featurize(sequence1, chain_poly_type)
+  feat2 = _featurize(sequence2, chain_poly_type)
+  return feat1 == feat2
+
+
+class Msa:
+  """Multiple Sequence Alignment container with methods for manipulating it."""
+
+  def __init__(
+      self,
+      query_sequence: str,
+      chain_poly_type: str,
+      sequences: Sequence[str],
+      descriptions: Sequence[str],
+      deduplicate: bool = True,
+  ):
+    """Raw constructor, prefer using the from_{a3m,multiple_msas} class methods.
+
+    The first sequence must be equal (in featurised form) to the query sequence.
+    If sequences/descriptions are empty, they will be initialised to the query.
+
+    Args:
+      query_sequence: The sequence that was used to search for MSA.
+      chain_poly_type: Polymer type of the query sequence, see mmcif_names.
+      sequences: The sequences returned by the MSA search tool.
+      descriptions: Metadata for the sequences returned by the MSA search tool.
+      deduplicate: If True, the MSA sequences will be deduplicated in the input
+        order. Lowercase letters (insertions) are ignored when deduplicating.
+    """
+    if len(sequences) != len(descriptions):
+      raise ValueError('The number of sequences and descriptions must match.')
+
+    self.query_sequence = query_sequence
+    self.chain_poly_type = chain_poly_type
+
+    if not deduplicate:
+      self.sequences = sequences
+      self.descriptions = descriptions
+    else:
+      self.sequences = []
+      self.descriptions = []
+      # A replacement table that removes all lowercase characters.
+      deletion_table = str.maketrans('', '', string.ascii_lowercase)
+      unique_sequences = set()
+      for seq, desc in zip(sequences, descriptions, strict=True):
+        # Using string.translate is faster than re.sub('[a-z]+', '').
+        sequence_no_deletions = seq.translate(deletion_table)
+        if sequence_no_deletions not in unique_sequences:
+          unique_sequences.add(sequence_no_deletions)
+          self.sequences.append(seq)
+          self.descriptions.append(desc)
+
+    # Make sure the MSA always has at least the query.
+    self.sequences = self.sequences or [query_sequence]
+    self.descriptions = self.descriptions or ['Original query']
+
+    # Check if the 1st MSA sequence matches the query sequence. Since it may be
+    # mutated by the search tool (jackhmmer) check using the featurized version.
+    if not sequences_are_feature_equivalent(
+        self.sequences[0], query_sequence, chain_poly_type
+    ):
+      raise ValueError(
+          f'First MSA sequence {self.sequences[0]} is not the {query_sequence=}'
+      )
+
+  @classmethod
+  def from_multiple_msas(
+      cls, msas: Sequence[Self], deduplicate: bool = True
+  ) -> Self:
+    """Initializes the MSA from multiple MSAs.
+
+    Args:
+      msas: A sequence of Msa objects representing individual MSAs produced by
+        different tools/dbs.
+      deduplicate: If True, the MSA sequences will be deduplicated in the input
+        order. Lowercase letters (insertions) are ignored when deduplicating.
+
+    Returns:
+      An Msa object created by merging multiple MSAs.
+    """
+    if not msas:
+      raise ValueError('At least one MSA must be provided.')
+
+    query_sequence = msas[0].query_sequence
+    chain_poly_type = msas[0].chain_poly_type
+    sequences = []
+    descriptions = []
+
+    for msa in msas:
+      if msa.query_sequence != query_sequence:
+        raise ValueError(
+            f'Query sequences must match: {[m.query_sequence for m in msas]}'
+        )
+      if msa.chain_poly_type != chain_poly_type:
+        raise ValueError(
+            f'Chain poly types must match: {[m.chain_poly_type for m in msas]}'
+        )
+      sequences.extend(msa.sequences)
+      descriptions.extend(msa.descriptions)
+
+    return cls(
+        query_sequence=query_sequence,
+        chain_poly_type=chain_poly_type,
+        sequences=sequences,
+        descriptions=descriptions,
+        deduplicate=deduplicate,
+    )
+
+  @classmethod
+  def from_multiple_a3ms(
+      cls, a3ms: Sequence[str], chain_poly_type: str, deduplicate: bool = True
+  ) -> Self:
+    """Initializes the MSA from multiple A3M strings.
+
+    Args:
+      a3ms: A sequence of A3M strings representing individual MSAs produced by
+        different tools/dbs.
+      chain_poly_type: Polymer type of the query sequence, see mmcif_names.
+      deduplicate: If True, the MSA sequences will be deduplicated in the input
+        order. Lowercase letters (insertions) are ignored when deduplicating.
+
+    Returns:
+      An Msa object created by merging multiple A3Ms.
+    """
+    if not a3ms:
+      raise ValueError('At least one A3M must be provided.')
+
+    query_sequence = None
+    all_sequences = []
+    all_descriptions = []
+
+    for a3m in a3ms:
+      sequences, descriptions = parsers.parse_fasta(a3m)
+      if query_sequence is None:
+        query_sequence = sequences[0]
+
+      if sequences[0] != query_sequence:
+        raise ValueError(
+            f'Query sequences must match: {sequences[0]=} != {query_sequence=}'
+        )
+      all_sequences.extend(sequences)
+      all_descriptions.extend(descriptions)
+
+    return cls(
+        query_sequence=query_sequence,
+        chain_poly_type=chain_poly_type,
+        sequences=all_sequences,
+        descriptions=all_descriptions,
+        deduplicate=deduplicate,
+    )
+
+  @classmethod
+  def from_a3m(
+      cls,
+      query_sequence: str,
+      chain_poly_type: str,
+      a3m: str,
+      max_depth: int | None = None,
+      deduplicate: bool = True,
+  ) -> Self:
+    """Parses the single A3M and builds the Msa object."""
+    sequences, descriptions = parsers.parse_fasta(a3m)
+
+    if max_depth is not None and 0 < max_depth < len(sequences):
+      logging.info(
+          'MSA cropped from depth of %d to %d for %s.',
+          len(sequences),
+          max_depth,
+          query_sequence,
+      )
+      sequences = sequences[:max_depth]
+      descriptions = descriptions[:max_depth]
+
+    return cls(
+        query_sequence=query_sequence,
+        chain_poly_type=chain_poly_type,
+        sequences=sequences,
+        descriptions=descriptions,
+        deduplicate=deduplicate,
+    )
+
+  @classmethod
+  def from_empty(cls, query_sequence: str, chain_poly_type: str) -> Self:
+    """Creates an empty Msa containing just the query sequence."""
+    return cls(
+        query_sequence=query_sequence,
+        chain_poly_type=chain_poly_type,
+        sequences=[],
+        descriptions=[],
+        deduplicate=False,
+    )
+
+  @property
+  def depth(self) -> int:
+    return len(self.sequences)
+
+  def __repr__(self) -> str:
+    return f'Msa({self.depth} sequences, {self.chain_poly_type})'
+
+  def to_a3m(self) -> str:
+    """Returns the MSA in the A3M format."""
+    a3m_lines = []
+    for desc, seq in zip(self.descriptions, self.sequences, strict=True):
+      a3m_lines.append(f'>{desc}')
+      a3m_lines.append(seq)
+    return '\n'.join(a3m_lines) + '\n'
+
+  def featurize(self) -> MutableMapping[str, np.ndarray]:
+    """Featurises the MSA and returns a map of feature names to features.
+
+    Returns:
+      A dictionary mapping feature names to values.
+
+    Raises:
+      msa.Error:
+        * If the sequences in the MSA don't have the same length after deletions
+          (lower case letters) are removed.
+        * If the MSA contains an unknown amino acid code.
+        * If there are no sequences after aligning.
+    """
+    try:
+      msa, deletion_matrix = msa_features.extract_msa_features(
+          msa_sequences=self.sequences, chain_poly_type=self.chain_poly_type
+      )
+    except ValueError as e:
+      raise Error(f'Error extracting MSA or deletion features: {e}') from e
+
+    if msa.shape == (0, 0):
+      raise Error(f'Empty MSA feature for {self}')
+
+    species_ids = msa_features.extract_species_ids(self.descriptions)
+
+    return {
+        'msa_species_identifiers': np.array(species_ids, dtype=object),
+        'num_alignments': np.array(self.depth, dtype=np.int32),
+        'msa': msa,
+        'deletion_matrix': deletion_matrix,
+    }
+
+
+def get_msa_tool(
+    msa_tool_config: msa_config.JackhmmerConfig | msa_config.NhmmerConfig | msa_config.MmseqsConfig,
+) -> msa_tool.MsaTool:
+  """Returns the requested MSA tool."""
+
+  match msa_tool_config:
+    case msa_config.JackhmmerConfig():
+      return jackhmmer.Jackhmmer(
+          binary_path=msa_tool_config.binary_path,
+          database_path=msa_tool_config.database_config.path,
+          n_cpu=msa_tool_config.n_cpu,
+          n_iter=msa_tool_config.n_iter,
+          e_value=msa_tool_config.e_value,
+          z_value=msa_tool_config.z_value,
+          max_sequences=msa_tool_config.max_sequences,
+          max_threads=msa_tool_config.max_threads,
+      )
+    case msa_config.NhmmerConfig():
+      return nhmmer.Nhmmer(
+          binary_path=msa_tool_config.binary_path,
+          hmmalign_binary_path=msa_tool_config.hmmalign_binary_path,
+          hmmbuild_binary_path=msa_tool_config.hmmbuild_binary_path,
+          database_path=msa_tool_config.database_config.path,
+          n_cpu=msa_tool_config.n_cpu,
+          e_value=msa_tool_config.e_value,
+          max_sequences=msa_tool_config.max_sequences,
+          max_threads=msa_tool_config.max_threads,
+          alphabet=msa_tool_config.alphabet,
+      )
+    case msa_config.MmseqsConfig():
+      return mmseqs.Mmseqs(
+          binary_path=msa_tool_config.binary_path,
+          database_path=msa_tool_config.database_config.path,
+          n_cpu=msa_tool_config.n_cpu,
+          use_gpu=msa_tool_config.use_gpu,
+          mmseqs_options=msa_tool_config.mmseqs_options,
+          result2msa_options=msa_tool_config.result2msa_options,
+      )
+    case _:
+      raise ValueError(f'Unknown MSA tool: {msa_tool_config}.')
+
+
+def get_msa(
+    target_sequence: str,
+    run_config: msa_config.RunConfig,
+    chain_poly_type: str,
+    deduplicate: bool = False,
+) -> Msa:
+  """Computes the MSA for a given query sequence.
+
+  Args:
+    target_sequence: The target amino-acid sequence.
+    run_config: MSA run configuration.
+    chain_poly_type: The type of chain for which to get an MSA.
+    deduplicate: If True, the MSA sequences will be deduplicated in the input
+      order. Lowercase letters (insertions) are ignored when deduplicating.
+
+  Returns:
+    Aligned MSA sequences.
+  """
+
+  return Msa.from_a3m(
+      query_sequence=target_sequence,
+      chain_poly_type=chain_poly_type,
+      a3m=get_msa_tool(run_config.config).query(target_sequence).a3m,
+      max_depth=run_config.crop_size,
+      deduplicate=deduplicate,
+  )

flax_model/alphafold3/data/msa_store.py

@@ -0,0 +1,58 @@
+
+
+"""Interface and implementations for fetching MSA data."""
+
+from collections.abc import Sequence
+from typing import Protocol, TypeAlias
+
+from flax_model.alphafold3.data import msa
+from flax_model.alphafold3.data import msa_config
+
+
+MsaErrors: TypeAlias = Sequence[tuple[msa_config.RunConfig, str]]
+
+
+class MsaProvider(Protocol):
+  """Interface for providing Multiple Sequence Alignments."""
+
+  def __call__(
+      self,
+      query_sequence: str,
+      chain_polymer_type: str,
+  ) -> tuple[msa.Msa, MsaErrors]:
+    """Retrieve MSA for the given polymer query_sequence.
+
+    Args:
+      query_sequence: The residue sequence of the polymer to search for.
+      chain_polymer_type: The polymer type of the query_sequence. This must
+        match the chain_polymer_type of the provider.
+
+    Returns:
+      A tuple containing the MSA and MsaErrors. MsaErrors is a Sequence
+      containing a tuple for each msa_query that failed. Each tuple contains
+      the failing query and the associated error message.
+    """
+
+
+class EmptyMsaProvider:
+  """MSA provider that returns just the query sequence, useful for testing."""
+
+  def __init__(self, chain_polymer_type: str):
+    self._chain_polymer_type = chain_polymer_type
+
+  def __call__(
+      self, query_sequence: str, chain_polymer_type: str
+  ) -> tuple[msa.Msa, MsaErrors]:
+    """Returns an MSA containing just the query sequence, never errors."""
+    if chain_polymer_type != self._chain_polymer_type:
+      raise ValueError(
+          f'EmptyMsaProvider of type {self._chain_polymer_type} called with '
+          f'sequence of {chain_polymer_type=}, {query_sequence=}.'
+      )
+    return (
+        msa.Msa.from_empty(
+            query_sequence=query_sequence,
+            chain_poly_type=self._chain_polymer_type,
+        ),
+        (),
+    )

flax_model/alphafold3/data/template_realign.py

@@ -0,0 +1,160 @@
+
+
+"""Realign sequences found in PDB seqres to the actual CIF sequences."""
+
+from collections.abc import Mapping
+
+
+class AlignmentError(Exception):
+  """Failed alignment between the hit sequence and the actual mmCIF sequence."""
+
+
+def realign_hit_to_structure(
+    *,
+    hit_sequence: str,
+    hit_start_index: int,
+    hit_end_index: int,
+    full_length: int,
+    structure_sequence: str,
+    query_to_hit_mapping: Mapping[int, int],
+) -> Mapping[int, int]:
+  """Realigns the hit sequence to the Structure sequence.
+
+  For example, for the given input:
+    query_sequence : ABCDEFGHIJKL
+    hit_sequence   : ---DEFGHIJK-
+    struc_sequence : XDEFGHKL
+  the mapping is {3: 0, 4: 1, 5: 2, 6: 3, 7: 4, 8: 5, 9: 6, 10: 7}. However, the
+  actual Structure sequence has an extra X at the start as well as no IJ. So the
+  alignment from the query to the Structure sequence will be:
+    hit_sequence   : ---DEFGHIJK-
+    struc_aligned  : --XDEFGH--KL
+  and the new mapping will therefore be: {3: 1, 4: 2, 5: 3, 6: 4, 7: 5, 10: 6}.
+
+  Args:
+    hit_sequence: The PDB seqres hit sequence obtained from Hmmsearch, but
+      without any gaps. This is not the full PDB seqres template sequence but
+      rather just its subsequence from hit_start_index to hit_end_index.
+    hit_start_index: The start index of the hit sequence in the full PDB seqres
+      template sequence (inclusive).
+    hit_end_index: The end index of the hit sequence in the full PDB seqres
+      template sequence (exclusive).
+    full_length: The length of the full PDB seqres template sequence.
+    structure_sequence: The actual sequence extracted from the Structure
+      corresponding to this template. In vast majority of cases this is the same
+      as the PDB seqres sequence, but this function handles the cases when not.
+    query_to_hit_mapping: The mapping from the query sequence to the
+      hit_sequence.
+
+  Raises:
+    AlignmentError: if the alignment between the sequence returned by Hmmsearch
+      differs from the actual sequence found in the mmCIF and can't be aligned
+      using the simple alignment algorithm.
+
+  Returns:
+    A mapping from the query sequence to the actual Structure sequence.
+  """
+  max_num_gaps = full_length - len(structure_sequence)
+  if max_num_gaps < 0:
+    raise AlignmentError(
+        f'The Structure sequence ({len(structure_sequence)}) '
+        f'must be shorter than the PDB seqres sequence ({full_length}):\n'
+        f'Structure sequence : {structure_sequence}\n'
+        f'PDB seqres sequence: {hit_sequence}'
+    )
+
+  if len(hit_sequence) != hit_end_index - hit_start_index:
+    raise AlignmentError(
+        f'The difference of {hit_end_index=} and {hit_start_index=} does not '
+        f'equal to the length of the {hit_sequence}: {len(hit_sequence)}'
+    )
+
+  best_score = -1
+  best_start = 0
+  best_query_to_hit_mapping = query_to_hit_mapping
+  max_num_gaps_before_subseq = min(hit_start_index, max_num_gaps)
+  # It is possible the gaps needed to align the PDB seqres subsequence and
+  # the Structure subsequence need to be inserted before the match region.
+  # Try and pick the alignment with the best number of aligned residues.
+  for num_gaps_before_subseq in range(0, max_num_gaps_before_subseq + 1):
+    start = hit_start_index - num_gaps_before_subseq
+    end = hit_end_index - num_gaps_before_subseq
+    structure_subseq = structure_sequence[start:end]
+
+    new_query_to_hit_mapping, score = _remap_to_struc_seq(
+        hit_seq=hit_sequence,
+        struc_seq=structure_subseq,
+        max_num_gaps=max_num_gaps - num_gaps_before_subseq,
+        mapping=query_to_hit_mapping,
+    )
+    if score >= best_score:
+      # Use >= to prefer matches with larger number of gaps before.
+      best_score = score
+      best_start = start
+      best_query_to_hit_mapping = new_query_to_hit_mapping
+
+  return {q: h + best_start for q, h in best_query_to_hit_mapping.items()}
+
+
+def _remap_to_struc_seq(
+    *,
+    hit_seq: str,
+    struc_seq: str,
+    max_num_gaps: int,
+    mapping: Mapping[int, int],
+) -> tuple[Mapping[int, int], int]:
+  """Remaps the query -> hit mapping to match the actual Structure sequence.
+
+  Args:
+    hit_seq: The hit sequence - a subsequence of the PDB seqres sequence without
+      any Hmmsearch modifications like inserted gaps or lowercased residues.
+    struc_seq: The actual sequence obtained from the corresponding Structure.
+    max_num_gaps: The maximum number of gaps that can be inserted in the
+      Structure sequence. In practice, this is the length difference between the
+      PDB seqres sequence and the actual Structure sequence.
+    mapping: The mapping from the query residues to the hit residues. This will
+      be remapped to point to the actual Structure sequence using a simple
+      realignment algorithm.
+
+  Returns:
+    A tuple of (mapping, score):
+      * Mapping from the query to the actual Structure sequence.
+      * Score which is the number of matching aligned residues.
+
+  Raises:
+    ValueError if the structure sequence isn't shorter than the seqres sequence.
+    ValueError if the alignment fails.
+  """
+  hit_seq_idx = 0
+  struc_seq_idx = 0
+  hit_to_struc_seq_mapping = {}
+  score = 0
+
+  # This while loop is guaranteed to terminate since we increase both
+  # struc_seq_idx and hit_seq_idx by at least 1 in each iteration.
+  remaining_num_gaps = max_num_gaps
+  while hit_seq_idx < len(hit_seq) and struc_seq_idx < len(struc_seq):
+    if hit_seq[hit_seq_idx] != struc_seq[struc_seq_idx]:
+      # Explore which alignment aligns the next residue (if present).
+      best_shift = 0
+      for shift in range(0, remaining_num_gaps + 1):
+        next_hit_res = hit_seq[hit_seq_idx + shift : hit_seq_idx + shift + 1]
+        next_struc_res = struc_seq[struc_seq_idx : struc_seq_idx + 1]
+        if next_hit_res == next_struc_res:
+          best_shift = shift
+          break
+      hit_seq_idx += best_shift
+      remaining_num_gaps -= best_shift
+
+    hit_to_struc_seq_mapping[hit_seq_idx] = struc_seq_idx
+    score += hit_seq[hit_seq_idx] == struc_seq[struc_seq_idx]
+    hit_seq_idx += 1
+    struc_seq_idx += 1
+
+  fixed_mapping = {}
+  for query_idx, original_hit_idx in mapping.items():
+    fixed_hit_idx = hit_to_struc_seq_mapping.get(original_hit_idx)
+    if fixed_hit_idx is not None:
+      fixed_mapping[query_idx] = fixed_hit_idx
+
+  return fixed_mapping, score

flax_model/alphafold3/data/templates.py

@@ -0,0 +1,978 @@
+
+
+"""API for retrieving and manipulating template search results."""
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+import dataclasses
+import datetime
+import functools
+import os
+import re
+from typing import Any, Final, Self, TypeAlias
+
+from absl import logging
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.common import resources
+from flax_model.alphafold3.constants import atom_types
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.data import msa_config
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data import structure_stores
+from flax_model.alphafold3.data import template_realign
+from flax_model.alphafold3.data.tools import hmmsearch
+from flax_model.alphafold3.structure import mmcif
+import numpy as np
+
+
+_POLYMER_FEATURES: Final[Mapping[str, np.float64 | np.int32 | object]] = {
+    'template_aatype': np.int32,
+    'template_all_atom_masks': np.float64,
+    'template_all_atom_positions': np.float64,
+    'template_domain_names': object,
+    'template_release_date': object,
+    'template_sequence': object,
+}
+
+_LIGAND_FEATURES: Final[Mapping[str, Any]] = {
+    'ligand_features': Mapping[str, Any]
+}
+
+
+TemplateFeatures: TypeAlias = Mapping[
+    str, np.ndarray | bytes | Mapping[str, np.ndarray | bytes]
+]
+_REQUIRED_METADATA_COLUMNS: Final[Sequence[str]] = (
+    'seq_release_date',
+    'seq_unresolved_res_num',
+    'seq_author_chain_id',
+    'seq_sequence',
+)
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True, slots=True)
+class _Polymer:
+  """Container for alphabet specific (dna, rna, protein) atom information."""
+
+  min_atoms: int
+  num_atom_types: int
+  atom_order: Mapping[str, int]
+
+
+_POLYMERS = {
+    mmcif_names.PROTEIN_CHAIN: _Polymer(
+        min_atoms=5,
+        num_atom_types=atom_types.ATOM37_NUM,
+        atom_order=atom_types.ATOM37_ORDER,
+    ),
+    mmcif_names.DNA_CHAIN: _Polymer(
+        min_atoms=21,
+        num_atom_types=atom_types.ATOM29_NUM,
+        atom_order=atom_types.ATOM29_ORDER,
+    ),
+    mmcif_names.RNA_CHAIN: _Polymer(
+        min_atoms=20,
+        num_atom_types=atom_types.ATOM29_NUM,
+        atom_order=atom_types.ATOM29_ORDER,
+    ),
+}
+
+
+def _encode_restype(
+    chain_poly_type: str,
+    sequence: str,
+) -> Sequence[int]:
+  """Encodes a sequence of residue names as a sequence of ints.
+
+  Args:
+    chain_poly_type: Polymer chain type to determine sequence encoding.
+    sequence: Polymer residues. Protein encoded by single letters. RNA and DNA
+      encoded by  multi-letter CCD codes.
+
+  Returns:
+    A sequence of integers encoding amino acid types for the given chain type.
+  """
+  if chain_poly_type == mmcif_names.PROTEIN_CHAIN:
+    return [
+        residue_names.PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP_TO_INT[
+            _STANDARDIZED_AA.get(res, res)
+        ]
+        for res in sequence
+    ]
+
+  unk_nucleic = residue_names.UNK_NUCLEIC_ONE_LETTER
+  unk_nucleic_idx = residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP[
+      unk_nucleic
+  ]
+  if chain_poly_type == mmcif_names.RNA_CHAIN:
+    return [
+        residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP.get(
+            res, unk_nucleic_idx
+        )
+        for res in sequence
+    ]
+  elif chain_poly_type == mmcif_names.DNA_CHAIN:
+    # Map UNK DNA to the generic nucleic UNK (N), which happens to also be the
+    # same as the RNA UNK.
+    return [
+        residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP.get(
+            residue_names.DNA_COMMON_ONE_TO_TWO.get(res, unk_nucleic),
+            unk_nucleic_idx,
+        )
+        for res in sequence
+    ]
+
+  raise NotImplementedError(f'"{chain_poly_type}" unsupported.')
+
+
+_DAYS_BEFORE_QUERY_DATE: Final[int] = 60
+_HIT_DESCRIPTION_REGEX = re.compile(
+    r'(?P<pdb_id>[a-z0-9]{4,})_(?P<chain_id>\w+)/(?P<start>\d+)-(?P<end>\d+) '
+    r'.* length:(?P<length>\d+)\b.*'
+)
+
+_STANDARDIZED_AA = {'B': 'D', 'J': 'X', 'O': 'X', 'U': 'C', 'Z': 'E'}
+
+
+class Error(Exception):
+  """Base class for exceptions."""
+
+
+class HitDateError(Error):
+  """An error indicating that invalid release date was detected."""
+
+
+class InvalidTemplateError(Error):
+  """An error indicating that template is invalid."""
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Hit:
+  """Template hit metrics derived from the MSA for filtering and featurising.
+
+  Attributes:
+    pdb_id: The PDB ID of the hit.
+    auth_chain_id: The author chain ID of the hit.
+    hmmsearch_sequence: Hit sequence as given in hmmsearch a3m output.
+    structure_sequence: Hit sequence as given in PDB structure.
+    unresolved_res_indices: Indices of unresolved residues in the structure
+      sequence. 0-based.
+    query_sequence: The query nucleotide/amino acid sequence.
+    start_index: The start index of the sequence relative to the full PDB seqres
+      sequence. Inclusive and uses 0-based indexing.
+    end_index: The end index of the sequence relative to the full PDB seqres
+      sequence. Exclusive and uses 0-based indexing.
+    full_length: Length of the full PDB seqres sequence. This can be different
+      from the length from the actual sequence we get from the mmCIF and we use
+      this to detect whether we need to realign or not.
+    release_date: The release date of the PDB corresponding to this hit.
+    chain_poly_type: The polymer type of the selected hit structure.
+  """
+
+  pdb_id: str
+  auth_chain_id: str
+  hmmsearch_sequence: str
+  structure_sequence: str
+  unresolved_res_indices: Sequence[int] | None
+  query_sequence: str
+  start_index: int
+  end_index: int
+  full_length: int
+  release_date: datetime.date
+  chain_poly_type: str
+
+  @functools.cached_property
+  def query_to_hit_mapping(self) -> Mapping[int, int]:
+    """0-based query index to hit index mapping."""
+    query_to_hit_mapping = {}
+    hit_index = 0
+    query_index = 0
+    for residue in self.hmmsearch_sequence:
+      # Gap inserted in the template
+      if residue == '-':
+        query_index += 1
+      # Deleted residue in the template (would be a gap in the query).
+      elif residue.islower():
+        hit_index += 1
+      # Normal aligned residue, in both query and template. Add to mapping.
+      elif residue.isupper():
+        query_to_hit_mapping[query_index] = hit_index
+        query_index += 1
+        hit_index += 1
+
+    structure_subseq = self.structure_sequence[
+        self.start_index : self.end_index
+    ]
+    if self.matching_sequence != structure_subseq:
+      # The seqres sequence doesn't match the structure sequence. Two cases:
+      # 1. The sequences have the same length. The sequences are different
+      #    because our 3->1 residue code mapping is different from the one PDB
+      #    uses. We don't do anything in this case as both sequences have the
+      #    same length, so the original query to hit mapping stays valid.
+      # 2. The sequences don't have the same length, the one in structure is
+      #    shorter. In this case we change the mapping to match the actual
+      #    structure sequence using a simple realignment algorithm.
+      # This procedure was validated on all PDB seqres (2023_01_12) sequences
+      # and handles all cases that can happen.
+      if self.full_length != len(self.structure_sequence):
+        return template_realign.realign_hit_to_structure(
+            hit_sequence=self.matching_sequence,
+            hit_start_index=self.start_index,
+            hit_end_index=self.end_index,
+            full_length=self.full_length,
+            structure_sequence=self.structure_sequence,
+            query_to_hit_mapping=query_to_hit_mapping,
+        )
+
+    # Hmmsearch returns a subsequence and so far indices have been relative to
+    # the subsequence. Add an offset to index relative to the full structure
+    # sequence.
+    return {q: h + self.start_index for q, h in query_to_hit_mapping.items()}
+
+  @property
+  def matching_sequence(self) -> str:
+    """Returns the matching hit sequence including insertions.
+
+    Make deleted residues uppercase and remove gaps ("-").
+    """
+    return self.hmmsearch_sequence.upper().replace('-', '')
+
+  @functools.cached_property
+  def output_templates_sequence(self) -> str:
+    """Returns the final template sequence."""
+    result_seq = ['-'] * len(self.query_sequence)
+    for query_index, template_index in self.query_to_hit_mapping.items():
+      result_seq[query_index] = self.structure_sequence[template_index]
+    return ''.join(result_seq)
+
+  @property
+  def length_ratio(self) -> float:
+    """Ratio of the length of the hit sequence to the query."""
+    return len(self.matching_sequence) / len(self.query_sequence)
+
+  @property
+  def align_ratio(self) -> float:
+    """Ratio of the number of aligned residues to the query length."""
+    return len(self.query_to_hit_mapping) / len(self.query_sequence)
+
+  @functools.cached_property
+  def is_valid(self) -> bool:
+    """Whether hit can be used as a template."""
+    if self.unresolved_res_indices is None:
+      return False
+
+    return bool(
+        set(self.query_to_hit_mapping.values())
+        - set(self.unresolved_res_indices)
+    )
+
+  @property
+  def full_name(self) -> str:
+    """A full name of the hit."""
+    return f'{self.pdb_id}_{self.auth_chain_id}'
+
+  def __post_init__(self):
+    if not self.pdb_id.islower() and not self.pdb_id.isdigit():
+      raise ValueError(f'pdb_id must be lowercase {self.pdb_id}')
+
+    if not (0 <= self.start_index <= self.end_index):
+      raise ValueError(
+          'Start must be non-negative and less than or equal to end index. '
+          f'Range: {self.start_index}-{self.end_index}'
+      )
+
+    if len(self.matching_sequence) != (self.end_index - self.start_index):
+      raise ValueError(
+          'Sequence length must be equal to end_index - start_index. '
+          f'{len(self.matching_sequence)} != {self.end_index} - '
+          f'{self.start_index}'
+      )
+
+    if self.full_length < 0:
+      raise ValueError(f'Full length must be non-negative: {self.full_length}')
+
+  def keep(
+      self,
+      *,
+      release_date_cutoff: datetime.date | None,
+      max_subsequence_ratio: float | None,
+      min_hit_length: int | None,
+      min_align_ratio: float | None,
+  ) -> bool:
+    """Returns whether the hit should be kept.
+
+    In addition to filtering on all of the provided parameters, this method also
+    excludes hits with unresolved residues.
+
+    Args:
+      release_date_cutoff: Maximum release date of the template.
+      max_subsequence_ratio: If set, excludes hits which are an exact
+        subsequence of the query sequence, and longer than this ratio. Useful to
+        avoid ground truth leakage.
+      min_hit_length: If set, excludes hits which have fewer residues than this.
+      min_align_ratio: If set, excludes hits where the number of residues
+        aligned to the query is less than this proportion of the template
+        length.
+    """
+    # Exclude hits which are too recent.
+    if (
+        release_date_cutoff is not None
+        and self.release_date > release_date_cutoff
+    ):
+      return False
+
+    # Exclude hits which are large duplicates of the query_sequence.
+    if (
+        max_subsequence_ratio is not None
+        and self.length_ratio > max_subsequence_ratio
+    ):
+      if self.matching_sequence in self.query_sequence:
+        return False
+
+    # Exclude hits which are too short.
+    if (
+        min_hit_length is not None
+        and len(self.matching_sequence) < min_hit_length
+    ):
+      return False
+
+    # Exclude hits with unresolved residues.
+    if not self.is_valid:
+      return False
+
+    # Exclude hits with too few alignments.
+    try:
+      if min_align_ratio is not None and self.align_ratio <= min_align_ratio:
+        return False
+    except template_realign.AlignmentError as e:
+      logging.warning('Failed to align %s: %s', self, str(e))
+      return False
+
+    return True
+
+
+def _filter_hits(
+    hits: Iterable[Hit],
+    release_date_cutoff: datetime.date,
+    max_subsequence_ratio: float | None,
+    min_align_ratio: float | None,
+    min_hit_length: int | None,
+    deduplicate_sequences: bool,
+    max_hits: int | None,
+) -> Sequence[Hit]:
+  """Filters hits based on the filter config."""
+  filtered_hits = []
+  seen_before = set()
+  for hit in hits:
+    if not hit.keep(
+        max_subsequence_ratio=max_subsequence_ratio,
+        min_align_ratio=min_align_ratio,
+        min_hit_length=min_hit_length,
+        release_date_cutoff=release_date_cutoff,
+    ):
+      continue
+
+    # Remove duplicate templates, keeping the first.
+    if deduplicate_sequences:
+      if hit.output_templates_sequence in seen_before:
+        continue
+      seen_before.add(hit.output_templates_sequence)
+
+    filtered_hits.append(hit)
+    if max_hits and len(filtered_hits) == max_hits:
+      break
+
+  return filtered_hits
+
+
+@dataclasses.dataclass(init=False)
+class Templates:
+  """A container for templates that were found for the given query sequence.
+
+  The structure_store is constructed from the config by default. Callers can
+  optionally supply a structure_store to the constructor to avoid the cost of
+  construction and metadata loading.
+  """
+
+  def __init__(
+      self,
+      *,
+      query_sequence: str,
+      hits: Sequence[Hit],
+      max_template_date: datetime.date,
+      structure_store: structure_stores.StructureStore,
+      query_release_date: datetime.date | None = None,
+  ):
+    self._query_sequence = query_sequence
+    self._hits = tuple(hits)
+    self._max_template_date = max_template_date
+    self._query_release_date = query_release_date
+    self._hit_structures = {}
+    self._structure_store = structure_store
+
+    if any(h.query_sequence != self._query_sequence for h in self.hits):
+      raise ValueError('All hits must match the query sequence.')
+
+    if self._hits:
+      chain_poly_type = self._hits[0].chain_poly_type
+      if any(h.chain_poly_type != chain_poly_type for h in self.hits):
+        raise ValueError('All hits must have the same chain_poly_type.')
+
+  @classmethod
+  def from_seq_and_a3m(
+      cls,
+      *,
+      query_sequence: str,
+      msa_a3m: str,
+      max_template_date: datetime.date,
+      database_path: os.PathLike[str] | str,
+      hmmsearch_config: msa_config.HmmsearchConfig,
+      max_a3m_query_sequences: int | None,
+      structure_store: structure_stores.StructureStore,
+      filter_config: msa_config.TemplateFilterConfig | None = None,
+      query_release_date: datetime.date | None = None,
+      chain_poly_type: str = mmcif_names.PROTEIN_CHAIN,
+  ) -> Self:
+    """Creates templates from a run of hmmsearch tool against a custom a3m.
+
+    Args:
+      query_sequence: The polymer sequence of the target query.
+      msa_a3m: An a3m of related polymers aligned to the query sequence, this is
+        used to create an HMM for the hmmsearch run.
+      max_template_date: This is used to filter templates for training, ensuring
+        that they do not leak ground truth information used in testing sets.
+      database_path: A path to the sequence database to search for templates.
+      hmmsearch_config: Config with Hmmsearch settings.
+      max_a3m_query_sequences: The maximum number of input MSA sequences to use
+        to construct the profile which is then used to search for templates.
+      structure_store: Structure store to fetch template structures from.
+      filter_config: Optional config that controls which and how many hits to
+        keep. More performant than constructing and then filtering. If not
+        provided, no filtering is done.
+      query_release_date: The release_date of the template query, this is used
+        to filter templates for training, ensuring that they do not leak
+        structure information from the future.
+      chain_poly_type: The polymer type of the templates.
+
+    Returns:
+      Templates object containing a list of Hits initialised from the
+      structure_store metadata and a3m alignments.
+    """
+    hmmsearch_a3m = run_hmmsearch_with_a3m(
+        database_path=database_path,
+        hmmsearch_config=hmmsearch_config,
+        max_a3m_query_sequences=max_a3m_query_sequences,
+        a3m=msa_a3m,
+    )
+    return cls.from_hmmsearch_a3m(
+        query_sequence=query_sequence,
+        a3m=hmmsearch_a3m,
+        max_template_date=max_template_date,
+        query_release_date=query_release_date,
+        chain_poly_type=chain_poly_type,
+        structure_store=structure_store,
+        filter_config=filter_config,
+    )
+
+  @classmethod
+  def from_hmmsearch_a3m(
+      cls,
+      *,
+      query_sequence: str,
+      a3m: str,
+      max_template_date: datetime.date,
+      structure_store: structure_stores.StructureStore,
+      filter_config: msa_config.TemplateFilterConfig | None = None,
+      query_release_date: datetime.date | None = None,
+      chain_poly_type: str = mmcif_names.PROTEIN_CHAIN,
+  ) -> Self:
+    """Creates Templates from a Hmmsearch A3M.
+
+    Args:
+      query_sequence: The polymer sequence of the target query.
+      a3m: Results of Hmmsearch in A3M format. This provides a list of potential
+        template alignments and pdb codes.
+      max_template_date: This is used to filter templates for training, ensuring
+        that they do not leak ground truth information used in testing sets.
+      structure_store: Structure store to fetch template structures from.
+      filter_config: Optional config that controls which and how many hits to
+        keep. More performant than constructing and then filtering. If not
+        provided, no filtering is done.
+      query_release_date: The release_date of the template query, this is used
+        to filter templates for training, ensuring that they do not leak
+        structure information from the future.
+      chain_poly_type: The polymer type of the templates.
+
+    Returns:
+      Templates object containing a list of Hits initialised from the
+      structure_store metadata and a3m alignments.
+    """
+
+    def hit_generator(a3m: str):
+      if not a3m:
+        return  # Hmmsearch could return an empty string if there are no hits.
+
+      for hit_seq, hit_desc in parsers.lazy_parse_fasta_string(a3m):
+        pdb_id, auth_chain_id, start, end, full_length = _parse_hit_description(
+            hit_desc
+        )
+
+        release_date, sequence, unresolved_res_ids = _parse_hit_metadata(
+            structure_store, pdb_id, auth_chain_id
+        )
+        if unresolved_res_ids is None:
+          continue
+
+        # seq_unresolved_res_num are 1-based, setting to 0-based indices.
+        unresolved_indices = [i - 1 for i in unresolved_res_ids]
+
+        yield Hit(
+            pdb_id=pdb_id,
+            auth_chain_id=auth_chain_id,
+            hmmsearch_sequence=hit_seq,
+            structure_sequence=sequence,
+            query_sequence=query_sequence,
+            unresolved_res_indices=unresolved_indices,
+            start_index=start - 1,  # Raw value is residue number, not index.
+            end_index=end,
+            full_length=full_length,
+            release_date=datetime.date.fromisoformat(release_date),
+            chain_poly_type=chain_poly_type,
+        )
+
+    if filter_config is None:
+      hits = tuple(hit_generator(a3m))
+    else:
+      hits = _filter_hits(
+          hit_generator(a3m),
+          release_date_cutoff=filter_config.max_template_date,
+          max_subsequence_ratio=filter_config.max_subsequence_ratio,
+          min_align_ratio=filter_config.min_align_ratio,
+          min_hit_length=filter_config.min_hit_length,
+          deduplicate_sequences=filter_config.deduplicate_sequences,
+          max_hits=filter_config.max_hits,
+      )
+
+    return Templates(
+        query_sequence=query_sequence,
+        query_release_date=query_release_date,
+        hits=hits,
+        max_template_date=max_template_date,
+        structure_store=structure_store,
+    )
+
+  @property
+  def query_sequence(self) -> str:
+    return self._query_sequence
+
+  @property
+  def hits(self) -> tuple[Hit, ...]:
+    return self._hits
+
+  @property
+  def query_release_date(self) -> datetime.date | None:
+    return self._query_release_date
+
+  @property
+  def num_hits(self) -> int:
+    return len(self._hits)
+
+  @functools.cached_property
+  def release_date_cutoff(self) -> datetime.date:
+    if self.query_release_date is None:
+      return self._max_template_date
+    return min(
+        self._max_template_date,
+        self.query_release_date
+        - datetime.timedelta(days=_DAYS_BEFORE_QUERY_DATE),
+    )
+
+  def __repr__(self) -> str:
+    return f'Templates({self.num_hits} hits)'
+
+  def filter(
+      self,
+      *,
+      max_subsequence_ratio: float | None,
+      min_align_ratio: float | None,
+      min_hit_length: int | None,
+      deduplicate_sequences: bool,
+      max_hits: int | None,
+  ) -> Self:
+    """Returns a new Templates object with only the hits that pass all filters.
+
+    This also filters on query_release_date and max_template_date.
+
+    Args:
+      max_subsequence_ratio: If set, excludes hits which are an exact
+        subsequence of the query sequence, and longer than this ratio. Useful to
+        avoid ground truth leakage.
+      min_align_ratio: If set, excludes hits where the number of residues
+        aligned to the query is less than this proportion of the template
+        length.
+      min_hit_length: If set, excludes hits which have fewer residues than this.
+      deduplicate_sequences: Whether to exclude duplicate template sequences,
+        keeping only the first. This can be useful in increasing the diversity
+        of hits especially in the case of homomer hits.
+      max_hits: If set, excludes any hits which exceed this count.
+    """
+    filtered_hits = _filter_hits(
+        hits=self._hits,
+        release_date_cutoff=self.release_date_cutoff,
+        max_subsequence_ratio=max_subsequence_ratio,
+        min_align_ratio=min_align_ratio,
+        min_hit_length=min_hit_length,
+        deduplicate_sequences=deduplicate_sequences,
+        max_hits=max_hits,
+    )
+    return Templates(
+        query_sequence=self.query_sequence,
+        query_release_date=self.query_release_date,
+        hits=filtered_hits,
+        max_template_date=self._max_template_date,
+        structure_store=self._structure_store,
+    )
+
+  def get_hits_with_structures(
+      self,
+  ) -> Sequence[tuple[Hit, structure.Structure]]:
+    """Returns hits + Structures, Structures filtered to the hit's chain."""
+    results = []
+    structures = {struc.name.lower(): struc for struc in self.structures}
+    for hit in self.hits:
+      if not hit.is_valid:
+        raise InvalidTemplateError(
+            'Hits must be filtered before calling get_hits_with_structures.'
+        )
+      struc = structures[hit.pdb_id]
+      label_chain_id = struc.polymer_auth_asym_id_to_label_asym_id().get(
+          hit.auth_chain_id
+      )
+      results.append((hit, struc.filter(chain_id=label_chain_id)))
+    return results
+
+  def featurize(
+      self,
+      include_ligand_features: bool = True,
+  ) -> TemplateFeatures:
+    """Featurises the templates and returns a map of feature names to features.
+
+    NB: If you don't do any prefiltering, this method might be slow to run
+    as it has to fetch many CIFs and featurize them all.
+
+    Args:
+      include_ligand_features: Whether to compute ligand features.
+
+    Returns:
+      Template features: A mapping of template feature labels to features, which
+        may be numpy arrays, bytes objects, or for the special case of label
+        `ligand_features` (if `include_ligand_features` is True), a nested
+        feature map of labels to numpy arrays.
+
+    Raises:
+      InvalidTemplateError: If hits haven't been filtered before featurization.
+    """
+    hits_by_pdb_id = {}
+    for idx, hit in enumerate(self.hits):
+      if not hit.is_valid:
+        raise InvalidTemplateError(
+            f'Hits must be filtered before featurizing, got unprocessed {hit=}'
+        )
+      hits_by_pdb_id.setdefault(hit.pdb_id, []).append((idx, hit))
+
+    unsorted_features = []
+    for struc in self.structures:
+      pdb_id = str(struc.name).lower()
+      for idx, hit in hits_by_pdb_id[pdb_id]:
+        try:
+          label_chain_id = struc.polymer_auth_asym_id_to_label_asym_id()[
+              hit.auth_chain_id
+          ]
+          hit_features = {
+              **get_polymer_features(
+                  chain=struc.filter(chain_id=label_chain_id),
+                  chain_poly_type=hit.chain_poly_type,
+                  query_sequence_length=len(hit.query_sequence),
+                  query_to_hit_mapping=hit.query_to_hit_mapping,
+              ),
+          }
+          if include_ligand_features:
+            hit_features['ligand_features'] = _get_ligand_features(struc)
+          unsorted_features.append((idx, hit_features))
+        except Error as e:
+          raise type(e)(f'Failed to featurise {hit=}') from e
+
+    sorted_features = sorted(unsorted_features, key=lambda x: x[0])
+    sorted_features = [feat for _, feat in sorted_features]
+    return package_template_features(
+        hit_features=sorted_features,
+        include_ligand_features=include_ligand_features,
+    )
+
+  @property
+  def structures(self) -> Iterator[structure.Structure]:
+    """Yields template structures for each unique PDB ID among hits.
+
+    If there are multiple hits in the same Structure, the Structure will be
+    included only once by this method.
+
+    Yields:
+      A Structure object for each unique PDB ID among hits.
+
+    Raises:
+      HitDateError: If template's release date exceeds max cutoff date.
+    """
+
+    for hit in self.hits:
+      if hit.release_date > self.release_date_cutoff:  # pylint: disable=comparison-with-callable
+        raise HitDateError(
+            f'Invalid release date for hit {hit.pdb_id=}, when release date '
+            f'cutoff is {self.release_date_cutoff}.'
+        )
+
+    # Get the set of pdbs to load. In particular, remove duplicate PDB IDs.
+    targets_to_load = tuple({hit.pdb_id for hit in self.hits})
+
+    for target_name in targets_to_load:
+      yield structure.from_mmcif(
+          mmcif_string=self._structure_store.get_mmcif_str(target_name),
+          fix_mse_residues=True,
+          fix_arginines=True,
+          include_water=False,
+          include_bonds=False,
+          include_other=True,  # For non-standard polymer chains.
+      )
+
+
+def _parse_hit_description(description: str) -> tuple[str, str, int, int, int]:
+  """Parses the hmmsearch A3M sequence description line."""
+  # Example lines (protein, nucleic, no description):
+  # >4pqx_A/2-217 [subseq from] mol:protein length:217  Free text
+  # >4pqx_A/2-217 [subseq from] mol:na length:217  Free text
+  # >5g3r_A/1-55 [subseq from] mol:protein length:352
+  if match := re.fullmatch(_HIT_DESCRIPTION_REGEX, description):
+    return (
+        match['pdb_id'],
+        match['chain_id'],
+        int(match['start']),
+        int(match['end']),
+        int(match['length']),
+    )
+  else:
+    raise ValueError(f'Could not parse description "{description}"')
+
+
+def _parse_hit_metadata(
+    structure_store: structure_stores.StructureStore,
+    pdb_id: str,
+    auth_chain_id: str,
+) -> tuple[Any, str | None, Sequence[int] | None]:
+  """Parse hit metadata by parsing mmCIF from structure store."""
+  try:
+    cif = mmcif.from_string(structure_store.get_mmcif_str(pdb_id))
+  except structure_stores.NotFoundError:
+    logging.warning(
+        'Failed to get mmCIF for %s (author chain %s).', pdb_id, auth_chain_id
+    )
+    return None, None, None
+  release_date = mmcif.get_release_date(cif)
+
+  try:
+    struc = structure.from_parsed_mmcif(
+        cif,
+        model_id=structure.ModelID.ALL,
+        include_water=True,
+        include_other=True,
+        include_bonds=False,
+    )
+  except ValueError:
+    struc = structure.from_parsed_mmcif(
+        cif,
+        model_id=structure.ModelID.FIRST,
+        include_water=True,
+        include_other=True,
+        include_bonds=False,
+    )
+
+  sequence = struc.polymer_author_chain_single_letter_sequence(
+      include_missing_residues=True,
+      protein=True,
+      dna=True,
+      rna=True,
+      other=True,
+  )[auth_chain_id]
+
+  unresolved_res_ids = struc.filter(
+      chain_auth_asym_id=auth_chain_id
+  ).unresolved_residues.id
+
+  return release_date, sequence, unresolved_res_ids
+
+
+def get_polymer_features(
+    *,
+    chain: structure.Structure,
+    chain_poly_type: str,
+    query_sequence_length: int,
+    query_to_hit_mapping: Mapping[int, int],
+) -> Mapping[str, Any]:
+  """Returns features for this polymer chain.
+
+  Args:
+    chain: Structure object representing the template. Must be already filtered
+      to a single chain.
+    chain_poly_type: The chain polymer type (protein, DNA, RNA).
+    query_sequence_length: The length of the query sequence.
+    query_to_hit_mapping: 0-based query index to hit index mapping.
+
+  Returns:
+    A dictionary with polymer features for template_chain_id in the struc.
+
+  Raises:
+    ValueError: If the input structure contains more than just a single chain.
+  """
+  # if len(chain.polymer_auth_asym_id_to_label_asym_id()) != 1:
+  #   raise ValueError('The structure must be filtered to a single chain.')
+
+  if chain.name is None:
+    raise ValueError('Template structure must have a name.')
+
+  if chain.release_date is None:
+    raise ValueError(
+        f'Template structure {chain.name} must have a release date. You can do'
+        ' this by setting "_pdbx_audit_revision_history.revision_date" in the'
+        ' template mmCIF to a date in the ISO-8601 format (e.g. 1989-11-17).'
+    )
+
+  num_polymer_chains = len(chain.polymer_auth_asym_id_to_label_asym_id())
+  if num_polymer_chains != 1:
+    raise ValueError(
+        f'Template structure {chain.name} must be filtered to a single polymer'
+        f' chain but got a structure with {num_polymer_chains} polymer chains.'
+    )
+
+  auth_chain_id, label_chain_id = next(
+      iter(chain.polymer_auth_asym_id_to_label_asym_id().items())
+  )
+  chain_sequence = chain.chain_single_letter_sequence()[label_chain_id]
+
+  polymer = _POLYMERS[chain_poly_type]
+  res_arrays = chain.to_res_arrays(
+      include_missing_residues=True, atom_order=polymer.atom_order
+  )
+  positions = res_arrays.atom_positions
+  positions_mask = res_arrays.atom_mask
+  template_all_atom_positions = np.zeros(
+      (query_sequence_length, polymer.num_atom_types, 3), dtype=np.float64
+  )
+  template_all_atom_masks = np.zeros(
+      (query_sequence_length, polymer.num_atom_types), dtype=np.int64
+  )
+
+  template_sequence = ['-'] * query_sequence_length
+  for query_index, template_index in query_to_hit_mapping.items():
+    template_all_atom_positions[query_index] = positions[template_index]
+    template_all_atom_masks[query_index] = positions_mask[template_index]
+    template_sequence[query_index] = chain_sequence[template_index]
+
+  template_sequence = ''.join(template_sequence)
+  template_aatype = _encode_restype(chain_poly_type, template_sequence)
+  template_name = f'{chain.name.lower()}_{auth_chain_id}'
+  release_date = chain.release_date.strftime('%Y-%m-%d')
+  return {
+      'template_all_atom_positions': template_all_atom_positions,
+      'template_all_atom_masks': template_all_atom_masks,
+      'template_sequence': template_sequence.encode(),
+      'template_aatype': np.array(template_aatype, dtype=np.int32),
+      'template_domain_names': np.array(template_name.encode(), dtype=object),
+      'template_release_date': np.array(release_date.encode(), dtype=object),
+  }
+
+
+def _get_ligand_features(
+    struc: structure.Structure,
+) -> Mapping[str, Mapping[str, np.ndarray | bytes]]:
+  """Returns features for the ligands in this structure."""
+  ligand_struc = struc.filter_to_entity_type(ligand=True)
+  assert ligand_struc.coords is not None
+  assert ligand_struc.atom_name is not None
+  assert ligand_struc.atom_occupancy is not None
+
+  ligand_features = {}
+  for ligand_chain_id in ligand_struc.chains:
+    idxs = np.where(ligand_struc.chain_id == ligand_chain_id)[0]
+    if idxs.shape[0]:
+      ligand_features[ligand_chain_id] = {
+          'ligand_atom_positions': (
+              ligand_struc.coords[idxs, :].astype(np.float32)
+          ),
+          'ligand_atom_names': ligand_struc.atom_name[idxs].astype(object),
+          'ligand_atom_occupancies': (
+              ligand_struc.atom_occupancy[idxs].astype(np.float32)
+          ),
+          'ccd_id': ligand_struc.res_name[idxs][0].encode(),
+      }
+  return ligand_features
+
+
+def package_template_features(
+    *,
+    hit_features: Sequence[Mapping[str, Any]],
+    include_ligand_features: bool,
+) -> Mapping[str, Any]:
+  """Stacks polymer features, adds empty and keeps ligand features unstacked."""
+
+  features_to_include = set(_POLYMER_FEATURES)
+  if include_ligand_features:
+    features_to_include.update(_LIGAND_FEATURES)
+
+  features = {
+      feat: [single_hit_features[feat] for single_hit_features in hit_features]
+      for feat in features_to_include
+  }
+
+  stacked_features = {}
+  for k, v in features.items():
+    if k in _POLYMER_FEATURES:
+      v = np.stack(v, axis=0) if v else np.array([], dtype=_POLYMER_FEATURES[k])
+    stacked_features[k] = v
+
+  return stacked_features
+
+
+def _resolve_path(path: os.PathLike[str] | str) -> str:
+  """Resolves path for data dep paths, stringifies otherwise."""
+  # Data dependency paths: db baked into the binary.
+  resolved_path = resources.filename(path)
+  if os.path.exists(resolved_path):
+    return resolved_path
+  else:
+    # Other paths, e.g. local.
+    return str(path)
+
+
+def run_hmmsearch_with_a3m(
+    *,
+    database_path: os.PathLike[str] | str,
+    hmmsearch_config: msa_config.HmmsearchConfig,
+    max_a3m_query_sequences: int | None,
+    a3m: str | None,
+) -> str:
+  """Runs Hmmsearch to get a3m string of hits."""
+  searcher = hmmsearch.Hmmsearch(
+      binary_path=hmmsearch_config.hmmsearch_binary_path,
+      hmmbuild_binary_path=hmmsearch_config.hmmbuild_binary_path,
+      database_path=_resolve_path(database_path),
+      e_value=hmmsearch_config.e_value,
+      inc_e=hmmsearch_config.inc_e,
+      dom_e=hmmsearch_config.dom_e,
+      incdom_e=hmmsearch_config.incdom_e,
+      alphabet=hmmsearch_config.alphabet,
+      filter_f1=hmmsearch_config.filter_f1,
+      filter_f2=hmmsearch_config.filter_f2,
+      filter_f3=hmmsearch_config.filter_f3,
+      filter_max=hmmsearch_config.filter_max,
+  )
+  # STO enables us to annotate query non-gap columns as reference columns.
+  sto = parsers.convert_a3m_to_stockholm(a3m, max_a3m_query_sequences)
+  return searcher.query_with_sto(sto, model_construction='hand')

flax_model/alphafold3/data/tools/hmmbuild.py

@@ -0,0 +1,136 @@
+
+
+"""A Python wrapper for hmmbuild - construct HMM profiles from MSA."""
+
+import os
+import re
+import tempfile
+from typing import Literal
+
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import subprocess_utils
+
+
+class Hmmbuild(object):
+  """Python wrapper of the hmmbuild binary."""
+
+  def __init__(
+      self,
+      *,
+      binary_path: str,
+      singlemx: bool = False,
+      alphabet: str | None = None,
+  ):
+    """Initializes the Python hmmbuild wrapper.
+
+    Args:
+      binary_path: The path to the hmmbuild executable.
+      singlemx: Whether to use --singlemx flag. If True, it forces HMMBuild to
+        just use a common substitution score matrix.
+      alphabet: The alphabet to assert when building a profile. Useful when
+        hmmbuild cannot guess the alphabet. If None, no alphabet is asserted.
+
+    Raises:
+      RuntimeError: If hmmbuild binary not found within the path.
+    """
+    self._binary_path = binary_path
+    self._singlemx = singlemx
+    self._alphabet = alphabet
+
+    subprocess_utils.check_binary_exists(path=self._binary_path, name='hmmbuild')
+
+  def build_profile_from_sto(self, sto: str, model_construction='fast') -> str:
+    """Builds a HHM for the aligned sequences given as an A3M string.
+
+    Args:
+      sto: A string with the aligned sequences in the Stockholm format.
+      model_construction: Whether to use reference annotation in the msa to
+        determine consensus columns ('hand') or default ('fast').
+
+    Returns:
+      A string with the profile in the HMM format.
+
+    Raises:
+      RuntimeError: If hmmbuild fails.
+    """
+    return self._build_profile(
+        sto, informat='stockholm', model_construction=model_construction
+    )
+
+  def build_profile_from_a3m(self, a3m: str) -> str:
+    """Builds a HHM for the aligned sequences given as an A3M string.
+
+    Args:
+      a3m: A string with the aligned sequences in the A3M format.
+
+    Returns:
+      A string with the profile in the HMM format.
+
+    Raises:
+      RuntimeError: If hmmbuild fails.
+    """
+    lines = []
+    for sequence, description in parsers.lazy_parse_fasta_string(a3m):
+      sequence = re.sub('[a-z]+', '', sequence)  # Remove inserted residues.
+      lines.append(f'>{description}\n{sequence}\n')
+    msa = ''.join(lines)
+    return self._build_profile(msa, informat='afa')
+
+  def _build_profile(
+      self,
+      msa: str,
+      informat: Literal['afa', 'stockholm'],
+      model_construction: str = 'fast',
+  ) -> str:
+    """Builds a HMM for the aligned sequences given as an MSA string.
+
+    Args:
+      msa: A string with the aligned sequences, in A3M or STO format.
+      informat: One of 'afa' (aligned FASTA) or 'sto' (Stockholm).
+      model_construction: Whether to use reference annotation in the msa to
+        determine consensus columns ('hand') or default ('fast').
+
+    Returns:
+      A string with the profile in the HMM format.
+
+    Raises:
+      RuntimeError: If hmmbuild fails.
+      ValueError: If unspecified arguments are provided.
+    """
+    if model_construction not in {'hand', 'fast'}:
+      raise ValueError(f'Bad {model_construction=}. Only hand or fast allowed.')
+
+    with tempfile.TemporaryDirectory() as query_tmp_dir:
+      input_msa_path = os.path.join(query_tmp_dir, 'query.msa')
+      output_hmm_path = os.path.join(query_tmp_dir, 'output.hmm')
+
+      with open(input_msa_path, 'w') as f:
+        f.write(msa)
+
+      # Specify the format as we don't specify the input file extension. See
+      # https://github.com/EddyRivasLab/hmmer/issues/321 for more details.
+      cmd_flags = ['--informat', informat]
+      # If adding flags, we have to do so before the output and input:
+      if model_construction == 'hand':
+        cmd_flags.append(f'--{model_construction}')
+      if self._singlemx:
+        cmd_flags.append('--singlemx')
+      if self._alphabet:
+        cmd_flags.append(f'--{self._alphabet}')
+
+      cmd_flags.extend([output_hmm_path, input_msa_path])
+
+      cmd = [self._binary_path, *cmd_flags]
+
+      subprocess_utils.run(
+          cmd=cmd,
+          cmd_name='Hmmbuild',
+          log_stdout=False,
+          log_stderr=True,
+          log_on_process_error=True,
+      )
+
+      with open(output_hmm_path) as f:
+        hmm = f.read()
+
+    return hmm

flax_model/alphafold3/data/tools/hmmsearch.py

@@ -0,0 +1,141 @@
+
+
+"""A Python wrapper for hmmsearch - search profile against a sequence db."""
+
+import os
+import tempfile
+
+from absl import logging
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import hmmbuild
+from flax_model.alphafold3.data.tools import subprocess_utils
+
+
+class Hmmsearch(object):
+  """Python wrapper of the hmmsearch binary."""
+
+  def __init__(
+      self,
+      *,
+      binary_path: str,
+      hmmbuild_binary_path: str,
+      database_path: str,
+      alphabet: str = 'amino',
+      filter_f1: float | None = None,
+      filter_f2: float | None = None,
+      filter_f3: float | None = None,
+      e_value: float | None = None,
+      inc_e: float | None = None,
+      dom_e: float | None = None,
+      incdom_e: float | None = None,
+      filter_max: bool = False,
+  ):
+    """Initializes the Python hmmsearch wrapper.
+
+    Args:
+      binary_path: The path to the hmmsearch executable.
+      hmmbuild_binary_path: The path to the hmmbuild executable. Used to build
+        an hmm from an input a3m.
+      database_path: The path to the hmmsearch database (FASTA format).
+      alphabet: Chain type e.g. amino, rna, dna.
+      filter_f1: MSV and biased composition pre-filter, set to >1.0 to turn off.
+      filter_f2: Viterbi pre-filter, set to >1.0 to turn off.
+      filter_f3: Forward pre-filter, set to >1.0 to turn off.
+      e_value: E-value criteria for inclusion in tblout.
+      inc_e: E-value criteria for inclusion in MSA/next round.
+      dom_e: Domain e-value criteria for inclusion in tblout.
+      incdom_e: Domain e-value criteria for inclusion of domains in MSA/next
+        round.
+      filter_max: Remove all filters, will ignore all filter_f* settings.
+
+    Raises:
+      RuntimeError: If hmmsearch binary not found within the path.
+    """
+    self._binary_path = binary_path
+    self._hmmbuild_runner = hmmbuild.Hmmbuild(
+        alphabet=alphabet, binary_path=hmmbuild_binary_path
+    )
+    self._database_path = database_path
+    flags = []
+    if filter_max:
+      flags.append('--max')
+    else:
+      if filter_f1 is not None:
+        flags.extend(('--F1', filter_f1))
+      if filter_f2 is not None:
+        flags.extend(('--F2', filter_f2))
+      if filter_f3 is not None:
+        flags.extend(('--F3', filter_f3))
+
+    if e_value is not None:
+      flags.extend(('-E', e_value))
+    if inc_e is not None:
+      flags.extend(('--incE', inc_e))
+    if dom_e is not None:
+      flags.extend(('--domE', dom_e))
+    if incdom_e is not None:
+      flags.extend(('--incdomE', incdom_e))
+
+    self._flags = tuple(map(str, flags))
+
+    subprocess_utils.check_binary_exists(
+        path=self._binary_path, name='hmmsearch'
+    )
+
+    if not os.path.exists(self._database_path):
+      logging.error('Could not find hmmsearch database %s', database_path)
+      raise ValueError(f'Could not find hmmsearch database {database_path}')
+
+  def query_with_hmm(self, hmm: str) -> str:
+    """Queries the database using hmmsearch using a given hmm."""
+    with tempfile.TemporaryDirectory() as query_tmp_dir:
+      hmm_input_path = os.path.join(query_tmp_dir, 'query.hmm')
+      sto_out_path = os.path.join(query_tmp_dir, 'output.sto')
+      with open(hmm_input_path, 'w') as f:
+        f.write(hmm)
+
+      cmd = [
+          self._binary_path,
+          '--noali',  # Don't include the alignment in stdout.
+          *('--cpu', '8'),
+      ]
+      # If adding flags, we have to do so before the output and input:
+      if self._flags:
+        cmd.extend(self._flags)
+      cmd.extend([
+          *('-A', sto_out_path),
+          hmm_input_path,
+          self._database_path,
+      ])
+
+      subprocess_utils.run(
+          cmd=cmd,
+          cmd_name=f'Hmmsearch ({os.path.basename(self._database_path)})',
+          log_stdout=False,
+          log_stderr=True,
+          log_on_process_error=True,
+      )
+
+      with open(sto_out_path) as f:
+        a3m_out = parsers.convert_stockholm_to_a3m(
+            f, remove_first_row_gaps=False, linewidth=60
+        )
+
+    return a3m_out
+
+  def query_with_a3m(self, a3m_in: str) -> str:
+    """Query the database using hmmsearch using a given a3m."""
+
+    # Only the "fast" model construction makes sense with A3M, as it doesn't
+    # have any way to annotate reference columns.
+    hmm = self._hmmbuild_runner.build_profile_from_a3m(a3m_in)
+    return self.query_with_hmm(hmm)
+
+  def query_with_sto(
+      self, msa_sto: str, model_construction: str = 'fast'
+  ) -> str:
+    """Queries the database using hmmsearch using a given stockholm msa."""
+    hmm = self._hmmbuild_runner.build_profile_from_sto(
+        msa_sto, model_construction=model_construction
+    )
+    return self.query_with_hmm(hmm)

flax_model/alphafold3/data/tools/jackhmmer.py

@@ -0,0 +1,326 @@
+
+
+"""Library to run Jackhmmer from Python."""
+
+from collections.abc import Iterable, Sequence
+from concurrent import futures
+import heapq
+import os
+import pathlib
+import shutil
+import tempfile
+import time
+
+from absl import logging
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import shards
+from flax_model.alphafold3.data.tools import msa_tool
+from flax_model.alphafold3.data.tools import subprocess_utils
+
+
+class Jackhmmer(msa_tool.MsaTool):
+  """Python wrapper of the Jackhmmer binary."""
+
+  def __init__(
+      self,
+      *,
+      binary_path: str,
+      database_path: str,
+      n_cpu: int = 8,
+      n_iter: int = 3,
+      e_value: float | None = 1e-3,
+      z_value: float | int | None = None,
+      dom_e: float | None = None,
+      dom_z_value: float | int | None = None,
+      max_sequences: int = 5000,
+      filter_f1: float = 5e-4,
+      filter_f2: float = 5e-5,
+      filter_f3: float = 5e-7,
+      max_threads: int | None = None,
+      **unused_kwargs,
+  ):
+    """Initializes the Python Jackhmmer wrapper.
+
+    NOTE: The MSA obtained by running against sharded dbs won't be always
+    exactly the same as the MSA obtained by running against an unsharded db.
+    This is because of Jackhmmer deduplication logic, which won't spot duplicate
+    hits across multiple shards. Usually this means that the sharded search
+    finds more hits (likely bounded by the number of shards), but this should
+    not pose an issue given how the results are used downstream. The problem is
+    more pronounced with deep MSAs and lower in the hit list (higher e-values).
+
+    Make sure to set the Z and domZ values when searching against a sharded
+    database, otherwise the results won't match the normal unsharded search.
+
+    Args:
+      binary_path: The path to the jackhmmer executable.
+      database_path: The path to the jackhmmer database (FASTA format). Sharded
+        file specs, e.g. `<db_path>@<num_shards>`, are supported.
+      n_cpu: The number of CPUs to give Jackhmmer.
+      n_iter: The number of Jackhmmer iterations.
+      e_value: The E-value, see Jackhmmer docs for more details.
+      z_value: The Z-value representing the number of comparisons done (i.e
+        correct database size) for E-value calculation. Make sure to set this
+        when searching against a sharded database, otherwise the e-values will
+        be incorrectly scaled.
+      dom_e: Domain e-value criteria for inclusion in tblout.
+      dom_z_value: Domain z-value representing the number of comparisons done
+        (i.e correct database size) for domain E-value calculation. Make sure to
+        set this when searching against a sharded database, otherwise the domain
+        e-values will be incorrectly scaled.
+      max_sequences: Maximum number of sequences to return in the MSA.
+      filter_f1: MSV and biased composition pre-filter, set to >1.0 to turn off.
+      filter_f2: Viterbi pre-filter, set to >1.0 to turn off.
+      filter_f3: Forward pre-filter, set to >1.0 to turn off.
+      max_threads: If given, the maximum number of threads used when running
+        sharded databases.
+
+    Raises:
+      RuntimeError: If Jackhmmer binary not found within the path.
+      ValueError: If an invalid configuration is provided in the args.
+    """
+    self._database_path = database_path
+
+    if shard_paths := shards.get_sharded_paths(self._database_path):
+      if n_iter != 1:
+        raise ValueError('For a sharded db, only n_iter=1 is supported.')
+      if z_value is None:
+        raise ValueError(
+            'The Z-value must be set when searching against a sharded database '
+            'to correctly scale e-values.'
+        )
+      if max_sequences <= 1:
+        raise ValueError(
+            'max_sequences must be greater than 1 when running in sharded '
+            'mode, because each shard would return only the query sequence.'
+        )
+
+      self._shard_paths = shard_paths
+      self._max_threads = len(self._shard_paths)
+      if max_threads is not None:
+        self._max_threads = min(max_threads, self._max_threads)
+      logging.info('Jackhmmer running with max_threads = %d', self._max_threads)
+    else:
+      self._shard_paths = None
+      self._max_threads = None
+
+    self._binary_path = binary_path
+    subprocess_utils.check_binary_exists(
+        path=self._binary_path, name='Jackhmmer'
+    )
+
+    self._n_cpu = n_cpu
+    self._n_iter = n_iter
+    self._e_value = e_value
+    self._z_value = z_value
+    self._dom_e = dom_e
+    self._dom_z_value = dom_z_value
+    self._max_sequences = max_sequences
+    self._filter_f1 = filter_f1
+    self._filter_f2 = filter_f2
+    self._filter_f3 = filter_f3
+
+    # If Jackhmmer supports the --seq_limit flag (via our patch), use it to
+    # prevent writing out redundant sequences and increasing peak memory usage.
+    # If not, the Jackhmmer will be run without the --seq_limit flag.
+    self._supports_seq_limit = subprocess_utils.jackhmmer_seq_limit_supported(
+        self._binary_path
+    )
+
+  def query(self, target_sequence: str) -> msa_tool.MsaToolResult:
+    """Query the database (sharded or unsharded) using Jackhmmer."""
+    if self._shard_paths:
+      # Sharded case, run the query against each database shard in parallel.
+      logging.info(
+          'Query sequence (sharded db): %s',
+          target_sequence
+          if len(target_sequence) <= 16
+          else f'{target_sequence[:16]}... (len {len(target_sequence)})',
+      )
+
+      global_temp_dir = tempfile.mkdtemp()
+
+      def _query_shard_fn(
+          shard_path: str,
+      ) -> tuple[msa_tool.MsaToolResult, float]:
+        t_start = time.time()
+        result = self._query_db_shard(
+            target_sequence=target_sequence,
+            db_shard_path=shard_path,
+            get_tblout=True,  # Tblout contains e-values needed for merging.
+            global_temp_dir=global_temp_dir,
+        )
+        return result, time.time() - t_start
+
+      with futures.ThreadPoolExecutor(max_workers=self._max_threads) as ex:
+        tool_outputs, timings = zip(*ex.map(_query_shard_fn, self._shard_paths))
+
+      logging.info(
+          'Finished query for %d shards, shard timings (seconds): %s',
+          len(tool_outputs),
+          ', '.join(f'{t:.1f}' for t in timings),
+      )
+
+      shutil.rmtree(global_temp_dir, ignore_errors=True)
+      return _merge_jackhmmer_results(tool_outputs, self._max_sequences)
+
+    else:
+      # Non-sharded case, run the query against the whole database.
+      logging.info(
+          'Query sequence (non-sharded db): %s',
+          target_sequence
+          if len(target_sequence) <= 16
+          else f'{target_sequence[:16]}... (len {len(target_sequence)})',
+      )
+      return self._query_db_shard(
+          target_sequence=target_sequence,
+          db_shard_path=self._database_path,
+          get_tblout=False,
+      )
+
+  def _query_db_shard(
+      self,
+      *,
+      target_sequence: str,
+      db_shard_path: str,
+      get_tblout: bool,
+      global_temp_dir: str | None = None,
+  ) -> msa_tool.MsaToolResult:
+    """Query the database shard using Jackhmmer."""
+
+    with tempfile.TemporaryDirectory(dir=global_temp_dir) as query_tmp_dir:
+      input_fasta_path = os.path.join(query_tmp_dir, 'query.fasta')
+      subprocess_utils.create_query_fasta_file(
+          sequence=target_sequence, path=input_fasta_path
+      )
+      output_sto_path = os.path.join(query_tmp_dir, 'output.sto')
+      pathlib.Path(output_sto_path).touch()
+
+      # The F1/F2/F3 are the expected proportion to pass each of the filtering
+      # stages (which get progressively more expensive), reducing these
+      # speeds up the pipeline at the expensive of sensitivity.  They are
+      # currently set very low to make querying Mgnify run in a reasonable
+      # amount of time.
+      cmd_flags = [
+          *('-o', '/dev/null'),  # Don't pollute stdout with Jackhmmer output.
+          *('-A', output_sto_path),
+          '--noali',
+          *('--F1', str(self._filter_f1)),
+          *('--F2', str(self._filter_f2)),
+          *('--F3', str(self._filter_f3)),
+          *('--cpu', str(self._n_cpu)),
+          *('-N', str(self._n_iter)),
+      ]
+
+      if get_tblout:
+        output_tblout_path = pathlib.Path(query_tmp_dir, 'tblout.txt')
+        output_tblout_path.touch()
+        cmd_flags.extend(['--tblout', str(output_tblout_path)])
+      else:
+        output_tblout_path = None
+
+      # Report only sequences with E-values <= x in per-sequence output.
+      if self._e_value is not None:
+        cmd_flags.extend(['-E', str(self._e_value)])
+
+        # Use the same value as the reporting e-value (`-E` flag).
+        cmd_flags.extend(['--incE', str(self._e_value)])
+
+      if self._z_value is not None:
+        cmd_flags.extend(['-Z', str(self._z_value)])
+
+      if self._dom_z_value is not None:
+        cmd_flags.extend(['--domZ', str(self._dom_z_value)])
+
+      if self._dom_e is not None:
+        cmd_flags.extend(['--domE', str(self._dom_e)])
+
+      if self._max_sequences is not None and self._supports_seq_limit:
+        cmd_flags.extend(['--seq_limit', str(self._max_sequences)])
+
+      # The input FASTA and the input db are the last two arguments.
+      cmd = [self._binary_path] + cmd_flags + [input_fasta_path, db_shard_path]
+
+      subprocess_utils.run(
+          cmd=cmd,
+          cmd_name=f'Jackhmmer ({os.path.basename(db_shard_path)})',
+          log_stdout=False,
+          log_stderr=True,
+          log_on_process_error=True,
+      )
+
+      with open(output_sto_path) as f:
+        a3m = parsers.convert_stockholm_to_a3m(
+            f, max_sequences=self._max_sequences
+        )
+
+      # Get the tabular output which has e.g. e-value for each target.
+      tbl = '' if output_tblout_path is None else output_tblout_path.read_text()
+
+      return msa_tool.MsaToolResult(
+          target_sequence=target_sequence,
+          a3m=a3m,
+          e_value=self._e_value,
+          tblout=tbl,
+      )
+
+
+def _merge_jackhmmer_results(
+    jh_results: Sequence[msa_tool.MsaToolResult], max_sequences: int
+) -> msa_tool.MsaToolResult:
+  """Merges Jackhmmer result protos into a single one."""
+  assert len(set(jh_res.target_sequence for jh_res in jh_results)) == 1
+  assert len(set(jh_res.e_value for jh_res in jh_results)) == 1
+
+  # Parse the TBL output, create a mapping from hit name to TBL line.
+  parsed_tbl = {}
+  for jh_result in jh_results:
+    assert jh_result.tblout is not None
+    for line in jh_result.tblout.splitlines():
+      if not line.startswith('#'):
+        parsed_tbl[line.partition(' ')[0]] = line
+
+  # Create an iterator and merge a3m info with tbl info.
+  def _merged_a3m_tbl_iter(a3m: str) -> Iterable[tuple[str, str, str, str]]:
+    # Don't parse the entire a3m, lazily parse only as many sequences as needed.
+    iterator = iter(parsers.lazy_parse_fasta_string(a3m))
+    next(iterator)  # Skip the query which isn't present in tblout.
+    for sequence, description in iterator:
+      name = description.partition(' ')[0].partition('/')[0]
+      if tbl_info := parsed_tbl.get(name):
+        # Skip sequences for which we don't have tbl information.
+        yield sequence, description, tbl_info, name
+
+  def sort_key(seq_data: tuple[str, str, str, str]) -> tuple[float, str]:
+    unused_seq, unused_description, tbl_info, name = seq_data
+    # Tblout lines have 19 whitespace delimited columns. "-" used if no value
+    # present. We want e-value in column with index 4, so do only 5 splits.
+    # Use the name in case of a e-value tie.
+    return float(tbl_info.split(maxsplit=5)[4]), name
+
+  # A3M/TBL is sorted by e-value and name, hence we can merge them efficiently.
+  merged_a3m_and_tblout = heapq.merge(
+      *[_merged_a3m_tbl_iter(res.a3m) for res in jh_results],
+      key=sort_key,
+  )
+
+  # Truncate the a3m to max_sequences. Do not truncate the tblout.
+  merged_tblout = []
+  merged_a3m = [f'>query\n{jh_results[0].target_sequence}']
+  for seq, description, tbl_info, _ in merged_a3m_and_tblout:
+    merged_tblout.append(tbl_info)
+    if len(merged_a3m) < max_sequences:
+      merged_a3m.append(f'>{description}\n{seq}')
+
+  logging.info(
+      'Limiting merged MSA depth from %d to %d',
+      len(merged_tblout),
+      max_sequences,
+  )
+
+  return msa_tool.MsaToolResult(
+      target_sequence=jh_results[0].target_sequence,
+      a3m='\n'.join(merged_a3m),
+      e_value=jh_results[0].e_value,
+      tblout=None,  # We no longer need the tblout.
+  )

flax_model/alphafold3/data/tools/mmseqs.py

@@ -0,0 +1,303 @@
+
+
+"""Library to run Mmseqs from Python."""
+
+import os
+import tempfile
+import subprocess
+import time
+from typing import List, Dict, Optional
+
+
+from absl import logging
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import msa_tool
+from flax_model.alphafold3.data.tools import subprocess_utils
+import shlex
+
+import shutil
+import re
+
+
+class Mmseqs(msa_tool.MsaTool):
+    
+    _database_types = ["mgnify", "uniprot_cluster_annot", "uniref90", "small_bfd"]
+
+    def __init__(
+        self,
+        *,
+        binary_path: str,
+        database_path: str,
+        n_cpu: int = 8,
+        e_value: float = 1e-4,
+        max_sequences: int = 5000,
+        use_gpu: int = 1,
+        msa_format_mode: int = 4,
+        mmseqs_options: str = "",
+        result2msa_options: str = "",
+    ):
+        self.binary_path = binary_path
+        self.database_path = database_path
+        
+        subprocess_utils.check_binary_exists(path=self.binary_path, name='MMseqs')
+        
+        if not os.path.exists(self.database_path):
+            raise ValueError(f'Database not found: {database_path}')
+        
+        self.n_cpu = n_cpu
+        self.e_value = e_value
+        self.max_sequences = max_sequences
+        self.use_gpu = use_gpu
+        self.msa_format_mode = msa_format_mode
+        self.mmseqs_options = mmseqs_options
+        self.result2msa_options = result2msa_options
+
+    def _get_gpu_memory_gb(self) -> float:
+        hip_devices = os.environ.get('HIP_VISIBLE_DEVICES', '')
+        cuda_devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')
+
+        hy_smi = shutil.which('hy-smi')
+        nvidia_smi = shutil.which('nvidia-smi')
+
+        prefer_hip = bool(hip_devices or hy_smi)
+
+        device_id: Optional[str] = None
+        if prefer_hip and hip_devices:
+            ids = [d.strip() for d in hip_devices.split(',') if d.strip()]
+            device_id = ids[0] if ids else None
+        elif cuda_devices:
+            ids = [d.strip() for d in cuda_devices.split(',') if d.strip()]
+            device_id = ids[0] if ids else None
+
+        # 1) hy-smi (HIP path)
+        if hy_smi:
+            try:
+                mem = self._get_gpu_memory_gb_via_hy_smi(hy_smi, device_id)
+                if mem is not None:
+                    logging.info(f"GPU memory (hy-smi): {mem:.1f}GB")
+                    return mem
+            except Exception as e:  # noqa: BLE001
+                logging.warning(f"hy-smi query failed: {e}")
+
+        # 2) NVIDIA fallback when CUDA devices are provided
+        if nvidia_smi and cuda_devices:
+            try:
+                ids = [d.strip() for d in cuda_devices.split(',') if d.strip()]
+                if ids:
+                    cmd = [
+                        nvidia_smi,
+                        '--query-gpu=memory.total',
+                        '--format=csv,noheader,nounits',
+                        f'--id={ids[0]}',
+                    ]
+                    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+                    memory_mb = float(result.stdout.strip())
+                    return memory_mb / 1024.0
+            except (subprocess.CalledProcessError, ValueError, FileNotFoundError) as e:
+                logging.warning(f"nvidia-smi query failed: {e}")
+
+        logging.info("Falling back to default GPU memory: 40GB")
+        return 40.0
+
+    def _parse_memory_gb_from_text(self, text: str) -> Optional[float]:
+        units = {
+            'gib': 1.0,
+            'gb': 1.0,
+            'mib': 1.0 / 1024.0,
+            'mb': 1.0 / 1024.0,
+            'kib': 1.0 / (1024.0 * 1024.0),
+            'kb': 1.0 / (1024.0 * 1024.0),
+            'b': 1.0 / (1024.0 * 1024.0 * 1024.0),
+        }
+        pattern = re.compile(r"(\d+(?:\.\d+)?)\s*(GiB|GB|MiB|MB|KiB|KB|B)", re.IGNORECASE)
+        candidates = []
+        for match in pattern.finditer(text):
+            value = float(match.group(1))
+            unit = match.group(2).lower()
+            factor = units.get(unit, None)
+            if factor is not None:
+                candidates.append(value * factor)
+        if candidates:
+            return max(candidates)
+        # rocm-smi often prints bytes as: Total Memory (B): 17163091968
+        bytes_pattern = re.compile(r"Total\s*Memory\s*\(B\)\s*[:=]\s*(\d{6,})", re.IGNORECASE)
+        m = bytes_pattern.search(text)
+        if m:
+            try:
+                return float(m.group(1)) / (1024.0 * 1024.0 * 1024.0)
+            except ValueError:
+                return None
+        # generic Total: <bytes>
+        generic_bytes = re.compile(r"Total\s*[:=]\s*(\d{6,})", re.IGNORECASE)
+        m2 = generic_bytes.search(text)
+        if m2:
+            try:
+                return float(m2.group(1)) / (1024.0 * 1024.0 * 1024.0)
+            except ValueError:
+                return None
+        return None
+
+    def _get_gpu_memory_gb_via_hy_smi(self, hy_smi_path: str, device_id: Optional[str]) -> Optional[float]:
+        candidate_cmds = []
+        # 1) Plain text output exactly as user provided; most reliable for MiB lines
+        if device_id is not None:
+            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '-d', str(device_id)])
+        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram'])
+        # 2) CSV/JSON (may not include units next to values, kept as fallback)
+        if device_id is not None:
+            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--csv', '-d', str(device_id)])
+        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--csv'])
+        if device_id is not None:
+            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--json', '-d', str(device_id)])
+        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--json'])
+        # 3) All-in-one info as a last resort
+        candidate_cmds.append([hy_smi_path, '-a'])
+
+        for cmd in candidate_cmds:
+            try:
+                result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+                # Prefer exact vram parser for hy-smi
+                parsed = self._parse_hysmi_vram_total_gb(result.stdout, device_id)
+                if parsed is None:
+                    parsed = self._parse_memory_gb_from_text(result.stdout)
+                if parsed is not None and parsed > 0:
+                    return parsed
+            except subprocess.CalledProcessError:
+                continue
+        return None
+
+    def _parse_hysmi_vram_total_gb(self, text: str, device_id: Optional[str]) -> Optional[float]:
+        # Match lines like: HCU[0]          : vram Total Memory (MiB): 65520
+        pattern = re.compile(r"(?:DCU|HCU)\[(\d+)\].*?vram\s+Total\s+Memory\s*\(MiB\)\s*:\s*(\d+)", re.IGNORECASE)
+        matches = pattern.findall(text)
+        if not matches:
+            return None
+        values_gb = []
+        for dev, mib_str in matches:
+            try:
+                if device_id is not None and str(device_id) != str(dev):
+                    continue
+                mib = float(mib_str)
+                values_gb.append(mib / 1024.0)
+            except ValueError:
+                continue
+        if not values_gb:
+            return None
+        return max(values_gb)
+
+    
+    
+    def _get_database_type(self, database_path: str) -> str:
+        path_lower = database_path.lower()
+        for db_type in self._database_types:
+            if db_type in path_lower:
+                return db_type
+        return "unknown"
+    
+    def _calculate_memory_allocation(self, database_path: str) -> int:
+        db_type = self._get_database_type(database_path)
+        if db_type == "unknown":
+            return 8
+        
+        if not self.use_gpu:
+            return 8
+        
+        gpu_memory_gb = self._get_gpu_memory_gb()
+        
+        available_memory = max(0, gpu_memory_gb)
+        
+        if db_type == "mgnify":
+            small_dbs_memory = 3 * 2
+            allocated_memory = int(available_memory - small_dbs_memory)
+            allocated_memory = max(2, allocated_memory)
+        else:
+            allocated_memory = 2
+        
+        logging.info(f"Database {db_type}: allocated {allocated_memory}GB GPU memory")
+        return allocated_memory
+
+    def query(self, target_sequence: str) -> msa_tool.MsaToolResult:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            input_fasta = os.path.join(tmp_dir, 'query.fasta')
+            subprocess_utils.create_query_fasta_file(target_sequence, input_fasta)
+            
+            query_db = os.path.join(tmp_dir, 'queryDB')
+            self._run_createdb(input_fasta, query_db)
+            
+            result_db = os.path.join(tmp_dir, 'resultDB')
+            self._run_search(query_db, result_db)  
+            
+            output_sto = os.path.join(tmp_dir, 'output.sto')
+            self._run_result2msa(query_db, result_db, output_sto)
+            
+            with open(output_sto) as f:
+                a3m = self._parse_output(f)
+        
+        return msa_tool.MsaToolResult(
+            target_sequence=target_sequence,
+            a3m=a3m,
+            e_value=self.e_value
+        )
+
+    def _run_search(self, input_db: str, output_db: str):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            cmd = [
+                self.binary_path,
+                'search',
+                input_db,
+                self.database_path,
+                output_db,
+                tmp_dir,
+                '--threads', str(self.n_cpu),
+                '-e', str(self.e_value),
+                '--gpu', str(self.use_gpu),
+                *shlex.split(self.mmseqs_options)
+            ]
+            
+            if self.use_gpu:
+                max_gpu_mem = self._calculate_memory_allocation(self.database_path)
+                cmd.extend(['--max-gpu-mem', f'{max_gpu_mem}G'])
+                logging.info(f"Running MMseqs search with max GPU memory: {max_gpu_mem}GB")
+            else:
+                logging.info("Running MMseqs search in CPU mode")
+            
+            subprocess_utils.run(
+                cmd=cmd,
+                cmd_name='MMseqs2 search',
+                log_stderr=True
+            )
+
+    def _run_createdb(self, input_fasta: str, output_db: str):
+        cmd = [self.binary_path, 'createdb', input_fasta, output_db]
+        subprocess_utils.run(cmd, 'MMseqs2 createdb')
+
+    def _run_result2msa(self, query_db: str, result_db: str, output_sto: str):
+        cmd = [
+            self.binary_path,
+            'result2msa',
+            query_db,
+            self.database_path,
+            result_db,
+            output_sto,
+            '--msa-format-mode', str(self.msa_format_mode),
+            *shlex.split(self.result2msa_options)
+        ]
+        subprocess_utils.run(cmd, 'MMseqs2 result2msa')
+
+    def _parse_output(self, file_handle):
+        try:
+            return parsers.convert_stockholm_to_a3m(file_handle, self.max_sequences)
+        except Exception as e:
+            logging.warning(f"Stockholm parse failed: {e}")
+            
+            file_handle.seek(0)
+            content = file_handle.read()
+            debug_file = f"/tmp/stockholm_debug_{int(time.time())}.sto"
+            with open(debug_file, 'w') as f:
+                f.write(content)
+            logging.info(f"Saved debug file: {debug_file}")
+            
+            file_handle.seek(0)
+            return parsers.convert_mmseqs_stockholm_to_a3m(file_handle, self.max_sequences)
+
+

flax_model/alphafold3/data/tools/rdkit_utils.py

@@ -0,0 +1,535 @@
+
+
+"""Tools for calculating features for ligands."""
+
+import collections
+from collections.abc import Mapping, Sequence
+
+from absl import logging
+from flax_model.alphafold3.cpp import cif_dict
+import numpy as np
+import rdkit.Chem as rd_chem
+from rdkit.Chem import AllChem as rd_all_chem
+
+
+_RDKIT_MMCIF_TO_BOND_TYPE: Mapping[str, rd_chem.BondType] = {
+    'SING': rd_chem.BondType.SINGLE,
+    'DOUB': rd_chem.BondType.DOUBLE,
+    'TRIP': rd_chem.BondType.TRIPLE,
+}
+
+_RDKIT_BOND_TYPE_TO_MMCIF: Mapping[rd_chem.BondType, str] = {
+    v: k for k, v in _RDKIT_MMCIF_TO_BOND_TYPE.items()
+}
+
+_RDKIT_BOND_STEREO_TO_MMCIF: Mapping[rd_chem.BondStereo, str] = {
+    rd_chem.BondStereo.STEREONONE: 'N',
+    rd_chem.BondStereo.STEREOE: 'E',
+    rd_chem.BondStereo.STEREOZ: 'Z',
+    rd_chem.BondStereo.STEREOCIS: 'Z',
+    rd_chem.BondStereo.STEREOTRANS: 'E',
+}
+
+
+class MolFromMmcifError(Exception):
+  """Raised when conversion from mmCIF to RDKit Mol fails."""
+
+
+class UnsupportedMolBondError(Exception):
+  """Raised when we try to handle unsupported RDKit bonds."""
+
+
+def _populate_atoms_in_mol(
+    mol: rd_chem.Mol,
+    atom_names: Sequence[str],
+    atom_types: Sequence[str],
+    atom_charges: Sequence[int],
+    implicit_hydrogens: bool,
+    ligand_name: str,
+    atom_leaving_flags: Sequence[str],
+):
+  """Populate the atoms of a Mol given atom features.
+
+  Args:
+    mol: Mol object.
+    atom_names: Names of the atoms.
+    atom_types: Types of the atoms.
+    atom_charges: Charges of the atoms.
+    implicit_hydrogens: Whether to mark the atoms to allow implicit Hs.
+    ligand_name: Name of the ligand which the atoms are in.
+    atom_leaving_flags: Whether the atom is possibly a leaving atom. Values from
+      the CCD column `_chem_comp_atom.pdbx_leaving_atom_flag`. The expected
+      values are 'Y' (yes), 'N' (no), '?' (unknown/unset, interpreted as no).
+
+  Raises:
+    ValueError: If atom type is invalid.
+  """
+  # Map atom names to the position they will take in the rdkit molecule.
+  atom_name_to_idx = {name: i for i, name in enumerate(atom_names)}
+
+  for atom_name, atom_type, atom_charge, atom_leaving_flag in zip(
+      atom_names, atom_types, atom_charges, atom_leaving_flags, strict=True
+  ):
+    try:
+      if atom_type == 'X':
+        atom_type = '*'
+      atom = rd_chem.Atom(atom_type)
+    except RuntimeError as e:
+      raise ValueError(f'Failed to use atom type: {str(e)}') from e
+
+    if not implicit_hydrogens:
+      atom.SetNoImplicit(True)
+
+    atom.SetProp('atom_name', atom_name)
+    atom.SetProp('atom_leaving_flag', atom_leaving_flag)
+    atom.SetFormalCharge(atom_charge)
+    residue_info = rd_chem.AtomPDBResidueInfo()
+    residue_info.SetName(_format_atom_name(atom_name, atom_type))
+    residue_info.SetIsHeteroAtom(True)
+    residue_info.SetResidueName(ligand_name)
+    residue_info.SetResidueNumber(1)
+    atom.SetPDBResidueInfo(residue_info)
+    atom_index = mol.AddAtom(atom)
+    assert atom_index == atom_name_to_idx[atom_name]
+
+
+def _populate_bonds_in_mol(
+    mol: rd_chem.Mol,
+    atom_names: Sequence[str],
+    bond_begins: Sequence[str],
+    bond_ends: Sequence[str],
+    bond_orders: Sequence[str],
+    bond_is_aromatics: Sequence[bool],
+):
+  """Populate the bonds of a Mol given bond features.
+
+  Args:
+    mol: Mol object.
+    atom_names: Names of atoms in the molecule.
+    bond_begins: Names of atoms at the beginning of the bond.
+    bond_ends: Names of atoms at the end of the bond.
+    bond_orders: What order the bonds are.
+    bond_is_aromatics: Whether the bonds are aromatic.
+  """
+  atom_name_to_idx = {name: i for i, name in enumerate(atom_names)}
+  for begin, end, bond_type, is_aromatic in zip(
+      bond_begins, bond_ends, bond_orders, bond_is_aromatics, strict=True
+  ):
+    begin_name, end_name = atom_name_to_idx[begin], atom_name_to_idx[end]
+    bond_idx = mol.AddBond(begin_name, end_name, bond_type)
+    mol.GetBondWithIdx(bond_idx - 1).SetIsAromatic(is_aromatic)
+
+
+def sanitize_mol(mol, sort_alphabetically, remove_hydrogens) -> rd_chem.Mol:
+  # https://www.rdkit.org/docs/source/rdkit.Chem.rdmolops.html#rdkit.Chem.rdmolops.SanitizeMol
+  # Kekulize, check valencies, set aromaticity, conjugation and hybridization.
+  # This can repair e.g. incorrect aromatic flags.
+  rd_chem.SanitizeMol(mol)
+  if sort_alphabetically:
+    mol = sort_atoms_by_name(mol)
+  if remove_hydrogens:
+    mol = rd_chem.RemoveHs(mol)
+  return mol
+
+
+def _add_conformer_to_mol(mol, conformer, force_parse) -> rd_chem.Mol:
+  # Create conformer and use it to assign stereochemistry.
+  if conformer is not None:
+    try:
+      mol.AddConformer(conformer)
+      rd_chem.AssignStereochemistryFrom3D(mol)
+    except ValueError as e:
+      logging.warning('Failed to parse conformer: %s', e)
+      if not force_parse:
+        raise
+
+
+def mol_from_ccd_cif(
+    mol_cif: cif_dict.CifDict,
+    *,
+    force_parse: bool = False,
+    sort_alphabetically: bool = True,
+    remove_hydrogens: bool = True,
+    implicit_hydrogens: bool = False,
+) -> rd_chem.Mol:
+  """Creates an rdkit Mol object from a CCD mmcif data block.
+
+  The atoms are renumbered so that their names are in alphabetical order and
+  these names are placed on the atoms under property 'atom_name'.
+  Only hydrogens which are not required to define the molecule are removed.
+  For example, hydrogens that define stereochemistry around a double bond are
+  retained.
+  See this link for more details.
+  https://www.rdkit.org/docs/source/rdkit.Chem.rdmolops.html#rdkit.Chem.rdmolops.RemoveHs
+
+  Args:
+     mol_cif: An mmcif object representing a molecule.
+     force_parse: If True, assumes missing aromatic flags are false, substitutes
+       deuterium for hydrogen, assumes missing charges are 0 and ignores missing
+       conformer / stereochemistry information.
+     sort_alphabetically: True: sort atom alphabetically; False: keep CCD order
+     remove_hydrogens: if True, remove non-important hydrogens
+     implicit_hydrogens: Sets a marker on the atom that allows implicit Hs.
+
+  Returns:
+     An rdkit molecule, with the atoms sorted by name.
+
+  Raises:
+    MolToMmcifError: If conversion from mmcif to rdkit Mol fails. More detailed
+      error is available as this error's cause.
+  """
+  # Read data fields.
+  try:
+    atom_names, atom_types, atom_charges, atom_leaving_flags = parse_atom_data(
+        mol_cif, force_parse
+    )
+    bond_begins, bond_ends, bond_orders, bond_is_aromatics = parse_bond_data(
+        mol_cif, force_parse
+    )
+    lig_name = mol_cif['_chem_comp.id'][0].rjust(3)
+  except (KeyError, ValueError) as e:
+    raise MolFromMmcifError from e
+
+  # Build Rdkit molecule.
+  mol = rd_chem.RWMol()
+
+  # Per atom features.
+  try:
+    _populate_atoms_in_mol(
+        mol=mol,
+        atom_names=atom_names,
+        atom_types=atom_types,
+        atom_charges=atom_charges,
+        implicit_hydrogens=implicit_hydrogens,
+        ligand_name=lig_name,
+        atom_leaving_flags=atom_leaving_flags,
+    )
+  except (ValueError, RuntimeError) as e:
+    raise MolFromMmcifError from e
+
+  _populate_bonds_in_mol(
+      mol, atom_names, bond_begins, bond_ends, bond_orders, bond_is_aromatics
+  )
+
+  try:
+    conformer = _parse_ideal_conformer(mol_cif)
+  except (KeyError, ValueError) as e:
+    logging.warning('Failed to parse ideal conformer: %s', e)
+    if not force_parse:
+      raise MolFromMmcifError from e
+    conformer = None
+
+  mol.UpdatePropertyCache(strict=False)
+
+  try:
+    _add_conformer_to_mol(mol, conformer, force_parse)
+    mol = sanitize_mol(mol, sort_alphabetically, remove_hydrogens)
+  except (
+      ValueError,
+      rd_chem.KekulizeException,
+      rd_chem.AtomValenceException,
+  ) as e:
+    raise MolFromMmcifError from e
+
+  return mol
+
+
+def mol_to_ccd_cif(
+    mol: rd_chem.Mol,
+    component_id: str,
+    pdbx_smiles: str | None = None,
+    include_hydrogens: bool = True,
+) -> cif_dict.CifDict:
+  """Creates a CCD-like mmcif data block from an rdkit Mol object.
+
+  Only a subset of associated mmcif fields is populated, but that is
+  sufficient for further usage, e.g. in featurization code.
+
+  Atom names can be specified via `atom_name` property. For atoms with
+  unspecified value of that property, the name is assigned based on element type
+  and the order in the Mol object.
+
+  If the Mol object has associated conformers, atom positions from the first of
+  them will be populated in the resulting mmcif file.
+
+  Args:
+     mol: An rdkit molecule.
+     component_id: Name of the molecule to use in the resulting mmcif. That is
+       equivalent to CCD code.
+     pdbx_smiles: If specified, the value will be used to populate
+       `_chem_comp.pdbx_smiles`.
+     include_hydrogens: Whether to include atom and bond data involving
+       hydrogens.
+
+  Returns:
+     An mmcif data block corresponding for the given rdkit molecule.
+
+  Raises:
+    UnsupportedMolBond: When a molecule contains a bond that can't be
+      represented with mmcif.
+  """
+  mol = rd_chem.Mol(mol)
+  if include_hydrogens:
+    mol = rd_chem.AddHs(mol)
+  rd_chem.Kekulize(mol)
+
+  if mol.GetNumConformers() > 0:
+    ideal_conformer = mol.GetConformer(0).GetPositions()
+    ideal_conformer = np.vectorize(lambda x: f'{x:.3f}')(ideal_conformer)
+  else:
+    # No data will be populated in the resulting mmcif if the molecule doesn't
+    # have any conformers attached to it.
+    ideal_conformer = None
+
+  mol_cif = collections.defaultdict(list)
+  mol_cif['data_'] = [component_id]
+  mol_cif['_chem_comp.id'] = [component_id]
+  if pdbx_smiles:
+    mol_cif['_chem_comp.pdbx_smiles'] = [pdbx_smiles]
+
+  mol = assign_atom_names_from_graph(mol, keep_existing_names=True)
+
+  for atom_idx, atom in enumerate(mol.GetAtoms()):
+    element = atom.GetSymbol()
+    if not include_hydrogens and element in ('H', 'D'):
+      continue
+
+    mol_cif['_chem_comp_atom.comp_id'].append(component_id)
+    mol_cif['_chem_comp_atom.atom_id'].append(atom.GetProp('atom_name'))
+    mol_cif['_chem_comp_atom.type_symbol'].append(atom.GetSymbol().upper())
+    mol_cif['_chem_comp_atom.charge'].append(str(atom.GetFormalCharge()))
+    if ideal_conformer is not None:
+      coords = ideal_conformer[atom_idx]
+      mol_cif['_chem_comp_atom.pdbx_model_Cartn_x_ideal'].append(coords[0])
+      mol_cif['_chem_comp_atom.pdbx_model_Cartn_y_ideal'].append(coords[1])
+      mol_cif['_chem_comp_atom.pdbx_model_Cartn_z_ideal'].append(coords[2])
+
+  for bond in mol.GetBonds():
+    atom1 = bond.GetBeginAtom()
+    atom2 = bond.GetEndAtom()
+    if not include_hydrogens and (
+        atom1.GetSymbol() in ('H', 'D') or atom2.GetSymbol() in ('H', 'D')
+    ):
+      continue
+    mol_cif['_chem_comp_bond.comp_id'].append(component_id)
+    mol_cif['_chem_comp_bond.atom_id_1'].append(
+        bond.GetBeginAtom().GetProp('atom_name')
+    )
+    mol_cif['_chem_comp_bond.atom_id_2'].append(
+        bond.GetEndAtom().GetProp('atom_name')
+    )
+    try:
+      bond_type = bond.GetBondType()
+      # Older versions of RDKit did not have a DATIVE bond type. Convert it to
+      # SINGLE to match the AF3 training setup.
+      if bond_type == rd_chem.BondType.DATIVE:
+        bond_type = rd_chem.BondType.SINGLE
+      mol_cif['_chem_comp_bond.value_order'].append(
+          _RDKIT_BOND_TYPE_TO_MMCIF[bond_type]
+      )
+      mol_cif['_chem_comp_bond.pdbx_stereo_config'].append(
+          _RDKIT_BOND_STEREO_TO_MMCIF[bond.GetStereo()]
+      )
+    except KeyError as e:
+      raise UnsupportedMolBondError from e
+    mol_cif['_chem_comp_bond.pdbx_aromatic_flag'].append(
+        'Y' if bond.GetIsAromatic() else 'N'
+    )
+
+  return cif_dict.CifDict(mol_cif)
+
+
+def _format_atom_name(atom_name: str, atom_type: str) -> str:
+  """Formats an atom name to fit in the four characters specified in PDB.
+
+  See for example the following note on atom name formatting in PDB files:
+  https://www.cgl.ucsf.edu/chimera/docs/UsersGuide/tutorials/pdbintro.html#note1
+
+  Args:
+    atom_name: The unformatted atom name.
+    atom_type: The atom element symbol.
+
+  Returns:
+    formatted_atom_name: The formatted 4-character atom name.
+  """
+  atom_name = atom_name.strip()
+  atom_type = atom_type.strip().upper()
+  if len(atom_name) == 1:
+    return atom_name.rjust(2).ljust(4)
+  elif len(atom_name) == 2:
+    if atom_name == atom_type:
+      return atom_name.ljust(4)
+    return atom_name.center(4)
+  elif len(atom_name) == 3:
+    if atom_name[:2] == atom_type:
+      return atom_name.ljust(4)
+    return atom_name.rjust(4)
+  elif len(atom_name) == 4:
+    return atom_name
+  else:
+    raise ValueError(
+        f'Atom name `{atom_name}` has more than four characters '
+        'or is an empty string.'
+    )
+
+
+def parse_atom_data(
+    mol_cif: cif_dict.CifDict | Mapping[str, Sequence[str]], force_parse: bool
+) -> tuple[Sequence[str], Sequence[str], Sequence[int], Sequence[str]]:
+  """Parses atoms. If force_parse is True, fix deuterium and missing charge."""
+  atom_types = [t.capitalize() for t in mol_cif['_chem_comp_atom.type_symbol']]
+  atom_names = mol_cif['_chem_comp_atom.atom_id']
+  atom_charges = mol_cif['_chem_comp_atom.charge']
+  atom_leaving_flags = ['?'] * len(atom_names)
+  if '_chem_comp_atom.pdbx_leaving_atom_flag' in mol_cif:
+    atom_leaving_flags = mol_cif['_chem_comp_atom.pdbx_leaving_atom_flag']
+
+  if force_parse:
+    # Replace missing charges with 0.
+    atom_charges = [charge if charge != '?' else '0' for charge in atom_charges]
+    # Deuterium for hydrogen.
+    atom_types = [type_ if type_ != 'D' else 'H' for type_ in atom_types]
+
+  atom_charges = [int(atom_charge) for atom_charge in atom_charges]
+  return atom_names, atom_types, atom_charges, atom_leaving_flags
+
+
+def parse_bond_data(
+    mol_cif: cif_dict.CifDict | Mapping[str, Sequence[str]], force_parse: bool
+) -> tuple[
+    Sequence[str], Sequence[str], Sequence[rd_chem.BondType], Sequence[bool]
+]:
+  """Parses bond data. If force_parse is True, ignore missing aromatic flags."""
+  # The bond table isn't present if there are no bonds. Use [] in that case.
+  begin_atoms = mol_cif.get('_chem_comp_bond.atom_id_1', [])
+  end_atoms = mol_cif.get('_chem_comp_bond.atom_id_2', [])
+  orders = mol_cif.get('_chem_comp_bond.value_order', [])
+  bond_types = [_RDKIT_MMCIF_TO_BOND_TYPE[order] for order in orders]
+
+  try:
+    aromatic_flags = mol_cif.get('_chem_comp_bond.pdbx_aromatic_flag', [])
+    is_aromatic = [{'Y': True, 'N': False}[flag] for flag in aromatic_flags]
+  except KeyError:
+    if force_parse:
+      # Set them all to not aromatic.
+      is_aromatic = [False for _ in begin_atoms]
+    else:
+      raise
+
+  return begin_atoms, end_atoms, bond_types, is_aromatic
+
+
+def _parse_ideal_conformer(mol_cif: cif_dict.CifDict) -> rd_chem.Conformer:
+  """Builds a conformer containing the ideal coordinates from the CCD.
+
+  Args:
+     mol_cif: An mmcif object representing a molecule.
+
+  Returns:
+     An rdkit conformer filled with the ideal positions from the mmcif.
+
+  Raises:
+     ValueError: if the positions can't be interpreted.
+  """
+  atom_x = [
+      float(x) for x in mol_cif['_chem_comp_atom.pdbx_model_Cartn_x_ideal']
+  ]
+  atom_y = [
+      float(y) for y in mol_cif['_chem_comp_atom.pdbx_model_Cartn_y_ideal']
+  ]
+  atom_z = [
+      float(z) for z in mol_cif['_chem_comp_atom.pdbx_model_Cartn_z_ideal']
+  ]
+  atom_positions = zip(atom_x, atom_y, atom_z, strict=True)
+
+  conformer = rd_chem.Conformer(len(atom_x))
+  for atom_index, atom_position in enumerate(atom_positions):
+    conformer.SetAtomPosition(atom_index, atom_position)
+
+  return conformer
+
+
+def sort_atoms_by_name(mol: rd_chem.Mol) -> rd_chem.Mol:
+  """Sorts the atoms in the molecule by their names."""
+  atom_names = {
+      atom.GetProp('atom_name'): atom.GetIdx() for atom in mol.GetAtoms()
+  }
+
+  # Sort the name, int tuples by the names.
+  sorted_atom_names = sorted(atom_names.items())
+
+  # Zip these tuples back together to the sorted indices.
+  _, new_order = zip(*sorted_atom_names, strict=True)
+
+  # Reorder the molecule.
+  # new_order is effectively an argsort of the names.
+  return rd_chem.RenumberAtoms(mol, new_order)
+
+
+def assign_atom_names_from_graph(
+    mol: rd_chem.Mol,
+    keep_existing_names: bool = False,
+) -> rd_chem.Mol:
+  """Assigns atom names from the molecular graph.
+
+  The atom name is stored as an atom property 'atom_name', accessible
+  with atom.GetProp('atom_name'). If the property is already specified, and
+  keep_existing_names is True we keep the original name.
+
+  We traverse the graph in the order of the rdkit atom index and give each atom
+  a name equal to '{ELEMENT_TYPE}{INDEX}'. E.g. C5 is the name for the fifth
+  unnamed carbon encountered.
+
+  NOTE: A new mol is returned, the original is not changed in place.
+
+  Args:
+    mol: Mol object.
+    keep_existing_names: If True, atoms that already have the atom_name property
+      will keep their assigned names.
+
+  Returns:
+    A new mol, with potentially new 'atom_name' properties.
+  """
+  mol = rd_chem.Mol(mol)
+
+  specified_atom_names = {
+      atom.GetProp('atom_name')
+      for atom in mol.GetAtoms()
+      if atom.HasProp('atom_name') and keep_existing_names
+  }
+
+  element_counts = collections.Counter()
+  for atom in mol.GetAtoms():
+    if not atom.HasProp('atom_name') or not keep_existing_names:
+      element = atom.GetSymbol()
+      while True:
+        element_counts[element] += 1
+        # Standardize names by using uppercase element type, as in CCD. Only
+        # effects elements with more than one letter, e.g. 'Cl' becomes 'CL'.
+        new_name = f'{element.upper()}{element_counts[element]}'
+        if new_name not in specified_atom_names:
+          break
+      atom.SetProp('atom_name', new_name)
+
+  return mol
+
+
+def get_random_conformer(
+    mol: rd_chem.Mol,
+    random_seed: int,
+    max_iterations: int | None,
+    logging_name: str,
+) -> rd_chem.Conformer | None:
+  """Stochastic conformer search method using V3 ETK."""
+  params = rd_all_chem.ETKDGv3()
+  params.randomSeed = random_seed
+  if max_iterations is not None:  # Override default value.
+    params.maxIterations = max_iterations
+  mol_copy = rd_chem.Mol(mol)
+  try:
+    conformer_id = rd_all_chem.EmbedMolecule(mol_copy, params)
+    conformer = mol_copy.GetConformer(conformer_id)
+  except ValueError:
+    logging.warning('Failed to generate conformer for: %s', logging_name)
+    conformer = None
+  return conformer

flax_model/alphafold3/data/tools/shards.py

@@ -0,0 +1,94 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+#
+# To request access to the AlphaFold 3 model parameters, follow the process set
+# out at https://github.com/google-deepmind/alphafold3. You may only use these
+# if received directly from Google. Use is subject to terms of use available at
+# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+"""A library to handle shards of the format file_path@NUM_SHARDS.
+
+For instance, /path/to/file@20 will generate the following shards:
+
+- /path/to/file-00000-of-00020
+- /path/to/file-00001-of-00020
+- ...
+- /path/to/file-00019-of-00020
+
+This also supports @* pattern, which will determine the number of shards based
+on the filesystem content.
+"""
+
+from collections.abc import Sequence
+import dataclasses
+import pathlib
+import re
+
+
+_MAX_NUM_SHARDS = 99_999
+_SHARD_RE = re.compile(
+    r"""
+    ^(?P<prefix>[^\?\],\*]+)@
+     (?P<shards>(\d{1,5})|\*)
+     (?P<suffix>[\._][^\?\]@\*\/]*)?
+    $""",
+    re.X,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class ShardSpec:
+  prefix: str
+  num_shards: int
+  suffix: str
+
+
+def parse_shard_spec(path: str) -> ShardSpec | None:
+  """Returns the shard spec or None if the path is not a shard spec.
+
+  For instance, if the shard spec is '/path/to/file@20', the output will be
+  ('/path/to/file', 20).
+
+  Args:
+    path: the path to parse, e.g. /path/to/file@20 or /path/to/file@*.
+  """
+  parsed = re.fullmatch(_SHARD_RE, path)
+  if not parsed:
+    return None
+  prefix = parsed.group('prefix')
+  shards = parsed.group('shards')
+  suffix = parsed.group('suffix') or ''
+
+  if shards != '*':
+    return ShardSpec(prefix=prefix, num_shards=int(shards), suffix=suffix)
+  shard_slice = slice(len(prefix) + 10, len(prefix) + 15)
+  shard_path = pathlib.Path(f'{prefix}-00000-of-?????{suffix}')
+  for shard in sorted(shard_path.parent.glob(shard_path.name), reverse=True):
+    try:
+      num_shards = int(str(shard)[shard_slice])
+      return ShardSpec(prefix=prefix, num_shards=num_shards, suffix=suffix)
+    except ValueError:
+      continue
+  return None
+
+
+def get_sharded_paths(shard_spec: str) -> Sequence[str] | None:
+  """Returns a list of file path or None if the input is not a shard spec.
+
+  Args:
+    shard_spec: the specifications of the shard, e.g. /path/to/file@20.
+  """
+  parsed_spec = parse_shard_spec(shard_spec)
+  if not parsed_spec:
+    return None
+
+  prefix = parsed_spec.prefix
+  num_shards = parsed_spec.num_shards
+  suffix = parsed_spec.suffix
+  if num_shards > _MAX_NUM_SHARDS:
+    raise ValueError(f'Shard count for {shard_spec} exceeds {_MAX_NUM_SHARDS}')
+  return [
+      f'{prefix}-{i:05d}-of-{num_shards:05d}{suffix}' for i in range(num_shards)
+  ]

flax_model/alphafold3/data/tools/subprocess_utils.py

@@ -0,0 +1,112 @@
+
+
+"""Helper functions for launching external tools."""
+
+from collections.abc import Sequence
+import os
+import subprocess
+import time
+from typing import Any
+
+from absl import logging
+
+
+def create_query_fasta_file(sequence: str, path: str, linewidth: int = 80):
+  """Creates a fasta file with the sequence with line width limit."""
+  with open(path, 'w') as f:
+    f.write('>query\n')
+
+    i = 0
+    while i < len(sequence):
+      f.write(f'{sequence[i:(i + linewidth)]}\n')
+      i += linewidth
+
+
+def check_binary_exists(path: str, name: str) -> None:
+  """Checks if a binary exists on the given path and raises otherwise."""
+  if not os.path.exists(path):
+    raise RuntimeError(f'{name} binary not found at {path}')
+
+
+def jackhmmer_seq_limit_supported(jackhmmer_path: str) -> bool:
+  """Checks if Jackhmmer supports the --seq-limit flag."""
+  try:
+    subprocess.run(
+        [jackhmmer_path, '-h', '--seq_limit', '1'],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=True,
+    )
+  except subprocess.CalledProcessError:
+    return False
+  return True
+
+
+def run(
+    cmd: Sequence[str],
+    cmd_name: str,
+    log_on_process_error: bool = False,
+    log_stderr: bool = False,
+    log_stdout: bool = False,
+    max_out_streams_len: int | None = 500_000,
+    **run_kwargs,
+) -> subprocess.CompletedProcess[Any]:
+  """Launches a subprocess, times it, and checks for errors.
+
+  Args:
+    cmd: Command to launch.
+    cmd_name: Human-readable command name to be used in logs.
+    log_on_process_error: Whether to use `logging.error` to log the process'
+      stderr on failure.
+    log_stderr: Whether to log the stderr of the command.
+    log_stdout: Whether to log the stdout of the command.
+    max_out_streams_len: Max length of prefix of stdout and stderr included in
+      the exception message. Set to `None` to disable truncation.
+    **run_kwargs: Any other kwargs for `subprocess.run`.
+
+  Returns:
+    The completed process object.
+
+  Raises:
+    RuntimeError: if the process completes with a non-zero return code.
+  """
+
+  logging.info('Launching subprocess "%s"', ' '.join(cmd))
+
+  start_time = time.time()
+  try:
+    completed_process = subprocess.run(
+        cmd,
+        check=True,
+        stderr=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+        **run_kwargs,
+    )
+  except subprocess.CalledProcessError as e:
+    if log_on_process_error:
+      # Logs have a 15k character limit, so log the error line by line.
+      logging.error('%s failed. %s stderr begin:', cmd_name, cmd_name)
+      for error_line in e.stderr.splitlines():
+        if stripped_error_line := error_line.strip():
+          logging.error(stripped_error_line)
+      logging.error('%s stderr end.', cmd_name)
+
+    error_msg = (
+        f'{cmd_name} failed'
+        f'\nstdout:\n{e.stdout[:max_out_streams_len]}\n'
+        f'\nstderr:\n{e.stderr[:max_out_streams_len]}'
+    )
+    raise RuntimeError(error_msg) from e
+  end_time = time.time()
+
+  logging.info('Finished %s in %.3f seconds', cmd_name, end_time - start_time)
+  stdout, stderr = completed_process.stdout, completed_process.stderr
+
+  if log_stdout and stdout:
+    logging.info('%s stdout:\n%s', cmd_name, stdout)
+
+  if log_stderr and stderr:
+    logging.info('%s stderr:\n%s', cmd_name, stderr)
+
+  return completed_process

flax_model/alphafold3/jax/attention/attention_base.py

@@ -0,0 +1,355 @@
+
+
+"""Common types and utilities for attention kernels."""
+
+import abc
+import dataclasses
+import enum
+import functools
+import math
+from typing import Any, Self
+
+from flax_model.alphafold3.jax.common import array_view
+from flax_model.alphafold3.jax.common import precision as precision_lib
+import jax
+import jax.numpy as jnp
+from jax.typing import DTypeLike  # pylint: disable=g-importing-member
+import jaxtyping
+from jaxtyping import Array, Bool, Float, Int  # pylint: disable=g-multiple-import,g-importing-member
+import typeguard
+
+
+class AUTO:  # Used as a sentinel value.
+  pass
+
+
+DotPrecisionLike = jax.lax.Precision | precision_lib.DotPrecision
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclasses.dataclass(frozen=True)
+class Mask:
+  """An attention mask.
+
+  `k_start` (inclusive) and `k_end` (exclusive) define range of enabled
+  k-sequence values for each row of logits.
+
+  For example, a local attention mask could be defined as follows:
+  ```
+  seq_len_q = seq_len_k = 4
+  window_size = 2
+  k_start = jnp.maximum(0, jnp.arange(seq_len_q) + 1 - window_size)
+  mask = Mask(k_start=k_start, is_causal=True)
+  assert mask.as_array(seq_len_q, seq_len_k) == jnp.array(
+      [[1, 0, 0, 0],
+       [1, 1, 0, 0],
+       [0, 1, 1, 0],
+       [0, 0, 1, 1]], dtype=bool)
+  ```
+  Or equivalently (but less efficiently):
+  ```
+  k_end = jnp.arange(seq_len_q) + 1
+  k_start = jnp.maximum(0, k_end - window_size)
+  mask = Mask(k_start=k_start, k_end=k_end)
+  assert mask.as_array(seq_len_q, seq_len_k) == jnp.array(
+      [[1, 0, 0, 0],
+       [1, 1, 0, 0],
+       [0, 1, 1, 0],
+       [0, 0, 1, 1]], dtype=bool)
+  ```
+
+  A mask for two independent causal sequences could be defined as follows:
+  ```
+  k_start = jnp.array([0, 0, 2, 2])
+  mask = Mask(k_start=k_start, is_causal=True)
+  assert mask.as_array(seq_len_q, seq_len_k) == jnp.array(
+      [[1, 0, 0, 0],
+       [1, 1, 0, 0],
+       [0, 0, 1, 0],
+       [0, 0, 1, 1]], dtype=bool)
+  ```
+  """
+
+  bool_mask: Bool[Array, "*#B #T #t"] | None = None
+  _: dataclasses.KW_ONLY
+  q_start: Int[Array, "*#B #t"] | None = None
+  q_end: Int[Array, "*#B #t"] | None = None
+  k_start: Int[Array, "*#B #T"] | None = None
+  k_end: Int[Array, "*#B #T"] | None = None
+  is_causal: bool = False
+
+  def tree_flatten(self):
+    return (
+        self.bool_mask,
+        self.q_start,
+        self.q_end,
+        self.k_start,
+        self.k_end,
+    ), (self.is_causal,)
+
+  @classmethod
+  def tree_unflatten(cls, aux, children) -> Self:
+    (is_causal,) = aux
+    bool_mask, q_start, q_end, k_start, k_end = children
+    return cls(
+        bool_mask,
+        q_start=q_start,
+        q_end=q_end,
+        k_start=k_start,
+        k_end=k_end,
+        is_causal=is_causal,
+    )
+
+  def as_array(
+      self,
+      q_len_or_indices: int | Int[Array, "*#B T"],
+      k_len_or_indices: int | Int[Array, "*#B t"],
+  ) -> Bool[Array, "*#B #T #t"] | None:
+    """Returns the mask as a boolean array."""
+    if isinstance(q_len_or_indices, int):
+      q_indices = jnp.arange(q_len_or_indices)
+    else:
+      q_indices = q_len_or_indices
+
+    if isinstance(k_len_or_indices, int):
+      k_indices = jnp.arange(k_len_or_indices)
+    else:
+      k_indices = k_len_or_indices
+
+    q_indices = q_indices[..., None]
+    k_indices = k_indices[..., None, :]
+
+    mask = []
+    if self.bool_mask is not None:
+      mask.append(self.bool_mask)
+      # Check `bool_mask` shape is compatible with `{q,kv}_indices`.
+      _ = jnp.broadcast_shapes(
+          q_indices.shape, k_indices.shape, self.bool_mask.shape
+      )
+
+    if self.q_start is not None:
+      mask.append(q_indices >= self.q_start[..., None, :])
+
+    if self.q_end is not None:
+      mask.append(q_indices < self.q_end[..., None, :])
+
+    if self.k_start is not None:
+      mask.append(k_indices >= self.k_start[..., None])
+
+    if self.k_end is not None:
+      mask.append(k_indices < self.k_end[..., None])
+
+    if self.is_causal:
+      mask.append(q_indices >= k_indices)
+
+    logical_and = functools.partial(functools.reduce, jnp.logical_and)
+    return jax.lax.broadcast_to_rank(logical_and(mask), 3) if mask else None
+
+  def take(self, *attrs: str) -> tuple[Any, ...]:
+    """Returns a mask with attrs removed and the removed attrs."""
+    default_mask = type(self)()
+    replacements = {attr: getattr(default_mask, attr) for attr in attrs}
+    values = (getattr(self, attr) for attr in attrs)
+    return dataclasses.replace(self, **replacements), *values
+
+  def __and__(self, other: "Bool[Array, '*#B #T #t'] | Mask") -> "Mask":  # pylint: disable=g-inconsistent-quotes
+    """Returns the intersection of two masks."""
+    if not isinstance(other, Mask):
+      other = Mask(other)
+
+    def combine(op):
+      return lambda a, b: b if a is None else a if b is None else op(a, b)
+
+    return Mask(
+        bool_mask=combine(jnp.logical_and)(self.bool_mask, other.bool_mask),
+        q_end=combine(jnp.minimum)(self.q_end, other.q_end),
+        k_start=combine(jnp.maximum)(self.k_start, other.k_start),
+        k_end=combine(jnp.minimum)(self.k_end, other.k_end),
+        is_causal=self.is_causal or other.is_causal,
+    )
+
+
+CAUSAL_MASK = Mask(is_causal=True)
+
+
+SoftmaxResidual = (
+    tuple[Float[Array, "*B H T"], Float[Array, "*B H T"]]
+    | Float[Array, "*B H T"]
+)
+
+
+@enum.unique
+class SoftmaxResidualMode(enum.Enum):
+  """The mode of storing softmax residuals for the backwards pass.
+
+  The stable softmax calculation performs two reductions calculating:
+    - the maximum input value (`x_max`),
+    - the sum of exponentiated values (`denom`).
+
+  We can store these values as residuals to avoid the need to recompute them
+  in the backwards pass.
+
+  It is also possible to combine the two residuals into a single residual,
+  `res = x_max + log(denom)`, as `exp(x - res) === exp(x - x_max - log(denom))
+  === exp(x - x_max) / denom`. Combining the residuals reduces the memory usage
+  of the residuals, but will reduce the accuracy of the backwards pass if
+  `abs(x_max) >> log(denom)`.
+  """
+
+  SEPARATE = "separate"
+  COMBINED = "combined"
+
+  def conform(self, aux: SoftmaxResidual) -> SoftmaxResidual | None:
+    match self, aux:
+      case None, _:
+        return None
+      case SoftmaxResidualMode.SEPARATE, (_, _):
+        return aux
+      case SoftmaxResidualMode.SEPARATE, _:  # pytype: disable=redundant-match  # b/300135240
+        raise ValueError("`aux` has been combined.")
+      case SoftmaxResidualMode.COMBINED, (x_max, denom):
+        return x_max + jnp.log(denom)
+      case SoftmaxResidualMode.COMBINED, _:  # pytype: disable=redundant-match  # b/300135240
+        return aux
+
+
+class DotProductAttention(abc.ABC):
+  """Dot product attention function."""
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  def __call__(
+      self,
+      query: Float[Array | array_view.ArrayView, "*B T H D"],
+      key: Float[Array | array_view.ArrayView, "*B t h D"],
+      value: Float[Array | array_view.ArrayView, "*B t h D"],
+      *,
+      precision: (
+          DotPrecisionLike | tuple[DotPrecisionLike, DotPrecisionLike]
+      ) = jax.lax.Precision.DEFAULT,
+      logits_dtype: DTypeLike | type[AUTO] = AUTO,
+      bias: Float[Array, "*#B #H #T #t"] | None = None,
+      mask: Bool[Array, "*#B #H #T #t"] | Mask | None = None,
+      q_indices: Int[Array, "*#B #H T"] | None = None,
+      k_indices: Int[Array, "*#B #H t"] | None = None,
+  ) -> Float[Array, "*B T H D"]:
+    """Performs scaled dot-product attention.
+
+    Scaled dot-product attention from "Attention is all you need"
+    https://arxiv.org/abs/1706.03762.
+
+    Computes self- or cross-attention. The following is computed:
+    softmax(qk_scale * query @ key^T + bias) @ value.
+
+    Supports both multi-head and multi-query attention
+    (https://arxiv.org/abs/1911.02150).
+
+    Arguments:
+      query: Query array of shape `[batch, seq_len_q, num_heads_q, head_dim]`.
+        It must be a multiple of num_heads_kv.
+        Here's an example of how q/kv heads are interleaved:
+          For 8 key/value heads and 4 query heads:
+          - key/value heads [0, 1] see query head 0
+          - key/value heads [2, 3] see query head 1
+          - key/value heads [4, 5] see query head 2
+      key: Key array of shape `[batch, seq_len_kv, num_heads_kv, head_dim]`. It
+        must be divisible by num_heads_q.
+      value: Value array of shape `[batch, seq_len_kv, num_heads_kv, head_dim]`.
+      precision: The precision for the dot products. Either a tuple `(
+        query_key_dot_precision, weights_value_dot_precision)` or a single
+        precision applied to both dot products.
+      logits_dtype: Data type for attention logits (`query @ key^T`). If `AUTO`
+        is passed (the default), the accumulator type from the `query @ key^T`
+        dot product will be used.
+      bias: Optional bias array, broadcastable to shape `[batch, num_heads,
+        seq_len_q, seq_len_kv]`.
+      mask: Optional boolean mask, broadcastable to `[batch, num_heads,
+        seq_len_q, seq_len_kv]`. Attention weights are masked out if the
+        corresponding mask value is `False`.
+      q_indices: Optional indices for each token in query sequence.
+      k_indices: Optional indices for each token in key/value sequence.
+
+    Returns:
+      An array with the same shape as `query`.
+    """  # fmt: skip
+    return self.fwd(
+        query,
+        key,
+        value,
+        precision=precision,
+        logits_dtype=logits_dtype,
+        bias=bias,
+        mask=mask,
+        q_indices=q_indices,
+        k_indices=k_indices,
+    )
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  def fwd(
+      self,
+      query: Float[Array | array_view.ArrayView, "*B T H D"],
+      key: Float[Array | array_view.ArrayView, "*B t h D"],
+      value: Float[Array | array_view.ArrayView, "*B t h D"],
+      *,
+      precision: (
+          DotPrecisionLike | tuple[DotPrecisionLike, DotPrecisionLike]
+      ) = jax.lax.Precision.DEFAULT,
+      logits_dtype: DTypeLike | type[AUTO] = AUTO,
+      bias: Float[Array, "*#B #H #T #t"] | None = None,
+      mask: Bool[Array, "*#B #H #T #t"] | Mask | None = None,
+      q_indices: Int[Array, "*#B #H T"] | None = None,
+      k_indices: Int[Array, "*#B #H t"] | None = None,
+  ) -> Float[Array, "*B T H D"]:
+    """Performs attention."""
+    if not isinstance(precision, tuple):
+      precision = (precision, precision)
+
+    q_k_dot_precision, weights_v_dot_precision = precision
+
+    if not isinstance(q_k_dot_precision, precision_lib.DotPrecision):
+      q_k_dot_precision = precision_lib.get_equivalent_dot_precision(
+          query.dtype, key.dtype, q_k_dot_precision
+      )
+
+    if not isinstance(weights_v_dot_precision, precision_lib.DotPrecision):
+      weights_v_dot_precision = precision_lib.get_equivalent_dot_precision(
+          value.dtype, value.dtype, weights_v_dot_precision
+      )
+
+    if logits_dtype is AUTO:
+      logits_dtype = q_k_dot_precision.accumulator_dtype
+
+    if not isinstance(mask, Mask):
+      mask = Mask(mask)
+
+    return self._fwd(
+        array_view.as_array_view(query),
+        array_view.as_array_view(key),
+        array_view.as_array_view(value),
+        q_k_dot_precision=q_k_dot_precision,
+        logits_dtype=jnp.dtype(logits_dtype),
+        logits_scale=1 / math.sqrt(query.shape[-1]),
+        bias=bias,
+        mask=mask,
+        weights_v_dot_precision=weights_v_dot_precision,
+        q_indices=q_indices,
+        k_indices=k_indices,
+    )
+
+  @abc.abstractmethod
+  def _fwd(
+      self,
+      q: Float[array_view.ArrayView, "*B T H D"],
+      k: Float[array_view.ArrayView, "*B t h D"],
+      v: Float[array_view.ArrayView, "*B t h D"],
+      *,
+      q_k_dot_precision: precision_lib.DotPrecision,
+      logits_dtype: jnp.dtype,
+      logits_scale: float,
+      bias: Float[Array, "*#B #H #T #t"] | None,
+      mask: Mask | None,
+      weights_v_dot_precision: precision_lib.DotPrecision,
+      q_indices: Int[Array, "*#B #H T"] | None = None,
+      k_indices: Int[Array, "*#B #H t"] | None = None,
+  ) -> Float[Array, "*B T H D"]:
+    """Performs attention."""
+    ...

flax_model/alphafold3/jax/attention/flash_attention.py

@@ -0,0 +1,697 @@
+
+
+"""Triton FlashAttention implementation."""
+
+import dataclasses
+import functools
+
+from flax_model.alphafold3.jax.attention import attention_base as base
+from flax_model.alphafold3.jax.common import array_view
+from flax_model.alphafold3.jax.common import precision as precision_lib
+from flax_model.alphafold3.jax.common import triton_utils
+import jax
+import jax.numpy as jnp
+import jax_triton as jt
+import jaxtyping
+from jaxtyping import Array, Bool, Float, Int  # pylint: disable=g-multiple-import,g-importing-member
+import triton
+import triton.language as tl
+import typeguard
+
+
+@triton.jit
+def _fwd_kernel_inner(
+    start_loop,
+    end_loop,
+    q,
+    span_q,
+    k_block_ptr,
+    v_block_ptr,
+    bias_block_ptr,
+    mask_block_ptr,
+    k_start,
+    k_end,
+    seq_len_k,
+    acc,
+    m_i,
+    l_i,
+    bias_advance: tl.constexpr,
+    mask_advance: tl.constexpr,
+    is_causal: tl.constexpr,
+    use_attention_mask: tl.constexpr,
+    use_k_start: tl.constexpr,
+    use_k_end: tl.constexpr,
+    use_bias: tl.constexpr,
+    block_k: tl.constexpr,
+    use_mask_k: tl.constexpr,
+    k_boundary_check: tl.constexpr,
+    v_boundary_check: tl.constexpr,
+    dot_fn_qk: tl.constexpr,
+    dot_fn_kv: tl.constexpr,
+):
+  """Triton MHA forward kernel's inner loop."""
+
+  for start_k in range(start_loop, end_loop, block_k):
+    start_k = tl.multiple_of(start_k, block_k)
+    span_k = start_k + tl.arange(0, block_k)
+
+    k = tl.load(
+        k_block_ptr,
+        boundary_check=k_boundary_check,
+        padding_option="zero" if len(k_boundary_check.value) else "",
+    )
+    v = tl.load(
+        v_block_ptr,
+        boundary_check=v_boundary_check,
+        padding_option="zero" if len(v_boundary_check.value) else "",
+    )
+
+    if use_bias:
+      bias = tl.load(bias_block_ptr)
+
+    qk = dot_fn_qk(q.to(k.dtype), k)  # [block_q, block_k]
+
+    if use_bias:
+      # Prevent dot accumulating into the bias tensor. It appears that Triton
+      # doesn't pipeline the bias load as it does the `k` load, so the bias load
+      # blocks the matmul if the add is merged.
+      qk = qk.to(tl.uint32, bitcast=True) & 0xFFFFFFFF
+      qk = qk.to(tl.float32, bitcast=True)
+      qk += bias
+
+    if use_attention_mask | use_k_start | use_k_end:
+      mask_value = float(jnp.finfo(jnp.float32).min)
+
+    if use_attention_mask:
+      mask = tl.load(mask_block_ptr)
+      qk = tl.where(mask, qk, mask_value)
+
+    if use_k_start:
+      # This check is there to work around a triton compiler bug, but it
+      # shouldn't be strictly needed.
+      if tl.sum(k_start) != 0:
+        qk = tl.where(k_start[:, None] <= span_k[None, :], qk, mask_value)
+    if is_causal:
+      qk = tl.where(span_q[:, None] >= span_k[None, :], qk, float("-inf"))
+    elif use_k_end:
+      # When called with k_end and is_causal=True, the causal mask gets folded
+      # into k_end and is_causal is set to False.
+      qk = tl.where(k_end[:, None] > span_k[None, :], qk, mask_value)
+
+    if use_mask_k:
+      qk = tl.where((span_k < seq_len_k)[None, :], qk, float("-inf"))
+
+    m_ij = tl.maximum(m_i, tl.max(qk, axis=1))  # Shape [block_q].
+    p = tl.exp(qk - m_ij[:, None])  # Shape [block_q, block_k].
+    alpha = tl.exp(m_i - m_ij)
+    m_i = m_ij
+    acc *= alpha[:, None]
+    l_i *= alpha
+    l_i += tl.sum(p, axis=1)
+
+    # Add the new block of attention weights.
+    acc += dot_fn_kv(p.to(v.dtype), v)
+
+    k_block_ptr = tl.advance(k_block_ptr, (0, block_k))
+    v_block_ptr = tl.advance(v_block_ptr, (block_k, 0))
+    bias_block_ptr = tl.advance(bias_block_ptr, bias_advance.value)
+    mask_block_ptr = tl.advance(mask_block_ptr, mask_advance.value)
+
+  return (
+      k_block_ptr,
+      v_block_ptr,
+      bias_block_ptr,
+      mask_block_ptr,
+      acc,
+      m_i,
+      l_i,
+  )
+
+
+# Based on Algorithm 1 of https://arxiv.org/abs/2205.14135.
+# Inspired by the official Triton tutorial implementation
+# https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
+@triton.jit
+def _fwd_kernel(
+    # Input arrays.
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    bias_ptr,
+    mask_ptr,
+    k_start_ptr,
+    k_end_ptr,
+    # Scalar inputs.
+    q_offset,
+    k_offset,
+    v_offset,
+    q_stride_b,
+    q_stride_s,
+    q_stride_h,
+    q_stride_d,
+    k_stride_b,
+    k_stride_s,
+    k_stride_h,
+    k_stride_d,
+    v_stride_b,
+    v_stride_s,
+    v_stride_h,
+    v_stride_d,
+    bias_stride_b,
+    bias_stride_h,
+    bias_stride_sq,
+    bias_stride_sk,
+    mask_stride_b,
+    mask_stride_h,
+    mask_stride_sq,
+    mask_stride_sk,
+    k_start_stride_b,
+    k_start_stride_h,
+    k_start_stride_sq,
+    k_end_stride_b,
+    k_end_stride_h,
+    k_end_stride_sq,
+    o_stride_b,
+    o_stride_s,
+    o_stride_h,
+    o_stride_d,
+    num_heads_q,
+    num_heads_k,
+    seq_len_q,
+    seq_len_k,
+    # Output arrays.
+    o_ptr,
+    # Compile-time constants.
+    is_causal: tl.constexpr,
+    use_attention_mask: tl.constexpr,
+    use_k_start: tl.constexpr,
+    use_k_end: tl.constexpr,
+    use_bias: tl.constexpr,
+    sm_scale: tl.constexpr,
+    block_q: tl.constexpr,
+    block_k: tl.constexpr,
+    head_dim: tl.constexpr,
+    use_mask_q: tl.constexpr,
+    use_mask_k: tl.constexpr,
+    bias_bcast_sq: tl.constexpr,
+    mask_bcast_sq: tl.constexpr,
+    dot_fn_qk: tl.constexpr,
+    dot_fn_kv: tl.constexpr,
+):
+  """Triton MHA forward kernel."""
+  # pytype: disable=annotation-type-mismatch,unsupported-operands
+  #block_d: tl.constexpr = jt.utils.next_power_of_2(head_dim.value)
+  a: tl.constexpr = jt.utils.next_power_of_2(head_dim.value)
+  block_d: tl.constexpr = a if a>16 else 32
+
+  # Each thread block processes one batch element (b) and one head (h).
+  start_q = tl.program_id(1) * block_q
+  off_h = tl.program_id(0)  # int in [0, num_heads_o).
+  off_b = tl.program_id(2)  # int in [0, batch_size)
+
+  off_h_k = off_h // (num_heads_q // num_heads_k)
+
+  q_ptr += off_h * q_stride_h + off_b * q_stride_b + q_offset
+  k_ptr += off_h_k * k_stride_h + off_b * k_stride_b + k_offset
+  v_ptr += off_h_k * v_stride_h + off_b * v_stride_b + v_offset
+  o_ptr += off_h * o_stride_h + off_b * o_stride_b
+
+  if use_bias:
+    bias_ptr += off_b * bias_stride_b + off_h * bias_stride_h
+  if use_attention_mask:
+    mask_ptr += off_b * mask_stride_b + off_h * mask_stride_h
+  if use_k_start:
+    k_start_ptr += off_b * k_start_stride_b + off_h * k_start_stride_h
+  if use_k_end:
+    k_end_ptr += off_b * k_end_stride_b + off_h * k_end_stride_h
+
+  q_block_ptr = tl.make_block_ptr(
+      q_ptr,
+      shape=(seq_len_q, head_dim),
+      strides=(q_stride_s, q_stride_d),
+      offsets=(start_q, 0),
+      block_shape=(block_q, block_d),
+      order=(1, 0),
+  )
+  k_block_ptr = tl.make_block_ptr(
+      k_ptr,
+      shape=(head_dim, seq_len_k),
+      strides=(k_stride_d, k_stride_s),
+      offsets=(0, 0),
+      block_shape=(block_d, block_k),
+      order=(0, 1),
+  )
+  v_block_ptr = tl.make_block_ptr(
+      v_ptr,
+      shape=(seq_len_k, head_dim),
+      strides=(v_stride_s, v_stride_d),
+      offsets=(0, 0),
+      block_shape=(block_k, block_d),
+      order=(1, 0),
+  )
+
+  q_boundary_check0: tl.constexpr = (0,) if use_mask_q else ()
+  q_boundary_check1: tl.constexpr = (1,) if head_dim != block_d else ()
+  q_boundary_check: tl.constexpr = q_boundary_check0 + q_boundary_check1
+  q_padding_option: tl.constexpr = "zero" if len(q_boundary_check.value) else ""
+  k_boundary_check: tl.constexpr = (0,) if head_dim != block_d else ()
+  v_boundary_check: tl.constexpr = (0,) if use_mask_k else ()
+
+  # If broadcasting in a given dim, use a 1D block (observed to be faster).
+  bias_start_dim: tl.constexpr = 1 if bias_bcast_sq else 0
+  bias_block_ptr = tl.make_block_ptr(
+      bias_ptr,
+      shape=(seq_len_q, seq_len_k)[bias_start_dim:],
+      strides=(bias_stride_sq, bias_stride_sk)[bias_start_dim:],
+      offsets=(start_q, 0)[bias_start_dim:],
+      block_shape=(block_q, block_k)[bias_start_dim:],
+      order=(1, 0)[bias_start_dim:],
+  )
+  bias_advance: tl.constexpr = (0, block_k)[bias_start_dim:]
+
+  mask_start_dim: tl.constexpr = 1 if mask_bcast_sq else 0
+  mask_block_ptr = tl.make_block_ptr(
+      mask_ptr,
+      shape=(seq_len_q, seq_len_k)[mask_start_dim:],
+      strides=(mask_stride_sq, mask_stride_sk)[mask_start_dim:],
+      offsets=(start_q, 0)[mask_start_dim:],
+      block_shape=(block_q, block_k)[mask_start_dim:],
+      order=(1, 0)[mask_start_dim:],
+  )
+  mask_advance: tl.constexpr = (0, block_k)[mask_start_dim:]
+
+  k_start_block_ptr = tl.make_block_ptr(
+      k_start_ptr,
+      shape=(seq_len_q,),
+      strides=(k_start_stride_sq,),
+      offsets=(start_q,),
+      block_shape=(block_q,),
+      order=(0,),
+  )
+  k_end_block_ptr = tl.make_block_ptr(
+      k_end_ptr,
+      shape=(seq_len_q,),
+      strides=(k_end_stride_sq,),
+      offsets=(start_q,),
+      block_shape=(block_q,),
+      order=(0,),
+  )
+  # pytype: enable=annotation-type-mismatch,unsupported-operands
+
+  # Each thread block processes a block of block_q queries.
+  span_q = start_q + tl.arange(0, block_q)
+
+  # m_i and l_i (see FlashAttention paper) are updated during the k,v loop.
+  m_i = tl.full([block_q], float("-inf"), dtype=tl.float32)
+  l_i = tl.zeros([block_q], dtype=tl.float32)
+  # acc is the buffer where we accumulate the output on sram.
+  acc = tl.zeros([block_q, block_d], dtype=tl.float32)
+
+  # Load q: it will stay in smem throughout. Indices form a matrix because we
+  # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index.
+  q = tl.load(
+      q_block_ptr,
+      boundary_check=q_boundary_check,
+      padding_option=q_padding_option,
+  )
+  q *= sm_scale
+
+  # In FlashAttention algorithm 1 there are 2 loops: slow over tiles of kv (size
+  # (Bc == block_k here), and fast over blocks of q (size Br == block_q here).
+  # Here we only loop over blocks of kv to process entire seq_len, the loop over
+  # blocks of q is carried out by the grid.
+  k_start = None
+  if use_k_start:
+    k_start = tl.load(k_start_block_ptr)
+    start_loop = tl.maximum(tl.min(k_start), 0)
+    blocks_to_skip = start_loop // block_k
+    start_loop = block_k * blocks_to_skip  # Floor to multiple of block_k.
+    for _ in range(blocks_to_skip):
+      # Advance all block pointers to the first valid block.
+      k_block_ptr = tl.advance(k_block_ptr, (0, block_k))
+      v_block_ptr = tl.advance(v_block_ptr, (block_k, 0))
+      bias_block_ptr = tl.advance(bias_block_ptr, bias_advance.value)
+      mask_block_ptr = tl.advance(mask_block_ptr, mask_advance.value)
+  else:
+    start_loop = 0
+
+  k_end = None
+  if is_causal:
+    end_loop = tl.minimum((start_q // block_k) * block_k, seq_len_k)
+  elif use_k_end:
+    k_end = tl.load(k_end_block_ptr)
+    end_loop = tl.minimum(tl.max(k_end), seq_len_k)
+  else:
+    end_loop = seq_len_k
+
+  (
+      k_block_ptr,
+      v_block_ptr,
+      bias_block_ptr,
+      mask_block_ptr,
+      acc,
+      m_i,
+      l_i,
+  ) = _fwd_kernel_inner(
+      start_loop,
+      end_loop,
+      q,
+      span_q,
+      k_block_ptr,
+      v_block_ptr,
+      bias_block_ptr,
+      mask_block_ptr,
+      k_start,
+      k_end,
+      seq_len_k,
+      acc,
+      m_i,
+      l_i,
+      bias_advance,
+      mask_advance,
+      False,  # is_causal
+      use_attention_mask,
+      use_k_start,
+      use_k_end,
+      use_bias,
+      block_k,
+      use_mask_k,
+      k_boundary_check,
+      v_boundary_check,
+      dot_fn_qk,
+      dot_fn_kv,
+  )
+
+  if is_causal:
+    tl.debug_barrier()  # Help compiler schedule loops independently.
+    start_loop, end_loop = end_loop, tl.minimum(end_loop + block_k, seq_len_k)
+
+    _, _, _, _, acc, _, l_i = _fwd_kernel_inner(
+        start_loop,
+        end_loop,
+        q,
+        span_q,
+        k_block_ptr,
+        v_block_ptr,
+        bias_block_ptr,
+        mask_block_ptr,
+        k_start,
+        k_end,
+        seq_len_k,
+        acc,
+        m_i,
+        l_i,
+        bias_advance,
+        mask_advance,
+        True,  # is_causal
+        use_attention_mask,
+        use_k_start,
+        use_k_end,
+        use_bias,
+        block_k,
+        use_mask_k,
+        k_boundary_check,
+        v_boundary_check,
+        dot_fn_qk,
+        dot_fn_kv,
+    )
+
+  # It is possible that every value in a row was masked to f32 min or that the
+  # main loop has been completely optimised out, and that `l_i` is `0` for that
+  # row. Add epsilon value to avoid NaNs from `0 / 0`.
+  l_i += float(jnp.finfo(jnp.float32).tiny)
+
+  acc /= l_i[:, None]
+
+  # Write output to dram.
+  o_block_ptr = tl.make_block_ptr(
+      o_ptr,
+      shape=(seq_len_q, head_dim),
+      strides=(o_stride_s, o_stride_d),
+      offsets=(start_q, 0),
+      block_shape=(block_q, block_d),
+      order=(1, 0),
+  )
+  acc = acc.to(o_ptr.dtype.element_ty)
+  tl.store(o_block_ptr, acc, boundary_check=q_boundary_check)
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def _fwd(
+    q: Float[array_view.ArrayView, "*B T H D"],
+    k: Float[array_view.ArrayView, "*B t h D"],
+    v: Float[array_view.ArrayView, "*B t h D"],
+    bias: Float[Array, "*#B #H #T #t"] | None,
+    mask: Bool[Array, "*#B #H #T #t"] | None,
+    k_start: Int[Array, "*#B #H #T"] | None,
+    k_end: Int[Array, "*#B #H #T"] | None,
+    *,
+    logits_scale: float,
+    is_causal: bool,
+    q_k_dot_precision: precision_lib.DotPrecision,
+    weights_v_dot_precision: precision_lib.DotPrecision,
+) -> Float[Array, "*B T H D"]:
+  """Forward pass of Triton FlashAttention."""
+
+  orig_q_shape = q.shape
+  q = q.collapse(0, -3, allow_copy=True)
+  batch_size, seq_len_q, num_heads_q, head_dim = q.shape
+  *_, seq_len_k, num_heads_kv, _ = k.shape
+  # Maybe broadcast `k`/`v` heads dimension.
+  kv_shape = (batch_size, seq_len_k, num_heads_kv, head_dim)
+  k = k.collapse(0, -3, allow_copy=True).broadcast_to(kv_shape)
+  v = v.collapse(0, -3, allow_copy=True).broadcast_to(kv_shape)
+
+  def get_bias_mask_view(x, dtype):
+    if x is None:
+      x = jnp.array([], dtype=dtype)
+      return array_view.ArrayView(x, shape=(0, 0, 0, 0), strides=(0, 0, 0, 0))
+
+    shape = orig_q_shape[:-3] + (num_heads_q, seq_len_q, seq_len_k)
+    return (
+        array_view.ArrayView(x)
+        .broadcast_to(shape)
+        .collapse(0, -3, allow_copy=True)
+    )
+
+  bias = get_bias_mask_view(bias, dtype=q.dtype)
+  mask = get_bias_mask_view(mask, dtype=jnp.bool_)
+
+  def get_range_view(x, seq_len):
+    if x is None:
+      x = jnp.array([], dtype=jnp.int32)
+      return array_view.ArrayView(x, shape=(0, 0, 0), strides=(0, 0, 0))
+
+    shape = orig_q_shape[:-3] + (num_heads_q, seq_len)
+    return (
+        array_view.ArrayView(x)
+        .broadcast_to(shape)
+        .collapse(0, -2, allow_copy=True)
+    )
+
+  k_start = get_range_view(k_start, seq_len_q)
+  k_end = get_range_view(k_end, seq_len_q)
+
+  block_q = 64
+  block_k = 64
+
+  return jt.triton_call(
+      q.base,
+      k.base,
+      v.base,
+      bias.base,
+      mask.base,
+      k_start.base,
+      k_end.base,
+      q.offset,
+      k.offset,
+      v.offset,
+      *q.strides,
+      *k.strides,
+      *v.strides,
+      *bias.strides,
+      *mask.strides,
+      k_start.strides,
+      k_end.strides,
+      *jt.utils.strides_from_shape(q.shape),  # out strides.
+      num_heads_q,
+      num_heads_kv,
+      seq_len_q,
+      seq_len_k,
+      kernel=_fwd_kernel,
+      name="triton_flash_attention",
+      out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
+      grid=(num_heads_q, triton.cdiv(seq_len_q, block_q), batch_size),
+      num_stages=1,
+      num_warps=4,
+      is_causal=is_causal,
+      use_attention_mask=(mask.size != 0),
+      use_k_start=(k_start.size != 0),
+      use_k_end=(k_end.size != 0),
+      use_bias=(bias.size != 0),
+      sm_scale=logits_scale,
+      block_q=block_q,
+      block_k=block_k,
+      head_dim=head_dim,
+      use_mask_q=(seq_len_q % block_q != 0),
+      use_mask_k=(seq_len_k % block_q != 0),
+      bias_bcast_sq=(bias.strides[-2] == 0),
+      mask_bcast_sq=(mask.strides[-2] == 0),
+      dot_fn_qk=triton_utils.get_tl_dot_fn(q_k_dot_precision),
+      dot_fn_kv=triton_utils.get_tl_dot_fn(weights_v_dot_precision),
+  ).reshape(orig_q_shape)
+
+
+def _as_batched_array_view(x, axis_size):
+  batched_shape = (axis_size,) + x.shape
+  batched_strides = (x.base.size // axis_size,) + x.strides
+  return dataclasses.replace(x, shape=batched_shape, strides=batched_strides)
+
+
+def _fwd_vmap_rule(
+    axis_size, in_batched, *args, fn: jax.custom_batching.custom_vmap
+):
+  """`vmap` rule for Triton FlashAttention forward op."""
+  q, k, v, bias, mask, k_start, k_end = args
+  (
+      q_batched,
+      k_batched,
+      v_batched,
+      bias_batched,
+      mask_batched,
+      k_start_batched,
+      k_end_batched,
+  ) = in_batched
+
+  if q_batched.base:
+    q = _as_batched_array_view(q, axis_size)
+  if k_batched.base:
+    k = _as_batched_array_view(k, axis_size)
+  if v_batched.base:
+    v = _as_batched_array_view(v, axis_size)
+
+  # Triton op requires `q`, `k`, `v` batch dims to be identical.
+  if q_batched.base and k_batched.base and v_batched.base:
+    if bias is not None and not bias_batched:
+      bias = jax.lax.broadcast_to_rank(bias, bias.ndim + 1)
+    if mask is not None and not mask_batched:
+      mask = jax.lax.broadcast_to_rank(mask, mask.ndim + 1)
+    if k_start is not None and not k_start_batched:
+      k_start = jax.lax.broadcast_to_rank(k_start, k_start.ndim + 1)
+    if k_end is not None and not k_end_batched:
+      k_end = jax.lax.broadcast_to_rank(k_end, k_end.ndim + 1)
+    out = fn(q, k, v, bias, mask, k_start, k_end)
+    out_batched = True
+    return out, out_batched
+
+  # Fallback to sequential loop.
+  q, k, v = map(jnp.asarray, (q, k, v))
+  in_batched = [
+      q_batched.base,
+      k_batched.base,
+      v_batched.base,
+      bias_batched,
+      mask_batched,
+      k_start_batched,
+      k_end_batched,
+  ]
+
+  def f(q, k, v, *args, **kwargs):
+    q, k, v = map(array_view.ArrayView, (q, k, v))
+    return fn.fun(q, k, v, *args, **kwargs)
+
+  sequential_vmap = jax.custom_batching.sequential_vmap(f)
+  return sequential_vmap.vmap_rule(axis_size, in_batched, q, k, v, *args[3:])
+
+
+def _decompose_mask(mask, q, k, q_indices, k_indices):
+  """Decomposes `mask` into a mask array, `is_causal`, `k_start` and `k_end`."""
+  if mask is None:
+    return None, False, None, None
+
+  is_causal = False
+  k_start = None
+  k_end = None
+  if q_indices is None and k_indices is None:
+    mask, is_causal, k_start, k_end = mask.take("is_causal", "k_start", "k_end")
+    if k_start is not None:
+      k_start = jax.lax.broadcast_to_rank(k_start, 2)
+    if k_end is not None:
+      k_end = jax.lax.broadcast_to_rank(k_end, 2)
+      if is_causal:  # Fold is_causal into k_end
+        k_end = jnp.minimum(k_end, jnp.arange(1, q.shape[-3] + 1))
+        is_causal = False
+
+  q_len_or_indices = q.shape[-3] if q_indices is None else q_indices
+  k_len_or_indices = k.shape[-3] if k_indices is None else k_indices
+  return (
+      mask.as_array(q_len_or_indices, k_len_or_indices),
+      is_causal,
+      k_start,
+      k_end,
+  )
+
+
+@dataclasses.dataclass(frozen=True)
+class TritonFlashAttention(base.DotProductAttention):
+  """Triton FlashAttention implementation."""
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  def _fwd(
+      self,
+      q: Float[array_view.ArrayView, "*B T H D"],
+      k: Float[array_view.ArrayView, "*B t h D"],
+      v: Float[array_view.ArrayView, "*B t h D"],
+      bias: Float[Array, "*#B #H #T #t"] | None,
+      *,
+      q_k_dot_precision: precision_lib.DotPrecision,
+      logits_dtype: jnp.dtype,
+      logits_scale: float,
+      mask: base.Mask | None,
+      weights_v_dot_precision: precision_lib.DotPrecision,
+      q_indices: Int[Array, "*#B #H T"] | None = None,
+      k_indices: Int[Array, "*#B #H t"] | None = None,
+  ) -> Float[Array, "*B T H D"]:
+    if logits_dtype != jnp.float32:
+      raise ValueError("`logits_dtype` must be float32.")
+
+    kwargs = dict(
+        logits_scale=logits_scale,
+        q_k_dot_precision=q_k_dot_precision,
+        weights_v_dot_precision=weights_v_dot_precision,
+    )
+
+    def attend_fwd(
+        q,
+        k,
+        v,
+        bias,
+        mask_,
+        q_indices,
+        k_indices,
+    ):
+
+      mask, is_causal, k_start, k_end = _decompose_mask(
+          mask_, q, k, q_indices, k_indices
+      )
+
+      fwd_closed_kwargs = dict(
+          is_causal=is_causal,
+          **kwargs,
+      )
+
+      fwd_closed = functools.partial(_fwd, **fwd_closed_kwargs)
+      fwd_closed = jax.custom_batching.custom_vmap(fwd_closed)
+      fwd_closed.def_vmap(functools.partial(_fwd_vmap_rule, fn=fwd_closed))
+
+      return fwd_closed(q, k, v, bias, mask, k_start, k_end)
+
+    return attend_fwd(
+        q,
+        k,
+        v,
+        bias,
+        mask,
+        q_indices,
+        k_indices,
+    )

flax_model/alphafold3/jax/attention/xla_attention.py

@@ -0,0 +1,132 @@
+
+
+"""XLA implementation of scaled dot-product attention."""
+
+import dataclasses
+
+from flax_model.alphafold3.jax.attention import attention_base as base
+from flax_model.alphafold3.jax.common import array_view
+from flax_model.alphafold3.jax.common import precision as precision_lib
+import jax
+import jax.numpy as jnp
+import jaxtyping
+from jaxtyping import Array, Float, Int  # pylint: disable=g-multiple-import,g-importing-member
+import typeguard
+
+
+def _get_precision(
+    backend: str, precision: precision_lib.DotPrecision
+) -> jax.lax.Precision:
+  if backend == "gpu" and precision == precision_lib.DotPrecision.F32_F32:
+    return jax.lax.Precision.HIGHEST
+  return jax.lax.Precision.DEFAULT
+
+
+def einsum_with_dot_precision(
+    subscript: str,
+    a: jax.Array,
+    b: jax.Array,
+    *,
+    precision: precision_lib.DotPrecision,
+) -> jax.Array:
+  """Evaluate `fn` with the given precision."""
+  result = jnp.einsum(
+      subscript,
+      a.astype(precision.operand_dtype),
+      b.astype(precision.operand_dtype),
+      precision=_get_precision(jax.default_backend().lower(), precision),
+      preferred_element_type=precision.accumulator_dtype,
+  )
+  assert result.dtype == precision.accumulator_dtype
+  return result
+
+
+def _softmax(x: jax.Array) -> jax.Array:
+  """Computes softmax."""
+  # Always perform reductions in at least f32 precision.
+  dtype = jnp.promote_types(x.dtype, jnp.float32)
+  x_max = jnp.max(x.astype(dtype), axis=-1, keepdims=True)
+  unnormalized = jnp.exp(x - x_max)
+  denom = jnp.sum(unnormalized, axis=-1, keepdims=True)
+  return unnormalized / denom
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def _attend(
+    q: Float[array_view.ArrayView, "*B T H D"],
+    k: Float[array_view.ArrayView, "*B t #H D"],
+    v: Float[array_view.ArrayView, "*B t #H D"],
+    *,
+    q_k_dot_precision: precision_lib.DotPrecision,
+    logits_dtype: jnp.dtype,
+    logits_scale: float,
+    bias: Float[Array, "*#B #H #T #t"] | None,
+    mask: base.Mask | None,
+    weights_v_dot_precision: precision_lib.DotPrecision,
+    q_indices: Int[Array, "*#B #H T"] | None,
+    k_indices: Int[Array, "*#B #H t"] | None,
+) -> Float[Array, "*B T H D"]:
+  """Computes attention."""
+  logits = einsum_with_dot_precision(
+      "...qhd,...khd->...hqk", q, k, precision=q_k_dot_precision
+  ).astype(logits_dtype)
+
+  logits *= logits_scale
+
+  if bias is not None:
+    logits += bias
+
+  if mask is not None:
+    q_len_or_indices = q.shape[-3] if q_indices is None else q_indices
+    k_len_or_indices = k.shape[-3] if k_indices is None else k_indices
+    mask = mask.as_array(q_len_or_indices, k_len_or_indices)
+
+  if mask is not None:
+    mask_value = float(jnp.finfo(logits.dtype).min)
+
+    logits = jnp.where(jnp.asarray(mask), logits, mask_value)
+
+  weights = _softmax(logits)
+
+  weights = weights.astype(v.dtype)
+  out = einsum_with_dot_precision(
+      "...hqk,...khd->...qhd", weights, v, precision=weights_v_dot_precision
+  ).astype(q.dtype)
+  return out
+
+
+@dataclasses.dataclass(frozen=True)
+class XlaDotProductAttention(base.DotProductAttention):
+  """XLA dot product attention function."""
+
+  _: dataclasses.KW_ONLY
+
+  def _fwd(
+      self,
+      q: Float[array_view.ArrayView, "*B T H D"],
+      k: Float[array_view.ArrayView, "*B t #H D"],
+      v: Float[array_view.ArrayView, "*B t #H D"],
+      *,
+      q_k_dot_precision: precision_lib.DotPrecision,
+      logits_dtype: jnp.dtype,
+      logits_scale: float,
+      bias: Float[Array, "*#B #H #T #t"] | None,
+      mask: base.Mask | None,
+      weights_v_dot_precision: precision_lib.DotPrecision,
+      q_indices: Int[Array, "*#B #H T"] | None = None,
+      k_indices: Int[Array, "*#B #H t"] | None = None,
+  ) -> Float[Array, "*B T H D"]:
+
+    return _attend(
+        q,
+        k,
+        v,
+        bias=bias,
+        mask=mask,
+        q_indices=q_indices,
+        k_indices=k_indices,
+        q_k_dot_precision=q_k_dot_precision,
+        logits_dtype=logits_dtype,
+        logits_scale=logits_scale,
+        weights_v_dot_precision=weights_v_dot_precision,
+    )

flax_model/alphafold3/jax/common/triton_utils.py

@@ -0,0 +1,118 @@
+
+
+"""Triton utils."""
+
+from collections.abc import Callable, Mapping
+
+from flax_model.alphafold3.jax.common import precision as precision_lib
+import jax
+import jax.numpy as jnp
+import triton
+import triton.language as tl
+
+
+_JNP_TO_TL_DTYPES: Mapping[jnp.dtype, tl.dtype] = {
+    jnp.bool_: tl.int1,
+    jnp.int8: tl.int8,
+    jnp.int16: tl.int16,
+    jnp.int32: tl.int32,
+    jnp.int64: tl.int64,
+    jnp.uint8: tl.uint8,
+    jnp.uint16: tl.uint16,
+    jnp.uint32: tl.uint32,
+    jnp.uint64: tl.uint64,
+    jnp.float16: tl.float16,
+    jnp.bfloat16: tl.bfloat16,
+    jnp.float32: tl.float32,
+    jnp.float64: tl.float64,
+}
+
+
+def jnp_to_tl_dtype(jnp_dtype: jnp.dtype) -> tl.dtype:
+  return _JNP_TO_TL_DTYPES[jnp_dtype]
+
+
+def get_tl_dot_fn(
+    precision: precision_lib.DotPrecision,
+) -> Callable[..., tl.tensor]:
+  """Returns a tl `dot` implementation with the specified precision.
+
+  Args:
+    precision: The `dot` precision.
+  """
+  if not is_precision_supported(precision):
+    raise ValueError(f'Unsupported dot precision: {precision}')
+
+  if precision == precision_lib.DotPrecision.TF32_F32_3X:
+    return _dot_tf32_f32_3x
+
+  in_dtype = jnp_to_tl_dtype(precision.operand_dtype)
+  out_dtype = jnp_to_tl_dtype(precision.accumulator_dtype)
+  allow_tf32 = precision == precision_lib.DotPrecision.TF32_F32
+
+  @tl.core.extern
+  def _dot_fn(
+      a: tl.core.tensor,
+      b: tl.core.tensor,
+      *,
+      trans_a: bool = False,
+      trans_b: bool = False,
+      _builder,
+  ):
+    if in_dtype == tl.float32:
+      tl.static_assert(a.dtype == tl.float32, _builder=_builder)
+      tl.static_assert(b.dtype == tl.float32, _builder=_builder)
+    else:
+      tl.static_assert(a.dtype.is_standard_floating(), _builder=_builder)
+      tl.static_assert(b.dtype.is_standard_floating(), _builder=_builder)
+    a = a.to(in_dtype, _builder=_builder)
+    b = b.to(in_dtype, _builder=_builder)
+    a = tl.trans(a, _builder=_builder) if trans_a else a
+    b = tl.trans(b, _builder=_builder) if trans_b else b
+    return tl.dot(
+        a, b, allow_tf32=allow_tf32, out_dtype=out_dtype, _builder=_builder
+    )
+
+  return _dot_fn
+
+
+def is_precision_supported(precision: precision_lib.DotPrecision) -> bool:
+  return precision in {
+      precision_lib.DotPrecision.F32_F32,
+      precision_lib.DotPrecision.TF32_F32,
+      precision_lib.DotPrecision.F16_F32,
+      precision_lib.DotPrecision.BF16_F32,
+      precision_lib.DotPrecision.TF32_F32_3X,
+  }
+
+
+@triton.jit
+def _dot_tf32_f32_3x(a, b, trans_a=False, trans_b=False):
+  """Perform the 3-pass tf32 dot function."""
+  tl.static_assert(a.dtype == tl.float32)
+  tl.static_assert(b.dtype == tl.float32)
+  a_ = (a.to(tl.uint32, bitcast=True) & 0xFFFFE000).to(tl.float32, bitcast=True)
+  b_ = (b.to(tl.uint32, bitcast=True) & 0xFFFFE000).to(tl.float32, bitcast=True)
+  a_err = a - a_
+  b_err = b - b_
+  if trans_a:
+    a_ = tl.trans(a_)
+    a_err = tl.trans(a_err)
+  if trans_b:
+    b_ = tl.trans(b_)
+    b_err = tl.trans(b_err)
+  # Add smallest terms first for better accuracy.
+  return tl.dot(a_, b_, out_dtype=tl.float32) + (
+      tl.dot(a_, b_err, out_dtype=tl.float32)
+      + tl.dot(a_err, b_, out_dtype=tl.float32)
+  )
+
+
+def has_triton_support() -> bool:
+  """Returns True if Triton is supported by the default JAX device."""
+  if jax.default_backend() != 'gpu':
+    return False
+
+  # Only currently supported for Ampere and above.
+  # return float(jax.devices()[0].compute_capability) >= 8.0
+  return True

flax_model/alphafold3/jax/gated_linear_unit/block.py

@@ -0,0 +1,118 @@
+
+
+"""Pallas block load / store utilities."""
+
+from collections.abc import Sequence
+from typing import Any, TypeAlias
+
+from flax_model.alphafold3.jax.common import array_view
+import jax
+import jax.experimental
+from jax.experimental import pallas as pl
+import jax.numpy as jnp
+import jaxtyping
+from jaxtyping import Int  # pylint: disable=g-importing-member
+import numpy as np
+import typeguard
+
+ArrayT: TypeAlias = Any
+ScalarInt: TypeAlias = (
+    Int[ArrayT, ""] | Int[np.generic, ""] | Int[jnp.generic, ""]
+)
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def load_block(
+    ref,
+    idx: Sequence[int | ScalarInt],
+    *,
+    block_shape: Sequence[int | None],
+    other=None,
+    **kwargs,
+) -> jax.Array:
+  """Loads a block from the given `ref`, masking where necessary."""
+  idx, mask = _get_block_indexer_and_mask(ref, idx, block_shape=block_shape)
+  if isinstance(ref, array_view.ArrayView):
+    idx = ref[idx].offsets
+    ref = ref.base
+  other = None if mask is None else other
+  with jax.experimental.enable_x64():
+    return pl.load(ref, idx, mask=mask, other=other, **kwargs)
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def store_block(
+    ref,
+    val: jax.Array,
+    idx: Sequence[int | ScalarInt],
+    *,
+    block_shape: Sequence[int | None] | None = None,
+    **kwargs,
+):
+  """Stores a block from the given `ref`, masking where necessary."""
+  if block_shape is None:
+    block_shape = val.shape
+  idx, mask = _get_block_indexer_and_mask(ref, idx, block_shape=block_shape)
+  if isinstance(ref, array_view.ArrayView):
+    idx = ref[idx].offsets
+    ref = ref.base
+  with jax.experimental.enable_x64():
+    pl.store(ref, idx, val.astype(ref.dtype), mask=mask, **kwargs)
+
+
+def in_bounds_mask(
+    idx: Sequence[int | slice | pl.Slice | jax.Array],
+    shape: Sequence[int],
+    *,
+    check: Sequence[bool] | None = None,
+) -> jax.Array | None:
+  """Returns a boolean mask denoting which indices are within bounds.
+
+  Args:
+    idx: Indices for each dimension.
+    shape: Shape designating the valid bounds.
+    check: Whether or not to check bounds in each dimension. Useful for ignoring
+      indices known to be in bounds. Defaults to all True.
+  """
+  if check is None:
+    check = [True] * len(shape)
+
+  # Remove `int` indexed dims (mask shape must match slice result shape).
+  shape = [dim for i, dim in enumerate(shape) if not isinstance(idx[i], int)]
+  check = [chk for i, chk in enumerate(check) if not isinstance(idx[i], int)]
+  idx = [idx for idx in idx if not isinstance(idx, int)]
+
+  mask = None
+  for i, (dim_idx, dim, chk) in enumerate(zip(idx, shape, check, strict=True)):
+    if not chk:
+      continue
+
+    if isinstance(dim_idx, slice):
+      dim_idx = pl.Slice.from_slice(dim_idx, dim)
+    if isinstance(dim_idx, pl.Slice):
+      dim_idx = dim_idx.start + dim_idx.stride * jnp.arange(dim_idx.size)
+    if dim_idx.ndim != 1:
+      raise NotImplementedError("Only one-dimensional indices are supported.")
+
+    bcast_axes = [a for a in range(len(shape)) if a != i]
+    dim_mask = jnp.expand_dims(dim_idx < dim, bcast_axes)
+    mask = dim_mask if mask is None else (mask & dim_mask)
+  return mask
+
+
+def _get_block_indexer_and_mask(
+    ref, idx: Sequence[int | ScalarInt], *, block_shape: Sequence[int | None]
+) -> tuple[tuple[int | slice | pl.Slice, ...], jax.Array | None]:
+  """Return indices and mask for loading / storing a block."""
+  shape = ref.shape
+  idxs = []
+  check = []
+  for dim, block_idx, block_dim in zip(shape, idx, block_shape, strict=True):
+    if block_dim is None:
+      idxs.append(block_idx)
+      check.append(False)
+    else:
+      idxs.append(pl.dslice(block_dim * block_idx, block_dim))
+      check.append(dim % block_dim != 0)
+
+  return tuple(idxs), in_bounds_mask(idxs, shape, check=check)

flax_model/alphafold3/jax/gated_linear_unit/gated_linear_unit.py

@@ -0,0 +1,116 @@
+
+
+"""Public API for gated linear unit functions."""
+
+from collections.abc import Callable
+import typing
+from typing import Literal, TypeAlias
+
+from flax_model.alphafold3.jax.common import array_view
+from flax_model.alphafold3.jax.common import triton_utils
+from flax_model.alphafold3.jax.gated_linear_unit import gated_linear_unit_base
+from flax_model.alphafold3.jax.gated_linear_unit import matmul_ext
+import jax
+import jaxtyping
+from jaxtyping import Array, Float  # pylint: disable=g-importing-member,g-multiple-import
+import typeguard
+
+Implementation: TypeAlias = Literal['xla', 'triton']
+
+
+class PallasGatedLinearUnit(gated_linear_unit_base.GatedLinearUnit):
+  """Pallas gated linear unit."""
+
+  def _fwd(self, x, weight, *, activation, precision):
+    weight_view = array_view.ArrayView(weight)
+    return self.apply_vmap_rule_forward(
+        matmul_ext.gated_linear_unit,
+        activation=activation,
+        precision=precision,
+    )(
+        x,
+        weight_view[:, 1],
+        weight_view[:, 0],
+    )
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def gated_linear_unit(
+    x: Float[Array, '*B M K'],
+    weight: Float[Array, 'K 2 N'],
+    *,
+    activation: Callable[[jax.Array], jax.Array] | None = None,
+    precision: jax.lax.Precision | None = None,
+    implementation: Implementation | None = None,
+) -> Float[Array, '*B M N']:
+  """Applies a gated linear unit (https://arxiv.org/abs/1612.08083).
+
+  Computes `activation(x @ weight[:, 0]) * x @ weight[:, 1]`.
+
+  This is SwiGLU when `activation=jax.nn.swish`, GEGLU when
+  `activation=jax.nn.gelu`, REGLU when `activation=jax.nn.relu`, and GLU when
+  `activation=jax.nn.sigmoid` (https://arxiv.org/abs/2002.05202).
+
+  Args:
+    x: the input array.
+    weight: the combined weight array.
+    activation: optional activation function.
+    precision: specifies the matrix multiplication precision. Either `None`
+      (default), which means the default precision for the backend, or a
+      `jax.lax.Precision` enum.
+    implementation: if `None` (default), an implementation is automatically
+      chosen. 'xla' will use standard XLA and work on any platform, and 'triton'
+      will use a fused Triton GPU kernel. Only a subset of data types, shapes
+      and GPUs are supported by 'triton', with an exception thrown in this case.
+
+  Raises:
+    NotImplementedError: if `implementation='triton'` does not support a given
+      input or device.
+    ValueError: if the arguments are invalid.
+
+  Returns:
+    The output array.
+  """
+
+  match implementation:
+    case 'triton':
+      if not triton_utils.has_triton_support():
+        raise NotImplementedError('Triton not supported on this platform.')
+    case _:
+      ...
+
+  if x.dtype.name != weight.dtype.name:
+    raise ValueError(
+        f'Input and weight must have the same dtype. {x.dtype} !='
+        f' {weight.dtype}'
+    )
+
+  if implementation is not None:
+    named_args = typing.get_args(Implementation)
+    if implementation not in named_args:
+      raise ValueError(
+          f'Unsupported named implementation. Must be one of {named_args}.'
+      )
+
+  # if implementation is None or implementation == 'triton':
+  #   try:
+  #     return PallasGatedLinearUnit()(
+  #         x=x,
+  #         weight=weight,
+  #         activation=activation,
+  #         precision=precision,
+  #     )
+  #     # When `implementation=None`, we must catch any exception, and use XLA
+  #     # as a fallback. As we rely on a third-party library (Triton), it might
+  #     # not be possible to enumerate all possible exceptions that could be
+  #     # thrown, hence catching the broadest possible one.
+  #   except Exception as e:  # pylint: disable=broad-exception-caught
+  #     if implementation == 'triton':
+  #       raise e
+
+  return gated_linear_unit_base.gated_linear_unit_xla(
+      x=x,
+      weight=weight,
+      activation=activation,
+      precision=precision,
+  )

flax_model/alphafold3/jax/gated_linear_unit/matmul_config.py

@@ -0,0 +1,70 @@
+
+
+"""Auto-tuned configs for matmul."""
+
+import dataclasses
+import functools
+import math
+
+import jax
+from jax.experimental import pallas as pl
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Config:
+  block_m: int
+  block_n: int
+  block_k: int
+  num_warps: int
+  num_stages: int
+
+
+@functools.cache
+def _get_best_block_size(
+    m: int, n: int, k: int, core_count: int
+) -> tuple[int, int, int]:
+  """Returns the best block size for the given shape."""
+  min_block_dim = 32
+  block_m = min(max(min_block_dim, pl.next_power_of_2(m)), 128)
+  block_n = min(max(min_block_dim, pl.next_power_of_2(n)), 256)
+  block_n = min(block_n, (128 * 128) // block_m)
+  block_k = 32
+  split_k = 1
+  num_blocks = pl.cdiv(m, block_m) * pl.cdiv(n, block_n)
+  while num_blocks < core_count:
+    if block_m > min_block_dim:
+      block_m //= 2
+      num_blocks = pl.cdiv(m, block_m) * pl.cdiv(n, block_n)
+    elif split_k * block_k < pl.next_power_of_2(k):
+      split_k *= 2
+      num_blocks *= 2
+    else:
+      break
+  return block_m, block_n, block_k
+
+
+def _abstractify(x):
+  return jax.api_util.shaped_abstractify(x) if isinstance(x, jax.Array) else x
+
+
+def get_config(
+    x: jax.Array, w: jax.Array, core_count: int | None = None
+) -> Config:
+  """Returns a config for the given args."""
+  if core_count is None:
+    core_count = jax.devices()[0].core_count
+  x = _abstractify(x)
+  w = _abstractify(w)
+  m, k = math.prod(x.shape[:-1]), x.shape[-1]
+  n = w.shape[1]
+  if n >= m:  # Prefer `block_n` > `block_m`.
+    block_m, block_n, block_k = _get_best_block_size(m, n, k, core_count)
+  else:
+    block_n, block_m, block_k = _get_best_block_size(n, m, k, core_count)
+  return Config(
+      block_m=block_m,
+      block_n=block_n // 2,  # Halve `block_n` as we read two `w` blocks.
+      block_k=block_k,
+      num_warps=4,
+      num_stages=4,
+  )

flax_model/alphafold3/jax/geometry/__init__.py

@@ -0,0 +1,20 @@
+
+
+"""Geometry Module."""
+
+from flax_model.alphafold3.jax.geometry import rigid_matrix_vector
+from flax_model.alphafold3.jax.geometry import rotation_matrix
+from flax_model.alphafold3.jax.geometry import struct_of_array
+from flax_model.alphafold3.jax.geometry import vector
+
+Rot3Array = rotation_matrix.Rot3Array
+Rigid3Array = rigid_matrix_vector.Rigid3Array
+
+StructOfArray = struct_of_array.StructOfArray
+
+Vec3Array = vector.Vec3Array
+square_euclidean_distance = vector.square_euclidean_distance
+euclidean_distance = vector.euclidean_distance
+dihedral_angle = vector.dihedral_angle
+dot = vector.dot
+cross = vector.cross

scripts/infer_jackhmmer.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+
+export PATH="${PROJECT_ROOT}/flax_model/alphafold3/_tools/hmmer/bin:${PATH}"
+export TF_CPP_MIN_LOG_LEVEL="${TF_CPP_MIN_LOG_LEVEL:-0}"
+export JAX_TRACEBACK_FILTERING="${JAX_TRACEBACK_FILTERING:-off}"
+export XLA_CLIENT_MEM_FRACTION="${XLA_CLIENT_MEM_FRACTION:-0.95}"
+export TRITON_ENABLE_GLOBAL_TO_LOCAL="${TRITON_ENABLE_GLOBAL_TO_LOCAL:-1}"
+export TRITON_USE_MAKE_BLOCK_PTR="${TRITON_USE_MAKE_BLOCK_PTR:-1}"
+export TRITON_DEFAULT_ENABLE_NUM_VGPRS512="${TRITON_DEFAULT_ENABLE_NUM_VGPRS512:-1}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+
+MODEL_ROOT="${ONESCIENCE_MODELS_DIR:-${PROJECT_ROOT}/weight}"
+MODEL_DIR="${ALPHAFOLD3_MODEL_DIR:-${MODEL_ROOT}}"
+DATASET_ROOT="${ALPHAFOLD3_DATASET_ROOT:-${ONESCIENCE_DATASETS_DIR:-${PROJECT_ROOT}/data}/alphafold3}"
+DB_DIRS="${ALPHAFOLD3_DB_DIR:-${DATASET_ROOT}/public_databases}"
+MMSEQS_DB_DIRS="${ALPHAFOLD3_MMSEQS_DB_DIR:-${DATASET_ROOT}/mmseqsDB}"
+JSON_PATH="${ALPHAFOLD3_JSON_PATH:-${PROJECT_ROOT}/inputs/t1119_search.json}"
+OUTPUT_DIR="${ALPHAFOLD3_OUTPUT_DIR:-${PROJECT_ROOT}/outputs}"
+RUN_INFERENCE="${ALPHAFOLD3_RUN_INFERENCE:-false}"
+FLASH_ATTENTION="${ALPHAFOLD3_FLASH_ATTENTION:-cutlass}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+python "${PROJECT_ROOT}/scripts/run_alphafold.py" \
+  --json_path="${JSON_PATH}" \
+  --model_dir="${MODEL_DIR}" \
+  --output_dir="${OUTPUT_DIR}" \
+  --run_data_pipeline=true \
+  --run_inference="${RUN_INFERENCE}" \
+  --flash_attention_implementation="${FLASH_ATTENTION}" \
+  --use_mmseqs=false \
+  --use_mmseqs_gpu=false \
+  --db_dir="${DB_DIRS}" \
+  --mmseqs_db_dir="${MMSEQS_DB_DIRS}" \
+  --small_bfd_database_path="${DATASET_ROOT}/jackhmmer_split/bfd-first_non_consensus_sequences.fasta@64" \
+  --small_bfd_z_value=65928866 \
+  --mgnify_database_path="${DATASET_ROOT}/jackhmmer_split/mgy_clusters_2022_05.fa@512" \
+  --mgnify_z_value=623796864 \
+  --uniprot_cluster_annot_database_path="${DATASET_ROOT}/jackhmmer_split/uniprot_cluster_annot_2021_04.fa@256" \
+  --uniprot_cluster_annot_z_value=225619586 \
+  --uniref90_database_path="${DATASET_ROOT}/jackhmmer_split/uniref90_2022_05.fa@128" \
+  --uniref90_z_value=153742194 \
+  --jackhmmer_n_cpu="${ALPHAFOLD3_JACKHMMER_N_CPU:-2}" \
+  --jackhmmer_max_parallel_shards="${ALPHAFOLD3_JACKHMMER_MAX_PARALLEL_SHARDS:-1}" \
+  --jackhmmer_max_threads="${ALPHAFOLD3_JACKHMMER_MAX_THREADS:-8}" \
+  --nhmmer_n_cpu="${ALPHAFOLD3_NHMMER_N_CPU:-2}" \
+  --nhmmer_max_parallel_shards="${ALPHAFOLD3_NHMMER_MAX_PARALLEL_SHARDS:-1}" \
+  --nhmmer_max_threads="${ALPHAFOLD3_NHMMER_MAX_THREADS:-8}"

scripts/preflight.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+
+try:
+    import yaml
+except ImportError as exc:
+    raise SystemExit("PyYAML is required for preflight: pip install pyyaml") from exc
+
+
+MODEL_FILES = {
+    "af3.bin": (1146811260, "a43670ea1fae790cd30ccc8c8bf836c4098014c45250fb310a5338407e321de7"),
+    "libflash_atten_c.so": (75882104, "2ea956e8d13317dd795ffce846079f20e502a39df66802f485d5f25900ef493b"),
+    "mmseqs/bin/mmseqs": (14417192, "d166c67ca79089ea5f89ec1133f67ab1c216e1734a0408bbc7205e588fe865ed"),
+    "mmseqs/lib/libmarv.so": (58277568, "fdc787467b6d05d36db1f034a2cb8944b5b072bbe6d46172013810a066583928"),
+}
+
+DATASET_REQUIRED = [
+    "infer_input_data/readme.md",
+    "infer_input_data/all_data/7r6r_data.json",
+    "infer_input_data/all_data/t1119_data.json",
+    "public_databases/pdb_seqres_2022_09_28.fasta",
+    "mmseqsDB/small_bfd_db",
+    "mmseqsDB/small_bfd_db.dbtype",
+    "jackhmmer_split/bfd-first_non_consensus_sequences.fasta-00000-of-00064",
+]
+
+
+def sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def check_json(path: Path) -> None:
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if data.get("dialect") != "alphafold3":
+        raise AssertionError(f"{path} dialect is not alphafold3")
+    if not data.get("name"):
+        raise AssertionError(f"{path} missing name")
+    if not isinstance(data.get("sequences"), list) or not data["sequences"]:
+        raise AssertionError(f"{path} missing sequences")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--root", default=".")
+    parser.add_argument("--dataset-root", default="data/alphafold3_dataset")
+    parser.add_argument("--model-dir", default="checkpoints/AlphaFold3")
+    parser.add_argument("--skip-hash", action="store_true")
+    args = parser.parse_args()
+
+    root = Path(args.root).resolve()
+    model_dir = (root / args.model_dir).resolve()
+    dataset_root = (root / args.dataset_root).resolve()
+
+    config = yaml.safe_load((root / "conf" / "alphafold3_paths.yaml").read_text(encoding="utf-8"))
+    if config["model"]["repo_id"] != "OneScience/AlphaFold3/":
+        raise AssertionError("model repo_id mismatch in conf/alphafold3_paths.yaml")
+    if config["dataset"]["repo_id"] != "OneScience/AlphaFold3_dataset":
+        raise AssertionError("dataset repo_id mismatch in conf/alphafold3_paths.yaml")
+
+    for rel, (size, digest) in MODEL_FILES.items():
+        path = model_dir / rel
+        if not path.is_file():
+            raise AssertionError(f"missing model file: {path}")
+        if path.stat().st_size != size:
+            raise AssertionError(f"model size mismatch: {path}")
+        if not args.skip_hash and sha256(path) != digest:
+            raise AssertionError(f"model sha256 mismatch: {path}")
+
+    for rel in DATASET_REQUIRED:
+        path = dataset_root / rel
+        if not path.exists():
+            raise AssertionError(f"missing dataset dependency: {path}")
+
+    check_json(root / "inputs" / "7r6r_data.json")
+    check_json(root / "inputs" / "t1119_search.json")
+    print("model_preflight_ok: true")
+
+
+if __name__ == "__main__":
+    main()

scripts/run_alphafold.py

@@ -0,0 +1,1194 @@
+
+
+"""AlphaFold 3 structure prediction script.
+
+AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+
+To request access to the AlphaFold 3 model parameters, follow the process set
+out at https://github.com/google-deepmind/alphafold3. You may only use these
+if received directly from Google. Use is subject to terms of use available at
+https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+"""
+
+from collections.abc import Callable, Sequence
+import csv
+import dataclasses
+import datetime
+import functools
+import multiprocessing
+import os
+import pathlib
+import shutil
+import string
+import sys
+import textwrap
+import time
+import typing
+from typing import overload
+
+_PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[1]
+if str(_PROJECT_ROOT) not in sys.path:
+  sys.path.insert(0, str(_PROJECT_ROOT))
+
+from absl import app
+from absl import flags
+from flax_model.alphafold3.common import folding_input
+from flax_model.alphafold3.common import resources
+from flax_model.alphafold3.constants import chemical_components
+import flax_model.alphafold3.cpp as af3_cpp
+from flax_model.alphafold3.data import featurisation
+from flax_model.alphafold3.data import pipeline
+from flax_model.alphafold3.jax.attention import attention
+from flax_model.alphafold3.model import features
+from flax_model.alphafold3.model import model
+from flax_model.alphafold3.model import params
+from flax_model.alphafold3.model import post_processing
+from flax_model.alphafold3.model.components import utils
+import haiku as hk
+import jax
+from jax import numpy as jnp
+import numpy as np
+
+
+_HOME_DIR = pathlib.Path(os.environ.get('HOME', _PROJECT_ROOT))
+_MODELS_ROOT = pathlib.Path(
+    os.environ.get('ONESCIENCE_MODELS_DIR', _PROJECT_ROOT / 'weight')
+)
+_DATASETS_ROOT = pathlib.Path(
+    os.environ.get('ONESCIENCE_DATASETS_DIR', _HOME_DIR)
+) / 'alphafold3'
+_DEFAULT_MODEL_DIR = pathlib.Path(
+    os.environ.get('ALPHAFOLD3_MODEL_DIR', _MODELS_ROOT / 'AlphaFold3')
+)
+_DEFAULT_DB_DIR = pathlib.Path(
+    os.environ.get('ALPHAFOLD3_DB_DIR', _DATASETS_ROOT / 'public_databases')
+)
+_DEFAULT_MMSEQS_DB_DIR = pathlib.Path(
+    os.environ.get('ALPHAFOLD3_MMSEQS_DB_DIR', _DATASETS_ROOT / 'mmseqsDB')
+)
+
+# Input and output paths.
+_JSON_PATH = flags.DEFINE_string(
+    'json_path',
+    None,
+    'Path to the input JSON file.',
+)
+_INPUT_DIR = flags.DEFINE_string(
+    'input_dir',
+    None,
+    'Path to the directory containing input JSON files.',
+)
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'Path to a directory where the results will be saved.',
+)
+MODEL_DIR = flags.DEFINE_string(
+    'model_dir',
+    _DEFAULT_MODEL_DIR.as_posix(),
+    'Path to the model to use for inference.',
+)
+
+# Control which stages to run.
+_RUN_DATA_PIPELINE = flags.DEFINE_bool(
+    'run_data_pipeline',
+    True,
+    'Whether to run the data pipeline on the fold inputs.',
+)
+_RUN_INFERENCE = flags.DEFINE_bool(
+    'run_inference',
+    True,
+    'Whether to run inference on the fold inputs.',
+)
+
+
+_DEFAULT_MMSEQS_OPTIONS ='--num-iterations 1 --db-load-mode 2 -a --max-seqs 10000 --prefilter-mode 1'
+_DEFAULT_R2MSA_OPTIONS ='--filter-msa 1 --filter-min-enable 1000 --diff 3000 --qid 0.0,0.2,0.4,0.6,0.8,1.0 --qsc 0 --max-seq-id 0.95'
+_AF3_DIR = _PROJECT_ROOT / 'flax_model' / 'alphafold3'
+_HMMER_BIN_DIR = _AF3_DIR / '_tools' / 'hmmer' / 'bin'
+
+
+def _resolve_af3_tool(binary_name: str) -> str | None:
+  local_binary = _HMMER_BIN_DIR / binary_name
+  if local_binary.exists():
+    return str(local_binary)
+  return shutil.which(binary_name)
+
+
+def _available_cpu_count() -> int:
+  if hasattr(os, 'sched_getaffinity'):
+    return len(os.sched_getaffinity(0))
+  return multiprocessing.cpu_count()
+
+
+_USE_MMSEQS = flags.DEFINE_bool(
+    'use_mmseqs',
+    False,
+    'Whether to use mmseqs for protein MSA search',
+)
+_USE_MMSEQS_GPU = flags.DEFINE_bool(
+    'use_mmseqs_gpu',
+    False,
+    'Whether to use mmseqs GPU for protein MSA search',
+)
+_MMSEQS_OPTIONS = flags.DEFINE_string(
+    'mmseqs_options',
+    _DEFAULT_MMSEQS_OPTIONS,
+    'mmseqs serach options',
+)
+_R2MSA_OPTIONS = flags.DEFINE_string(
+    'result2msa_options',
+    _DEFAULT_R2MSA_OPTIONS,
+    'mmseqs result2msa options',
+)
+
+
+# Binary paths.
+_JACKHMMER_BINARY_PATH = flags.DEFINE_string(
+    'jackhmmer_binary_path',
+    _resolve_af3_tool('jackhmmer'),
+    'Path to the Jackhmmer binary.',
+)
+_NHMMER_BINARY_PATH = flags.DEFINE_string(
+    'nhmmer_binary_path',
+    _resolve_af3_tool('nhmmer'),
+    'Path to the Nhmmer binary.',
+)
+_HMMALIGN_BINARY_PATH = flags.DEFINE_string(
+    'hmmalign_binary_path',
+    _resolve_af3_tool('hmmalign'),
+    'Path to the Hmmalign binary.',
+)
+_HMMSEARCH_BINARY_PATH = flags.DEFINE_string(
+    'hmmsearch_binary_path',
+    _resolve_af3_tool('hmmsearch'),
+    'Path to the Hmmsearch binary.',
+)
+_HMMBUILD_BINARY_PATH = flags.DEFINE_string(
+    'hmmbuild_binary_path',
+    _resolve_af3_tool('hmmbuild'),
+    'Path to the Hmmbuild binary.',
+)
+_MMSEQS_BINARY_PATH = flags.DEFINE_string(
+    'mmseqs_binary_path',
+    shutil.which('mmseqs'),
+    'Path to the mmseqs binary.',
+)
+
+# Database paths.
+DB_DIR = flags.DEFINE_multi_string(
+    'db_dir',
+    (_DEFAULT_DB_DIR.as_posix(),),
+    'Path to the directory containing the databases. Can be specified multiple'
+    ' times to search multiple directories in order.',
+)
+MMSEQS_DB_DIR = flags.DEFINE_multi_string(
+    'mmseqs_db_dir',
+    (_DEFAULT_MMSEQS_DB_DIR.as_posix(),),
+    'Path to the directory containing the mmseqs databases. Can be specified multiple'
+    ' times to search multiple directories in order.',  
+)
+
+_SMALL_BFD_DATABASE_PATH = flags.DEFINE_string(
+    'small_bfd_database_path',
+    '${DB_DIR}/bfd-first_non_consensus_sequences.fasta',
+    'Small BFD database path, used for protein MSA search.',
+)
+_SMALL_BFD_Z_VALUE = flags.DEFINE_integer(
+    'small_bfd_z_value',
+    None,
+    'The Z-value representing the database size in number of sequences for'
+    ' E-value calculation. Must be set for sharded databases.',
+    lower_bound=0,
+)
+_MGNIFY_DATABASE_PATH = flags.DEFINE_string(
+    'mgnify_database_path',
+    '${DB_DIR}/mgy_clusters_2022_05.fa',
+    'Mgnify database path, used for protein MSA search.',
+)
+_MGNIFY_Z_VALUE = flags.DEFINE_integer(
+    'mgnify_z_value',
+    None,
+    'The Z-value representing the database size in number of sequences for'
+    ' E-value calculation. Must be set for sharded databases.',
+    lower_bound=0,
+)
+_UNIPROT_CLUSTER_ANNOT_DATABASE_PATH = flags.DEFINE_string(
+    'uniprot_cluster_annot_database_path',
+    '${DB_DIR}/uniprot_all_2021_04.fa',
+    'UniProt database path, used for protein paired MSA search.',
+)
+_UNIPROT_CLUSTER_ANNOT_Z_VALUE = flags.DEFINE_integer(
+    'uniprot_cluster_annot_z_value',
+    None,
+    'The Z-value representing the database size in number of sequences for'
+    ' E-value calculation. Must be set for sharded databases.',
+    lower_bound=0,
+)
+_UNIREF90_DATABASE_PATH = flags.DEFINE_string(
+    'uniref90_database_path',
+    '${DB_DIR}/uniref90_2022_05.fa',
+    'UniRef90 database path, used for MSA search. The MSA obtained by '
+    'searching it is used to construct the profile for template search.',
+)
+_UNIREF90_Z_VALUE = flags.DEFINE_integer(
+    'uniref90_z_value',
+    None,
+    'The Z-value representing the database size in number of sequences for'
+    ' E-value calculation. Must be set for sharded databases.',
+    lower_bound=0,
+)
+_NTRNA_DATABASE_PATH = flags.DEFINE_string(
+    'ntrna_database_path',
+    '${DB_DIR}/nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fasta',
+    'NT-RNA database path, used for RNA MSA search.',
+)
+_NTRNA_Z_VALUE = flags.DEFINE_float(
+    'ntrna_z_value',
+    None,
+    'The Z-value representing the database size in megabases for E-value'
+    ' calculation. Must be set for sharded databases.',
+    lower_bound=0.0,
+)
+_RFAM_DATABASE_PATH = flags.DEFINE_string(
+    'rfam_database_path',
+    '${DB_DIR}/rfam_14_9_clust_seq_id_90_cov_80_rep_seq.fasta',
+    'Rfam database path, used for RNA MSA search.',
+)
+_RFAM_Z_VALUE = flags.DEFINE_float(
+    'rfam_z_value',
+    None,
+    'The Z-value representing the database size in megabases for E-value'
+    ' calculation. Must be set for sharded databases.',
+    lower_bound=0.0,
+)
+_RNA_CENTRAL_DATABASE_PATH = flags.DEFINE_string(
+    'rna_central_database_path',
+    '${DB_DIR}/rnacentral_active_seq_id_90_cov_80_linclust.fasta',
+    'RNAcentral database path, used for RNA MSA search.',
+)
+_RNA_CENTRAL_Z_VALUE = flags.DEFINE_float(
+    'rna_central_z_value',
+    None,
+    'The Z-value representing the database size in megabases for E-value'
+    ' calculation. Must be set for sharded databases.',
+    lower_bound=0.0,
+)
+_PDB_DATABASE_PATH = flags.DEFINE_string(
+    'pdb_database_path',
+    '${DB_DIR}/mmcif_files',
+    'PDB database directory with mmCIF files path, used for template search.',
+)
+_SEQRES_DATABASE_PATH = flags.DEFINE_string(
+    'seqres_database_path',
+    '${DB_DIR}/pdb_seqres_2022_09_28.fasta',
+    'PDB sequence database path, used for template search.',
+)
+
+# MMSEQS Database paths.
+_MMSEQS_SMALL_BFD_DATABASE_PATH = flags.DEFINE_string(
+    'mmseqs_small_bfd_database_path',
+    '${MMSEQS_DB_DIR}/small_bfd_db',
+    'Small BFD database path, used for protein MSA search.',
+) 
+_MMSEQS_MGNIFY_DATABASE_PATH = flags.DEFINE_string(
+    'mmseqs_mgnify_database_path',
+    '${MMSEQS_DB_DIR}/mgnify_db',
+    'Mgnify database path, used for protein MSA search.',
+)
+_MMSEQS_UNIPROT_CLUSTER_ANNOT_DATABASE_PATH = flags.DEFINE_string(
+    'mmseqs_uniprot_cluster_annot_database_path',
+    '${MMSEQS_DB_DIR}/uniprot_cluster_annot_db',
+    'UniProt database path, used for protein paired MSA search.',
+)
+_MMSEQS_UNIREF90_DATABASE_PATH = flags.DEFINE_string(
+    'mmseqs_uniref90_database_path',
+    '${MMSEQS_DB_DIR}/uniref90_db',
+    'UniRef90 database path, used for MSA search. ',
+)
+
+_JACKHMMER_MAX_THREADS = flags.DEFINE_integer(
+    'jackhmmer_max_threads',
+    None,
+    'Maximum number of threads used when running sharded databases. If unset,'
+    ' defaults to None (no limit).',
+    lower_bound=1,
+)
+# Number of CPUs to use for MSA tools.
+_JACKHMMER_N_CPU = flags.DEFINE_integer(
+    'jackhmmer_n_cpu',
+    # Unfortunately, os.process_cpu_count() is only available in Python 3.13+.
+    min(_available_cpu_count(), 8),
+    'Number of CPUs to use for Jackhmmer. Defaults to min(cpu_count, 8). Going'
+    ' above 8 CPUs provides very little additional speedup.',
+    lower_bound=0,
+)
+_JACKHMMER_MAX_PARALLEL_SHARDS = flags.DEFINE_integer(
+    'jackhmmer_max_parallel_shards',
+    None,
+    'Maximum number of shards to search against in parallel. If unset, one'
+    ' Jackhmmer instance will be run per shard. Only applicable if the'
+    ' database is sharded.',
+    lower_bound=1,
+)
+_NHMMER_N_CPU = flags.DEFINE_integer(
+    'nhmmer_n_cpu',
+    # Unfortunately, os.process_cpu_count() is only available in Python 3.13+.
+    min(_available_cpu_count(), 8),
+    'Number of CPUs to use for Nhmmer. Defaults to min(cpu_count, 8). Going'
+    ' above 8 CPUs provides very little additional speedup.',
+    lower_bound=0,
+)
+_NHMMER_MAX_PARALLEL_SHARDS = flags.DEFINE_integer(
+    'nhmmer_max_parallel_shards',
+    None,
+    'Maximum number of shards to search against in parallel. If unset, one'
+    ' Nhmmer instance will be run per shard. Only applicable if the'
+    ' database is sharded.',
+    lower_bound=1,
+)
+_NHMMER_MAX_THREADS = flags.DEFINE_integer(
+    'nhmmer_max_threads',
+    None,
+    'Maximum number of threads used when running sharded databases. If unset,'
+    ' defaults to None (no limit).',
+    lower_bound=1,
+)
+# Data pipeline configuration.
+_RESOLVE_MSA_OVERLAPS = flags.DEFINE_bool(
+    'resolve_msa_overlaps',
+    True,
+    'Whether to deduplicate unpaired MSA against paired MSA. The default'
+    ' behaviour matches the method described in the AlphaFold 3 paper. Set this'
+    ' to false if providing custom paired MSA using the unpaired MSA field to'
+    ' keep it exactly as is as deduplication against the paired MSA could break'
+    ' the manually crafted pairing between MSA sequences.',
+)
+_MMSEQS_N_CPU = flags.DEFINE_integer(
+    'mmseqs_n_cpu',
+    min(multiprocessing.cpu_count(), 8),
+    'Number of CPUs to use for MMseqs. Default to min(cpu_count, 8). Going'
+    ' beyond 8 CPUs provides very little additional speedup.',
+)
+
+# Template search configuration.
+_MAX_TEMPLATE_DATE = flags.DEFINE_string(
+    'max_template_date',
+    '2021-09-30',  # By default, use the date from the AlphaFold 3 paper.
+    'Maximum template release date to consider. Format: YYYY-MM-DD. All'
+    ' templates released after this date will be ignored. Controls also whether'
+    ' to allow use of model coordinates for a chemical component from the CCD'
+    ' if RDKit conformer generation fails and the component does not have ideal'
+    ' coordinates set. Only for components that have been released before this'
+    ' date the model coordinates can be used as a fallback.',
+)
+
+_CONFORMER_MAX_ITERATIONS = flags.DEFINE_integer(
+    'conformer_max_iterations',
+    None,  # Default to RDKit default parameters value.
+    'Optional override for maximum number of iterations to run for RDKit '
+    'conformer search.',
+    lower_bound=0,
+)
+
+# JAX inference performance tuning.
+_JAX_COMPILATION_CACHE_DIR = flags.DEFINE_string(
+    'jax_compilation_cache_dir',
+    None,
+    'Path to a directory for the JAX compilation cache.',
+)
+_GPU_DEVICE = flags.DEFINE_integer(
+    'gpu_device',
+    0,
+    'Optional override for the GPU device to use for inference, uses zero-based'
+    ' indexing. Defaults to the 0th GPU on the system. Useful on multi-GPU'
+    ' systems to pin each run to a specific GPU. Note that if GPUs are already'
+    ' pre-filtered by the environment (e.g. by using CUDA_VISIBLE_DEVICES),'
+    ' this flag refers to the GPU index after the filtering has been done.',
+)
+_BUCKETS = flags.DEFINE_list(
+    'buckets',
+    # pyformat: disable
+    ['256', '512', '768', '1024', '1280', '1536', '2048', '2560', '3072',
+     '3584', '4096', '4608', '5120'],
+    # pyformat: enable
+    'Strictly increasing order of token sizes for which to cache compilations.'
+    ' For any input with more tokens than the largest bucket size, a new bucket'
+    ' is created for exactly that number of tokens.',
+)
+_FLASH_ATTENTION_IMPLEMENTATION = flags.DEFINE_enum(
+    'flash_attention_implementation',
+    default='triton',
+    enum_values=['triton', 'cudnn', 'xla', 'cutlass',],
+    help=(
+        "Flash attention implementation to use. 'triton' and 'cudnn' uses a"
+        ' Triton and cuDNN flash attention implementation, respectively. The'
+        ' Triton kernel is fastest and has been tested more thoroughly. The'
+        " Triton and cuDNN kernels require Ampere GPUs or later. 'xla' uses an"
+        ' XLA attention implementation (no flash attention) and is portable'
+        ' across GPU devices.'
+    ),
+)
+_NUM_RECYCLES = flags.DEFINE_integer(
+    'num_recycles',
+    10,
+    'Number of recycles to use during inference.',
+    lower_bound=1,
+)
+_NUM_DIFFUSION_SAMPLES = flags.DEFINE_integer(
+    'num_diffusion_samples',
+    5,
+    'Number of diffusion samples to generate.',
+    lower_bound=1,
+)
+_NUM_SEEDS = flags.DEFINE_integer(
+    'num_seeds',
+    None,
+    'Number of seeds to use for inference. If set, only a single seed must be'
+    ' provided in the input JSON. AlphaFold 3 will then generate random seeds'
+    ' in sequence, starting from the single seed specified in the input JSON.'
+    ' The full input JSON produced by AlphaFold 3 will include the generated'
+    ' random seeds. If not set, AlphaFold 3 will use the seeds as provided in'
+    ' the input JSON.',
+    lower_bound=1,
+)
+
+# Output controls.
+_SAVE_EMBEDDINGS = flags.DEFINE_bool(
+    'save_embeddings',
+    False,
+    'Whether to save the final trunk single and pair embeddings in the output.'
+    ' Note that the embeddings are large float16 arrays: num_tokens * 384'
+    ' + num_tokens * num_tokens * 128.',
+)
+_SAVE_DISTOGRAM = flags.DEFINE_bool(
+    'save_distogram',
+    False,
+    'Whether to save the final distogram in the output. Note that the distogram'
+    ' is a large float16 array: num_tokens * num_tokens * 64.',
+)
+_FORCE_OUTPUT_DIR = flags.DEFINE_bool(
+    'force_output_dir',
+    False,
+    'Whether to force the output directory to be used even if it already exists'
+    ' and is non-empty. Useful to set this to True to run the data pipeline and'
+    ' the inference separately, but use the same output directory.',
+)
+
+
+def make_model_config(
+    *,
+    flash_attention_implementation: attention.Implementation = 'triton',
+    num_diffusion_samples: int = 5,
+    num_recycles: int = 10,
+    return_embeddings: bool = False,
+    return_distogram: bool = False,
+) -> model.Model.Config:
+  """Returns a model config with some defaults overridden."""
+  config = model.Model.Config()
+  config.global_config.flash_attention_implementation = (
+      flash_attention_implementation
+  )
+  config.heads.diffusion.eval.num_samples = num_diffusion_samples
+  config.num_recycles = num_recycles
+  config.return_embeddings = return_embeddings
+  config.return_distogram = return_distogram
+  return config
+
+
+class ModelRunner:
+  """Helper class to run structure prediction stages."""
+
+  def __init__(
+      self,
+      config: model.Model.Config,
+      device: jax.Device,
+      model_dir: pathlib.Path,
+  ):
+    self._model_config = config
+    self._device = device
+    self._model_dir = model_dir
+
+  @functools.cached_property
+  def model_params(self) -> hk.Params:
+    """Loads model parameters from the model directory."""
+    return params.get_model_haiku_params(model_dir=self._model_dir)
+
+  @functools.cached_property
+  def _model(
+      self,
+  ) -> Callable[[jnp.ndarray, features.BatchDict], model.ModelResult]:
+    """Loads model parameters and returns a jitted model forward pass."""
+
+    @hk.transform
+    def forward_fn(batch):
+      return model.Model(self._model_config)(batch)
+
+    return functools.partial(
+        jax.jit(forward_fn.apply, device=self._device), self.model_params
+    )
+
+  def run_inference(
+      self, featurised_example: features.BatchDict, rng_key: jnp.ndarray
+  ) -> model.ModelResult:
+    """Computes a forward pass of the model on a featurised example."""
+    featurised_example = jax.device_put(
+        jax.tree_util.tree_map(
+            jnp.asarray, utils.remove_invalidly_typed_feats(featurised_example)
+        ),
+        self._device,
+    )
+
+    result = self._model(rng_key, featurised_example)
+    result = jax.tree.map(np.asarray, result)
+    result = jax.tree.map(
+        lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x,
+        result,
+    )
+    result = dict(result)
+    identifier = self.model_params['__meta__']['__identifier__'].tobytes()
+    result['__identifier__'] = identifier
+    return result
+
+  def extract_inference_results(
+      self,
+      batch: features.BatchDict,
+      result: model.ModelResult,
+      target_name: str,
+  ) -> list[model.InferenceResult]:
+    """Extracts inference results from model outputs."""
+    return list(
+        model.Model.get_inference_result(
+            batch=batch, result=result, target_name=target_name
+        )
+    )
+
+  def extract_embeddings(
+      self, result: model.ModelResult, num_tokens: int
+  ) -> dict[str, np.ndarray] | None:
+    """Extracts embeddings from model outputs."""
+    embeddings = {}
+    if 'single_embeddings' in result:
+      embeddings['single_embeddings'] = result['single_embeddings'][
+          :num_tokens
+      ].astype(np.float16)
+    if 'pair_embeddings' in result:
+      embeddings['pair_embeddings'] = result['pair_embeddings'][
+          :num_tokens, :num_tokens
+      ].astype(np.float16)
+    return embeddings or None
+
+  def extract_distogram(
+      self, result: model.ModelResult, num_tokens: int
+  ) -> np.ndarray | None:
+    """Extracts distogram from model outputs."""
+    if 'distogram' not in result['distogram']:
+      return None
+    distogram = result['distogram']['distogram'][:num_tokens, :num_tokens, :]
+    return distogram
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ResultsForSeed:
+  """Stores the inference results (diffusion samples) for a single seed.
+
+  Attributes:
+    seed: The seed used to generate the samples.
+    inference_results: The inference results, one per sample.
+    full_fold_input: The fold input that must also include the results of
+      running the data pipeline - MSA and templates.
+    embeddings: The final trunk single and pair embeddings, if requested.
+    distogram: The token distance histogram, if requested.
+  """
+
+  seed: int
+  inference_results: Sequence[model.InferenceResult]
+  full_fold_input: folding_input.Input
+  embeddings: dict[str, np.ndarray] | None = None
+  distogram: np.ndarray | None = None
+
+
+def predict_structure(
+    fold_input: folding_input.Input,
+    model_runner: ModelRunner,
+    buckets: Sequence[int] | None = None,
+    ref_max_modified_date: datetime.date | None = None,
+    conformer_max_iterations: int | None = None,
+    resolve_msa_overlaps: bool = True,
+) -> Sequence[ResultsForSeed]:
+  """Runs the full inference pipeline to predict structures for each seed."""
+
+  print(f'Featurising data with {len(fold_input.rng_seeds)} seed(s)...')
+  featurisation_start_time = time.time()
+  ccd = chemical_components.Ccd(user_ccd=fold_input.user_ccd)
+  featurised_examples = featurisation.featurise_input(
+      fold_input=fold_input,
+      buckets=buckets,
+      ccd=ccd,
+      verbose=True,
+      ref_max_modified_date=ref_max_modified_date,
+      conformer_max_iterations=conformer_max_iterations,
+      resolve_msa_overlaps=resolve_msa_overlaps,
+  )
+  print(
+      f'Featurising data with {len(fold_input.rng_seeds)} seed(s) took'
+      f' {time.time() - featurisation_start_time:.2f} seconds.'
+  )
+  print(
+      'Running model inference and extracting output structure samples with'
+      f' {len(fold_input.rng_seeds)} seed(s)...'
+  )
+  all_inference_start_time = time.time()
+  all_inference_results = []
+  for seed, example in zip(fold_input.rng_seeds, featurised_examples):
+    print(f'Running model inference with seed {seed}...')
+    inference_start_time = time.time()
+    rng_key = jax.random.PRNGKey(seed)
+    result = model_runner.run_inference(example, rng_key)
+    print(
+        f'Running model inference with seed {seed} took'
+        f' {time.time() - inference_start_time:.2f} seconds.'
+    )
+    print(f'Extracting inference results with seed {seed}...')
+    extract_structures = time.time()
+    inference_results = model_runner.extract_inference_results(
+        batch=example, result=result, target_name=fold_input.name
+    )
+    num_tokens = len(inference_results[0].metadata['token_chain_ids'])
+    embeddings = model_runner.extract_embeddings(
+        result=result, num_tokens=num_tokens
+    )
+    distogram = model_runner.extract_distogram(
+        result=result, num_tokens=num_tokens
+    )
+    print(
+        f'Extracting {len(inference_results)} inference samples with'
+        f' seed {seed} took {time.time() - extract_structures:.2f} seconds.'
+    )
+
+    all_inference_results.append(
+        ResultsForSeed(
+            seed=seed,
+            inference_results=inference_results,
+            full_fold_input=fold_input,
+            embeddings=embeddings,
+            distogram=distogram,
+        )
+    )
+  print(
+      'Running model inference and extracting output structures with'
+      f' {len(fold_input.rng_seeds)} seed(s) took'
+      f' {time.time() - all_inference_start_time:.2f} seconds.'
+  )
+  return all_inference_results
+
+
+def write_fold_input_json(
+    fold_input: folding_input.Input,
+    output_dir: os.PathLike[str] | str,
+) -> None:
+  """Writes the input JSON to the output directory."""
+  os.makedirs(output_dir, exist_ok=True)
+  path = os.path.join(output_dir, f'{fold_input.sanitised_name()}_data.json')
+  print(f'Writing model input JSON to {path}')
+  with open(path, 'wt') as f:
+    f.write(fold_input.to_json())
+
+
+def write_outputs(
+    all_inference_results: Sequence[ResultsForSeed],
+    output_dir: os.PathLike[str] | str,
+    job_name: str,
+) -> None:
+  """Writes outputs to the specified output directory."""
+  ranking_scores = []
+  max_ranking_score = None
+  max_ranking_result = None
+  try:
+    output_terms = (
+        pathlib.Path(af3_cpp.__file__).parent / 'OUTPUT_TERMS_OF_USE.md'
+    ).read_text()
+  except FileNotFoundError:
+    output_terms = None
+  os.makedirs(output_dir, exist_ok=True)
+  for results_for_seed in all_inference_results:
+    seed = results_for_seed.seed
+    for sample_idx, result in enumerate(results_for_seed.inference_results):
+      sample_dir = os.path.join(output_dir, f'seed-{seed}_sample-{sample_idx}')
+      os.makedirs(sample_dir, exist_ok=True)
+      post_processing.write_output(
+          inference_result=result,
+          output_dir=sample_dir,
+          name=f'{job_name}_seed-{seed}_sample-{sample_idx}',
+      )
+      ranking_score = float(result.metadata['ranking_score'])
+      ranking_scores.append((seed, sample_idx, ranking_score))
+      if max_ranking_score is None or ranking_score > max_ranking_score:
+        max_ranking_score = ranking_score
+        max_ranking_result = result
+
+    if embeddings := results_for_seed.embeddings:
+      embeddings_dir = os.path.join(output_dir, f'seed-{seed}_embeddings')
+      os.makedirs(embeddings_dir, exist_ok=True)
+      post_processing.write_embeddings(
+          embeddings=embeddings,
+          output_dir=embeddings_dir,
+          name=f'{job_name}_seed-{seed}',
+      )
+
+    if (distogram := results_for_seed.distogram) is not None:
+      distogram_dir = os.path.join(output_dir, f'seed-{seed}_distogram')
+      os.makedirs(distogram_dir, exist_ok=True)
+      distogram_path = os.path.join(
+          distogram_dir, f'{job_name}_seed-{seed}_distogram.npz'
+      )
+      with open(distogram_path, 'wb') as f:
+        np.savez_compressed(f, distogram=distogram.astype(np.float16))
+
+  if max_ranking_result is not None:  # True iff ranking_scores non-empty.
+    post_processing.write_output(
+        inference_result=max_ranking_result,
+        output_dir=output_dir,
+        # The output terms of use are the same for all seeds/samples.
+        terms_of_use=output_terms,
+        name=job_name,
+    )
+    # Save csv of ranking scores with seeds and sample indices, to allow easier
+    # comparison of ranking scores across different runs.
+    with open(
+        os.path.join(output_dir, f'{job_name}_ranking_scores.csv'), 'wt'
+    ) as f:
+      writer = csv.writer(f)
+      writer.writerow(['seed', 'sample', 'ranking_score'])
+      writer.writerows(ranking_scores)
+
+
+def replace_db_dir(path_with_db_dir: str, db_dirs: Sequence[str]) -> str:
+  """Replaces the DB_DIR placeholder in a path with the given DB_DIR."""
+  template = string.Template(path_with_db_dir)
+  if 'DB_DIR' in template.get_identifiers():
+    for db_dir in db_dirs:
+      path = template.substitute(DB_DIR=db_dir)
+      if os.path.exists(path):
+        return path
+    raise FileNotFoundError(
+        f'{path_with_db_dir} with ${{DB_DIR}} not found in any of {db_dirs}.'
+    )
+  if not os.path.exists(path_with_db_dir):
+    raise FileNotFoundError(f'{path_with_db_dir} does not exist.')
+  return path_with_db_dir
+
+
+def replace_mmseqs_db_dir(
+    path_with_db_dir: str, 
+    db_dirs: Sequence[str],
+    mmseqs_db_dirs: Sequence[str]
+) -> str:
+  """Replaces the MMSEQS_DB_DIR placeholder in a path with the given MMSEQS_DB_DIR.
+  
+  Args:
+    path_with_db_dir: Path containing MMSEQS_DB_DIR placeholder
+    db_dirs: List of database directories
+    mmseqs_db_dirs: List of mmseqs database directories
+    use_gpu: Whether to use GPU version of mmseqs databases
+    
+  Returns:
+    The expanded path if found, otherwise raises FileNotFoundError
+  """
+  template = string.Template(path_with_db_dir)
+  
+  is_jackhmmer_db = any(
+      path_with_db_dir.endswith(db) for db in [
+          'bfd-first_non_consensus_sequences.fasta',
+          'mgy_clusters_2022_05.fa',
+          'uniprot_all_2021_04.fa',
+          'uniref90_2022_05.fa'
+      ]
+  )
+  
+  if 'MMSEQS_DB_DIR' in template.get_identifiers():
+    db_suffixes = [
+        'small_bfd_db',
+        'mgnify_db',
+        'uniprot_cluster_annot_db',
+        'uniref90_db'
+    ]
+    
+    is_mmseqs_db = any(
+        path_with_db_dir.endswith(suffix) for suffix in db_suffixes
+    )
+    
+    for mmseqs_db_dir in mmseqs_db_dirs:
+      path = template.substitute(MMSEQS_DB_DIR=mmseqs_db_dir)
+      if is_mmseqs_db:
+        if os.path.exists(path):
+          return path
+      else:
+        return path
+    
+    if is_mmseqs_db:
+      raise FileNotFoundError(
+          f'{path_with_db_dir} with ${{MMSEQS_DB_DIR}} not found in any of {mmseqs_db_dirs}.'
+      )
+    return template.substitute(MMSEQS_DB_DIR=mmseqs_db_dirs[0])
+      
+  if 'DB_DIR' in template.get_identifiers():
+    for db_dir in db_dirs:
+      path = template.substitute(DB_DIR=db_dir)
+      if is_jackhmmer_db:
+        return path
+      if os.path.exists(path):
+        return path
+    if is_jackhmmer_db:
+      return template.substitute(DB_DIR=db_dirs[0])
+    raise FileNotFoundError(
+        f'{path_with_db_dir} with ${{DB_DIR}} not found in any of {db_dirs}.'
+    )
+  if not is_jackhmmer_db and not os.path.exists(path_with_db_dir):
+    raise FileNotFoundError(f'{path_with_db_dir} does not exist.')
+  return path_with_db_dir
+
+
+@overload
+def process_fold_input(
+    fold_input: folding_input.Input,
+    data_pipeline_config: pipeline.DataPipelineConfig | None,
+    model_runner: None,
+    output_dir: os.PathLike[str] | str,
+    buckets: Sequence[int] | None = None,
+    ref_max_modified_date: datetime.date | None = None,
+    conformer_max_iterations: int | None = None,
+    resolve_msa_overlaps: bool = True,
+    force_output_dir: bool = False,
+) -> folding_input.Input:
+  ...
+
+
+@overload
+def process_fold_input(
+    fold_input: folding_input.Input,
+    data_pipeline_config: pipeline.DataPipelineConfig | None,
+    model_runner: ModelRunner,
+    output_dir: os.PathLike[str] | str,
+    buckets: Sequence[int] | None = None,
+    ref_max_modified_date: datetime.date | None = None,
+    conformer_max_iterations: int | None = None,
+    resolve_msa_overlaps: bool = True,
+    force_output_dir: bool = False,
+) -> Sequence[ResultsForSeed]:
+  ...
+
+
+def process_fold_input(
+    fold_input: folding_input.Input,
+    data_pipeline_config: pipeline.DataPipelineConfig | None,
+    model_runner: ModelRunner | None,
+    output_dir: os.PathLike[str] | str,
+    buckets: Sequence[int] | None = None,
+    ref_max_modified_date: datetime.date | None = None,
+    conformer_max_iterations: int | None = None,
+    resolve_msa_overlaps: bool = True,
+    force_output_dir: bool = False,
+) -> folding_input.Input | Sequence[ResultsForSeed]:
+  """Runs data pipeline and/or inference on a single fold input.
+
+  Args:
+    fold_input: Fold input to process.
+    data_pipeline_config: Data pipeline config to use. If None, skip the data
+      pipeline.
+    model_runner: Model runner to use. If None, skip inference.
+    output_dir: Output directory to write to.
+    buckets: Bucket sizes to pad the data to, to avoid excessive re-compilation
+      of the model. If None, calculate the appropriate bucket size from the
+      number of tokens. If not None, must be a sequence of at least one integer,
+      in strictly increasing order. Will raise an error if the number of tokens
+      is more than the largest bucket size.
+    ref_max_modified_date: Optional maximum date that controls whether to allow
+      use of model coordinates for a chemical component from the CCD if RDKit
+      conformer generation fails and the component does not have ideal
+      coordinates set. Only for components that have been released before this
+      date the model coordinates can be used as a fallback.
+    conformer_max_iterations: Optional override for maximum number of iterations
+      to run for RDKit conformer search.
+    resolve_msa_overlaps: Whether to deduplicate unpaired MSA against paired
+      MSA. The default behaviour matches the method described in the AlphaFold 3
+      paper. Set this to false if providing custom paired MSA using the unpaired
+      MSA field to keep it exactly as is as deduplication against the paired MSA
+      could break the manually crafted pairing between MSA sequences.
+    force_output_dir: If True, do not create a new output directory even if the
+      existing one is non-empty. Instead use the existing output directory and
+      potentially overwrite existing files. If False, create a new timestamped
+      output directory instead if the existing one is non-empty.
+
+  Returns:
+    The processed fold input, or the inference results for each seed.
+
+  Raises:
+    ValueError: If the fold input has no chains.
+  """
+  print(f'\nRunning fold job {fold_input.name}...')
+
+  if not fold_input.chains:
+    raise ValueError('Fold input has no chains.')
+
+  if (
+      not force_output_dir
+      and os.path.exists(output_dir)
+      and os.listdir(output_dir)
+  ):
+    new_output_dir = (
+        f'{output_dir}_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}'
+    )
+    print(
+        f'Output will be written in {new_output_dir} since {output_dir} is'
+        ' non-empty.'
+    )
+    output_dir = new_output_dir
+  else:
+    print(f'Output will be written in {output_dir}')
+
+   # Create model loader callback function
+  def load_model_callback():
+    if model_runner is not None:
+      _ = model_runner.model_params
+
+  if data_pipeline_config is None:
+    print('Skipping data pipeline...')
+    # Load model immediately when skipping pipeline
+    if model_runner is not None:
+      print('Loading model parameters (no pipeline to wait for)...')
+      load_model_callback()
+  else:	
+    print('Running data pipeline...')	
+    fold_input = pipeline.DataPipeline(data_pipeline_config, load_model_callback).process(fold_input)	
+
+  write_fold_input_json(fold_input, output_dir)
+  if model_runner is None:
+    print('Skipping model inference...')
+    output = fold_input
+  else:
+    print(
+        f'Predicting 3D structure for {fold_input.name} with'
+        f' {len(fold_input.rng_seeds)} seed(s)...'
+    )
+    all_inference_results = predict_structure(
+        fold_input=fold_input,
+        model_runner=model_runner,
+        buckets=buckets,
+        ref_max_modified_date=ref_max_modified_date,
+        conformer_max_iterations=conformer_max_iterations,
+        resolve_msa_overlaps=resolve_msa_overlaps,
+    )
+    print(f'Writing outputs with {len(fold_input.rng_seeds)} seed(s)...')
+    write_outputs(
+        all_inference_results=all_inference_results,
+        output_dir=output_dir,
+        job_name=fold_input.sanitised_name(),
+    )
+    output = all_inference_results
+
+  print(f'Fold job {fold_input.name} done, output written to {output_dir}\n')
+  return output
+
+
+def main(_):
+  if _JAX_COMPILATION_CACHE_DIR.value is not None:
+    jax.config.update(
+        'jax_compilation_cache_dir', _JAX_COMPILATION_CACHE_DIR.value
+    )
+
+  if _JSON_PATH.value is None == _INPUT_DIR.value is None:
+    raise ValueError(
+        'Exactly one of --json_path or --input_dir must be specified.'
+    )
+
+  if not _RUN_INFERENCE.value and not _RUN_DATA_PIPELINE.value:
+    raise ValueError(
+        'At least one of --run_inference or --run_data_pipeline must be'
+        ' set to true.'
+    )
+
+  if _INPUT_DIR.value is not None:
+    fold_inputs = folding_input.load_fold_inputs_from_dir(
+        pathlib.Path(_INPUT_DIR.value)
+    )
+  elif _JSON_PATH.value is not None:
+    fold_inputs = folding_input.load_fold_inputs_from_path(
+        pathlib.Path(_JSON_PATH.value)
+    )
+  else:
+    raise AssertionError(
+        'Exactly one of --json_path or --input_dir must be specified.'
+    )
+
+  # Make sure we can create the output directory before running anything.
+  try:
+    os.makedirs(_OUTPUT_DIR.value, exist_ok=True)
+  except OSError as e:
+    print(f'Failed to create output directory {_OUTPUT_DIR.value}: {e}')
+    raise
+
+  # if _RUN_INFERENCE.value:
+  #   # Fail early on incompatible devices, but only if we're running inference.
+  #   gpu_devices = jax.local_devices(backend='gpu')
+  #   if gpu_devices:
+  #     compute_capability = float(
+  #         gpu_devices[_GPU_DEVICE.value].compute_capability
+  #     )
+  #     if compute_capability < 6.0:
+  #       raise ValueError(
+  #           'AlphaFold 3 requires at least GPU compute capability 6.0 (see'
+  #           ' https://developer.nvidia.com/cuda-gpus).'
+  #       )
+  #     elif 7.0 <= compute_capability < 8.0:
+  #       xla_flags = os.environ.get('XLA_FLAGS')
+  #       required_flag = '--xla_disable_hlo_passes=custom-kernel-fusion-rewriter'
+  #       if not xla_flags or required_flag not in xla_flags:
+  #         raise ValueError(
+  #             'For devices with GPU compute capability 7.x (see'
+  #             ' https://developer.nvidia.com/cuda-gpus) the ENV XLA_FLAGS must'
+  #             f' include "{required_flag}".'
+  #         )
+  #       if _FLASH_ATTENTION_IMPLEMENTATION.value != 'xla':
+  #         raise ValueError(
+  #             'For devices with GPU compute capability 7.x (see'
+  #             ' https://developer.nvidia.com/cuda-gpus) the'
+  #             ' --flash_attention_implementation must be set to "xla".'
+  #         )
+
+  notice = textwrap.wrap(
+      'Running AlphaFold 3. Please note that standard AlphaFold 3 model'
+      ' parameters are only available under terms of use provided at'
+      ' https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md.'
+      ' If you do not agree to these terms and are using AlphaFold 3 derived'
+      ' model parameters, cancel execution of AlphaFold 3 inference with'
+      ' CTRL-C, and do not use the model parameters.',
+      break_long_words=False,
+      break_on_hyphens=False,
+      width=80,
+  )
+  print('\n' + '\n'.join(notice) + '\n')
+
+  max_template_date = datetime.date.fromisoformat(_MAX_TEMPLATE_DATE.value)
+  if _RUN_DATA_PIPELINE.value:
+    if not _USE_MMSEQS.value:
+      expand_path = lambda x: replace_db_dir(x, DB_DIR.value)
+      data_pipeline_config = pipeline.DataPipelineConfig(
+          jackhmmer_binary_path=_JACKHMMER_BINARY_PATH.value,
+          nhmmer_binary_path=_NHMMER_BINARY_PATH.value,
+          hmmalign_binary_path=_HMMALIGN_BINARY_PATH.value,
+          hmmsearch_binary_path=_HMMSEARCH_BINARY_PATH.value,
+          hmmbuild_binary_path=_HMMBUILD_BINARY_PATH.value,
+          mmseqs_binary_path=_MMSEQS_BINARY_PATH.value,
+          small_bfd_database_path=expand_path(_SMALL_BFD_DATABASE_PATH.value),
+          small_bfd_z_value=_SMALL_BFD_Z_VALUE.value,
+          mgnify_database_path=expand_path(_MGNIFY_DATABASE_PATH.value),
+          mgnify_z_value=_MGNIFY_Z_VALUE.value,
+          uniprot_cluster_annot_database_path=expand_path(
+              _UNIPROT_CLUSTER_ANNOT_DATABASE_PATH.value
+          ),
+          uniprot_cluster_annot_z_value=_UNIPROT_CLUSTER_ANNOT_Z_VALUE.value,
+          uniref90_database_path=expand_path(_UNIREF90_DATABASE_PATH.value),
+          uniref90_z_value=_UNIREF90_Z_VALUE.value,
+          ntrna_database_path=expand_path(_NTRNA_DATABASE_PATH.value),
+          ntrna_z_value=_NTRNA_Z_VALUE.value,
+          rfam_database_path=expand_path(_RFAM_DATABASE_PATH.value),
+          rfam_z_value=_RFAM_Z_VALUE.value,
+          rna_central_database_path=expand_path(_RNA_CENTRAL_DATABASE_PATH.value),
+          rna_central_z_value=_RNA_CENTRAL_Z_VALUE.value,
+          pdb_database_path=expand_path(_PDB_DATABASE_PATH.value),
+          seqres_database_path=expand_path(_SEQRES_DATABASE_PATH.value),
+          jackhmmer_n_cpu=_JACKHMMER_N_CPU.value,
+          jackhmmer_max_parallel_shards=_JACKHMMER_MAX_PARALLEL_SHARDS.value,
+          jackhmmer_max_threads=_JACKHMMER_MAX_THREADS.value,
+          nhmmer_n_cpu=_NHMMER_N_CPU.value,
+          nhmmer_max_parallel_shards=_NHMMER_MAX_PARALLEL_SHARDS.value,
+          nhmmer_max_threads=_NHMMER_MAX_THREADS.value,
+          max_template_date=max_template_date,
+          use_mmseqs=_USE_MMSEQS.value,
+          mmseqs_options=_MMSEQS_OPTIONS.value,
+          result2msa_options=_R2MSA_OPTIONS.value,
+      )
+    else:
+      expand_path = lambda x: replace_mmseqs_db_dir(x, DB_DIR.value, MMSEQS_DB_DIR.value)
+      data_pipeline_config = pipeline.DataPipelineConfig(
+          jackhmmer_binary_path=_JACKHMMER_BINARY_PATH.value,
+          nhmmer_binary_path=_NHMMER_BINARY_PATH.value,
+          hmmalign_binary_path=_HMMALIGN_BINARY_PATH.value,
+          hmmsearch_binary_path=_HMMSEARCH_BINARY_PATH.value,
+          hmmbuild_binary_path=_HMMBUILD_BINARY_PATH.value,
+          mmseqs_binary_path=_MMSEQS_BINARY_PATH.value,
+          small_bfd_database_path=expand_path(_MMSEQS_SMALL_BFD_DATABASE_PATH.value),
+          mgnify_database_path=expand_path(_MMSEQS_MGNIFY_DATABASE_PATH.value),
+          uniprot_cluster_annot_database_path=expand_path(
+              _MMSEQS_UNIPROT_CLUSTER_ANNOT_DATABASE_PATH.value
+          ),
+          uniref90_database_path=expand_path(_MMSEQS_UNIREF90_DATABASE_PATH.value),
+          ntrna_database_path=expand_path(_NTRNA_DATABASE_PATH.value),
+          rfam_database_path=expand_path(_RFAM_DATABASE_PATH.value),
+          rna_central_database_path=expand_path(_RNA_CENTRAL_DATABASE_PATH.value),
+          pdb_database_path=expand_path(_PDB_DATABASE_PATH.value),
+          seqres_database_path=expand_path(_SEQRES_DATABASE_PATH.value),
+          mmseqs_n_cpu=_MMSEQS_N_CPU.value,
+          nhmmer_n_cpu=_NHMMER_N_CPU.value,
+          max_template_date=max_template_date,
+          use_mmseqs=_USE_MMSEQS.value,
+          use_mmseqs_gpu=_USE_MMSEQS_GPU.value,
+          mmseqs_options=_MMSEQS_OPTIONS.value,
+          result2msa_options=_R2MSA_OPTIONS.value,
+      )      
+  else:
+    data_pipeline_config = None
+
+  if _RUN_INFERENCE.value:
+    devices = jax.local_devices(backend='gpu')
+    print(
+        f'Found local devices: {devices}, using device {_GPU_DEVICE.value}:'
+        f' {devices[_GPU_DEVICE.value]}'
+    )
+
+    print('Building model from scratch...')
+    model_runner = ModelRunner(
+        config=make_model_config(
+            flash_attention_implementation=typing.cast(
+                attention.Implementation, _FLASH_ATTENTION_IMPLEMENTATION.value
+            ),
+            num_diffusion_samples=_NUM_DIFFUSION_SAMPLES.value,
+            num_recycles=_NUM_RECYCLES.value,
+            return_embeddings=_SAVE_EMBEDDINGS.value,
+            return_distogram=_SAVE_DISTOGRAM.value,
+        ),
+        device=devices[_GPU_DEVICE.value],
+        model_dir=pathlib.Path(MODEL_DIR.value),
+    )
+    # Check we can load the model parameters before launching anything.
+    print('Checking that model parameters can be loaded...')
+    _ = model_runner.model_params
+  else:
+    model_runner = None
+
+  num_fold_inputs = 0
+  for fold_input in fold_inputs:
+    if _NUM_SEEDS.value is not None:
+      print(f'Expanding fold job {fold_input.name} to {_NUM_SEEDS.value} seeds')
+      fold_input = fold_input.with_multiple_seeds(_NUM_SEEDS.value)
+    process_fold_input(
+        fold_input=fold_input,
+        data_pipeline_config=data_pipeline_config,
+        model_runner=model_runner,
+        output_dir=os.path.join(_OUTPUT_DIR.value, fold_input.sanitised_name()),
+        buckets=tuple(int(bucket) for bucket in _BUCKETS.value),
+        ref_max_modified_date=max_template_date,
+        conformer_max_iterations=_CONFORMER_MAX_ITERATIONS.value,
+        resolve_msa_overlaps=_RESOLVE_MSA_OVERLAPS.value,
+        force_output_dir=_FORCE_OUTPUT_DIR.value,
+    )
+    num_fold_inputs += 1
+
+  print(f'Done running {num_fold_inputs} fold jobs.')
+
+
+if __name__ == '__main__':
+  flags.mark_flags_as_required(['output_dir'])
+  app.run(main)

weight/_dep/abseil-cpp/.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: false
+contact_links:
+    - name: Question
+      url: https://github.com/abseil/abseil-cpp/discussions
+      about: Have a question? Ask us anything! :-)

weight/_dep/abseil-cpp/.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,8 @@
+Thank you for your contribution to Abseil!
+
+Before submitting this PR, please be sure to read our [contributing
+guidelines](https://github.com/abseil/abseil-cpp/blob/master/CONTRIBUTING.md).
+
+If you are a Googler, please also note that it is required that you send us a
+Piper CL instead of using the GitHub pull-request process. The code propagation
+process will deliver the change to GitHub.

weight/_dep/abseil-cpp/BUILD.bazel

@@ -0,0 +1,25 @@
+#
+# Copyright 2020 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+# Expose license for external usage through bazel.
+exports_files([
+    "AUTHORS",
+    "LICENSE",
+])

weight/_dep/abseil-cpp/CMake/Googletest/CMakeLists.txt.in

@@ -0,0 +1,14 @@
+cmake_minimum_required(VERSION 3.10)
+
+project(googletest-external NONE)
+
+include(ExternalProject)
+ExternalProject_Add(googletest
+  URL               "${absl_gtest_download_url}"  # May be empty
+  SOURCE_DIR        "${absl_gtest_src_dir}"
+  BINARY_DIR        "${absl_gtest_build_dir}"
+  CONFIGURE_COMMAND ""
+  BUILD_COMMAND     ""
+  INSTALL_COMMAND   ""
+  TEST_COMMAND      ""
+)

weight/_dep/abseil-cpp/CMake/Googletest/DownloadGTest.cmake

@@ -0,0 +1,41 @@
+# Integrates googletest at configure time.  Based on the instructions at
+# https://github.com/google/googletest/tree/master/googletest#incorporating-into-an-existing-cmake-project
+
+# Set up the external googletest project, downloading the latest from Github
+# master if requested.
+configure_file(
+  ${CMAKE_CURRENT_LIST_DIR}/CMakeLists.txt.in
+  ${CMAKE_BINARY_DIR}/googletest-external/CMakeLists.txt
+)
+
+set(ABSL_SAVE_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})
+set(ABSL_SAVE_CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})
+if (BUILD_SHARED_LIBS)
+  set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DGTEST_CREATE_SHARED_LIBRARY=1")
+endif()
+
+# Configure and build the googletest source.
+execute_process(COMMAND ${CMAKE_COMMAND} -G "${CMAKE_GENERATOR}" .
+  RESULT_VARIABLE result
+  WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/googletest-external )
+if(result)
+  message(FATAL_ERROR "CMake step for googletest failed: ${result}")
+endif()
+
+execute_process(COMMAND ${CMAKE_COMMAND} --build .
+  RESULT_VARIABLE result
+  WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/googletest-external)
+if(result)
+  message(FATAL_ERROR "Build step for googletest failed: ${result}")
+endif()
+
+set(CMAKE_CXX_FLAGS ${ABSL_SAVE_CMAKE_CXX_FLAGS})
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${ABSL_SAVE_CMAKE_RUNTIME_OUTPUT_DIRECTORY})
+
+# Prevent overriding the parent project's compiler/linker settings on Windows
+set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
+
+# Add googletest directly to our build. This defines the gtest and gtest_main
+# targets.
+add_subdirectory(${absl_gtest_src_dir} ${absl_gtest_build_dir} EXCLUDE_FROM_ALL)

weight/_dep/abseil-cpp/CMake/README.md

@@ -0,0 +1,188 @@
+# Abseil CMake Build Instructions
+
+Abseil comes with a CMake build script ([CMakeLists.txt](../CMakeLists.txt))
+that can be used on a wide range of platforms ("C" stands for cross-platform.).
+If you don't have CMake installed already, you can download it for free from
+<https://www.cmake.org/>.
+
+CMake works by generating native makefiles or build projects that can
+be used in the compiler environment of your choice.
+
+For API/ABI compatibility reasons, we strongly recommend building Abseil in a
+subdirectory of your project or as an embedded dependency.
+
+## Incorporating Abseil Into a CMake Project
+
+The recommendations below are similar to those for using CMake within the
+googletest framework
+(<https://github.com/google/googletest/blob/master/googletest/README.md#incorporating-into-an-existing-cmake-project>)
+
+### Step-by-Step Instructions
+
+1. If you want to build the Abseil tests, integrate the Abseil dependency
+[Google Test](https://github.com/google/googletest) into your CMake
+project. To disable Abseil tests, you have to pass either
+`-DBUILD_TESTING=OFF` or `-DABSL_BUILD_TESTING=OFF` when configuring your
+project with CMake.
+
+2. Download Abseil and copy it into a subdirectory in your CMake project or add
+Abseil as a [git submodule](https://git-scm.com/docs/git-submodule) in your
+CMake project.
+
+3. You can then use the CMake command
+[`add_subdirectory()`](https://cmake.org/cmake/help/latest/command/add_subdirectory.html)
+to include Abseil directly in your CMake project.
+
+4. Add the **absl::** target you wish to use to the
+[`target_link_libraries()`](https://cmake.org/cmake/help/latest/command/target_link_libraries.html)
+section of your executable or of your library.<br>
+Here is a short CMakeLists.txt example of an application project using Abseil.
+
+```cmake
+cmake_minimum_required(VERSION 3.10)
+project(my_app_project)
+
+# Pick the C++ standard to compile with.
+# Abseil currently supports C++14, C++17, and C++20.
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+add_subdirectory(abseil-cpp)
+
+add_executable(my_exe source.cpp)
+target_link_libraries(my_exe absl::base absl::synchronization absl::strings)
+```
+
+Note that if you are developing a library designed for use by other clients, you
+should instead leave `CMAKE_CXX_STANDARD` unset (or only set if being built as
+the current top-level CMake project) and configure the minimum required C++
+standard at the target level. If you require a later minimum C++ standard than
+Abseil does, it's a good idea to also enforce that `CMAKE_CXX_STANDARD` (which
+will control Abseil library targets) is set to at least that minimum. For
+example:
+
+```cmake
+cmake_minimum_required(VERSION 3.10)
+project(my_lib_project)
+
+# Leave C++ standard up to the root application, so set it only if this is the
+# current top-level CMake project.
+if(CMAKE_SOURCE_DIR STREQUAL my_lib_project_SOURCE_DIR)
+  set(CMAKE_CXX_STANDARD 17)
+  set(CMAKE_CXX_STANDARD_REQUIRED ON)
+endif()
+
+add_subdirectory(abseil-cpp)
+
+add_library(my_lib source.cpp)
+target_link_libraries(my_lib absl::base absl::synchronization absl::strings)
+
+# Enforce that my_lib requires C++17. Important to document for clients that they
+# must set CMAKE_CXX_STANDARD to 17 or higher for proper Abseil ABI compatibility
+# (since otherwise, Abseil library targets could be compiled with a lower C++
+# standard than my_lib).
+target_compile_features(my_lib PUBLIC cxx_std_17)
+if(CMAKE_CXX_STANDARD LESS 17)
+  message(FATAL_ERROR
+      "my_lib_project requires CMAKE_CXX_STANDARD >= 17 (got: ${CMAKE_CXX_STANDARD})")
+endif()
+```
+
+Then the top-level application project that uses your library is responsible for
+setting a consistent `CMAKE_CXX_STANDARD` that is sufficiently high.
+
+### Running Abseil Tests with CMake
+
+Use the `-DABSL_BUILD_TESTING=ON` flag to run Abseil tests.  Note that
+BUILD_TESTING must also be on (the default).
+
+You will need to provide Abseil with a Googletest dependency.  There are two
+options for how to do this:
+
+* Use `-DABSL_USE_GOOGLETEST_HEAD`.  This will automatically download the latest
+Googletest source into the build directory at configure time.  Googletest will
+then be compiled directly alongside Abseil's tests.
+* Manually integrate Googletest with your build.  See
+https://github.com/google/googletest/blob/master/googletest/README.md#using-cmake
+for more information on using Googletest in a CMake project.
+
+For example, to run just the Abseil tests, you could use this script:
+
+```
+cd path/to/abseil-cpp
+mkdir build
+cd build
+cmake -DABSL_BUILD_TESTING=ON -DABSL_USE_GOOGLETEST_HEAD=ON ..
+make -j
+ctest
+```
+
+Currently, we only run our tests with CMake in a Linux environment, but we are
+working on the rest of our supported platforms. See
+https://github.com/abseil/abseil-cpp/projects/1 and
+https://github.com/abseil/abseil-cpp/issues/109 for more information.
+
+### Available Abseil CMake Public Targets
+
+Here's a non-exhaustive list of Abseil CMake public targets:
+
+```cmake
+absl::algorithm
+absl::base
+absl::debugging
+absl::flat_hash_map
+absl::flags
+absl::memory
+absl::meta
+absl::numeric
+absl::random_random
+absl::strings
+absl::synchronization
+absl::time
+absl::utility
+```
+
+## Traditional CMake Set-Up
+
+For larger projects, it may make sense to use the traditional CMake set-up where you build and install projects separately.
+
+First, you'd need to build and install Google Test:
+```
+cmake -S /source/googletest -B /build/googletest -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/installation/dir -DBUILD_GMOCK=ON
+cmake --build /build/googletest --target install
+```
+
+Then you need to configure and build Abseil. Make sure you enable `ABSL_USE_EXTERNAL_GOOGLETEST` and `ABSL_FIND_GOOGLETEST`. You also need to enable `ABSL_ENABLE_INSTALL` so that you can install Abseil itself.
+```
+cmake -S /source/abseil-cpp -B /build/abseil-cpp -DCMAKE_PREFIX_PATH=/installation/dir -DCMAKE_INSTALL_PREFIX=/installation/dir -DABSL_ENABLE_INSTALL=ON -DABSL_USE_EXTERNAL_GOOGLETEST=ON -DABSL_FIND_GOOGLETEST=ON
+cmake --build /temporary/build/abseil-cpp
+```
+
+(`CMAKE_PREFIX_PATH` is where you already have Google Test installed; `CMAKE_INSTALL_PREFIX` is where you want to have Abseil installed; they can be different.)
+
+Run the tests:
+```
+ctest --test-dir /temporary/build/abseil-cpp
+```
+
+And finally install:
+```
+cmake --build /temporary/build/abseil-cpp --target install
+```
+
+# CMake Option Synopsis
+
+## Enable Standard CMake Installation
+
+`-DABSL_ENABLE_INSTALL=ON`
+
+## Google Test Options
+
+`-DABSL_BUILD_TESTING=ON` must be set to enable testing
+
+- Have Abseil download and build Google Test for you: `-DABSL_USE_EXTERNAL_GOOGLETEST=OFF` (default)
+  - Download and build latest Google Test: `-DABSL_USE_GOOGLETEST_HEAD=ON`
+  - Download specific Google Test version (ZIP archive): `-DABSL_GOOGLETEST_DOWNLOAD_URL=https://.../version.zip`
+  - Use Google Test from specific local directory: `-DABSL_LOCAL_GOOGLETEST_DIR=/path/to/googletest`
+- Use Google Test included elsewhere in your project: `-DABSL_USE_EXTERNAL_GOOGLETEST=ON`
+- Use standard CMake `find_package(CTest)` to find installed Google Test: `-DABSL_USE_EXTERNAL_GOOGLETEST=ON -DABSL_FIND_GOOGLETEST=ON`

weight/_dep/abseil-cpp/CMake/install_test_project/CMakeLists.txt

@@ -0,0 +1,25 @@
+#
+# Copyright 2019 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# A simple CMakeLists.txt for testing cmake installation
+
+cmake_minimum_required(VERSION 3.10)
+project(absl_cmake_testing CXX)
+
+add_executable(simple simple.cc)
+
+find_package(absl REQUIRED)
+
+target_link_libraries(simple absl::strings absl::config)

weight/_dep/abseil-cpp/CMake/install_test_project/simple.cc

@@ -0,0 +1,32 @@
+//
+// Copyright 2019 The Abseil Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include "absl/base/config.h"
+#include "absl/strings/substitute.h"
+
+#if !defined(ABSL_LTS_RELEASE_VERSION) || ABSL_LTS_RELEASE_VERSION != 99998877
+#error ABSL_LTS_RELEASE_VERSION is not set correctly.
+#endif
+
+#if !defined(ABSL_LTS_RELEASE_PATCH_LEVEL) || ABSL_LTS_RELEASE_PATCH_LEVEL != 0
+#error ABSL_LTS_RELEASE_PATCH_LEVEL is not set correctly.
+#endif
+
+int main(int argc, char** argv) {
+  for (int i = 0; i < argc; ++i) {
+    std::cout << absl::Substitute("Arg $0: $1\n", i, argv[i]);
+  }
+}

weight/_dep/abseil-cpp/CMake/install_test_project/test.sh

@@ -0,0 +1,112 @@
+#!/bin/bash
+#
+# Copyright 2019 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Unit and integration tests for Abseil LTS CMake installation
+
+# Fail on any error. Treat unset variables an error. Print commands as executed.
+set -euox pipefail
+
+absl_dir=/abseil-cpp
+absl_build_dir=/buildfs
+googletest_builddir=/googletest_builddir
+project_dir="${absl_dir}"/CMake/install_test_project
+project_build_dir=/buildfs/project-build
+
+build_shared_libs="OFF"
+if [ "${LINK_TYPE:-}" = "DYNAMIC" ]; then
+  build_shared_libs="ON"
+fi
+
+# Build and install GoogleTest
+mkdir "${googletest_builddir}"
+pushd "${googletest_builddir}"
+curl -L "${ABSL_GOOGLETEST_DOWNLOAD_URL}" --output "${ABSL_GOOGLETEST_COMMIT}".zip
+unzip "${ABSL_GOOGLETEST_COMMIT}".zip
+pushd "googletest-${ABSL_GOOGLETEST_COMMIT}"
+mkdir build
+pushd build
+cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS="${build_shared_libs}" ..
+make -j $(nproc)
+make install
+ldconfig
+popd
+popd
+popd
+
+# Run the LTS transformations
+./create_lts.py 99998877
+
+# Build and install Abseil
+pushd "${absl_build_dir}"
+cmake "${absl_dir}" \
+  -DABSL_USE_EXTERNAL_GOOGLETEST=ON \
+  -DABSL_FIND_GOOGLETEST=ON  \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DABSL_BUILD_TESTING=ON \
+  -DBUILD_SHARED_LIBS="${build_shared_libs}"
+make -j $(nproc)
+ctest -j $(nproc) --output-on-failure
+make install
+ldconfig
+popd
+
+# Test the project against the installed Abseil
+mkdir -p "${project_build_dir}"
+pushd "${project_build_dir}"
+cmake "${project_dir}"
+cmake --build . --target simple
+
+output="$(${project_build_dir}/simple "printme" 2>&1)"
+if [[ "${output}" != *"Arg 1: printme"* ]]; then
+  echo "Faulty output on simple project:"
+  echo "${output}"
+  exit 1
+fi
+
+popd
+
+if ! grep absl::strings "/usr/local/lib/cmake/absl/abslTargets.cmake"; then
+  cat "/usr/local/lib/cmake/absl/abslTargets.cmake"
+  echo "CMake targets named incorrectly"
+  exit 1
+fi
+
+pushd "${HOME}"
+cat > hello-abseil.cc << EOF
+#include <cstdlib>
+
+#include "absl/strings/str_format.h"
+
+int main(int argc, char **argv) {
+  absl::PrintF("Hello Abseil!\n");
+  return EXIT_SUCCESS;
+}
+EOF
+
+if [ "${LINK_TYPE:-}" != "DYNAMIC" ]; then
+  pc_args=($(pkg-config --cflags --libs --static absl_str_format))
+  g++ -static -o hello-abseil hello-abseil.cc "${pc_args[@]}"
+else
+  pc_args=($(pkg-config --cflags --libs absl_str_format))
+  g++ -o hello-abseil hello-abseil.cc "${pc_args[@]}"
+fi
+hello="$(./hello-abseil)"
+[[ "${hello}" == "Hello Abseil!" ]]
+
+popd
+
+echo "Install test complete!"
+exit 0

weight/_dep/abseil-cpp/CMakeLists.txt

@@ -0,0 +1,270 @@
+#
+# Copyright 2017 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# https://github.com/google/oss-policies-info/blob/main/foundational-cxx-support-matrix.md
+# As of 2022-09-06, CMake 3.10 is the minimum supported version.
+cmake_minimum_required(VERSION 3.10)
+
+# Compiler id for Apple Clang is now AppleClang.
+if (POLICY CMP0025)
+  cmake_policy(SET CMP0025 NEW)
+endif (POLICY CMP0025)
+
+# if command can use IN_LIST
+if (POLICY CMP0057)
+  cmake_policy(SET CMP0057 NEW)
+endif (POLICY CMP0057)
+
+# Project version variables are the empty string if version is unspecified
+if (POLICY CMP0048)
+  cmake_policy(SET CMP0048 NEW)
+endif (POLICY CMP0048)
+
+# Honor the GTest_ROOT variable if specified
+if (POLICY CMP0074)
+  cmake_policy(SET CMP0074 NEW)
+endif (POLICY CMP0074)
+
+# option() honor variables
+if (POLICY CMP0077)
+  cmake_policy(SET CMP0077 NEW)
+endif (POLICY CMP0077)
+
+# Allow the user to specify the MSVC runtime
+if (POLICY CMP0091)
+  cmake_policy(SET CMP0091 NEW)
+endif (POLICY CMP0091)
+
+# try_compile() honors the CMAKE_CXX_STANDARD value
+if (POLICY CMP0067)
+  cmake_policy(SET CMP0067 NEW)
+endif (POLICY CMP0067)
+
+# Allow the user to specify the CMAKE_MSVC_DEBUG_INFORMATION_FORMAT
+if (POLICY CMP0141)
+  cmake_policy(SET CMP0141 NEW)
+endif (POLICY CMP0141)
+
+project(absl LANGUAGES CXX VERSION 20240116)
+include(CTest)
+
+# Output directory is correct by default for most build setups. However, when
+# building Abseil as a DLL, it is important to have the DLL in the same
+# directory as the executable using it. Thus, we put all executables in a single
+# /bin directory.
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
+
+# when absl is included as subproject (i.e. using add_subdirectory(abseil-cpp))
+# in the source tree of a project that uses it, install rules are disabled.
+if(NOT CMAKE_SOURCE_DIR STREQUAL PROJECT_SOURCE_DIR)
+  option(ABSL_ENABLE_INSTALL "Enable install rule" OFF)
+else()
+  option(ABSL_ENABLE_INSTALL "Enable install rule" ON)
+endif()
+
+option(ABSL_PROPAGATE_CXX_STD
+  "Use CMake C++ standard meta features (e.g. cxx_std_14) that propagate to targets that link to Abseil"
+  OFF)  # TODO: Default to ON for CMake 3.8 and greater.
+if(NOT ABSL_PROPAGATE_CXX_STD)
+  message(WARNING "A future Abseil release will default ABSL_PROPAGATE_CXX_STD to ON for CMake 3.8 and up. We recommend enabling this option to ensure your project still builds correctly.")
+endif()
+
+option(ABSL_USE_SYSTEM_INCLUDES
+  "Silence warnings in Abseil headers by marking them as SYSTEM includes"
+  OFF)
+
+list(APPEND CMAKE_MODULE_PATH
+  ${CMAKE_CURRENT_LIST_DIR}/CMake
+  ${CMAKE_CURRENT_LIST_DIR}/absl/copts
+)
+
+include(CMakePackageConfigHelpers)
+include(GNUInstallDirs)
+include(AbseilDll)
+include(AbseilHelpers)
+
+
+##
+## Using absl targets
+##
+## all public absl targets are
+## exported with the absl:: prefix
+##
+## e.g absl::base absl::synchronization absl::strings ....
+##
+## DO NOT rely on the internal targets outside of the prefix
+
+
+# include current path
+list(APPEND ABSL_COMMON_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR})
+
+if("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
+  set(ABSL_USING_CLANG ON)
+else()
+  set(ABSL_USING_CLANG OFF)
+endif()
+
+# find dependencies
+## pthread
+find_package(Threads REQUIRED)
+
+include(CMakeDependentOption)
+
+option(ABSL_BUILD_TESTING
+  "If ON, Abseil will build all of Abseil's own tests." OFF)
+
+option(ABSL_BUILD_TEST_HELPERS
+  "If ON, Abseil will build libraries that you can use to write tests against Abseil code. This option requires that Abseil is configured to use GoogleTest."
+  OFF)
+
+option(ABSL_USE_EXTERNAL_GOOGLETEST
+  "If ON, Abseil will assume that the targets for GoogleTest are already provided by the including project. This makes sense when Abseil is used with add_subdirectory." OFF)
+
+cmake_dependent_option(ABSL_FIND_GOOGLETEST
+  "If ON, Abseil will use find_package(GTest) rather than assuming that GoogleTest is already provided by the including project."
+  ON
+  "ABSL_USE_EXTERNAL_GOOGLETEST"
+  OFF)
+
+
+option(ABSL_USE_GOOGLETEST_HEAD
+  "If ON, abseil will download HEAD from GoogleTest at config time." OFF)
+
+set(ABSL_GOOGLETEST_DOWNLOAD_URL "" CACHE STRING "If set, download GoogleTest from this URL")
+
+set(ABSL_LOCAL_GOOGLETEST_DIR "/usr/src/googletest" CACHE PATH
+  "If ABSL_USE_GOOGLETEST_HEAD is OFF and ABSL_GOOGLETEST_URL is not set, specifies the directory of a local GoogleTest checkout."
+  )
+
+if((BUILD_TESTING AND ABSL_BUILD_TESTING) OR ABSL_BUILD_TEST_HELPERS)
+  if (ABSL_USE_EXTERNAL_GOOGLETEST)
+    if (ABSL_FIND_GOOGLETEST)
+      find_package(GTest REQUIRED)
+    elseif(NOT TARGET GTest::gtest)
+      if(TARGET gtest)
+        # When Google Test is included directly rather than through find_package, the aliases are missing.
+        add_library(GTest::gtest ALIAS gtest)
+        add_library(GTest::gtest_main ALIAS gtest_main)
+        add_library(GTest::gmock ALIAS gmock)
+        add_library(GTest::gmock_main ALIAS gmock_main)
+      else()
+        message(FATAL_ERROR "ABSL_USE_EXTERNAL_GOOGLETEST is ON and ABSL_FIND_GOOGLETEST is OFF, which means that the top-level project must build the Google Test project. However, the target gtest was not found.")
+      endif()
+    endif()
+  else()
+    set(absl_gtest_build_dir ${CMAKE_BINARY_DIR}/googletest-build)
+    if(ABSL_USE_GOOGLETEST_HEAD AND ABSL_GOOGLETEST_DOWNLOAD_URL)
+      message(FATAL_ERROR "Do not set both ABSL_USE_GOOGLETEST_HEAD and ABSL_GOOGLETEST_DOWNLOAD_URL")
+    endif()
+    if(ABSL_USE_GOOGLETEST_HEAD)
+      set(absl_gtest_download_url "https://github.com/google/googletest/archive/main.zip")
+    elseif(ABSL_GOOGLETEST_DOWNLOAD_URL)
+      set(absl_gtest_download_url ${ABSL_GOOGLETEST_DOWNLOAD_URL})
+    endif()
+    if(absl_gtest_download_url)
+      set(absl_gtest_src_dir ${CMAKE_BINARY_DIR}/googletest-src)
+    else()
+      set(absl_gtest_src_dir ${ABSL_LOCAL_GOOGLETEST_DIR})
+    endif()
+    include(CMake/Googletest/DownloadGTest.cmake)
+  endif()
+endif()
+
+add_subdirectory(absl)
+
+if(ABSL_ENABLE_INSTALL)
+  
+
+  # install as a subdirectory only
+  install(EXPORT ${PROJECT_NAME}Targets
+    NAMESPACE absl::
+    DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}"
+  )
+
+  configure_package_config_file(
+    CMake/abslConfig.cmake.in
+    "${PROJECT_BINARY_DIR}/${PROJECT_NAME}Config.cmake"
+    INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}"
+  )
+  install(FILES "${PROJECT_BINARY_DIR}/${PROJECT_NAME}Config.cmake"
+    DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}"
+  )
+
+  # Abseil only has a version in LTS releases.  This mechanism is accomplished
+  # Abseil's internal Copybara (https://github.com/google/copybara) workflows and
+  # isn't visible in the CMake buildsystem itself.
+  if(absl_VERSION)
+    write_basic_package_version_file(
+      "${PROJECT_BINARY_DIR}/${PROJECT_NAME}ConfigVersion.cmake"
+      COMPATIBILITY ExactVersion
+    )
+
+    install(FILES "${PROJECT_BINARY_DIR}/${PROJECT_NAME}ConfigVersion.cmake"
+      DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}"
+    )
+  endif()  # absl_VERSION
+
+  install(DIRECTORY absl
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+    FILES_MATCHING
+      PATTERN "*.inc"
+      PATTERN "*.h"
+      PATTERN "copts" EXCLUDE
+      PATTERN "testdata" EXCLUDE
+    )
+
+  # Rewrite options.h to use the compiled ABI.
+  file(READ "absl/base/options.h" ABSL_INTERNAL_OPTIONS_H_CONTENTS)
+
+  # Handle features that require at least C++20.
+  if (ABSL_INTERNAL_AT_LEAST_CXX20)
+    foreach(FEATURE "ORDERING")
+      string(REPLACE
+      "#define ABSL_OPTION_USE_STD_${FEATURE} 2"
+      "#define ABSL_OPTION_USE_STD_${FEATURE} 1"
+      ABSL_INTERNAL_OPTIONS_H_PINNED
+      "${ABSL_INTERNAL_OPTIONS_H_CONTENTS}")
+      set(ABSL_INTERNAL_OPTIONS_H_CONTENTS "${ABSL_INTERNAL_OPTIONS_H_PINNED}")
+    endforeach()
+  endif()
+
+  # Handle features that require at least C++17.
+  if (ABSL_INTERNAL_AT_LEAST_CXX17)
+    foreach(FEATURE "ANY" "OPTIONAL" "STRING_VIEW" "VARIANT")
+      string(REPLACE
+      "#define ABSL_OPTION_USE_STD_${FEATURE} 2"
+      "#define ABSL_OPTION_USE_STD_${FEATURE} 1"
+      ABSL_INTERNAL_OPTIONS_H_PINNED
+      "${ABSL_INTERNAL_OPTIONS_H_CONTENTS}")
+      set(ABSL_INTERNAL_OPTIONS_H_CONTENTS "${ABSL_INTERNAL_OPTIONS_H_PINNED}")
+    endforeach()
+  endif()
+
+  # Any feature that still has the value of 2 (because it was not handled above)
+  # should be set to 0.
+  string(REGEX REPLACE
+    "#define ABSL_OPTION_USE_STD_([^ ]*) 2"
+    "#define ABSL_OPTION_USE_STD_\\1 0"
+    ABSL_INTERNAL_OPTIONS_H_PINNED
+    "${ABSL_INTERNAL_OPTIONS_H_CONTENTS}")
+
+  file(WRITE "${CMAKE_BINARY_DIR}/options-pinned.h" "${ABSL_INTERNAL_OPTIONS_H_PINNED}")
+
+  install(FILES "${CMAKE_BINARY_DIR}/options-pinned.h"
+         DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/absl/base
+         RENAME "options.h")
+
+endif()  # ABSL_ENABLE_INSTALL

weight/_dep/abseil-cpp/CONTRIBUTING.md

@@ -0,0 +1,141 @@
+# How to Contribute to Abseil
+
+We'd love to accept your patches and contributions to this project. There are
+just a few small guidelines you need to follow.
+
+NOTE: If you are new to GitHub, please start by reading [Pull Request
+howto](https://help.github.com/articles/about-pull-requests/)
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement. You (or your employer) retain the copyright to your contribution,
+this simply gives us permission to use and redistribute your contributions as
+part of the project. Head over to <https://cla.developers.google.com/> to see
+your current agreements on file or to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Contribution Guidelines
+
+Potential contributors sometimes ask us if the Abseil project is the appropriate
+home for their utility library code or for specific functions implementing
+missing portions of the standard. Often, the answer to this question is "no".
+We’d like to articulate our thinking on this issue so that our choices can be
+understood by everyone and so that contributors can have a better intuition
+about whether Abseil might be interested in adopting a new library.
+
+### Priorities
+
+Although our mission is to augment the C++ standard library, our goal is not to
+provide a full forward-compatible implementation of the latest standard. For us
+to consider a library for inclusion in Abseil, it is not enough that a library
+is useful. We generally choose to release a library when it meets at least one
+of the following criteria:
+
+*   **Widespread usage** - Using our internal codebase to help gauge usage, most
+    of the libraries we've released have tens of thousands of users.
+*   **Anticipated widespread usage** - Pre-adoption of some standard-compliant
+    APIs may not have broad adoption initially but can be expected to pick up
+    usage when it replaces legacy APIs. `absl::from_chars`, for example,
+    replaces existing code that converts strings to numbers and will therefore
+    likely see usage growth.
+*   **High impact** - APIs that provide a key solution to a specific problem,
+    such as `absl::FixedArray`, have higher impact than usage numbers may signal
+    and are released because of their importance.
+*   **Direct support for a library that falls under one of the above** - When we
+    want access to a smaller library as an implementation detail for a
+    higher-priority library we plan to release, we may release it, as we did
+    with portions of `absl/meta/type_traits.h`. One consequence of this is that
+    the presence of a library in Abseil does not necessarily mean that other
+    similar libraries would be a high priority.
+
+### API Freeze Consequences
+
+Via the
+[Abseil Compatibility Guidelines](https://abseil.io/about/compatibility), we
+have promised a large degree of API stability. In particular, we will not make
+backward-incompatible changes to released APIs without also shipping a tool or
+process that can upgrade our users' code. We are not yet at the point of easily
+releasing such tools. Therefore, at this time, shipping a library establishes an
+API contract which is borderline unchangeable. (We can add new functionality,
+but we cannot easily change existing behavior.) This constraint forces us to
+very carefully review all APIs that we ship.
+
+
+## Coding Style
+
+To keep the source consistent, readable, diffable and easy to merge, we use a
+fairly rigid coding style, as defined by the
+[google-styleguide](https://github.com/google/styleguide) project. All patches
+will be expected to conform to the style outlined
+[here](https://google.github.io/styleguide/cppguide.html).
+
+## Guidelines for Pull Requests
+
+*   If you are a Googler, it is required that you send us a Piper CL instead of
+    using the GitHub pull-request process. The code propagation process will
+    deliver the change to GitHub.
+
+*   Create **small PRs** that are narrowly focused on **addressing a single
+    concern**. We often receive PRs that are trying to fix several things at a
+    time, but if only one fix is considered acceptable, nothing gets merged and
+    both author's & review's time is wasted. Create more PRs to address
+    different concerns and everyone will be happy.
+
+*   For speculative changes, consider opening an [Abseil
+    issue](https://github.com/abseil/abseil-cpp/issues) and discussing it first.
+    If you are suggesting a behavioral or API change, consider starting with an
+    [Abseil proposal template](ABSEIL_ISSUE_TEMPLATE.md).
+
+*   Provide a good **PR description** as a record of **what** change is being
+    made and **why** it was made. Link to a GitHub issue if it exists.
+
+*   Don't fix code style and formatting unless you are already changing that
+    line to address an issue. Formatting of modified lines may be done using
+   `git clang-format`. PRs with irrelevant changes won't be merged. If
+    you do want to fix formatting or style, do that in a separate PR.
+
+*   Unless your PR is trivial, you should expect there will be reviewer comments
+    that you'll need to address before merging. We expect you to be reasonably
+    responsive to those comments, otherwise the PR will be closed after 2-3
+    weeks of inactivity.
+
+*   Maintain **clean commit history** and use **meaningful commit messages**.
+    PRs with messy commit history are difficult to review and won't be merged.
+    Use `rebase -i upstream/master` to curate your commit history and/or to
+    bring in latest changes from master (but avoid rebasing in the middle of a
+    code review).
+
+*   Keep your PR up to date with upstream/master (if there are merge conflicts,
+    we can't really merge your change).
+
+*   **All tests need to be passing** before your change can be merged. We
+    recommend you **run tests locally** (see below)
+
+*   Exceptions to the rules can be made if there's a compelling reason for doing
+    so. That is - the rules are here to serve us, not the other way around, and
+    the rules need to be serving their intended purpose to be valuable.
+
+*   All submissions, including submissions by project members, require review.
+
+## Running Tests
+
+If you have [Bazel](https://bazel.build/) installed, use `bazel test
+--test_tag_filters="-benchmark" ...` to run the unit tests.
+
+If you are running the Linux operating system and have
+[Docker](https://www.docker.com/) installed, you can also run the `linux_*.sh`
+scripts under the `ci/`(https://github.com/abseil/abseil-cpp/tree/master/ci)
+directory to test Abseil under a variety of conditions.
+
+## Abseil Committers
+
+The current members of the Abseil engineering team are the only committers at
+present.
+
+## Release Process
+
+Abseil lives at head, where latest-and-greatest code can be found.

weight/_dep/abseil-cpp/FAQ.md

@@ -0,0 +1,167 @@
+# Abseil FAQ
+
+## Is Abseil the right home for my utility library?
+
+Most often the answer to the question is "no." As both the [About
+Abseil](https://abseil.io/about/) page and our [contributing
+guidelines](https://github.com/abseil/abseil-cpp/blob/master/CONTRIBUTING.md#contribution-guidelines)
+explain, Abseil contains a variety of core C++ library code that is widely used
+at [Google](https://www.google.com/). As such, Abseil's primary purpose is to be
+used as a dependency by Google's open source C++ projects. While we do hope that
+Abseil is also useful to the C++ community at large, this added constraint also
+means that we are unlikely to accept a contribution of utility code that isn't
+already widely used by Google.
+
+## How to I set the C++ dialect used to build Abseil?
+
+The short answer is that whatever mechanism you choose, you need to make sure
+that you set this option consistently at the global level for your entire
+project. If, for example, you want to set the C++ dialect to C++17, with
+[Bazel](https://bazel/build/) as the build system and `gcc` or `clang` as the
+compiler, there several ways to do this:
+* Pass `--cxxopt=-std=c++17` on the command line (for example, `bazel build
+  --cxxopt=-std=c++17 ...`)
+* Set the environment variable `BAZEL_CXXOPTS` (for example,
+  `BAZEL_CXXOPTS=-std=c++17`)
+* Add `build --cxxopt=-std=c++17` to your [`.bazelrc`
+  file](https://docs.bazel.build/versions/master/guide.html#bazelrc)
+
+If you are using CMake as the build system, you'll need to add a line like
+`set(CMAKE_CXX_STANDARD 17)` to your top level `CMakeLists.txt` file. If you
+are developing a library designed to be used by other clients, you should
+instead leave `CMAKE_CXX_STANDARD` unset and configure the minimum C++ standard
+required by each of your library targets via `target_compile_features`. See the
+[CMake build
+instructions](https://github.com/abseil/abseil-cpp/blob/master/CMake/README.md)
+for more information.
+
+For a longer answer to this question and to understand why some other approaches
+don't work, see the answer to ["What is ABI and why don't you recommend using a
+pre-compiled version of
+Abseil?"](#what-is-abi-and-why-dont-you-recommend-using-a-pre-compiled-version-of-abseil)
+
+## What is ABI and why don't you recommend using a pre-compiled version of Abseil?
+
+For the purposes of this discussion, you can think of
+[ABI](https://en.wikipedia.org/wiki/Application_binary_interface) as the
+compiled representation of the interfaces in code. This is in contrast to
+[API](https://en.wikipedia.org/wiki/Application_programming_interface), which
+you can think of as the interfaces as defined by the code itself. [Abseil has a
+strong promise of API compatibility, but does not make any promise of ABI
+compatibility](https://abseil.io/about/compatibility). Let's take a look at what
+this means in practice.
+
+You might be tempted to do something like this in a
+[Bazel](https://bazel.build/) `BUILD` file:
+
+```
+# DON'T DO THIS!!!
+cc_library(
+    name = "my_library",
+    srcs = ["my_library.cc"],
+    copts = ["-std=c++17"],  # May create a mixed-mode compile!
+    deps = ["@com_google_absl//absl/strings"],
+)
+```
+
+Applying `-std=c++17` to an individual target in your `BUILD` file is going to
+compile that specific target in C++17 mode, but it isn't going to ensure the
+Abseil library is built in C++17 mode, since the Abseil library itself is a
+different build target. If your code includes an Abseil header, then your
+program may contain conflicting definitions of the same
+class/function/variable/enum, etc. As a rule, all compile options that affect
+the ABI of a program need to be applied to the entire build on a global basis.
+
+C++ has something called the [One Definition
+Rule](https://en.wikipedia.org/wiki/One_Definition_Rule) (ODR). C++ doesn't
+allow multiple definitions of the same class/function/variable/enum, etc. ODR
+violations sometimes result in linker errors, but linkers do not always catch
+violations. Uncaught ODR violations can result in strange runtime behaviors or
+crashes that can be hard to debug.
+
+If you build the Abseil library and your code using different compile options
+that affect ABI, there is a good chance you will run afoul of the One Definition
+Rule. Examples of GCC compile options that affect ABI include (but aren't
+limited to) language dialect (e.g. `-std=`), optimization level (e.g. `-O2`),
+code generation flags (e.g. `-fexceptions`), and preprocessor defines
+(e.g. `-DNDEBUG`).
+
+If you use a pre-compiled version of Abseil, (for example, from your Linux
+distribution package manager or from something like
+[vcpkg](https://github.com/microsoft/vcpkg)) you have to be very careful to
+ensure ABI compatibility across the components of your program. The only way you
+can be sure your program is going to be correct regarding ABI is to ensure
+you've used the exact same compile options as were used to build the
+pre-compiled library. This does not mean that Abseil cannot work as part of a
+Linux distribution since a knowledgeable binary packager will have ensured that
+all packages have been built with consistent compile options. This is one of the
+reasons we warn against - though do not outright reject - using Abseil as a
+pre-compiled library.
+
+Another possible way that you might afoul of ABI issues is if you accidentally
+include two versions of Abseil in your program. Multiple versions of Abseil can
+end up within the same binary if your program uses the Abseil library and
+another library also transitively depends on Abseil (resulting in what is
+sometimes called the diamond dependency problem). In cases such as this you must
+structure your build so that all libraries use the same version of Abseil.
+[Abseil's strong promise of API compatibility between
+releases](https://abseil.io/about/compatibility) means the latest "HEAD" release
+of Abseil is almost certainly the right choice if you are doing as we recommend
+and building all of your code from source.
+
+For these reasons we recommend you avoid pre-compiled code and build the Abseil
+library yourself in a consistent manner with the rest of your code.
+
+## What is "live at head" and how do I do it?
+
+From Abseil's point-of-view, "live at head" means that every Abseil source
+release (which happens on an almost daily basis) is either API compatible with
+the previous release, or comes with an automated tool that you can run over code
+to make it compatible. In practice, the need to use an automated tool is
+extremely rare. This means that upgrading from one source release to another
+should be a routine practice that can and should be performed often.
+
+We recommend you update to the [latest commit in the `master` branch of
+Abseil](https://github.com/abseil/abseil-cpp/commits/master) as often as
+possible. Not only will you pick up bug fixes more quickly, but if you have good
+automated testing, you will catch and be able to fix any [Hyrum's
+Law](https://www.hyrumslaw.com/) dependency problems on an incremental basis
+instead of being overwhelmed by them and having difficulty isolating them if you
+wait longer between updates.
+
+If you are using the [Bazel](https://bazel.build/) build system and its
+[external dependencies](https://docs.bazel.build/versions/master/external.html)
+feature, updating the
+[`http_archive`](https://docs.bazel.build/versions/master/repo/http.html#http_archive)
+rule in your
+[`WORKSPACE`](https://docs.bazel.build/versions/master/be/workspace.html) for
+`com_google_abseil` to point to the [latest commit in the `master` branch of
+Abseil](https://github.com/abseil/abseil-cpp/commits/master) is all you need to
+do. For example, on February 11, 2020, the latest commit to the master branch
+was `98eb410c93ad059f9bba1bf43f5bb916fc92a5ea`. To update to this commit, you
+would add the following snippet to your `WORKSPACE` file:
+
+```
+http_archive(
+  name = "com_google_absl",
+  urls = ["https://github.com/abseil/abseil-cpp/archive/98eb410c93ad059f9bba1bf43f5bb916fc92a5ea.zip"],  # 2020-02-11T18:50:53Z
+  strip_prefix = "abseil-cpp-98eb410c93ad059f9bba1bf43f5bb916fc92a5ea",
+  sha256 = "aabf6c57e3834f8dc3873a927f37eaf69975d4b28117fc7427dfb1c661542a87",
+)
+```
+
+To get the `sha256` of this URL, run `curl -sL --output -
+https://github.com/abseil/abseil-cpp/archive/98eb410c93ad059f9bba1bf43f5bb916fc92a5ea.zip
+| sha256sum -`.
+
+You can commit the updated `WORKSPACE` file to your source control every time
+you update, and if you have good automated testing, you might even consider
+automating this.
+
+One thing we don't recommend is using GitHub's `master.zip` files (for example
+[https://github.com/abseil/abseil-cpp/archive/master.zip](https://github.com/abseil/abseil-cpp/archive/master.zip)),
+which are always the latest commit in the `master` branch, to implement live at
+head. Since these `master.zip` URLs are not versioned, you will lose build
+reproducibility. In addition, some build systems, including Bazel, will simply
+cache this file, which means you won't actually be updating to the latest
+release until your cache is cleared or invalidated.

weight/_dep/abseil-cpp/MODULE.bazel

@@ -0,0 +1,39 @@
+# Copyright 2024 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# https://bazel.build/external/overview#bzlmod
+
+module(
+    name = "abseil-cpp",
+    version = "20240116.2",
+    compatibility_level = 1,
+)
+
+# Only direct dependencies need to be listed below.
+# Please keep the versions in sync with the versions in the WORKSPACE file.
+
+bazel_dep(name = "bazel_skylib",
+          version = "1.5.0")
+
+bazel_dep(name = "google_benchmark",
+          version = "1.8.3",
+          repo_name = "com_github_google_benchmark",
+          dev_dependency = True)
+
+bazel_dep(name = "googletest",
+          version = "1.14.0.bcr.1",
+          repo_name = "com_google_googletest")
+
+bazel_dep(name = "platforms",
+          version = "0.0.8")

weight/_dep/abseil-cpp/PrivacyInfo.xcprivacy

@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>NSPrivacyTracking</key>
+	<false/>
+	<key>NSPrivacyCollectedDataTypes</key>
+	<array/>
+	<key>NSPrivacyTrackingDomains</key>
+	<array/>
+	<key>NSPrivacyAccessedAPITypes</key>
+	<array/>
+</dict>
+</plist>