Add files using upload-large-folder tool

flax_model/alphafold3/_tools/hmmer/.built

@@ -0,0 +1 @@
+built

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

@@ -0,0 +1,134 @@
+.TH "esl\-compstruct" 1 "Aug 2023" "Easel 0.49" "Easel Manual"
+
+.SH NAME
+esl\-compstruct \- calculate accuracy of RNA secondary structure predictions
+
+
+.SH SYNOPSIS
+.B esl\-compstruct
+[\fIoptions\fR]
+.I trusted_file
+.I test_file
+
+
+
+.SH DESCRIPTION
+
+.PP
+.B esl\-compstruct 
+evaluates the accuracy of RNA secondary structure predictions
+on a per-base-pair basis.
+The 
+.I trusted_file 
+contains one or more sequences with trusted (known) RNA
+secondary structure annotation. The
+.I test_file
+contains the same sequences, in the same order, with
+predicted RNA secondary structure annotation.
+.B esl\-compstruct 
+reads the structures and compares them,
+and calculates both the sensitivity (the number
+of true base pairs that are correctly predicted)
+and the positive predictive value (PPV;
+the number of predicted base pairs that are true).
+Results are reported for each individual sequence,
+and in summary for all sequences together.
+
+.PP
+Both files must contain secondary structure annotation in
+WUSS notation. Only SELEX and Stockholm formats support 
+structure markup at present.
+
+.PP
+The default definition of a correctly predicted base pair
+is that a true pair (i,j) must exactly match a predicted
+pair (i,j).
+
+.PP
+Mathews and colleagues (Mathews et al., JMB 288:911-940, 1999) use a
+more relaxed definition. Mathews defines "correct" as follows: a true
+pair (i,j) is correctly predicted if any of the following pairs are
+predicted: (i,j), (i+1,j), (i\-1,j), (i,j+1), or (i,j\-1). This rule
+allows for "slipped helices" off by one base.  The
+.B \-m
+option activates this rule for both sensitivity and for
+specificity. For specificity, the rule is reversed: predicted pair
+(i,j) is considered to be true if the true structure contains one of
+the five pairs (i,j), (i+1,j), (i\-1,j), (i,j+1), or (i,j\-1).
+
+
+
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Print brief help; includes version number and summary of
+all options, including expert options.
+
+.TP
+.B \-m
+Use the Mathews relaxed accuracy rule (see above), instead
+of requiring exact prediction of base pairs.
+
+.TP
+.B \-p 
+Count pseudoknotted base pairs towards the accuracy, in either trusted
+or predicted structures. By default, pseudoknots are ignored. 
+.IP
+Normally, only the 
+.I trusted_file 
+would have pseudoknot annotation, since most RNA secondary structure
+prediction programs do not predict pseudoknots. Using the
+.B \-p
+option allows you to penalize the prediction program for not
+predicting known pseudoknots. In a case where both the
+.I trusted_file 
+and the
+.I test_file 
+have pseudoknot annotation,  the
+.B \-p
+option lets you count pseudoknots in evaluating
+the prediction accuracy. Beware, however, the case where you
+use a pseudoknot-capable prediction program to generate the
+.IR test_file ,
+but the 
+.I trusted_file 
+does not have pseudoknot annotation; in this case,
+.B \-p
+will penalize any predicted pseudoknots when it calculates
+specificity, even if they're right, because they don't appear in the
+trusted annotation. This is probably not what you'd want to do.
+
+
+
+.SH EXPERT OPTIONS
+
+.TP
+.B \-\-quiet
+Don't print any verbose header information. (Used by regression test
+scripts, for example, to suppress version/date information.)
+
+
+
+.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/esl-construct.1

@@ -0,0 +1,196 @@
+.TH "esl\-construct" 1 "Aug 2023" "Easel 0.49" "Easel Manual"
+
+.SH NAME
+esl\-construct \- describe or create a consensus secondary structure
+
+.SH SYNOPSIS
+.B esl\-construct
+[\fIoptions\fR]
+.I msafile
+
+.SH DESCRIPTION
+
+.B esl\-construct 
+reports information on existing consensus secondary structure annotation
+of an alignment or derives new consensus secondary structures based on 
+structure annotation for individual aligned sequences.
+
+.PP
+The alignment file must contain either individual sequence secondary
+structure annotation (Stockholm #=GR SS), consensus secondary
+structure annotation (Stockohlm #=GC SS_cons), or both. All structure
+annotation must be in WUSS notation (Vienna dot paranetheses notation will be
+correctly interpreted). At present, the alignment file must be in
+Stockholm format and contain RNA or DNA sequences.
+
+.PP
+By default, 
+.B esl\-construct
+generates lists the sequences in the alignment that have structure
+annotation and the number of basepairs in those structures. If the
+alignment also contains consensus structure annotation, the default output
+will list how many of the individual basepairs overlap with the
+consensus basepairs and how many conflict with a consensus basepair.
+
+.PP
+For the purposes of this miniapp, a basepair 'conflict' exists between
+two basepairs in different structures, one between columns i and j and
+the other between columns k and l, if (i == k and j != l) or (j == l
+and i != k).
+
+.PP
+.B esl\-construct
+can also be used to derive a new consensus structure based on
+structure annotation for individual sequences in the alignment by
+using any of the following options: 
+.BR \-x ,
+.BR \-r ,
+.BR \-c ,
+\fB\-\-indi \fR\fI<s>\fR,
+\fB\-\-ffreq \fR\fI<x>\fR,
+.BR \-\-fmin .
+These are described below. All of these options require the 
+.BI \-o " <f>"
+option be used as well to specify that a new alignment file 
+.I <f> 
+be created. Differences between the new alignment(s) and the input
+alignment(s) will be limited to the the consensus secondary structure (#=GC
+SS_cons) annotation and possibly reference (#=GC RF) annotation.
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Print brief help; includes version number and summary of
+all options, including expert options.
+
+.TP
+.B \-a
+List all alignment positions that are involved in at least one
+conflicting basepair in at least one sequence to the screen, and then
+exit.
+
+.TP
+.B \-v
+Be verbose; with no other options, list individual sequence basepair
+conflicts as well as summary statistics.
+
+.TP
+.B \-x
+Compute a new consensus structure as the maximally sized set of
+basepairs (greatest number of basepairs) chosen from all individual
+structures that contains 0 conflicts. Output the alignment with the
+new SS_cons annotation. This option must be used in combination with
+the
+.B \-o 
+option. 
+
+.TP
+.B \-r
+Remove any consensus basepairs that conflict with >= 1 individual
+basepair and output the alignment with the new SS_cons annotation. 
+This option must be used in combination with the 
+.B \-o 
+option. 
+
+.TP
+.B \-c
+Define a new consensus secondary structure as the individual structure
+annotation that has the maximum number of consistent basepairs with
+the existing consensus secondary structure annotation.
+This option must be used in combination with the 
+.B \-o 
+option. 
+
+.TP
+.B \-\-rfc
+With 
+.BR \-c , 
+set the reference annotation (#=GC RF) as the sequence whose 
+individual structure becomes the consensus structure. 
+
+.TP
+.BI \-\-indi " <s>"
+Define a new consensus secondary structure as the individual structure
+annotation from sequence named
+.IR <s> .
+This option must be used in combination with
+the
+.B \-o 
+option. 
+
+.TP
+.B \-\-rfindi
+With 
+\fB\-\-indi \fR\fI<s>\fR, 
+set the reference annotation (#=GC RF) as the sequence named
+.BR <s> .
+
+.TP
+.BI \-\-ffreq " <x>"
+Define a new consensus structure as the set of basepairs between
+columns i:j that are paired in more than 
+.I <x>
+fraction of the individual sequence structures.
+This option must be used in combination with
+the
+.B \-o 
+option. 
+
+.TP
+.B \-\-fmin
+Same as
+.BI \-\-ffreq " <x>" 
+except find the maximal 
+.I <x>
+that gives a consistent consensus structure. A consistent structure
+has each base (alignment position) as a member of at most 1 basepair.  
+
+.TP
+.BI \-o " <s>", 
+Output the alignment(s) with new consensus structure annotation to file
+.IR <f> .
+
+.TP
+.B \-\-pfam 
+With
+.BR -o ,
+specify that the alignment output format be Pfam format, a
+special type of non-interleaved Stockholm on which each sequence
+appears on a single line.
+
+.TP
+.BI \-l " <f>"
+Create a new file 
+.I <f> 
+that lists the sequences that have at least one
+basepair that conflicts with a consensus basepair.
+
+.TP
+.BI \-\-lmax " <n>"
+With
+.BR \-l ,
+only list sequences that have more than 
+.I <n>
+basepairs that conflict with the consensus structure to the list file.
+
+
+.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/esl-seqrange.1

@@ -0,0 +1,99 @@
+.TH "esl\-seqrange" 1 "Aug 2023" "Easel 0.49" "Easel Manual"
+
+.SH NAME
+esl\-seqrange \- determine a range of sequences for one of many parallel processes
+
+.SH SYNOPSIS
+.B esl\-sfetch
+[\fIoptions\fR]
+.I seqfile
+.I procidx
+.I nproc
+
+.SH DESCRIPTION
+
+.PP
+.B esl\-seqrange
+reads an SSI-indexed 
+.I seqfile
+and determines the range of sequence indices in that file that process
+number 
+.I procidx
+out of
+.I nproc
+total processes should operate on during a parallel processing of 
+.IR seqfile .
+
+.PP
+The 
+.I seqfile 
+must be indexed first using 
+.B esl\-sfetch \-\-index
+.IR seqfile .
+This creates an SSI index file
+.IR seqfile .ssi.
+An SSI file is required in order for
+.B esl\-seqrange
+to work.
+
+.PP
+Sequence index ranges are calculated using a simple rule: the number
+of sequences for each process should be identical, or as close as
+possible to identical, across all processes. The lengths of the sequences
+are not considered (even though they probably should be).
+
+.SH OPTIONS
+
+.TP
+.B \-h
+Print brief help; includes version number and summary of
+all options, including expert 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).
+
+
+
+.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/data/tools/hmmalign.py

@@ -0,0 +1,134 @@
+
+
+"""A Python wrapper for hmmalign from the HMMER Suite."""
+
+from collections.abc import Mapping, Sequence
+import os
+import tempfile
+
+from flax_model.alphafold3.data import parsers
+from flax_model.alphafold3.data.tools import subprocess_utils
+
+
+def _to_a3m(sequences: Sequence[str], name_prefix: str = 'sequence') -> str:
+  a3m = ''
+  for i, sequence in enumerate(sequences, 1):
+    a3m += f'> {name_prefix} {i}\n{sequence}\n'
+  return a3m
+
+
+class Hmmalign:
+  """Python wrapper of the hmmalign binary."""
+
+  def __init__(self, binary_path: str):
+    """Initializes the Python hmmalign wrapper.
+
+    Args:
+      binary_path: Path to the hmmalign binary.
+
+    Raises:
+      RuntimeError: If hmmalign binary not found within the path.
+    """
+    self._binary_path = binary_path
+
+    subprocess_utils.check_binary_exists(path=self._binary_path, name='hmmalign')
+
+  def align_sequences(
+      self,
+      sequences: Sequence[str],
+      profile: str,
+      extra_flags: Mapping[str, str] | None = None,
+  ) -> str:
+    """Aligns sequence list to the profile and returns the alignment in A3M."""
+    return self.align(
+        a3m_str=_to_a3m(sequences, name_prefix='query'),
+        profile=profile,
+        extra_flags=extra_flags,
+    )
+
+  def align(
+      self,
+      a3m_str: str,
+      profile: str,
+      extra_flags: Mapping[str, str] | None = None,
+  ) -> str:
+    """Aligns sequences in A3M to the profile and returns the alignment in A3M.
+
+    Args:
+      a3m_str: A list of sequence strings.
+      profile: A hmm file with the hmm profile to align the sequences to.
+      extra_flags: Dictionary with extra flags, flag_name: flag_value, that are
+        added to hmmalign.
+
+    Returns:
+      An A3M string with the aligned sequences.
+
+    Raises:
+      RuntimeError: If hmmalign fails.
+    """
+    with tempfile.TemporaryDirectory() as query_tmp_dir:
+      input_profile = os.path.join(query_tmp_dir, 'profile.hmm')
+      input_sequences = os.path.join(query_tmp_dir, 'sequences.a3m')
+      output_a3m_path = os.path.join(query_tmp_dir, 'output.a3m')
+
+      with open(input_profile, 'w') as f:
+        f.write(profile)
+
+      with open(input_sequences, 'w') as f:
+        f.write(a3m_str)
+
+      cmd = [
+          self._binary_path,
+          *('-o', output_a3m_path),
+          *('--outformat', 'A2M'),  # A2M is A3M in the HMMER suite.
+      ]
+      if extra_flags:
+        for flag_name, flag_value in extra_flags.items():
+          cmd.extend([flag_name, flag_value])
+      cmd.extend([input_profile, input_sequences])
+
+      subprocess_utils.run(
+          cmd=cmd,
+          cmd_name='hmmalign',
+          log_stdout=False,
+          log_stderr=True,
+          log_on_process_error=True,
+      )
+
+      with open(output_a3m_path, encoding='utf-8') as f:
+        a3m = f.read()
+
+    return a3m
+
+  def align_sequences_to_profile(self, profile: str, sequences_a3m: str) -> str:
+    """Aligns the sequences to profile and returns the alignment in A3M string.
+
+    Uses hmmalign to align the sequences to the profile, then ouputs the
+    sequence contatenated at the beginning of the sequences in the A3M format.
+    As the sequences are represented by an alignment with possible gaps ('-')
+    and insertions (lowercase characters), the method first removes the gaps,
+    then uppercases the insertions to prepare the sequences for realignment.
+    Sequences with gaps cannot be aligned, as '-'s are not a valid symbol to
+    align; lowercase characters must be uppercased to preserve the original
+    sequences before realignment.
+
+    Args:
+      profile: The Hmmbuild profile to align the sequences to.
+      sequences_a3m: Sequences in A3M format to align to the profile.
+
+    Returns:
+      An A3M string with the aligned sequences.
+
+    Raises:
+      RuntimeError: If hmmalign fails.
+    """
+    deletion_table = str.maketrans('', '', '-')
+    sequences_no_gaps_a3m = []
+    for seq, desc in parsers.lazy_parse_fasta_string(sequences_a3m):
+      sequences_no_gaps_a3m.append(f'>{desc}')
+      sequences_no_gaps_a3m.append(seq.translate(deletion_table))
+    sequences_no_gaps_a3m = '\n'.join(sequences_no_gaps_a3m)
+
+    aligned_sequences = self.align(sequences_no_gaps_a3m, profile)
+
+    return aligned_sequences

flax_model/alphafold3/jax/attention/attention.py

@@ -0,0 +1,150 @@
+
+
+"""Scaled dot-product attention."""
+
+import typing
+from typing import Literal, TypeAlias
+
+from flax_model.alphafold3.jax.attention import attention_base as base
+from flax_model.alphafold3.jax.attention import flash_attention as attention_triton
+from flax_model.alphafold3.jax.attention import xla_attention
+from flax_model.alphafold3.jax.common import triton_utils
+import jax
+from jax.typing import DTypeLike  # pylint: disable=g-importing-member
+import jaxtyping
+from jaxtyping import Array  # pylint: disable=g-importing-member
+from jaxtyping import Bool  # pylint: disable=g-importing-member
+from jaxtyping import Float  # pylint: disable=g-importing-member
+import typeguard
+
+Implementation: TypeAlias = Literal["cudnn", "xla", "triton", "cutlass"]
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def dot_product_attention(
+    query: Float[Array, "*B T H D"],
+    key: Float[Array, "*B t #H D"],
+    value: Float[Array, "*B t #H D"],
+    *,
+    bias: Float[Array, "*#B #H #T #t"] | None = None,
+    mask: Bool[Array, "*#B #H #T #t"] | None = None,
+    implementation: Implementation | None = None,
+    logits_dtype: DTypeLike | None = None,
+    precision: (
+        jax.lax.Precision | tuple[jax.lax.Precision, jax.lax.Precision] | 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, head_dim]`.
+    key: Key array of shape `[batch, seq_len_kv, num_heads, head_dim]`.
+      `num_heads` can be 1 for multi-query attention.
+    value: Value array of shape `[batch, seq_len_kv, num_heads, head_dim]`.
+      `num_heads` can be 1 for multi-query attention.
+    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`.
+    implementation: if `None` (default), an implementation is automatically
+      chosen. 'xla' will use standard XLA and work on any platform, 'triton'
+      will use a fused Triton GPU kernel, and 'cudnn' a cuDNN FlashAttention
+      kernel. Only a subset of data types, shapes and GPUs are supported by
+      'triton' and 'cudnn', with an exception thrown in this case.
+    logits_dtype: Data type for attention logits (`query @ key^T`). If `None` is
+      passed (the default), the accumulator type from the `query @ key^T` dot
+      product will be used, which is FP32 for BF16/FP16/FP32 inputs. Note that
+      this default increases the memory usage for BF16/FP16 inputs when using
+      `implementation='xla'`, but does not increase memory usage when using
+      `implementation='triton'`.
+    precision: The precision for the dot products. Either `None` (default) which
+      uses the default JAX precision for a backend; a tuple `(
+      query_key_dot_precision, weights_value_dot_precision)` of
+      `jax.lax.Precision` objects; or a single `jax.lax.Precision` object
+      applied to both dot products.
+
+  Returns:
+    An array with the same shape as `query`.
+  """
+
+  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 == "cudnn":
+    if logits_dtype is not None:
+      raise ValueError(
+          "logits_dtype is not supported for cudnn implementation."
+      )
+    if precision is not None:
+      raise NotImplementedError(
+          "precision is not supported for cudnn implementation."
+      )
+
+    return jax.nn.dot_product_attention(
+        query=query,
+        key=key,
+        value=value,
+        bias=bias,
+        mask=mask,
+        implementation="cudnn",
+    )
+
+  if implementation == "cutlass":
+    if logits_dtype is not None:
+      raise ValueError(
+          "logits_dtype is not supported for cutlass implementation."
+      )
+    if precision is not None:
+      raise NotImplementedError(
+          "precision is not supported for cutlass implementation."
+      )
+
+    return jax.nn.dot_product_attention(
+        query=query,
+        key=key,
+        value=value,
+        bias=bias,
+        mask=mask,
+        implementation="cutlass",
+    )
+
+  logits_dtype = base.AUTO if logits_dtype is None else logits_dtype
+  precision = jax.lax.Precision.DEFAULT if precision is None else precision
+
+  args = (query, key, value)
+  kwargs = dict(
+      precision=precision,
+      logits_dtype=logits_dtype,
+      bias=bias,
+      mask=mask,
+  )
+
+  if implementation == "triton":
+    if not triton_utils.has_triton_support():
+      raise ValueError(
+          "implementation='triton' for FlashAttention is unsupported on this"
+          " GPU generation. Please use implementation='xla' instead."
+      )
+    return attention_triton.TritonFlashAttention()(*args, **kwargs)
+
+  if implementation is None and triton_utils.has_triton_support():
+    try:
+      return attention_triton.TritonFlashAttention()(*args, **kwargs)
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass  # Fallback to XLA.
+
+  return xla_attention.XlaDotProductAttention()(*args, **kwargs)

flax_model/alphafold3/jax/common/array_view.py

@@ -0,0 +1,397 @@
+
+
+"""Array view class and utilities."""
+
+from collections.abc import Sequence
+import dataclasses
+import math
+import operator
+from types import EllipsisType  # pylint: disable=g-importing-member
+from typing import Any, Self, TypeAlias, TypeVar
+
+import jax
+import jax.experimental
+from jax.experimental import pallas as pl
+import jax.numpy as jnp
+from jax.typing import ArrayLike  # pylint: disable=g-importing-member
+from jaxtyping import Int  # pylint: disable=g-importing-member
+import numpy as np
+
+ArrayT: TypeAlias = Any
+ScalarInt: TypeAlias = (
+    Int[ArrayT, ""] | Int[np.generic, ""] | Int[jnp.generic, ""]
+)
+
+Indexer: TypeAlias = int | ScalarInt | slice | pl.Slice | EllipsisType
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclasses.dataclass(frozen=True)
+class ArrayView:
+  """A strided view of a JAX array."""
+
+  base: jax.Array
+  _: dataclasses.KW_ONLY
+  # These are set by `__post_init__` so `None` value is never seen after init.
+  shape: tuple[int, ...] = None  # type: ignore
+  strides: tuple[int, ...] = None  # type: ignore
+  offset: int | ScalarInt = 0
+  flatten_base: bool = True
+
+  def __post_init__(self):
+    if self.shape is None:
+      object.__setattr__(self, "shape", self.base.shape)
+
+    if self.strides is None:
+      object.__setattr__(self, "strides", pl.strides_from_shape(self.shape))
+
+    if len(self.shape) != len(self.strides):
+      raise ValueError("`shape` and `strides` must have the same length.")
+
+    # Within `jax.vjp`, we can get non-`Array` values here (such as `object`).
+    if isinstance(self.base, jax.Array):
+      if isinstance(self.offset, int):
+        if not (0 <= self.offset < max(self.base.size, 1)):
+          raise ValueError("Invalid `offset`.")
+
+      if self.flatten_base:
+        if len(self.base.shape) != 1:
+          object.__setattr__(self, "base", self.base.reshape((-1,)))
+
+  def tree_flatten(self):
+    if isinstance(self.offset, int):
+      return (self.base,), (self.offset, self.shape, self.strides)
+    return (self.base, self.offset), (self.shape, self.strides)
+
+  @classmethod
+  def tree_unflatten(cls, aux, children) -> Self:
+    base, offset, shape, strides = (*children, *aux)
+    return cls(base, shape=shape, strides=strides, offset=offset)
+
+  @property
+  def dtype(self) -> jnp.dtype:
+    return self.base.dtype
+
+  @property
+  def size(self) -> int:
+    return math.prod(self.shape)
+
+  @property
+  def ndim(self) -> int:
+    return len(self.shape)
+
+  @property
+  def T(self) -> Self:  # pylint: disable=invalid-name
+    return self.transpose()
+
+  @property
+  def _index_dtype(self) -> jax.typing.DTypeLike:
+    i32_max = jnp.iinfo(jnp.int32).max
+    return jnp.int32 if (self.base.size <= i32_max) else jnp.int64
+
+  @property
+  def offsets(self) -> jax.Array:
+    """Returns array of offsets into `base` for each element."""
+    with jax.experimental.enable_x64():
+      idxs = jnp.indices(self.shape, sparse=True, dtype=self._index_dtype)
+      return self.offset + sum(s * idx for s, idx in zip(self.strides, idxs))
+
+  def astype(self, dtype: jax.typing.DTypeLike) -> Self:
+    return self._replace(base=self.base.astype(dtype))
+
+  def broadcast_to_rank(self, rank: int) -> Self:
+    """Returns a new view with the specified rank."""
+    if rank < self.ndim:
+      raise ValueError(f"Cannot broadcast to lower rank: {rank} < {self.ndim}.")
+
+    shape = (1,) * (rank - self.ndim) + self.shape
+    strides = (0,) * (rank - self.ndim) + self.strides
+    return self._replace(shape=shape, strides=strides)
+
+  def broadcast_to(self, shape: tuple[int, ...]) -> Self:
+    """Returns a new view with the specified shape."""
+    view = self.broadcast_to_rank(len(shape))
+    strides = []
+    for dim_size, stride, target_size in zip(
+        view.shape, view.strides, shape, strict=True
+    ):
+      if dim_size == target_size:
+        strides.append(stride)
+      elif dim_size == 1:
+        strides.append(0)
+      else:
+        raise ValueError(f"Cannot broadcast {self.shape} to {shape}.")
+    return self._replace(shape=shape, strides=strides)
+
+  def collapse(
+      self, start: int, stop: int | None = None, *, allow_copy: bool = False
+  ) -> Self:
+    """Returns a new view with the axis range collapsed into one axis."""
+    lo, hi, _ = slice(start, stop).indices(self.ndim)
+    if hi < lo:
+      raise ValueError(
+          "Invalid dimension range passed to collapse: "
+          f"{self.shape} [{start}:{stop}]"
+      )
+    shape = self.shape[:lo] + (-1,) + self.shape[hi:]
+    return self.reshape(shape, allow_copy=allow_copy)
+
+  def reshape(self, shape: Sequence[int], *, allow_copy: bool = False) -> Self:
+    """Returns a new view with the specified shape."""
+    try:
+      return self._reshape(tuple(shape))
+    except ValueError:
+      if not allow_copy:
+        raise
+    return type(self)(jnp.array(self)).reshape(shape)
+
+  def _reshape(self, shape: tuple[int, ...]) -> Self:
+    """Returns a new view with the specified shape."""
+
+    if (num_minus_one_dims := shape.count(-1)) > 0:
+      if num_minus_one_dims > 1:
+        raise ValueError("`shape` may only contain a single `-1` dimension.")
+      pos = shape.index(-1)
+      shape = list(shape)
+      shape[pos] = self.size // math.prod(d for d in shape if d != -1)
+
+    if math.prod(shape) != self.size:
+      raise ValueError("Mismatched number of elements.")
+
+    # Logic copied from `numpy` C++ code.
+    # Remove axes with length 1, to simplify logic below.
+    old_shape = [d for d in self.shape if d != 1]
+    old_strides = [s for i, s in enumerate(self.strides) if self.shape[i] != 1]
+    strides = [0] * len(shape)
+
+    # Axes currently being worked upon.
+    old_start, old_stop = 0, 1
+    new_start, new_stop = 0, 1
+
+    while (old_start < len(old_shape)) and (new_start < len(shape)):
+      old_axes_prod = old_shape[old_start]
+      new_axes_prod = shape[new_start]
+      while old_axes_prod != new_axes_prod:
+        if old_axes_prod < new_axes_prod:
+          old_axes_prod *= old_shape[old_stop]
+          old_stop += 1
+        else:
+          new_axes_prod *= shape[new_stop]
+          new_stop += 1
+
+      # Check if original axes can be combined.
+      for i in range(old_start, old_stop - 1):
+        if old_strides[i] != old_shape[i + 1] * old_strides[i + 1]:
+          raise ValueError("Cannot combine axes non-contiguous in memory.")
+
+      # Calculate new strides.
+      strides[new_stop - 1] = old_strides[old_stop - 1]
+      for i in range(new_stop - 1, new_start, -1):
+        strides[i - 1] = strides[i] * shape[i]
+
+      old_start, old_stop = old_stop, old_stop + 1
+      new_start, new_stop = new_stop, new_stop + 1
+
+    return self._replace(shape=shape, strides=strides)
+
+  def split(
+      self, indices_or_sections: int | Sequence[int], axis: int = 0
+  ) -> tuple[Self, ...]:
+    """Splits the view into multiple slice views."""
+    if isinstance(indices_or_sections, int):
+      if self.shape[axis] % indices_or_sections != 0:
+        raise ValueError("Axis size is not divisible by number of sections.")
+
+      chunk = self.shape[axis] // indices_or_sections
+      indices_or_sections = [i * chunk for i in range(1, indices_or_sections)]
+
+    los = (0, *indices_or_sections)
+    his = (*indices_or_sections, None)
+    slice_prefix = (slice(None),) * _canonicalize_axis(axis, self.ndim)
+    return tuple(self[*slice_prefix, slice(lo, hi)] for lo, hi in zip(los, his))
+
+  def swapaxes(self, axis1: int, axis2: int) -> Self:
+    """Returns a new view with the specified axis swapped."""
+    axes = list(range(self.ndim))
+    axes[axis1], axes[axis2] = axes[axis2], axes[axis1]
+    return self.transpose(axes)
+
+  def moveaxis(self, source: int, destination: int) -> Self:
+    """Returns a new view with the specified axis moved."""
+    source, destination = source % self.ndim, destination % self.ndim
+    axes = list(range(self.ndim))
+    del axes[source]
+    axes.insert(destination, source)
+    return self.transpose(axes)
+
+  def transpose(self, axes: Sequence[int] | None = None) -> Self:
+    """Returns a new view with the specified axes order."""
+    if axes is None:
+      axes = tuple(reversed(range(self.ndim)))
+    if len(axes) != self.ndim:
+      raise ValueError("`axes` must have the same dimensionality as the array.")
+    shape = tuple(self.shape[a] for a in axes)
+    strides = tuple(self.strides[a] for a in axes)
+    return self._replace(shape=shape, strides=strides)
+
+  def __getitem__(self, idxs: Indexer | tuple[Indexer, ...]) -> Self:
+    if not isinstance(idxs, tuple):
+      idxs = (idxs,)
+
+    if len(idxs) > self.ndim:
+      raise ValueError("Too many slice indices.")
+
+    num_ellipses = idxs.count(Ellipsis)
+    if num_ellipses > 1:
+      raise ValueError("Multiple `...` are not supported.")
+    elif num_ellipses == 0:
+      idxs += (Ellipsis,)  # `[a:b]` is equivalent to `[a:b, ...]`.
+
+    # Replace `...` with slices that take the entirety of the missing axes.
+    ellipsis_idx = idxs.index(Ellipsis)
+    ellipsis_slices = (slice(None),) * (self.ndim - len(idxs) + 1)
+    idxs = idxs[:ellipsis_idx] + ellipsis_slices + idxs[ellipsis_idx + 1 :]
+
+    shape = []
+    strides = []
+    with jax.experimental.enable_x64():
+
+      def as_index(x):
+        return x.astype(self._index_dtype) if isinstance(x, jax.Array) else x
+
+      offset = as_index(self.offset)
+
+      for idx, dim, stride in zip(idxs, self.shape, self.strides, strict=True):
+        if isinstance(idx, int):
+          if not (-dim <= idx < dim):
+            raise ValueError("Slice index out of range.")
+          offset += stride * (idx % dim)
+        elif isinstance(idx, ScalarInt):
+          offset += stride * as_index(idx)
+        elif isinstance(idx, slice):
+          start, stop, step = idx.indices(dim)
+          if step >= 0:
+            shape.append(pl.cdiv(stop - start, step))
+          else:
+            shape.append(pl.cdiv(start - stop, -step))
+          strides.append(stride * step)
+          offset += stride * start
+        elif isinstance(idx, pl.Slice):
+          shape.append(idx.size)
+          strides.append(stride * idx.stride)
+          offset += stride * as_index(idx.start)
+        else:
+          raise ValueError(f"Unexpected indexer: {idx}")
+
+    return self._replace(shape=shape, strides=strides, offset=offset)
+
+  def _replace(self, **kwargs) -> Self:
+    if "shape" in kwargs:
+      kwargs["shape"] = tuple(kwargs["shape"])
+    if "strides" in kwargs:
+      kwargs["strides"] = tuple(kwargs["strides"])
+    return dataclasses.replace(self, **kwargs)
+
+  def set(self, value: ArrayLike | "ArrayView") -> Self:
+    """Returns a new view with the views values set to `value`."""
+    if any(s == 0 for s in self.strides):
+      raise ValueError("Cannot set values on a broadcasted array.")
+
+    # Try to just transpose the value, if possible.
+    major_to_minor = np.argsort(-np.array(self.strides), kind="stable")
+    value = jnp.array(value)
+    value_transposed = value.transpose(major_to_minor)
+    if (
+        self.transpose(major_to_minor).strides
+        == ArrayView(value_transposed).strides
+    ):
+      base = jax.lax.dynamic_update_slice(
+          self.base, value_transposed.flatten(), (self.offset,)
+      )
+    else:
+      base = self.base.at[self.offsets].set(value)
+    return self._replace(base=base)
+
+  def __jax_array__(self) -> jax.Array:
+    """Returns values as a dense array."""
+    # Try to express using transpose, slice, and reshape, to encourage XLA to
+    # fuse into other ops, rather than materialising the values. Otherwise,
+    # fall back to using a gather.
+    if (self.ndim == 0) or any(s < 0 for s in self.strides):
+      return self.base[self.offsets]
+
+    major_to_minor = np.argsort(-np.array(self.strides), kind="stable")
+
+    # Construct a shape that gives us the correct strides.
+    bcast_axes = []
+    shape = []
+    for axis in major_to_minor[::-1]:  # minor to major
+      stride = self.strides[axis]
+      if stride == 0:
+        bcast_axes.append(axis)
+        shape.append(1)
+        continue
+
+      if stride % math.prod(shape) != 0:
+        raise ValueError("Cannot express as a reshape, then slice.")
+      shape.append(stride // math.prod(shape))
+
+    if self.base.size % math.prod(shape) != 0:
+      return self.base[self.offsets]
+
+    shape = [self.base.size // math.prod(shape), *reversed(shape)]
+    slice_sizes = [
+        *(1 if a in bcast_axes else self.shape[a] for a in major_to_minor),
+        1,
+    ]
+
+    if shape[0] == self.shape[major_to_minor[0]]:
+      needs_offset_slice = False
+    elif not isinstance(self.offset, int):
+      needs_offset_slice = True
+    else:
+      start_indices = np.unravel_index(self.offset, shape)
+      end_indices = [s + size for s, size in zip(start_indices, slice_sizes)]
+      needs_offset_slice = any(e > dim for e, dim in zip(end_indices, shape))
+
+    if needs_offset_slice:
+      shape[0] = self.shape[major_to_minor[0]]
+      size = math.prod(shape)
+      # The pad is necessary to ensure that the dynamic slice is in range.
+      vals = jnp.pad(self.base, (0, size))
+      vals = jax.lax.dynamic_slice(vals, (self.offset,), (size,))
+      start_indices = [0] * len(shape)
+    else:
+      vals = self.base
+      start_indices = jnp.unravel_index(self.offset, shape)
+
+    vals = vals.reshape(shape)
+    vals = jax.lax.dynamic_slice(vals, start_indices, slice_sizes)[..., 0]
+    # Move axes from their physical ordering to their logical ordering.
+    vals = vals.transpose(np.argsort(major_to_minor))
+    return jnp.broadcast_to(vals, self.shape)
+
+
+def as_array_view(x: jax.Array | ArrayView) -> ArrayView:
+  return x if isinstance(x, ArrayView) else ArrayView(x)
+
+
+T = TypeVar("T", jax.Array, ArrayView)
+
+
+def zeros_like(x: T) -> T:
+  if isinstance(x, ArrayView):
+    return x._replace(base=jnp.zeros_like(x.base))
+  return jnp.zeros_like(x)
+
+
+def _canonicalize_axis(axis, num_dims) -> int:
+  """Canonicalize an axis in [-num_dims, num_dims) to [0, num_dims)."""
+  axis = operator.index(axis)
+  if not -num_dims <= axis < num_dims:
+    raise ValueError(
+        f"axis {axis} is out of bounds for array of dimension {num_dims}"
+    )
+  if axis < 0:
+    axis = axis + num_dims
+  return axis

flax_model/alphafold3/jax/common/precision.py

@@ -0,0 +1,84 @@
+
+
+"""Precision classes and utilities."""
+
+import enum
+
+import jax
+import jax.numpy as jnp
+
+
+@enum.unique
+class DotPrecision(enum.Enum):
+  """Precision for `dot` operation.
+
+  Naming scheme: {OPERAND_DTYPE}_{ACCUMULATOR_DTYPE}[_{NUM_PASSES}x]
+  """
+
+  BF16_F32 = "bf16_f32"
+
+  # GPU only precisions.
+  F32_F32 = "f32_f32"  # Full f32 precision (doesn't use TensorCores).
+  TF32_F32 = "tf32_f32"  # Equivalent to `DEFAULT`/`HIGH` on GPU.
+  TF32_F32_3X = "tf32_f32_3x"
+  F16_F16 = "f16_f16"
+  F16_F32 = "f16_f32"
+
+  @property
+  def operand_dtype(self) -> jnp.dtype:
+    match self:
+      case DotPrecision.BF16_F32:
+        return jnp.bfloat16
+      case DotPrecision.F16_F16 | DotPrecision.F16_F32:
+        return jnp.float16
+      case _:
+        return jnp.float32
+
+  @property
+  def accumulator_dtype(self) -> jnp.dtype:
+    return jnp.float16 if (self == DotPrecision.F16_F16) else jnp.float32
+
+
+_JAX_GPU_PRECISION_MAP = {
+    (jnp.float16, jax.lax.Precision.DEFAULT): DotPrecision.F16_F32,
+    (jnp.bfloat16, jax.lax.Precision.DEFAULT): DotPrecision.BF16_F32,
+    (jnp.float32, jax.lax.Precision.DEFAULT): DotPrecision.TF32_F32,
+    (jnp.float32, jax.lax.Precision.HIGH): DotPrecision.TF32_F32,
+    (jnp.float32, jax.lax.Precision.HIGHEST): DotPrecision.F32_F32,
+}
+
+_JAX_CPU_PRECISION_MAP = {
+    (jnp.float16, jax.lax.Precision.DEFAULT): DotPrecision.F16_F32,
+    (jnp.bfloat16, jax.lax.Precision.DEFAULT): DotPrecision.F32_F32,
+    (jnp.float32, jax.lax.Precision.DEFAULT): DotPrecision.F32_F32,
+    (jnp.float32, jax.lax.Precision.HIGH): DotPrecision.F32_F32,
+    (jnp.float32, jax.lax.Precision.HIGHEST): DotPrecision.F32_F32,
+}
+
+
+def _create_jax_precision_map():
+  precision_map = {}
+  for (dtype, jax_precision), dot_precision in _JAX_GPU_PRECISION_MAP.items():
+    precision_map[("gpu", jnp.dtype(dtype), jax_precision)] = dot_precision
+  for (dtype, jax_precision), dot_precision in _JAX_CPU_PRECISION_MAP.items():
+    precision_map[("cpu", jnp.dtype(dtype), jax_precision)] = dot_precision
+  return precision_map
+
+
+_JAX_PRECISION_MAP = _create_jax_precision_map()
+
+
+def get_equivalent_dot_precision(
+    a_dtype: jnp.dtype, b_dtype: jnp.dtype, jax_precision: jax.lax.Precision
+) -> DotPrecision:
+  """Returns `DotPrecision` replicating default XLA behaviour."""
+  if a_dtype != b_dtype:
+    raise ValueError("Cannot infer precision if operand types differ.")
+
+  backend = jax.default_backend().lower()
+  if (jax_precision != jax.lax.Precision.DEFAULT) and (a_dtype != jnp.float32):
+    raise ValueError(
+        "`jax.lax.Precision` values other than `DEFAULT` only have an effect if"
+        " the operand type is `float32`."
+    )
+  return _JAX_PRECISION_MAP[(backend, a_dtype, jax_precision)]

flax_model/alphafold3/jax/gated_linear_unit/gated_linear_unit_base.py

@@ -0,0 +1,122 @@
+
+
+"""Common types for gated linear unit kernels."""
+
+import abc
+from collections.abc import Callable
+import functools
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import jaxtyping
+from jaxtyping import Array, Float  # pylint: disable=g-importing-member,g-multiple-import
+import typeguard
+
+
+class GatedLinearUnit(abc.ABC):
+  """Gated linear unit."""
+
+  def __call__(
+      self,
+      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,
+      **kwargs,
+  ) -> Float[Array, '*B M N']:
+    """Applies a gated linear unit (https://arxiv.org/abs/1612.08083).
+
+    Computes `activation(x @ weight[:, 0]) * x @ weight[:, 1]`.
+
+    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.
+      **kwargs: additional keyword arguments.
+
+    Returns:
+      The output array.
+    """
+    return self._fwd(
+        x, weight, activation=activation, precision=precision, **kwargs
+    )
+
+  # Default vmap rule.
+  @property
+  def vmap_rule_forward(self) -> Callable[..., Any]:
+    def _vmap_rule(
+        axis_size, in_batched, *args, fn: jax.custom_batching.custom_vmap
+    ):
+      sequential_vmap = jax.custom_batching.sequential_vmap(fn.fun)
+      return sequential_vmap.vmap_rule(axis_size, in_batched, *args)
+
+    return _vmap_rule
+
+  def apply_vmap_rule_forward(
+      self, fn: Callable[..., Any], **kwargs
+  ) -> jax.custom_batching.custom_vmap:
+    fn_closed = functools.partial(fn, **kwargs)
+    fn_closed = jax.custom_batching.custom_vmap(fn_closed)
+    vmap_rule = functools.partial(self.vmap_rule_forward, fn=fn_closed)
+    fn_closed.def_vmap(vmap_rule)
+    return fn_closed
+
+  @abc.abstractmethod
+  def _fwd(
+      self,
+      x: Float[Array, '*B M K'],
+      weight: Float[Array, 'K 2 N'],
+      *,
+      activation: Callable[[jax.Array], jax.Array] | None,
+      precision: jax.lax.Precision | None,
+  ) -> Float[Array, '*B M N']:
+    """Gated linear unit."""
+    ...
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def gated_linear_unit_xla(
+    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,
+) -> 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.
+
+  Returns:
+    The output array.
+  """
+
+  weight_reshaped = jax.lax.collapse(
+      weight, start_dimension=-2, stop_dimension=None
+  )
+  assert weight_reshaped.ndim == 2
+
+  y = jnp.dot(x, weight_reshaped, precision=precision)
+
+  # Apply activation and compute product of FP8/FP16/BF16 in FP32.
+  y = y.astype(jnp.promote_types(x.dtype, jnp.float32))
+  a, b = jnp.split(y, 2, axis=-1)
+  out = a * b if activation is None else activation(a) * b
+  out = out.astype(x.dtype)
+  return out

flax_model/alphafold3/jax/gated_linear_unit/matmul_ext.py

@@ -0,0 +1,265 @@
+
+
+"""Extended matmul ops."""
+
+from collections.abc import Callable
+import functools
+from typing import Any, 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 block
+from flax_model.alphafold3.jax.gated_linear_unit import matmul_config
+import jax
+from jax._src.state import discharge
+from jax.experimental import pallas as pl
+import jax.numpy as jnp
+import jaxtyping
+from jaxtyping import Array, Float, Int  # pylint: disable=g-importing-member,g-multiple-import
+import numpy as np
+import typeguard
+
+ArrayView = array_view.ArrayView
+PyTree: TypeAlias = Any
+ArrayT: TypeAlias = Any
+ScalarInt: TypeAlias = (
+    Int[ArrayT, ''] | Int[np.generic, ''] | Int[jnp.generic, '']
+)
+
+
+def _get_group_cache_usage(
+    group_size_m, num_blocks_m, num_blocks_n, block_m_bytes, block_n_bytes
+) -> int:
+  """Returns the cache usage in bytes for the given group size."""
+  num_live_progs = jax.devices()[0].core_count
+  num_live_blocks_n = min(pl.cdiv(num_live_progs, group_size_m), num_blocks_n)
+  num_live_groups = pl.cdiv(num_live_progs, group_size_m * num_live_blocks_n)
+  num_live_blocks_m = min(num_live_groups * group_size_m, num_blocks_m)
+  return num_live_blocks_m * block_m_bytes + num_live_blocks_n * block_n_bytes
+
+
+def _get_pids(
+    pid, num_blocks_m, num_blocks_n, group_size_m
+) -> tuple[ScalarInt, ScalarInt]:
+  """Returns the program IDs in each grid axis."""
+  # Use `floor_divide` and `remainder` (instead of lax.div and lax.rem)
+  # to handle dtypes: pid (int32) vs. num_blocks_n (int64) when `jax_enable_x64`
+  # is set.
+  if group_size_m == 1:
+    return jnp.floor_divide(pid, num_blocks_n), jnp.remainder(pid, num_blocks_n)
+
+  num_progs_in_group = group_size_m * num_blocks_n
+  group_start_m = jnp.floor_divide(pid, num_progs_in_group) * group_size_m
+  group_size_m = jnp.minimum(num_blocks_m - group_start_m, group_size_m)
+  pid_m = group_start_m + jnp.remainder(pid, group_size_m)
+  pid_n = jnp.floor_divide(jnp.remainder(pid, num_progs_in_group), group_size_m)
+  return pid_m, pid_n
+
+
+def _get_best_pids(
+    pid, *, m, n, block_m, block_n, a_dtype_bytes, b_dtype_bytes
+) -> tuple[ScalarInt, ScalarInt]:
+  """Returns the grouped program IDs that minimize cache usage."""
+  num_blocks_m = pl.cdiv(m, block_m)
+  num_blocks_n = pl.cdiv(n, block_n)
+  block_m_bytes = block_m * a_dtype_bytes
+  block_n_bytes = block_n * b_dtype_bytes
+
+  num_live_progs = jax.devices()[0].core_count
+
+  def group_size_m_usage(group_size_m):
+    return _get_group_cache_usage(
+        group_size_m, num_blocks_m, num_blocks_n, block_m_bytes, block_n_bytes
+    )
+
+  group_size_m = min(
+      range(1, min(num_live_progs, num_blocks_m) + 1), key=group_size_m_usage
+  )
+
+  def group_size_n_usage(group_size_n):
+    return _get_group_cache_usage(
+        group_size_n, num_blocks_n, num_blocks_m, block_n_bytes, block_m_bytes
+    )
+
+  group_size_n = min(
+      range(1, min(num_live_progs, num_blocks_n) + 1), key=group_size_n_usage
+  )
+
+  if group_size_m_usage(group_size_m) <= group_size_n_usage(group_size_n):
+    pid_m, pid_n = _get_pids(pid, num_blocks_m, num_blocks_n, group_size_m)
+  else:
+    pid_n, pid_m = _get_pids(pid, num_blocks_n, num_blocks_m, group_size_n)
+  return pid_m, pid_n
+
+
+def _apply_epilogue(
+    epilogue: Callable[..., jax.Array], x: jax.Array, args: PyTree
+) -> jax.Array:
+  """Applies the epilogue to the output."""
+  # Convert array view arguments to JAX arrays. This means that we can use the
+  # array view slices, rather than the gather that discharging state gives us.
+  is_leaf = lambda x: isinstance(x, ArrayView)
+  args_flat, args_tree = jax.tree.flatten((x, args), is_leaf=is_leaf)
+  args_flat = tuple(map(jnp.array, args_flat))
+
+  def epilogue_wrapper(refs):
+    x_ref, arg_refs = args_tree.unflatten(refs)
+    x_ref[:] = epilogue(x_ref[:], arg_refs, 0, 0)
+
+  return discharge.run_state_reference(epilogue_wrapper)(args_flat)[0]
+
+
+def _gated_linear_unit_kernel(
+    x_ref,
+    w_ref,
+    v_ref,
+    _,  # Destination, aliased with `out_ref`.
+    epilogue_in_refs,
+    out_ref,
+    *,
+    block_m,
+    block_n,
+    block_k,
+    activation,
+    precision,
+    epilogue,
+):
+  """Pallas GLU kernel."""
+  m = x_ref.shape[0]
+  n = w_ref.shape[1]
+  pid_m, pid_n = _get_best_pids(
+      pl.program_id(0),
+      m=m,
+      n=n,
+      block_m=block_m,
+      block_n=block_n,
+      a_dtype_bytes=jnp.dtype(x_ref.dtype).itemsize,
+      b_dtype_bytes=jnp.dtype(w_ref.dtype).itemsize * 2,  # Two blocks.
+  )
+
+  def body(i, acc):
+    x = block.load_block(x_ref, (pid_m, i), block_shape=(block_m, block_k))
+    w = block.load_block(w_ref, (i, pid_n), block_shape=(block_k, block_n))
+    v = block.load_block(v_ref, (i, pid_n), block_shape=(block_k, block_n))
+    acc[0] += pl.dot(x, w.astype(x.dtype), precision=precision)
+    acc[1] += pl.dot(x, v.astype(x.dtype), precision=precision)
+    return acc
+
+  num_iters = pl.cdiv(x_ref.shape[-1], block_k)
+  acc0 = jnp.zeros((block_m, block_n), dtype=jnp.float32)
+  acc1 = jnp.zeros((block_m, block_n), dtype=jnp.float32)
+  proj, gates = jax.lax.fori_loop(0, num_iters, body, init_val=[acc0, acc1])
+
+  proj = proj.astype(x_ref.dtype).astype(jnp.float32)
+  gates = gates.astype(x_ref.dtype).astype(jnp.float32)
+
+  out = proj * (gates if activation is None else activation(gates))
+
+  if epilogue is not None:
+    out = epilogue(out, epilogue_in_refs, pid_m, pid_n)
+
+  block.store_block(out_ref, out, (pid_m, pid_n))
+
+
+def _gated_linear_unit(
+    x: Float[Array | ArrayView, 'M K'],
+    weights_projection: Float[Array | ArrayView, 'K N'],
+    weights_gate: Float[Array | ArrayView, 'K N'],
+    *,
+    dst: Float[ArrayView, 'M N'] | None = None,
+    activation: Callable[[jax.Array], jax.Array] | None,
+    epilogue: Any,  # Callable[..., Any] | None - breaks `typed`.
+    epilogue_args: PyTree,
+    precision: jax.lax.Precision | None,
+) -> jax.Array:  # Float[Array, 'M N'] | Float[Array, 'N M']
+  """Applies a gated linear unit (arxiv.org/abs/1612.08083)."""
+  if epilogue is None and epilogue_args is not None:
+    raise ValueError('`epilogue_args` is specified but `epilogue` is None.')
+
+  name = 'pallas_glu'
+  if activation is not None:
+    name += f'_{getattr(activation, "__name__", repr(activation))}'
+  if epilogue is not None:
+    name += f'_{getattr(epilogue, "__name__", repr(epilogue))}'
+
+  w = weights_projection
+  config = matmul_config.get_config(x, w)
+
+  m = x.shape[0]
+  n = w.shape[1]
+  kernel = functools.partial(
+      _gated_linear_unit_kernel,
+      block_m=config.block_m,
+      block_n=config.block_n,
+      block_k=config.block_k,
+      activation=activation,
+      precision=precision,
+      epilogue=epilogue,
+  )
+
+  if dst is None:
+    input_output_aliases = {}
+  else:
+    input_output_aliases = {3: 0}
+
+  compiler_params = dict(
+      triton=dict(num_warps=config.num_warps, num_stages=config.num_stages)
+  )
+
+  return pl.pallas_call(
+      kernel,
+      name=name,
+      grid=(pl.cdiv(m, config.block_m) * pl.cdiv(n, config.block_n),),
+      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype) if dst is None else dst,
+      input_output_aliases=input_output_aliases,
+      compiler_params=compiler_params,
+      backend='triton',
+  )(x, weights_projection, weights_gate, dst, epilogue_args)
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+def gated_linear_unit(
+    x: Float[Array | ArrayView, '*B M K'],
+    weights_projection: Float[Array | ArrayView, 'K N'],
+    weights_gate: Float[Array | ArrayView, 'K N'],
+    *,
+    activation: Callable[[jax.Array], jax.Array] | None = None,
+    precision: jax.lax.Precision | None = None,
+) -> Float[Array | ArrayView, '*B M N']:
+  """Applies a gated linear unit (arxiv.org/abs/1612.08083).
+
+  Args:
+    x: Input activations.
+    weights_projection: Weights for linear projection.
+    weights_gate: Weights for gates.
+    activation: Optional activation function.
+    precision: Specifies the precision of the matmuls.
+
+  Returns:
+    `(x @ weights_projection) * activation(x @ weights_gate)`
+  """
+
+  supported_dtypes = {'float16', 'bfloat16', 'float32'}
+  if x.dtype.name not in supported_dtypes:
+    raise NotImplementedError(
+        f'Triton kernel does not support input datatype {x.dtype.name}. Must be'
+        f' one of {supported_dtypes}.'
+    )
+
+  if not triton_utils.has_triton_support():
+    raise NotImplementedError('Triton kernel not supported on current device.')
+
+  *batch, m, _ = x.shape
+  n = weights_projection.shape[1]
+  x = array_view.as_array_view(x).collapse(start=0, stop=-1)
+
+  return _gated_linear_unit(
+      x,
+      weights_projection,
+      weights_gate,
+      dst=None,
+      activation=activation,
+      precision=precision,
+      epilogue=None,
+      epilogue_args=None,
+  ).reshape(batch + [m, n])

flax_model/alphafold3/jax/geometry/rigid_matrix_vector.py

@@ -0,0 +1,216 @@
+
+
+"""Rigid3Array Transformations represented by a Matrix and a Vector."""
+
+from typing import Any, Final, Self, TypeAlias
+
+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 utils
+from flax_model.alphafold3.jax.geometry import vector
+import jax
+import jax.numpy as jnp
+
+
+Float: TypeAlias = float | jnp.ndarray
+
+VERSION: Final[str] = '0.1'
+
+
+# Disabling name in pylint, since the relevant variable in math are typically
+# referred to as X, Y in mathematical literature.
+def _compute_covariance_matrix(
+    row_values: vector.Vec3Array,
+    col_values: vector.Vec3Array,
+    weights: jnp.ndarray,
+    epsilon=1e-6,
+) -> jnp.ndarray:
+  """Compute covariance matrix.
+
+  The quantity computes is
+  cov_xy = weighted_avg_i(row_values[i, x] col_values[j, y]).
+  Here x and y run over the xyz coordinates.
+  This is used to construct frames when aligning points.
+
+  Args:
+    row_values: Values used for rows of covariance matrix, shape [..., n_point]
+    col_values: Values used for columns of covariance matrix, shape [...,
+      n_point]
+    weights: weights to weight points by, shape broacastable to [...]
+    epsilon: small value to add to denominator to avoid Nan's when all weights
+      are 0.
+
+  Returns:
+    Covariance Matrix as [..., 3, 3] array.
+  """
+  weights = jnp.asarray(weights)
+  weights = jnp.broadcast_to(weights, row_values.shape)
+
+  out = []
+
+  normalized_weights = weights / (weights.sum(axis=-1, keepdims=True) + epsilon)
+
+  weighted_average = lambda x: jnp.sum(normalized_weights * x, axis=-1)
+
+  out.append(
+      jnp.stack(
+          (
+              weighted_average(row_values.x * col_values.x),
+              weighted_average(row_values.x * col_values.y),
+              weighted_average(row_values.x * col_values.z),
+          ),
+          axis=-1,
+      )
+  )
+
+  out.append(
+      jnp.stack(
+          (
+              weighted_average(row_values.y * col_values.x),
+              weighted_average(row_values.y * col_values.y),
+              weighted_average(row_values.y * col_values.z),
+          ),
+          axis=-1,
+      )
+  )
+
+  out.append(
+      jnp.stack(
+          (
+              weighted_average(row_values.z * col_values.x),
+              weighted_average(row_values.z * col_values.y),
+              weighted_average(row_values.z * col_values.z),
+          ),
+          axis=-1,
+      )
+  )
+
+  return jnp.stack(out, axis=-2)
+
+
+@struct_of_array.StructOfArray(same_dtype=True)
+class Rigid3Array:
+  """Rigid Transformation, i.e. element of special euclidean group."""
+
+  rotation: rotation_matrix.Rot3Array
+  translation: vector.Vec3Array
+
+  def __matmul__(self, other: Self) -> Self:
+    new_rotation = self.rotation @ other.rotation
+    new_translation = self.apply_to_point(other.translation)
+    return Rigid3Array(new_rotation, new_translation)
+
+  def inverse(self) -> Self:
+    """Return Rigid3Array corresponding to inverse transform."""
+    inv_rotation = self.rotation.inverse()
+    inv_translation = inv_rotation.apply_to_point(-self.translation)
+    return Rigid3Array(inv_rotation, inv_translation)
+
+  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+    """Apply Rigid3Array transform to point."""
+    return self.rotation.apply_to_point(point) + self.translation
+
+  def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+    """Apply inverse Rigid3Array transform to point."""
+    new_point = point - self.translation
+    return self.rotation.apply_inverse_to_point(new_point)
+
+  def compose_rotation(self, other_rotation: rotation_matrix.Rot3Array) -> Self:
+    rot = self.rotation @ other_rotation
+    trans = jax.tree.map(
+        lambda x: jnp.broadcast_to(x, rot.shape), self.translation
+    )
+    return Rigid3Array(rot, trans)
+
+  @classmethod
+  def identity(cls, shape: Any, dtype: jnp.dtype = jnp.float32) -> Self:
+    """Return identity Rigid3Array of given shape."""
+    return cls(
+        rotation_matrix.Rot3Array.identity(shape, dtype=dtype),
+        vector.Vec3Array.zeros(shape, dtype=dtype),
+    )  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  def scale_translation(self, factor: Float) -> Self:
+    """Scale translation in Rigid3Array by 'factor'."""
+    return Rigid3Array(self.rotation, self.translation * factor)
+
+  def to_array(self):
+    rot_array = self.rotation.to_array()
+    vec_array = self.translation.to_array()
+    return jnp.concatenate([rot_array, vec_array[..., None]], axis=-1)
+
+  @classmethod
+  def from_array(cls, array):
+    rot = rotation_matrix.Rot3Array.from_array(array[..., :3])
+    vec = vector.Vec3Array.from_array(array[..., -1])
+    return cls(rot, vec)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  @classmethod
+  def from_array4x4(cls, array: jnp.ndarray) -> Self:
+    """Construct Rigid3Array from homogeneous 4x4 array."""
+    if array.shape[-2:] != (4, 4):
+      raise ValueError(f'array.shape({array.shape}) must be [..., 4, 4]')
+    rotation = rotation_matrix.Rot3Array(
+        *(array[..., 0, 0], array[..., 0, 1], array[..., 0, 2]),
+        *(array[..., 1, 0], array[..., 1, 1], array[..., 1, 2]),
+        *(array[..., 2, 0], array[..., 2, 1], array[..., 2, 2]),
+    )
+    translation = vector.Vec3Array(
+        array[..., 0, 3], array[..., 1, 3], array[..., 2, 3]
+    )
+    return cls(rotation, translation)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  @classmethod
+  def from_point_alignment(
+      cls,
+      points_to: vector.Vec3Array,
+      points_from: vector.Vec3Array,
+      weights: Float | None = None,
+      epsilon: float = 1e-6,
+  ) -> Self:
+    """Constructs Rigid3Array by finding transform aligning points.
+
+    This constructs the optimal Rigid Transform taking points_from to the
+    arrangement closest to points_to.
+
+    Args:
+      points_to: Points to align to.
+      points_from: Points to align from.
+      weights: weights for points.
+      epsilon: epsilon used to regularize covariance matrix.
+
+    Returns:
+      Rigid Transform.
+    """
+    if weights is None:
+      weights = 1.0
+
+    def compute_center(value):
+      return utils.weighted_mean(value=value, weights=weights, axis=-1)
+
+    points_to_center = jax.tree.map(compute_center, points_to)
+    points_from_center = jax.tree.map(compute_center, points_from)
+    centered_points_to = points_to - points_to_center[..., None]
+    centered_points_from = points_from - points_from_center[..., None]
+    cov_mat = _compute_covariance_matrix(
+        centered_points_to,
+        centered_points_from,
+        weights=weights,
+        epsilon=epsilon,
+    )
+    rots = rotation_matrix.Rot3Array.from_svd(
+        jnp.reshape(cov_mat, cov_mat.shape[:-2] + (9,))
+    )
+
+    translations = points_to_center - rots.apply_to_point(points_from_center)
+
+    return cls(rots, translations)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  def __getstate__(self):
+    return (VERSION, (self.rotation, self.translation))
+
+  def __setstate__(self, state):
+    version, (rot, trans) = state
+    del version
+    object.__setattr__(self, 'rotation', rot)
+    object.__setattr__(self, 'translation', trans)

flax_model/alphafold3/jax/geometry/rotation_matrix.py

@@ -0,0 +1,283 @@
+
+
+"""Rot3Array Matrix Class."""
+
+import dataclasses
+from typing import Any, Final, Self
+
+from flax_model.alphafold3.jax.geometry import struct_of_array
+from flax_model.alphafold3.jax.geometry import utils
+from flax_model.alphafold3.jax.geometry import vector
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+COMPONENTS: Final[tuple[str, ...]] = (
+    *('xx', 'xy', 'xz'),
+    *('yx', 'yy', 'yz'),
+    *('zx', 'zy', 'zz'),
+)
+VERSION: Final[str] = '0.1'
+
+
+def make_matrix_svd_factors() -> np.ndarray:
+  """Generates factors for converting 3x3 matrix to symmetric 4x4 matrix."""
+  factors = np.zeros((16, 9), dtype=np.float32)
+
+  factors[0, [0, 4, 8]] = 1.0
+
+  factors[[1, 4], 5] = 1.0
+  factors[[1, 4], 7] = -1.0
+
+  factors[[2, 8], 6] = 1.0
+  factors[[2, 8], 2] = -1.0
+
+  factors[[3, 12], 1] = 1.0
+  factors[[3, 12], 3] = -1.0
+
+  factors[5, 0] = 1.0
+  factors[5, [4, 8]] = -1.0
+
+  factors[[6, 9], 1] = 1.0
+  factors[[6, 9], 3] = 1.0
+
+  factors[[7, 13], 2] = 1.0
+  factors[[7, 13], 6] = 1.0
+
+  factors[10, 4] = 1.0
+  factors[10, [0, 8]] = -1.0
+
+  factors[[11, 14], 5] = 1.0
+  factors[[11, 14], 7] = 1.0
+
+  factors[15, 8] = 1.0
+  factors[15, [0, 4]] = -1.0
+
+  return factors
+
+
+@jax.custom_jvp
+def largest_evec(m):
+  _, eigvecs = jnp.linalg.eigh(m)
+  return eigvecs[..., -1]
+
+
+def largest_evec_jvp(primals, tangents):
+  """jvp for largest eigenvector."""
+  (m,) = primals
+  (t,) = tangents
+  eigvals, eigvecs = jnp.linalg.eigh(m)
+  large_eigvec = eigvecs[..., -1]
+  large_eigval = eigvals[..., -1]
+  other_eigvals = eigvals[..., :-1]
+  other_eigvecs = eigvecs[..., :-1]
+  other_ev_times_tangent = jnp.einsum(
+      '...aj,...ab -> ...bj',
+      other_eigvecs,
+      t,
+      precision=jax.lax.Precision.HIGHEST,
+  )
+  nominator = jnp.einsum(
+      '...bj,...b -> ...j',
+      other_ev_times_tangent,
+      large_eigvec,
+      precision=jax.lax.Precision.HIGHEST,
+  )
+  prefactor = nominator / jnp.maximum(
+      large_eigval[..., None] - other_eigvals, 1e-6
+  )
+  grad = jnp.sum(prefactor[..., None, :] * other_eigvecs, axis=-1)
+  return large_eigvec, grad
+
+
+largest_evec.defjvp(largest_evec_jvp)
+
+
+MATRIX_SVD_QUAT_FACTORS = make_matrix_svd_factors()
+
+
+@struct_of_array.StructOfArray(same_dtype=True)
+class Rot3Array:
+  """Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays."""
+
+  xx: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
+  xy: jnp.ndarray
+  xz: jnp.ndarray
+  yx: jnp.ndarray
+  yy: jnp.ndarray
+  yz: jnp.ndarray
+  zx: jnp.ndarray
+  zy: jnp.ndarray
+  zz: jnp.ndarray
+
+  __array_ufunc__ = None
+
+  def inverse(self) -> Self:
+    """Returns inverse of Rot3Array."""
+    return Rot3Array(
+        *(self.xx, self.yx, self.zx),
+        *(self.xy, self.yy, self.zy),
+        *(self.xz, self.yz, self.zz),
+    )
+
+  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+    """Applies Rot3Array to point."""
+    return vector.Vec3Array(
+        self.xx * point.x + self.xy * point.y + self.xz * point.z,
+        self.yx * point.x + self.yy * point.y + self.yz * point.z,
+        self.zx * point.x + self.zy * point.y + self.zz * point.z,
+    )
+
+  def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+    """Applies inverse Rot3Array to point."""
+    return self.inverse().apply_to_point(point)
+
+  def __matmul__(self, other: Self) -> Self:
+    """Composes two Rot3Arrays."""
+    c0 = self.apply_to_point(vector.Vec3Array(other.xx, other.yx, other.zx))
+    c1 = self.apply_to_point(vector.Vec3Array(other.xy, other.yy, other.zy))
+    c2 = self.apply_to_point(vector.Vec3Array(other.xz, other.yz, other.zz))
+    return Rot3Array(c0.x, c1.x, c2.x, c0.y, c1.y, c2.y, c0.z, c1.z, c2.z)
+
+  @classmethod
+  def identity(cls, shape: Any, dtype: jnp.dtype = jnp.float32) -> Self:
+    """Returns identity of given shape."""
+    ones = jnp.ones(shape, dtype=dtype)
+    zeros = jnp.zeros(shape, dtype=dtype)
+    return cls(ones, zeros, zeros, zeros, ones, zeros, zeros, zeros, ones)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  @classmethod
+  def from_two_vectors(cls, e0: vector.Vec3Array, e1: vector.Vec3Array) -> Self:
+    """Construct Rot3Array from two Vectors.
+
+    Rot3Array is constructed such that in the corresponding frame 'e0' lies on
+    the positive x-Axis and 'e1' lies in the xy plane with positive sign of y.
+
+    Args:
+      e0: Vector
+      e1: Vector
+
+    Returns:
+      Rot3Array
+    """
+    # Normalize the unit vector for the x-axis, e0.
+    e0 = e0.normalized()
+    # make e1 perpendicular to e0.
+    c = e1.dot(e0)
+    e1 = (e1 - c * e0).normalized()
+    # Compute e2 as cross product of e0 and e1.
+    e2 = e0.cross(e1)
+    return cls(e0.x, e1.x, e2.x, e0.y, e1.y, e2.y, e0.z, e1.z, e2.z)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  @classmethod
+  def from_array(cls, array: jnp.ndarray) -> Self:
+    """Construct Rot3Array Matrix from array of shape. [..., 3, 3]."""
+    unstacked = utils.unstack(array, axis=-2)
+    unstacked = sum([utils.unstack(x, axis=-1) for x in unstacked], [])
+    return cls(*unstacked)
+
+  def to_array(self) -> jnp.ndarray:
+    """Convert Rot3Array to array of shape [..., 3, 3]."""
+    return jnp.stack(
+        [
+            jnp.stack([self.xx, self.xy, self.xz], axis=-1),
+            jnp.stack([self.yx, self.yy, self.yz], axis=-1),
+            jnp.stack([self.zx, self.zy, self.zz], axis=-1),
+        ],
+        axis=-2,
+    )
+
+  @classmethod
+  def from_quaternion(
+      cls,
+      w: jnp.ndarray,
+      x: jnp.ndarray,
+      y: jnp.ndarray,
+      z: jnp.ndarray,
+      normalize: bool = True,
+      epsilon: float = 1e-6,
+  ) -> Self:
+    """Construct Rot3Array from components of quaternion."""
+    if normalize:
+      inv_norm = jax.lax.rsqrt(jnp.maximum(epsilon, w**2 + x**2 + y**2 + z**2))
+      w *= inv_norm
+      x *= inv_norm
+      y *= inv_norm
+      z *= inv_norm
+    xx = 1 - 2 * (jnp.square(y) + jnp.square(z))
+    xy = 2 * (x * y - w * z)
+    xz = 2 * (x * z + w * y)
+    yx = 2 * (x * y + w * z)
+    yy = 1 - 2 * (jnp.square(x) + jnp.square(z))
+    yz = 2 * (y * z - w * x)
+    zx = 2 * (x * z - w * y)
+    zy = 2 * (y * z + w * x)
+    zz = 1 - 2 * (jnp.square(x) + jnp.square(y))
+    return cls(xx, xy, xz, yx, yy, yz, zx, zy, zz)  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  @classmethod
+  def from_svd(cls, mat: jnp.ndarray, use_quat_formula: bool = True) -> Self:
+    """Constructs Rot3Array from arbitrary array of shape [3 * 3] using SVD.
+
+    The case when 'use_quat_formula' is False rephrases the problem of
+    projecting the matrix to a rotation matrix as a problem of finding the
+    largest eigenvector of a certain 4x4 matrix. This has the advantage of
+    having fewer numerical issues.
+    This approach follows:
+    https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.65.971&rep=rep1&type=pdf
+    In the other case we construct it via svd following
+    https://arxiv.org/pdf/2006.14616.pdf
+    In that case [∂L/∂M] is large if the two smallest singular values are close
+    to each other, or if they are close to 0.
+
+    Args:
+      mat: Array of shape [..., 3 * 3]
+      use_quat_formula: Whether to construct matrix via 4x4 eigenvalue problem.
+
+    Returns:
+      Rot3Array of shape [...]
+    """
+    assert mat.shape[-1] == 9
+    if use_quat_formula:
+      symmetric_4by4 = jnp.einsum(
+          'ji, ...i -> ...j',
+          MATRIX_SVD_QUAT_FACTORS,
+          mat,
+          precision=jax.lax.Precision.HIGHEST,
+      )
+      symmetric_4by4 = jnp.reshape(symmetric_4by4, mat.shape[:-1] + (4, 4))
+      largest_eigvec = largest_evec(symmetric_4by4)
+      return cls.from_quaternion(
+          *utils.unstack(largest_eigvec, axis=-1)
+      ).inverse()
+
+    else:
+      mat = jnp.reshape(mat, mat.shape[:-1] + (3, 3))
+      u, _, v_t = jnp.linalg.svd(mat, full_matrices=False)
+      det_uv_t = jnp.linalg.det(
+          jnp.matmul(u, v_t, precision=jax.lax.Precision.HIGHEST)
+      )
+      ones = jnp.ones_like(det_uv_t)
+      diag_array = jnp.stack([ones, ones, det_uv_t], axis=-1)
+      # This is equivalent to making diag_array into a diagonal array and matrix
+      # multiplying
+      diag_times_v_t = diag_array[..., None] * v_t
+      out = jnp.matmul(u, diag_times_v_t, precision=jax.lax.Precision.HIGHEST)
+      return cls.from_array(out)
+
+  @classmethod
+  def random_uniform(cls, key, shape, dtype=jnp.float32) -> Self:
+    """Samples uniform random Rot3Array according to Haar Measure."""
+    quat_array = jax.random.normal(key, tuple(shape) + (4,), dtype=dtype)
+    quats = utils.unstack(quat_array)
+    return cls.from_quaternion(*quats)
+
+  def __getstate__(self):
+    return (VERSION, [np.asarray(getattr(self, field)) for field in COMPONENTS])
+
+  def __setstate__(self, state):
+    version, state = state
+    del version
+    for i, field in enumerate(COMPONENTS):
+      object.__setattr__(self, field, state[i])

flax_model/alphafold3/jax/geometry/struct_of_array.py

@@ -0,0 +1,223 @@
+
+
+"""Class decorator to represent (nested) struct of arrays."""
+
+import dataclasses
+
+import jax
+
+
+def get_item(instance, key):
+  sliced = {}
+  for field in get_array_fields(instance):
+    num_trailing_dims = field.metadata.get('num_trailing_dims', 0)
+    this_key = key
+    if isinstance(key, tuple) and Ellipsis in this_key:
+      this_key += (slice(None),) * num_trailing_dims
+    sliced[field.name] = jax.tree.map(
+        lambda x: x[this_key],  # pylint: disable=cell-var-from-loop
+        getattr(instance, field.name),
+    )
+  return dataclasses.replace(instance, **sliced)
+
+
+@property
+def get_shape(instance):
+  """Returns Shape for given instance of dataclass."""
+  first_field = dataclasses.fields(instance)[0]
+  num_trailing_dims = first_field.metadata.get('num_trailing_dims', None)
+  value = getattr(instance, first_field.name)
+  if num_trailing_dims:
+    return value.shape[:-num_trailing_dims]
+  else:
+    return value.shape
+
+
+def get_len(instance):
+  """Returns length for given instance of dataclass."""
+  shape = instance.shape
+  if shape:
+    return shape[0]
+  else:
+    raise TypeError('len() of unsized object')  # Match jax.numpy behavior.
+
+
+@property
+def get_dtype(instance):
+  """Returns Dtype for given instance of dataclass."""
+  fields = dataclasses.fields(instance)
+  sets_dtype = [
+      field.name for field in fields if field.metadata.get('sets_dtype', False)
+  ]
+  if sets_dtype:
+    assert len(sets_dtype) == 1, 'at most field can set dtype'
+    field_value = getattr(instance, sets_dtype[0])
+  elif instance.same_dtype:
+    field_value = getattr(instance, fields[0].name)
+  else:
+    # Should this be Value Error?
+    raise AttributeError(
+        'Trying to access Dtype on Struct of Array without'
+        'either "same_dtype" or field setting dtype'
+    )
+
+  if hasattr(field_value, 'dtype'):
+    return field_value.dtype
+  else:
+    # Should this be Value Error?
+    raise AttributeError(f'field_value {field_value} does not have dtype')
+
+
+def replace(instance, **kwargs):
+  return dataclasses.replace(instance, **kwargs)
+
+
+def post_init(instance):
+  """Validate instance has same shapes & dtypes."""
+  array_fields = get_array_fields(instance)
+  arrays = list(get_array_fields(instance, return_values=True).values())
+  first_field = array_fields[0]
+  # These slightly weird constructions about checking whether the leaves are
+  # actual arrays is since e.g. vmap internally relies on being able to
+  # construct pytree's with object() as leaves, this would break the checking
+  # as such we are only validating the object when the entries in the dataclass
+  # Are arrays or other dataclasses of arrays.
+  try:
+    dtype = instance.dtype
+  except AttributeError:
+    dtype = None
+  if dtype is not None:
+    first_shape = instance.shape
+    for array, field in zip(arrays, array_fields, strict=True):
+      num_trailing_dims = field.metadata.get('num_trailing_dims', None)
+      if num_trailing_dims:
+        array_shape = array.shape
+        field_shape = array_shape[:-num_trailing_dims]
+        msg = (
+            f'field {field} should have number of trailing dims'
+            ' {num_trailing_dims}'
+        )
+        assert len(array_shape) == len(first_shape) + num_trailing_dims, msg
+      else:
+        field_shape = array.shape
+
+      shape_msg = (
+          f"Stripped Shape {field_shape} of field {field} doesn't "
+          f'match shape {first_shape} of field {first_field}'
+      )
+      assert field_shape == first_shape, shape_msg
+
+      field_dtype = array.dtype
+
+      allowed_metadata_dtypes = field.metadata.get('allowed_dtypes', [])
+      if allowed_metadata_dtypes:
+        msg = f'Dtype is {field_dtype} but must be in {allowed_metadata_dtypes}'
+        assert field_dtype in allowed_metadata_dtypes, msg
+
+      if 'dtype' in field.metadata:
+        target_dtype = field.metadata['dtype']
+      else:
+        target_dtype = dtype
+
+      msg = f'Dtype is {field_dtype} but must be {target_dtype}'
+      assert field_dtype == target_dtype, msg
+
+
+def flatten(instance):
+  """Flatten Struct of Array instance."""
+  array_likes = get_array_fields(instance, return_values=True).values()
+  flat_array_likes = []
+  inner_treedefs = []
+  num_arrays = []
+  for array_like in array_likes:
+    flat_array_like, inner_treedef = jax.tree_util.tree_flatten(array_like)
+    inner_treedefs.append(inner_treedef)
+    flat_array_likes += flat_array_like
+    num_arrays.append(len(flat_array_like))
+  metadata = get_metadata_fields(instance, return_values=True)
+  metadata = type(instance).metadata_cls(**metadata)
+  return flat_array_likes, (inner_treedefs, metadata, num_arrays)
+
+
+def make_metadata_class(cls):
+  metadata_fields = get_fields(
+      cls, lambda x: x.metadata.get('is_metadata', False)
+  )
+  metadata_cls = dataclasses.make_dataclass(
+      cls_name='Meta' + cls.__name__,
+      fields=[(field.name, field.type, field) for field in metadata_fields],
+      frozen=True,
+      eq=True,
+  )
+  return metadata_cls
+
+
+def get_fields(cls_or_instance, filterfn, return_values=False):
+  fields = dataclasses.fields(cls_or_instance)
+  fields = [field for field in fields if filterfn(field)]
+  if return_values:
+    return {
+        field.name: getattr(cls_or_instance, field.name) for field in fields
+    }
+  else:
+    return fields
+
+
+def get_array_fields(cls, return_values=False):
+  return get_fields(
+      cls,
+      lambda x: not x.metadata.get('is_metadata', False),
+      return_values=return_values,
+  )
+
+
+def get_metadata_fields(cls, return_values=False):
+  return get_fields(
+      cls,
+      lambda x: x.metadata.get('is_metadata', False),
+      return_values=return_values,
+  )
+
+
+class StructOfArray:
+  """Class Decorator for Struct Of Arrays."""
+
+  def __init__(self, same_dtype=True):
+    self.same_dtype = same_dtype
+
+  def __call__(self, cls):
+    cls.__array_ufunc__ = None
+    cls.replace = replace
+    cls.same_dtype = self.same_dtype
+    cls.dtype = get_dtype
+    cls.shape = get_shape
+    cls.__len__ = get_len
+    cls.__getitem__ = get_item
+    cls.__post_init__ = post_init
+    new_cls = dataclasses.dataclass(cls, frozen=True, eq=False)  # pytype: disable=wrong-keyword-args
+    # pytree claims to require metadata to be hashable, not sure why,
+    # But making derived dataclass that can just hold metadata
+    new_cls.metadata_cls = make_metadata_class(new_cls)
+
+    def unflatten(aux, data):
+      inner_treedefs, metadata, num_arrays = aux
+      array_fields = [field.name for field in get_array_fields(new_cls)]
+      value_dict = {}
+      array_start = 0
+      for num_array, inner_treedef, array_field in zip(
+          num_arrays, inner_treedefs, array_fields, strict=True
+      ):
+        value_dict[array_field] = jax.tree_util.tree_unflatten(
+            inner_treedef, data[array_start : array_start + num_array]
+        )
+        array_start += num_array
+      metadata_fields = get_metadata_fields(new_cls)
+      for field in metadata_fields:
+        value_dict[field.name] = getattr(metadata, field.name)
+
+      return new_cls(**value_dict)
+
+    jax.tree_util.register_pytree_node(
+        nodetype=new_cls, flatten_func=flatten, unflatten_func=unflatten
+    )
+    return new_cls

flax_model/alphafold3/jax/geometry/utils.py

@@ -0,0 +1,135 @@
+
+
+"""Utils for geometry library."""
+
+from collections.abc import Iterable
+import numbers
+
+import jax
+from jax import lax
+import jax.numpy as jnp
+
+
+def safe_select(condition, true_fn, false_fn):
+  """Safe version of selection (i.e. `where`).
+
+  This applies the double-where trick.
+  Like jnp.where, this function will still execute both branches and is
+  expected to be more lightweight than lax.cond.  Other than NaN-semantics,
+  safe_select(condition, true_fn, false_fn) is equivalent to
+
+    jax.tree.map(lambda x, y: jnp.where(condition, x, y),
+                 true_fn(),
+                 false_fn()),
+
+  Compared to the naive implementation above, safe_select provides the
+  following guarantee: in either the forward or backward pass, a NaN produced
+  *during the execution of true_fn()* will not propagate to the rest of the
+  computation and similarly for false_fn.  It is very important to note that
+  while true_fn and false_fn will typically close over other tensors (i.e. they
+  use values computed prior to the safe_select function), there is no NaN-safety
+  for the backward pass of closed over values.  It is important than any NaN's
+  are produced within the branch functions and not before them.  For example,
+
+    safe_select(x < eps, lambda: 0., lambda: jnp.sqrt(x))
+
+  will not produce NaN on the backward pass even if x == 0. since sqrt happens
+  within the false_fn, but the very similar
+
+    y = jnp.sqrt(x)
+    safe_select(x < eps, lambda: 0., lambda: y)
+
+  will produce a NaN on the backward pass if x == 0 because the sqrt happens
+  prior to the false_fn.
+
+  Args:
+    condition: Boolean array to use in where
+    true_fn: Zero-argument function to construct the values used in the True
+      condition.  Tensors that this function closes over will be extracted
+      automatically to implement the double-where trick to suppress spurious NaN
+      propagation.
+    false_fn: False branch equivalent of true_fn
+
+  Returns:
+    Resulting PyTree equivalent to tree_map line above.
+  """
+  true_fn, true_args = jax.closure_convert(true_fn)
+  false_fn, false_args = jax.closure_convert(false_fn)
+
+  true_args = jax.tree.map(
+      lambda x: jnp.where(condition, x, lax.stop_gradient(x)), true_args
+  )
+
+  false_args = jax.tree.map(
+      lambda x: jnp.where(condition, lax.stop_gradient(x), x), false_args
+  )
+
+  return jax.tree.map(
+      lambda x, y: jnp.where(condition, x, y),
+      true_fn(*true_args),
+      false_fn(*false_args),
+  )
+
+
+def unstack(value: jnp.ndarray, axis: int = -1) -> list[jnp.ndarray]:
+  return [
+      jnp.squeeze(v, axis=axis)
+      for v in jnp.split(value, value.shape[axis], axis=axis)
+  ]
+
+
+def angdiff(alpha: jnp.ndarray, beta: jnp.ndarray) -> jnp.ndarray:
+  """Compute absolute difference between two angles."""
+  d = alpha - beta
+  d = (d + jnp.pi) % (2 * jnp.pi) - jnp.pi
+  return d
+
+
+def safe_arctan2(
+    x1: jnp.ndarray, x2: jnp.ndarray, eps: float = 1e-8
+) -> jnp.ndarray:
+  """Safe version of arctan2 that avoids NaN gradients when x1=x2=0."""
+
+  return safe_select(
+      jnp.abs(x1) + jnp.abs(x2) < eps,
+      lambda: jnp.zeros_like(jnp.arctan2(x1, x2)),
+      lambda: jnp.arctan2(x1, x2),
+  )
+
+
+def weighted_mean(
+    *,
+    weights: jnp.ndarray,
+    value: jnp.ndarray,
+    axis: int | Iterable[int] | None = None,
+    eps: float = 1e-10,
+) -> jnp.ndarray:
+  """Computes weighted mean in a safe way that avoids NaNs.
+
+  This is equivalent to jnp.average for the case eps=0.0, but adds a small
+  constant to the denominator of the weighted average to avoid NaNs.
+  'weights' should be broadcastable to the shape of value.
+
+  Args:
+    weights: Weights to weight value by.
+    value: Values to average
+    axis: Axes to average over.
+    eps: Epsilon to add to the denominator.
+
+  Returns:
+    Weighted average.
+  """
+
+  weights = jnp.asarray(weights, dtype=value.dtype)
+  weights = jnp.broadcast_to(weights, value.shape)
+
+  weights_shape = weights.shape
+
+  if isinstance(axis, numbers.Integral):
+    axis = [axis]
+  elif axis is None:
+    axis = list(range(len(weights_shape)))
+
+  return jnp.sum(weights * value, axis=axis) / (
+      jnp.sum(weights, axis=axis) + eps
+  )

flax_model/alphafold3/jax/geometry/vector.py

@@ -0,0 +1,214 @@
+
+
+"""Vec3Array Class."""
+
+import dataclasses
+from typing import Final, Self, TypeAlias
+
+from flax_model.alphafold3.jax.geometry import struct_of_array
+from flax_model.alphafold3.jax.geometry import utils
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+Float: TypeAlias = float | jnp.ndarray
+
+VERSION: Final[str] = '0.1'
+
+
+@struct_of_array.StructOfArray(same_dtype=True)
+class Vec3Array:
+  """Vec3Array in 3 dimensional Space implemented as struct of arrays.
+
+  This is done in order to improve performance and precision.
+  On TPU small matrix multiplications are very suboptimal and will waste large
+  compute ressources, furthermore any matrix multiplication on TPU happens in
+  mixed bfloat16/float32 precision, which is often undesirable when handling
+  physical coordinates.
+
+  In most cases this will also be faster on CPUs/GPUs since it allows for easier
+  use of vector instructions.
+  """
+
+  x: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
+  y: jnp.ndarray
+  z: jnp.ndarray
+
+  def __post_init__(self):
+    if hasattr(self.x, 'dtype'):
+      if not self.x.dtype == self.y.dtype == self.z.dtype:
+        raise ValueError(
+            f'Type mismatch: {self.x.dtype}, {self.y.dtype}, {self.z.dtype}'
+        )
+      if not self.x.shape == self.y.shape == self.z.shape:
+        raise ValueError(
+            f'Shape mismatch: {self.x.shape}, {self.y.shape}, {self.z.shape}'
+        )
+
+  def __add__(self, other: Self) -> Self:
+    return jax.tree.map(lambda x, y: x + y, self, other)
+
+  def __sub__(self, other: Self) -> Self:
+    return jax.tree.map(lambda x, y: x - y, self, other)
+
+  def __mul__(self, other: Float) -> Self:
+    return jax.tree.map(lambda x: x * other, self)
+
+  def __rmul__(self, other: Float) -> Self:
+    return self * other
+
+  def __truediv__(self, other: Float) -> Self:
+    return jax.tree.map(lambda x: x / other, self)
+
+  def __neg__(self) -> Self:
+    return jax.tree.map(lambda x: -x, self)
+
+  def __pos__(self) -> Self:
+    return jax.tree.map(lambda x: x, self)
+
+  def cross(self, other: Self) -> Self:
+    """Compute cross product between 'self' and 'other'."""
+    new_x = self.y * other.z - self.z * other.y
+    new_y = self.z * other.x - self.x * other.z
+    new_z = self.x * other.y - self.y * other.x
+    return Vec3Array(new_x, new_y, new_z)
+
+  def dot(self, other: Self) -> Float:
+    """Compute dot product between 'self' and 'other'."""
+    return self.x * other.x + self.y * other.y + self.z * other.z
+
+  def norm(self, epsilon: float = 1e-6) -> Float:
+    """Compute Norm of Vec3Array, clipped to epsilon."""
+    # To avoid NaN on the backward pass, we must use maximum before the sqrt
+    norm2 = self.dot(self)
+    if epsilon:
+      norm2 = jnp.maximum(norm2, epsilon**2)
+    return jnp.sqrt(norm2)
+
+  def norm2(self):
+    return self.dot(self)
+
+  def normalized(self, epsilon: float = 1e-6) -> Self:
+    """Return unit vector with optional clipping."""
+    return self / self.norm(epsilon)
+
+  @classmethod
+  def zeros(cls, shape, dtype=jnp.float32):
+    """Return Vec3Array corresponding to zeros of given shape."""
+    return cls(
+        jnp.zeros(shape, dtype),
+        jnp.zeros(shape, dtype),
+        jnp.zeros(shape, dtype),
+    )  # pytype: disable=wrong-arg-count  # trace-all-classes
+
+  def to_array(self) -> jnp.ndarray:
+    return jnp.stack([self.x, self.y, self.z], axis=-1)
+
+  @classmethod
+  def from_array(cls, array):
+    return cls(*utils.unstack(array))
+
+  def __getstate__(self):
+    return (
+        VERSION,
+        [np.asarray(self.x), np.asarray(self.y), np.asarray(self.z)],
+    )
+
+  def __setstate__(self, state):
+    version, state = state
+    del version
+    for i, letter in enumerate('xyz'):
+      object.__setattr__(self, letter, state[i])
+
+
+def square_euclidean_distance(
+    vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-6
+) -> Float:
+  """Computes square of euclidean distance between 'vec1' and 'vec2'.
+
+  Args:
+    vec1: Vec3Array to compute  distance to
+    vec2: Vec3Array to compute  distance from, should be broadcast compatible
+      with 'vec1'
+    epsilon: distance is clipped from below to be at least epsilon
+
+  Returns:
+    Array of square euclidean distances;
+    shape will be result of broadcasting 'vec1' and 'vec2'
+  """
+  difference = vec1 - vec2
+  distance = difference.dot(difference)
+  if epsilon:
+    distance = jnp.maximum(distance, epsilon)
+  return distance
+
+
+def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
+  return vector1.dot(vector2)
+
+
+def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
+  return vector1.cross(vector2)
+
+
+def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
+  return vector.norm(epsilon)
+
+
+def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
+  return vector.normalized(epsilon)
+
+
+def euclidean_distance(
+    vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-6
+) -> Float:
+  """Computes euclidean distance between 'vec1' and 'vec2'.
+
+  Args:
+    vec1: Vec3Array to compute euclidean distance to
+    vec2: Vec3Array to compute euclidean distance from, should be broadcast
+      compatible with 'vec1'
+    epsilon: distance is clipped from below to be at least epsilon
+
+  Returns:
+    Array of euclidean distances;
+    shape will be result of broadcasting 'vec1' and 'vec2'
+  """
+  distance_sq = square_euclidean_distance(vec1, vec2, epsilon**2)
+  distance = jnp.sqrt(distance_sq)
+  return distance
+
+
+def dihedral_angle(
+    a: Vec3Array, b: Vec3Array, c: Vec3Array, d: Vec3Array
+) -> Float:
+  """Computes torsion angle for a quadruple of points.
+
+  For points (a, b, c, d), this is the angle between the planes defined by
+  points (a, b, c) and (b, c, d). It is also known as the dihedral angle.
+
+  Arguments:
+    a: A Vec3Array of coordinates.
+    b: A Vec3Array of coordinates.
+    c: A Vec3Array of coordinates.
+    d: A Vec3Array of coordinates.
+
+  Returns:
+    A tensor of angles in radians: [-pi, pi].
+  """
+  v1 = a - b
+  v2 = b - c
+  v3 = d - c
+
+  c1 = v1.cross(v2)
+  c2 = v3.cross(v2)
+  c3 = c2.cross(c1)
+
+  v2_mag = v2.norm()
+  return jnp.arctan2(c3.dot(v2), v2_mag * c1.dot(c2))
+
+
+def random_gaussian_vector(shape, key, dtype=jnp.float32) -> Vec3Array:
+  vec_array = jax.random.normal(key, shape + (3,), dtype)
+  return Vec3Array.from_array(vec_array)

flax_model/alphafold3/model/atom_layout/atom_layout.py

@@ -0,0 +1,1087 @@
+
+
+"""Helper functions for different atom layouts and conversion between them."""
+
+import collections
+from collections.abc import Mapping, Sequence
+import dataclasses
+import types
+from typing import Any, TypeAlias
+
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.constants import atom_types
+from flax_model.alphafold3.constants import chemical_component_sets
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.data.tools import rdkit_utils
+from flax_model.alphafold3.structure import chemical_components as struc_chem_comps
+import jax.numpy as jnp
+import numpy as np
+from rdkit import Chem
+
+
+xnp_ndarray: TypeAlias = np.ndarray | jnp.ndarray  # pylint: disable=invalid-name
+NumpyIndex: TypeAlias = Any
+
+
+@dataclasses.dataclass(frozen=True)
+class AtomLayout:
+  """Atom layout in a fixed shape (usually 1-dim or 2-dim).
+
+  Examples for atom layouts are atom37, atom14, and similar.
+  All members are np.ndarrays with the same shape, e.g.
+  - [num_atoms]
+  - [num_residues, max_atoms_per_residue]
+  - [num_fragments, max_fragments_per_residue]
+  All string arrays should have dtype=object to avoid pitfalls with Numpy's
+  fixed-size strings
+
+  Attributes:
+    atom_name: np.ndarray of str: atom names (e.g. 'CA', 'NE2'), padding
+      elements have an empty string (''), None or any other value, that maps to
+      False for .astype(bool). mmCIF field: _atom_site.label_atom_id.
+    res_id: np.ndarray of int: residue index (usually starting from 1) padding
+      elements can have an arbitrary value. mmCIF field:
+      _atom_site.label_seq_id.
+    chain_id: np.ndarray of str: chain names (e.g. 'A', 'B') padding elements
+      can have an arbitrary value. mmCIF field: _atom_site.label_seq_id.
+    atom_element: np.ndarray of str: atom elements (e.g. 'C', 'N', 'O'), padding
+      elements have an empty string (''), None or any other value, that maps to
+      False for .astype(bool). mmCIF field: _atom_site.type_symbol.
+    res_name: np.ndarray of str: residue names (e.g. 'ARG', 'TRP') padding
+      elements can have an arbitrary value. mmCIF field:
+      _atom_site.label_comp_id.
+    chain_type: np.ndarray of str: chain types (e.g. 'polypeptide(L)'). padding
+      elements can have an arbitrary value. mmCIF field: _entity_poly.type OR
+      _entity.type (for non-polymers).
+    shape: shape of the layout (just returns atom_name.shape)
+  """
+
+  atom_name: np.ndarray
+  res_id: np.ndarray
+  chain_id: np.ndarray
+  atom_element: np.ndarray | None = None
+  res_name: np.ndarray | None = None
+  chain_type: np.ndarray | None = None
+
+  def __post_init__(self):
+    """Assert all arrays have the same shape."""
+    attribute_names = (
+        'atom_name',
+        'atom_element',
+        'res_name',
+        'res_id',
+        'chain_id',
+        'chain_type',
+    )
+    _assert_all_arrays_have_same_shape(
+        obj=self,
+        expected_shape=self.atom_name.shape,
+        attribute_names=attribute_names,
+    )
+    # atom_name must have dtype object, such that we can convert it to bool to
+    # obtain the mask
+    if self.atom_name.dtype != object:
+      raise ValueError(
+          'atom_name must have dtype object, such that it can '
+          'be converted converted to bool to obtain the mask'
+      )
+
+  def __getitem__(self, key: NumpyIndex) -> 'AtomLayout':
+    return AtomLayout(
+        atom_name=self.atom_name[key],
+        res_id=self.res_id[key],
+        chain_id=self.chain_id[key],
+        atom_element=(
+            self.atom_element[key] if self.atom_element is not None else None
+        ),
+        res_name=(self.res_name[key] if self.res_name is not None else None),
+        chain_type=(
+            self.chain_type[key] if self.chain_type is not None else None
+        ),
+    )
+
+  def __eq__(self, other: 'AtomLayout') -> bool:
+    if not np.array_equal(self.atom_name, other.atom_name):
+      return False
+
+    mask = self.atom_name.astype(bool)
+    # Check essential fields.
+    for field in ('res_id', 'chain_id'):
+      my_arr = getattr(self, field)
+      other_arr = getattr(other, field)
+      if not np.array_equal(my_arr[mask], other_arr[mask]):
+        return False
+
+    # Check optional fields.
+    for field in ('atom_element', 'res_name', 'chain_type'):
+      my_arr = getattr(self, field)
+      other_arr = getattr(other, field)
+      if (
+          my_arr is not None
+          and other_arr is not None
+          and not np.array_equal(my_arr[mask], other_arr[mask])
+      ):
+        return False
+
+    return True
+
+  def copy_and_pad_to(self, shape: tuple[int, ...]) -> 'AtomLayout':
+    """Copies and pads the layout to the requested shape.
+
+    Args:
+      shape: new shape for the atom layout
+
+    Returns:
+      a copy of the atom layout padded to the requested shape
+
+    Raises:
+      ValueError: incompatible shapes.
+    """
+    if len(shape) != len(self.atom_name.shape):
+      raise ValueError(
+          f'Incompatible shape {shape}. Current layout has shape {self.shape}.'
+      )
+    if any(new < old for old, new in zip(self.atom_name.shape, shape)):
+      raise ValueError(
+          "Can't pad to a smaller shape. Current layout has shape "
+          f'{self.shape} and you requested shape {shape}.'
+      )
+    pad_width = [
+        (0, new - old) for old, new in zip(self.atom_name.shape, shape)
+    ]
+    pad_val = np.array('', dtype=object)
+    return AtomLayout(
+        atom_name=np.pad(self.atom_name, pad_width, constant_values=pad_val),
+        res_id=np.pad(self.res_id, pad_width, constant_values=0),
+        chain_id=np.pad(self.chain_id, pad_width, constant_values=pad_val),
+        atom_element=(
+            np.pad(self.atom_element, pad_width, constant_values=pad_val)
+            if self.atom_element is not None
+            else None
+        ),
+        res_name=(
+            np.pad(self.res_name, pad_width, constant_values=pad_val)
+            if self.res_name is not None
+            else None
+        ),
+        chain_type=(
+            np.pad(self.chain_type, pad_width, constant_values=pad_val)
+            if self.chain_type is not None
+            else None
+        ),
+    )
+
+  def to_array(self) -> np.ndarray:
+    """Stacks the fields to a numpy array with shape (6, <layout_shape>).
+
+    Creates a pure numpy array of type `object` by stacking the 6 fields of the
+    AtomLayout, i.e. (atom_name, atom_element, res_name, res_id, chain_id,
+    chain_type). This method together with from_array() provides an easy way to
+    apply pure numpy methods like np.concatenate() to `AtomLayout`s.
+
+    Returns:
+      np.ndarray of object with shape (6, <layout_shape>), e.g.
+      array([['N', 'CA', 'C', ..., 'CB', 'CG', 'CD'],
+       ['N', 'C', 'C', ..., 'C', 'C', 'C'],
+       ['LEU', 'LEU', 'LEU', ..., 'PRO', 'PRO', 'PRO'],
+       [1, 1, 1, ..., 403, 403, 403],
+       ['A', 'A', 'A', ..., 'D', 'D', 'D'],
+       ['polypeptide(L)', 'polypeptide(L)', ..., 'polypeptide(L)']],
+      dtype=object)
+    """
+    if (
+        self.atom_element is None
+        or self.res_name is None
+        or self.chain_type is None
+    ):
+      raise ValueError('All optional fields need to be present.')
+
+    return np.stack(dataclasses.astuple(self), axis=0)
+
+  @classmethod
+  def from_array(cls, arr: np.ndarray) -> 'AtomLayout':
+    """Creates an AtomLayout object from a numpy array with shape (6, ...).
+
+    see also to_array()
+    Args:
+      arr: np.ndarray of object with shape (6, <layout_shape>)
+
+    Returns:
+      AtomLayout object with shape (<layout_shape>)
+    """
+    if arr.shape[0] != 6:
+      raise ValueError(
+          'Given array must have shape (6, ...) to match the 6 fields of '
+          'AtomLayout (atom_name, atom_element, res_name, res_id, chain_id, '
+          f'chain_type). Your array has {arr.shape=}'
+      )
+    return cls(*arr)
+
+  @property
+  def shape(self) -> tuple[int, ...]:
+    return self.atom_name.shape
+
+
+@dataclasses.dataclass(frozen=True)
+class Residues:
+  """List of residues with meta data.
+
+  Attributes:
+    res_name: np.ndarray of str [num_res], e.g. 'ARG', 'TRP'
+    res_id: np.ndarray of int [num_res]
+    chain_id: np.ndarray of str [num_res], e.g. 'A', 'B'
+    chain_type: np.ndarray of str [num_res], e.g. 'polypeptide(L)'
+    is_start_terminus: np.ndarray of bool [num_res]
+    is_end_terminus: np.ndarray of bool [num_res]
+    deprotonation: (optional) np.ndarray of set() [num_res], e.g. {'HD1', 'HE2'}
+    smiles_string: (optional) np.ndarray of str [num_res], e.g. 'Cc1ccccc1'
+    shape: shape of the layout (just returns res_name.shape)
+  """
+
+  res_name: np.ndarray
+  res_id: np.ndarray
+  chain_id: np.ndarray
+  chain_type: np.ndarray
+  is_start_terminus: np.ndarray
+  is_end_terminus: np.ndarray
+  deprotonation: np.ndarray | None = None
+  smiles_string: np.ndarray | None = None
+
+  def __post_init__(self):
+    """Assert all arrays are 1D have the same shape."""
+    attribute_names = (
+        'res_name',
+        'res_id',
+        'chain_id',
+        'chain_type',
+        'is_start_terminus',
+        'is_end_terminus',
+        'deprotonation',
+        'smiles_string',
+    )
+    _assert_all_arrays_have_same_shape(
+        obj=self,
+        expected_shape=(self.res_name.shape[0],),
+        attribute_names=attribute_names,
+    )
+
+  def __getitem__(self, key: NumpyIndex) -> 'Residues':
+    return Residues(
+        res_name=self.res_name[key],
+        res_id=self.res_id[key],
+        chain_id=self.chain_id[key],
+        chain_type=self.chain_type[key],
+        is_start_terminus=self.is_start_terminus[key],
+        is_end_terminus=self.is_end_terminus[key],
+        deprotonation=(
+            self.deprotonation[key] if self.deprotonation is not None else None
+        ),
+        smiles_string=(
+            self.smiles_string[key] if self.smiles_string is not None else None
+        ),
+    )
+
+  def __eq__(self, other: 'Residues') -> bool:
+    return all(
+        np.array_equal(getattr(self, field.name), getattr(other, field.name))
+        for field in dataclasses.fields(self)
+    )
+
+  @property
+  def shape(self) -> tuple[int, ...]:
+    return self.res_name.shape
+
+
+@dataclasses.dataclass(frozen=True)
+class GatherInfo:
+  """Gather indices to translate from one atom layout to another.
+
+  All members are np or jnp ndarray (usually 1-dim or 2-dim) with the same
+  shape, e.g.
+  - [num_atoms]
+  - [num_residues, max_atoms_per_residue]
+  - [num_fragments, max_fragments_per_residue]
+
+  Attributes:
+    gather_idxs: np or jnp ndarray of int: gather indices into a flattened array
+    gather_mask: np or jnp ndarray of bool: mask for resulting array
+    input_shape: np or jnp ndarray of int: the shape of the unflattened input
+      array
+    shape: output shape. Just returns gather_idxs.shape
+  """
+
+  gather_idxs: xnp_ndarray
+  gather_mask: xnp_ndarray
+  input_shape: xnp_ndarray
+
+  def __post_init__(self):
+    if self.gather_mask.shape != self.gather_idxs.shape:
+      raise ValueError(
+          'All arrays must have the same shape. Got\n'
+          f'gather_idxs.shape = {self.gather_idxs.shape}\n'
+          f'gather_mask.shape = {self.gather_mask.shape}\n'
+      )
+
+  def __getitem__(self, key: NumpyIndex) -> 'GatherInfo':
+    return GatherInfo(
+        gather_idxs=self.gather_idxs[key],
+        gather_mask=self.gather_mask[key],
+        input_shape=self.input_shape,
+    )
+
+  @property
+  def shape(self) -> tuple[int, ...]:
+    return self.gather_idxs.shape
+
+  def as_np_or_jnp(self, xnp: types.ModuleType) -> 'GatherInfo':
+    return GatherInfo(
+        gather_idxs=xnp.array(self.gather_idxs),
+        gather_mask=xnp.array(self.gather_mask),
+        input_shape=xnp.array(self.input_shape),
+    )
+
+  def as_dict(
+      self,
+      key_prefix: str | None = None,
+  ) -> dict[str, xnp_ndarray]:
+    prefix = f'{key_prefix}:' if key_prefix else ''
+    return {
+        prefix + 'gather_idxs': self.gather_idxs,
+        prefix + 'gather_mask': self.gather_mask,
+        prefix + 'input_shape': self.input_shape,
+    }
+
+  @classmethod
+  def from_dict(
+      cls,
+      d: Mapping[str, xnp_ndarray],
+      key_prefix: str | None = None,
+  ) -> 'GatherInfo':
+    """Creates GatherInfo from a given dictionary."""
+    prefix = f'{key_prefix}:' if key_prefix else ''
+    return cls(
+        gather_idxs=d[prefix + 'gather_idxs'],
+        gather_mask=d[prefix + 'gather_mask'],
+        input_shape=d[prefix + 'input_shape'],
+    )
+
+
+def fill_in_optional_fields(
+    minimal_atom_layout: AtomLayout,
+    reference_atoms: AtomLayout,
+) -> AtomLayout:
+  """Fill in the optional fields (atom_element, res_name, chain_type).
+
+  Extracts the optional fields (atom_element, res_name, chain_type) from a
+  flat reference layout and fills them into the fields from this layout.
+
+  Args:
+    minimal_atom_layout: An AtomLayout that only contains the essential fields
+      (atom_name, res_id, chain_id).
+    reference_atoms: A flat layout that contains all fields for all atoms.
+
+  Returns:
+    An AtomLayout that contains all fields.
+
+  Raises:
+    ValueError: Reference atoms layout is not flat.
+    ValueError: Missing atoms in reference.
+  """
+  if len(reference_atoms.shape) > 1:
+    raise ValueError('Only flat layouts are supported as reference.')
+  ref_to_self = compute_gather_idxs(
+      source_layout=reference_atoms, target_layout=minimal_atom_layout
+  )
+  atom_mask = minimal_atom_layout.atom_name.astype(bool)
+  missing_atoms_mask = atom_mask & ~ref_to_self.gather_mask
+  if np.any(missing_atoms_mask):
+    raise ValueError(
+        f'{np.sum(missing_atoms_mask)} missing atoms in reference: '
+        f'{minimal_atom_layout[missing_atoms_mask]}'
+    )
+
+  def _convert_str_array(gather: GatherInfo, arr: np.ndarray):
+    output = arr[gather.gather_idxs]
+    output[~gather.gather_mask] = ''
+    return output
+
+  return dataclasses.replace(
+      minimal_atom_layout,
+      atom_element=_convert_str_array(
+          ref_to_self, reference_atoms.atom_element
+      ),
+      res_name=_convert_str_array(ref_to_self, reference_atoms.res_name),
+      chain_type=_convert_str_array(ref_to_self, reference_atoms.chain_type),
+  )
+
+
+def guess_deprotonation(residues: Residues) -> Residues:
+  """Convenience function to create a plausible deprotonation field.
+
+  Assumes a pH of 7 and always prefers HE2 over HD1 for HIS.
+  Args:
+    residues: a Residues object without a depronotation field
+
+  Returns:
+    a Residues object with a depronotation field
+  """
+  num_residues = residues.res_name.shape[0]
+  deprotonation = np.empty(num_residues, dtype=object)
+  deprotonation_at_ph7 = {
+      'ASP': 'HD2',
+      'GLU': 'HE2',
+      'HIS': 'HD1',
+  }
+  for idx, res_name in enumerate(residues.res_name):
+    deprotonation[idx] = set()
+    if res_name in deprotonation_at_ph7:
+      deprotonation[idx].add(deprotonation_at_ph7[res_name])
+    if residues.is_end_terminus[idx]:
+      deprotonation[idx].add('HXT')
+
+  return dataclasses.replace(residues, deprotonation=deprotonation)
+
+
+def atom_layout_from_structure(
+    struct: structure.Structure,
+    *,
+    fix_non_standard_polymer_res: bool = False,
+) -> AtomLayout:
+  """Extract AtomLayout from a Structure."""
+
+  if not fix_non_standard_polymer_res:
+    return AtomLayout(
+        atom_name=np.array(struct.atom_name, dtype=object),
+        atom_element=np.array(struct.atom_element, dtype=object),
+        res_name=np.array(struct.res_name, dtype=object),
+        res_id=np.array(struct.res_id, dtype=int),
+        chain_id=np.array(struct.chain_id, dtype=object),
+        chain_type=np.array(struct.chain_type, dtype=object),
+    )
+
+  # Target lists.
+  target_atom_names = []
+  target_atom_elements = []
+  target_res_ids = []
+  target_res_names = []
+  target_chain_ids = []
+  target_chain_types = []
+
+  for atom in struct.iter_atoms():
+    target_atom_names.append(atom['atom_name'])
+    target_atom_elements.append(atom['atom_element'])
+    target_res_ids.append(atom['res_id'])
+    target_chain_ids.append(atom['chain_id'])
+    target_chain_types.append(atom['chain_type'])
+    if mmcif_names.is_standard_polymer_type(atom['chain_type']):
+      fixed_res_name = mmcif_names.fix_non_standard_polymer_res(
+          res_name=atom['res_name'], chain_type=atom['chain_type']
+      )
+      target_res_names.append(fixed_res_name)
+    else:
+      target_res_names.append(atom['res_name'])
+
+  return AtomLayout(
+      atom_name=np.array(target_atom_names, dtype=object),
+      atom_element=np.array(target_atom_elements, dtype=object),
+      res_name=np.array(target_res_names, dtype=object),
+      res_id=np.array(target_res_ids, dtype=int),
+      chain_id=np.array(target_chain_ids, dtype=object),
+      chain_type=np.array(target_chain_types, dtype=object),
+  )
+
+
+def residues_from_structure(
+    struct: structure.Structure,
+    *,
+    include_missing_residues: bool = True,
+    fix_non_standard_polymer_res: bool = False,
+) -> Residues:
+  """Create a Residues object from a Structure object."""
+
+  def _get_smiles(res_name):
+    """Get SMILES string from chemical components."""
+    smiles = None
+    if (
+        struct.chemical_components_data is not None
+        and struct.chemical_components_data.chem_comp is not None
+        and struct.chemical_components_data.chem_comp.get(res_name)
+    ):
+      smiles = struct.chemical_components_data.chem_comp[res_name].pdbx_smiles
+    return smiles
+
+  res_names_per_chain = struct.chain_res_name_sequence(
+      include_missing_residues=include_missing_residues,
+      fix_non_standard_polymer_res=fix_non_standard_polymer_res,
+  )
+  res_name = []
+  res_id = []
+  chain_id = []
+  chain_type = []
+  smiles = []
+  is_start_terminus = []
+  for c in struct.iter_chains():
+    if include_missing_residues:
+      this_res_ids = [id for (_, id) in struct.all_residues[c['chain_id']]]
+    else:
+      this_res_ids = [
+          r['res_id']
+          for r in struct.iter_residues()
+          if r['chain_id'] == c['chain_id']
+      ]
+    fixed_res_names = res_names_per_chain[c['chain_id']]
+    assert len(this_res_ids) == len(
+        fixed_res_names
+    ), f'{len(this_res_ids)} != {len(fixed_res_names)}'
+    this_start_res_id = min(min(this_res_ids), 1)
+    this_is_start_terminus = [r == this_start_res_id for r in this_res_ids]
+    smiles.extend([_get_smiles(res_name) for res_name in fixed_res_names])
+    num_res = len(fixed_res_names)
+    res_name.extend(fixed_res_names)
+    res_id.extend(this_res_ids)
+    chain_id.extend([c['chain_id']] * num_res)
+    chain_type.extend([c['chain_type']] * num_res)
+    is_start_terminus.extend(this_is_start_terminus)
+  res_name = np.array(res_name, dtype=object)
+  res_id = np.array(res_id, dtype=int)
+  chain_id = np.array(chain_id, dtype=object)
+  chain_type = np.array(chain_type, dtype=object)
+  smiles = np.array(smiles, dtype=object)
+  is_start_terminus = np.array(is_start_terminus, dtype=bool)
+
+  res_uid_to_idx = {
+      uid: idx for idx, uid in enumerate(zip(chain_id, res_id, strict=True))
+  }
+
+  # Start terminus indicates whether residue index is 1 and chain is polymer.
+  is_polymer = np.isin(chain_type, tuple(mmcif_names.POLYMER_CHAIN_TYPES))
+  is_start_terminus = is_start_terminus & is_polymer
+
+  # Start also indicates whether amino acid is attached to H2 or proline to H.
+  start_terminus_atom_index = np.nonzero(
+      (struct.chain_type == mmcif_names.PROTEIN_CHAIN)
+      & (
+          (struct.atom_name == 'H2')
+          | ((struct.atom_name == 'H') & (struct.res_name == 'PRO'))
+      )
+  )[0]
+
+  # Translate atom idx to residue idx to assign start terminus.
+  for atom_idx in start_terminus_atom_index:
+    res_uid = (struct.chain_id[atom_idx], struct.res_id[atom_idx])
+    res_idx = res_uid_to_idx[res_uid]
+    is_start_terminus[res_idx] = True
+
+  # Infer end terminus: Check for OXT, or in case of
+  # include_missing_residues==True for the last residue of the chain.
+  num_all_residues = res_name.shape[0]
+  is_end_terminus = np.zeros(num_all_residues, dtype=bool)
+  end_term_atom_idxs = np.nonzero(struct.atom_name == 'OXT')[0]
+  for atom_idx in end_term_atom_idxs:
+    res_uid = (struct.chain_id[atom_idx], struct.res_id[atom_idx])
+    res_idx = res_uid_to_idx[res_uid]
+    is_end_terminus[res_idx] = True
+
+  if include_missing_residues:
+    for idx in range(num_all_residues - 1):
+      if is_polymer[idx] and chain_id[idx] != chain_id[idx + 1]:
+        is_end_terminus[idx] = True
+    if (num_all_residues > 0) and is_polymer[-1]:
+      is_end_terminus[-1] = True
+
+  # Infer (de-)protonation: Only if hydrogens are given.
+  num_hydrogens = np.sum(
+      (struct.atom_element == 'H') & (struct.chain_type == 'polypeptide(L)')
+  )
+  if num_hydrogens > 0:
+    deprotonation = np.empty(num_all_residues, dtype=object)
+    all_atom_uids = set(
+        zip(struct.chain_id, struct.res_id, struct.atom_name, strict=True)
+    )
+    for idx in range(num_all_residues):
+      deprotonation[idx] = set()
+      check_hydrogens = set()
+      if is_end_terminus[idx]:
+        check_hydrogens.add('HXT')
+      if res_name[idx] in atom_types.PROTONATION_HYDROGENS:
+        check_hydrogens.update(atom_types.PROTONATION_HYDROGENS[res_name[idx]])
+      for hydrogen in check_hydrogens:
+        if (chain_id[idx], res_id[idx], hydrogen) not in all_atom_uids:
+          deprotonation[idx].add(hydrogen)
+  else:
+    deprotonation = None
+
+  return Residues(
+      res_name=res_name,
+      res_id=res_id,
+      chain_id=chain_id,
+      chain_type=chain_type,
+      is_start_terminus=is_start_terminus.astype(bool),
+      is_end_terminus=is_end_terminus,
+      deprotonation=deprotonation,
+      smiles_string=smiles,
+  )
+
+
+def get_link_drop_atoms(
+    res_name: str,
+    chain_type: str,
+    *,
+    is_start_terminus: bool,
+    is_end_terminus: bool,
+    bonded_atoms: set[str],
+    drop_ligand_leaving_atoms: bool = False,
+) -> set[str]:
+  """Returns set of atoms that are dropped when this res_name gets linked.
+
+  Args:
+    res_name: residue name, e.g. 'ARG'
+    chain_type: chain_type, e.g. 'polypeptide(L)'
+    is_start_terminus: whether the residue is the n-terminus
+    is_end_terminus: whether the residue is the c-terminus
+    bonded_atoms: Names of atoms coming off this residue.
+    drop_ligand_leaving_atoms: Flag to switch on/off leaving atoms for ligands.
+
+  Returns:
+    Set of atoms that are dropped when this amino acid gets linked.
+  """
+  drop_atoms = set()
+  if chain_type == mmcif_names.PROTEIN_CHAIN:
+    if res_name == 'PRO':
+      if not is_start_terminus:
+        drop_atoms.update({'H', 'H2', 'H3'})
+      if not is_end_terminus:
+        drop_atoms.update({'OXT', 'HXT'})
+    else:
+      if not is_start_terminus:
+        drop_atoms.update({'H2', 'H3'})
+      if not is_end_terminus:
+        drop_atoms.update({'OXT', 'HXT'})
+  elif chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES:
+    if not is_start_terminus:
+      drop_atoms.update({'OP3'})
+  elif (
+      drop_ligand_leaving_atoms and chain_type in mmcif_names.LIGAND_CHAIN_TYPES
+  ):
+    if res_name in {
+        *chemical_component_sets.GLYCAN_OTHER_LIGANDS,
+        *chemical_component_sets.GLYCAN_LINKING_LIGANDS,
+    }:
+      if 'O1' not in bonded_atoms:
+        drop_atoms.update({'O1'})
+  return drop_atoms
+
+
+def get_bonded_atoms(
+    polymer_ligand_bonds: AtomLayout,
+    ligand_ligand_bonds: AtomLayout,
+    res_id: int,
+    chain_id: str,
+) -> set[str]:
+  """Finds the res_name on the opposite end of the bond, if a bond exists.
+
+  Args:
+    polymer_ligand_bonds: Bond information for polymer-ligand pairs.
+    ligand_ligand_bonds: Bond information for ligand-ligand pairs.
+    res_id: residue id in question.
+    chain_id: chain id of residue in question.
+
+  Returns:
+    res_name of bonded atom.
+  """
+  bonded_atoms = set()
+  if polymer_ligand_bonds:
+    # Filter before searching to speed this up.
+    bond_idx = np.logical_and(
+        polymer_ligand_bonds.res_id == res_id,
+        polymer_ligand_bonds.chain_id == chain_id,
+    ).any(axis=1)
+    relevant_polymer_bonds = polymer_ligand_bonds[bond_idx]
+    for atom_names, res_ids, chain_ids in zip(
+        relevant_polymer_bonds.atom_name,
+        relevant_polymer_bonds.res_id,
+        relevant_polymer_bonds.chain_id,
+    ):
+      if (res_ids[0], chain_ids[0]) == (res_id, chain_id):
+        bonded_atoms.add(atom_names[0])
+      elif (res_ids[1], chain_ids[1]) == (res_id, chain_id):
+        bonded_atoms.add(atom_names[1])
+  if ligand_ligand_bonds:
+    bond_idx = np.logical_and(
+        ligand_ligand_bonds.res_id == res_id,
+        ligand_ligand_bonds.chain_id == chain_id,
+    ).any(axis=1)
+    relevant_ligand_bonds = ligand_ligand_bonds[bond_idx]
+    for atom_names, res_ids, chain_ids in zip(
+        relevant_ligand_bonds.atom_name,
+        relevant_ligand_bonds.res_id,
+        relevant_ligand_bonds.chain_id,
+    ):
+      if (res_ids[0], chain_ids[0]) == (res_id, chain_id):
+        bonded_atoms.add(atom_names[0])
+      elif (res_ids[1], chain_ids[1]) == (res_id, chain_id):
+        bonded_atoms.add(atom_names[1])
+  return bonded_atoms
+
+
+def make_flat_atom_layout(
+    residues: Residues,
+    ccd: chemical_components.Ccd,
+    polymer_ligand_bonds: AtomLayout | None = None,
+    ligand_ligand_bonds: AtomLayout | None = None,
+    *,
+    with_hydrogens: bool = False,
+    skip_unk_residues: bool = True,
+    drop_ligand_leaving_atoms: bool = False,
+) -> AtomLayout:
+  """Make a flat atom layout for given residues.
+
+  Create a flat layout from a `Residues` object. The required atoms for each
+  amino acid type are taken from the CCD, hydrogens and oxygens are dropped to
+  make the linked residues. Terminal OXT's and protonation state for the
+  hydrogens come from the `Residues` object.
+
+  Args:
+    residues: a `Residues` object.
+    ccd: The chemical components dictionary.
+    polymer_ligand_bonds: Bond information for polymer-ligand pairs.
+    ligand_ligand_bonds: Bond information for ligand-ligand pairs.
+    with_hydrogens: whether to create hydrogens
+    skip_unk_residues: whether to skip 'UNK' resides -- default is True to be
+      compatible with the rest of AlphaFold that does not predict atoms for
+      unknown residues
+    drop_ligand_leaving_atoms: Flag to switch on/ off leaving atoms for ligands.
+
+  Returns:
+    an `AtomLayout` object
+  """
+  num_res = residues.res_name.shape[0]
+
+  # Target lists.
+  target_atom_names = []
+  target_atom_elements = []
+  target_res_ids = []
+  target_res_names = []
+  target_chain_ids = []
+  target_chain_types = []
+
+  for idx in range(num_res):
+    # skip 'UNK' residues if requested
+    if (
+        skip_unk_residues
+        and residues.res_name[idx] in residue_names.UNKNOWN_TYPES
+    ):
+      continue
+
+    # Get the atoms for this residue type from CCD.
+    if ccd.get(residues.res_name[idx]):
+      res_atoms = struc_chem_comps.get_all_atoms_in_entry(
+          ccd=ccd, res_name=residues.res_name[idx]
+      )
+      atom_names_elements = list(
+          zip(
+              res_atoms['_chem_comp_atom.atom_id'],
+              res_atoms['_chem_comp_atom.type_symbol'],
+              strict=True,
+          )
+      )
+    elif residues.smiles_string[idx]:
+      # Get atoms from RDKit via SMILES.
+      mol = Chem.MolFromSmiles(residues.smiles_string[idx])
+      if mol is None:
+        raise ValueError(
+            f'Failed to construct RDKit Mol for {residues.res_name[idx]} from'
+            f' SMILES string: {residues.smiles_string[idx]} . This is likely'
+            ' due to an issue with the SMILES string. Note that the userCCD'
+            ' input format provides an alternative way to define custom'
+            ' molecules directly without RDKit or SMILES.'
+        )
+      mol = rdkit_utils.assign_atom_names_from_graph(mol)
+      atom_names_elements = [
+          (a.GetProp('atom_name'), a.GetSymbol()) for a in mol.GetAtoms()
+      ]
+    else:
+      raise ValueError(
+          f'{residues.res_name[idx]} not found in CCD and no SMILES string'
+      )
+
+    # Remove hydrogens if requested.
+    if not with_hydrogens:
+      atom_names_elements = [
+          (n, e) for n, e in atom_names_elements if (e != 'H' and e != 'D')
+      ]
+    bonded_atoms = get_bonded_atoms(
+        polymer_ligand_bonds,
+        ligand_ligand_bonds,
+        residues.res_id[idx],
+        residues.chain_id[idx],
+    )
+    # Connect the amino-acids, i.e. remove OXT, HXT and H2.
+    drop_atoms = get_link_drop_atoms(
+        res_name=residues.res_name[idx],
+        chain_type=residues.chain_type[idx],
+        is_start_terminus=residues.is_start_terminus[idx],
+        is_end_terminus=residues.is_end_terminus[idx],
+        bonded_atoms=bonded_atoms,
+        drop_ligand_leaving_atoms=drop_ligand_leaving_atoms,
+    )
+
+    # If deprotonation info is available, remove the specific atoms.
+    if residues.deprotonation is not None:
+      drop_atoms.update(residues.deprotonation[idx])
+
+    atom_names_elements = [
+        (n, e) for n, e in atom_names_elements if n not in drop_atoms
+    ]
+
+    # Append the found atoms to the target lists.
+    target_atom_names.extend([n for n, _ in atom_names_elements])
+    target_atom_elements.extend([e for _, e in atom_names_elements])
+    num_atoms = len(atom_names_elements)
+    target_res_names.extend([residues.res_name[idx]] * num_atoms)
+    target_res_ids.extend([residues.res_id[idx]] * num_atoms)
+    target_chain_ids.extend([residues.chain_id[idx]] * num_atoms)
+    target_chain_types.extend([residues.chain_type[idx]] * num_atoms)
+
+  return AtomLayout(
+      atom_name=np.array(target_atom_names, dtype=object),
+      atom_element=np.array(target_atom_elements, dtype=object),
+      res_name=np.array(target_res_names, dtype=object),
+      res_id=np.array(target_res_ids, dtype=int),
+      chain_id=np.array(target_chain_ids, dtype=object),
+      chain_type=np.array(target_chain_types, dtype=object),
+  )
+
+
+def compute_gather_idxs(
+    *,
+    source_layout: AtomLayout,
+    target_layout: AtomLayout,
+    fill_value: int = 0,
+) -> GatherInfo:
+  """Produce gather indices and mask to convert from source layout to target."""
+  source_uid_to_idx = {
+      uid: idx
+      for idx, uid in enumerate(
+          zip(
+              source_layout.chain_id.ravel(),
+              source_layout.res_id.ravel(),
+              source_layout.atom_name.ravel(),
+              strict=True,
+          )
+      )
+  }
+  gather_idxs = []
+  gather_mask = []
+  for uid in zip(
+      target_layout.chain_id.ravel(),
+      target_layout.res_id.ravel(),
+      target_layout.atom_name.ravel(),
+      strict=True,
+  ):
+    if uid in source_uid_to_idx:
+      gather_idxs.append(source_uid_to_idx[uid])
+      gather_mask.append(True)
+    else:
+      gather_idxs.append(fill_value)
+      gather_mask.append(False)
+  target_shape = target_layout.atom_name.shape
+  return GatherInfo(
+      gather_idxs=np.array(gather_idxs, dtype=int).reshape(target_shape),
+      gather_mask=np.array(gather_mask, dtype=bool).reshape(target_shape),
+      input_shape=np.array(source_layout.atom_name.shape),
+  )
+
+
+def convert(
+    gather_info: GatherInfo,
+    arr: xnp_ndarray,
+    *,
+    layout_axes: tuple[int, ...] = (0,),
+) -> xnp_ndarray:
+  """Convert an array from one atom layout to another."""
+  # Translate negative indices to the corresponding positives.
+  layout_axes = tuple(i if i >= 0 else i + arr.ndim for i in layout_axes)
+
+  # Ensure that layout_axes are continuous.
+  layout_axes_begin = layout_axes[0]
+  layout_axes_end = layout_axes[-1] + 1
+
+  if layout_axes != tuple(range(layout_axes_begin, layout_axes_end)):
+    raise ValueError(f'layout_axes must be continuous. Got {layout_axes}.')
+  layout_shape = arr.shape[layout_axes_begin:layout_axes_end]
+
+  # Ensure that the layout shape is compatible
+  # with the gather_info. I.e. the first axis size must be equal or greater
+  # than the gather_info.input_shape, and all subsequent axes sizes must match.
+  if (len(layout_shape) != gather_info.input_shape.size) or (
+      isinstance(gather_info.input_shape, np.ndarray)
+      and (
+          (layout_shape[0] < gather_info.input_shape[0])
+          or (np.any(layout_shape[1:] != gather_info.input_shape[1:]))
+      )
+  ):
+    raise ValueError(
+        'Input array layout axes are incompatible. You specified layout '
+        f'axes {layout_axes} with an input array of shape {arr.shape}, but '
+        f'the gather info expects shape {gather_info.input_shape}. '
+        'Your first axis size must be equal or greater than the '
+        'gather_info.input_shape, and all subsequent axes sizes must '
+        'match.'
+    )
+
+  # Compute the shape of the input array with flattened layout.
+  batch_shape = arr.shape[:layout_axes_begin]
+  features_shape = arr.shape[layout_axes_end:]
+  arr_flattened_shape = batch_shape + (np.prod(layout_shape),) + features_shape
+
+  # Flatten input array and perform the gather.
+  arr_flattened = arr.reshape(arr_flattened_shape)
+  if layout_axes_begin == 0:
+    out_arr = arr_flattened[gather_info.gather_idxs, ...]
+  elif layout_axes_begin == 1:
+    out_arr = arr_flattened[:, gather_info.gather_idxs, ...]
+  elif layout_axes_begin == 2:
+    out_arr = arr_flattened[:, :, gather_info.gather_idxs, ...]
+  elif layout_axes_begin == 3:
+    out_arr = arr_flattened[:, :, :, gather_info.gather_idxs, ...]
+  elif layout_axes_begin == 4:
+    out_arr = arr_flattened[:, :, :, :, gather_info.gather_idxs, ...]
+  else:
+    raise ValueError(
+        'Only 4 batch axes supported. If you need more, the code '
+        'is easy to extend.'
+    )
+
+  # Broadcast the mask and apply it.
+  broadcasted_mask_shape = (
+      (1,) * len(batch_shape)
+      + gather_info.gather_mask.shape
+      + (1,) * len(features_shape)
+  )
+  out_arr *= gather_info.gather_mask.reshape(broadcasted_mask_shape)
+  return out_arr
+
+
+def make_structure(
+    flat_layout: AtomLayout,
+    atom_coords: np.ndarray,
+    name: str,
+    *,
+    atom_b_factors: np.ndarray | None = None,
+    all_physical_residues: Residues | None = None,
+) -> structure.Structure:
+  """Returns a Structure from a flat layout and atom coordinates.
+
+  The provided flat_layout must be 1-dim and must not contain any padding
+  elements. The flat_layout.atom_name must conform to the OpenMM/CCD standard
+  and must not contain deuterium.
+
+  Args:
+    flat_layout: flat 1-dim AtomLayout without pading elements
+    atom_coords: np.ndarray of float, shape (num_atoms, 3)
+    name: str: the name (usually PDB id), e.g. '1uao'
+    atom_b_factors: np.ndarray of float, shape (num_atoms,) or None. If None,
+      they will be set to all zeros.
+    all_physical_residues: a Residues object that contains all physically
+      existing residues, i.e. also those residues that have no resolved atoms.
+      This is common in experimental structures, but also appears in predicted
+      structures for 'UNK' or other non-standard residue types, where the model
+      does not predict coordinates. This will be used to create the
+      `all_residues` field of the structure object.
+  """
+
+  if flat_layout.atom_name.ndim != 1 or not np.all(
+      flat_layout.atom_name.astype(bool)
+  ):
+    raise ValueError(
+        'flat_layout must be 1-dim and must not contain anypadding element'
+    )
+  if (
+      flat_layout.atom_element is None
+      or flat_layout.res_name is None
+      or flat_layout.chain_type is None
+  ):
+    raise ValueError('All optional fields must be present.')
+
+  if atom_b_factors is None:
+    atom_b_factors = np.zeros(atom_coords.shape[:-1])
+
+  if all_physical_residues is not None:
+    # Create the all_residues field from a Residues object
+    # (unfortunately there is no central place to keep the chain_types in
+    # the structure class, so we drop it here)
+    all_residues = collections.defaultdict(list)
+    for chain_id, res_id, res_name in zip(
+        all_physical_residues.chain_id,
+        all_physical_residues.res_id,
+        all_physical_residues.res_name,
+        strict=True,
+    ):
+      all_residues[chain_id].append((res_name, res_id))
+  else:
+    # Create the all_residues field from the flat_layout
+    all_residues = collections.defaultdict(list)
+    if flat_layout.chain_id.shape[0] > 0:
+      all_residues[flat_layout.chain_id[0]].append(
+          (flat_layout.res_name[0], flat_layout.res_id[0])
+      )
+      for i in range(1, flat_layout.shape[0]):
+        if (
+            flat_layout.chain_id[i] != flat_layout.chain_id[i - 1]
+            or flat_layout.res_name[i] != flat_layout.res_name[i - 1]
+            or flat_layout.res_id[i] != flat_layout.res_id[i - 1]
+        ):
+          all_residues[flat_layout.chain_id[i]].append(
+              (flat_layout.res_name[i], flat_layout.res_id[i])
+          )
+
+  return structure.from_atom_arrays(
+      name=name,
+      all_residues=dict(all_residues),
+      chain_id=flat_layout.chain_id,
+      chain_type=flat_layout.chain_type,
+      res_id=flat_layout.res_id.astype(np.int32),
+      res_name=flat_layout.res_name,
+      atom_name=flat_layout.atom_name,
+      atom_element=flat_layout.atom_element,
+      atom_x=atom_coords[..., 0],
+      atom_y=atom_coords[..., 1],
+      atom_z=atom_coords[..., 2],
+      atom_b_factor=atom_b_factors,
+  )
+
+
+def _assert_all_arrays_have_same_shape(
+    *,
+    obj: AtomLayout | Residues | GatherInfo,
+    expected_shape: tuple[int, ...],
+    attribute_names: Sequence[str],
+) -> None:
+  """Checks that given attributes of the object have the expected shape."""
+  attribute_shapes_description = []
+  all_shapes_are_valid = True
+
+  for attribute_name in attribute_names:
+    attribute = getattr(obj, attribute_name)
+
+    if attribute is None:
+      attribute_shape = None
+    else:
+      attribute_shape = attribute.shape
+
+    if attribute_shape is not None and expected_shape != attribute_shape:
+      all_shapes_are_valid = False
+
+    attribute_shape_name = attribute_name + '.shape'
+    attribute_shapes_description.append(
+        f'{attribute_shape_name:25} = {attribute_shape}'
+    )
+
+  if not all_shapes_are_valid:
+    raise ValueError(
+        f'All arrays must have the same shape ({expected_shape=}). Got\n'
+        + '\n'.join(attribute_shapes_description)
+    )

flax_model/alphafold3/model/components/haiku_modules.py

@@ -0,0 +1,328 @@
+
+
+"""Common Haiku modules."""
+
+from collections.abc import Sequence
+import contextlib
+import numbers
+from typing import TypeAlias
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+PRECISION: TypeAlias = (
+    None
+    | str
+    | jax.lax.Precision
+    | tuple[str, str]
+    | tuple[jax.lax.Precision, jax.lax.Precision]
+)
+
+# Useful for mocking in tests.
+DEFAULT_PRECISION = None
+
+# Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
+TRUNCATED_NORMAL_STDDEV_FACTOR = np.asarray(
+    0.87962566103423978, dtype=np.float32
+)
+
+
+class LayerNorm(hk.LayerNorm):
+  """LayerNorm module.
+
+  Equivalent to hk.LayerNorm but with an extra 'upcast' option that casts
+  (b)float16 inputs to float32 before computing the layer norm, and then casts
+  the output back to the input type.
+
+  The learnable parameter shapes are also different from Haiku: they are always
+  vectors rather than possibly higher-rank tensors. This makes it easier
+  to change the layout whilst keep the model weight-compatible.
+  """
+
+  def __init__(
+      self,
+      *,
+      axis: int = -1,
+      create_scale: bool = True,
+      create_offset: bool = True,
+      eps: float = 1e-5,
+      scale_init: hk.initializers.Initializer | None = None,
+      offset_init: hk.initializers.Initializer | None = None,
+      use_fast_variance: bool = True,
+      name: str,
+      param_axis: int | None = None,
+      upcast: bool = True,
+  ):
+    super().__init__(
+        axis=axis,
+        create_scale=False,
+        create_offset=False,
+        eps=eps,
+        scale_init=None,
+        offset_init=None,
+        use_fast_variance=use_fast_variance,
+        name=name,
+        param_axis=param_axis,
+    )
+    self.upcast = upcast
+    self._temp_create_scale = create_scale
+    self._temp_create_offset = create_offset
+
+  def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+    dtype = x.dtype
+    is_16bit = x.dtype in [jnp.bfloat16, jnp.float16]
+    if self.upcast and is_16bit:
+      x = x.astype(jnp.float32)
+
+    param_axis = self.param_axis[0] if self.param_axis else -1
+    param_shape = (x.shape[param_axis],)
+
+    param_broadcast_shape = [1] * x.ndim
+    param_broadcast_shape[param_axis] = x.shape[param_axis]
+    scale = None
+    offset = None
+    if self._temp_create_scale:
+      scale = hk.get_parameter(
+          'scale', param_shape, x.dtype, init=self.scale_init
+      )
+      scale = scale.reshape(param_broadcast_shape)
+
+    if self._temp_create_offset:
+      offset = hk.get_parameter(
+          'offset', param_shape, x.dtype, init=self.offset_init
+      )
+      offset = offset.reshape(param_broadcast_shape)
+
+    out = super().__call__(x, scale=scale, offset=offset)
+
+    if self.upcast and is_16bit:
+      out = out.astype(dtype)
+
+    return out
+
+
+def haiku_linear_get_params(
+    inputs: jax.Array | jax.ShapeDtypeStruct,
+    *,
+    num_output: int | Sequence[int],
+    use_bias: bool = False,
+    num_input_dims: int = 1,
+    initializer: str = 'linear',
+    bias_init: float = 0.0,
+    transpose_weights: bool = False,
+    name: str | None = None,
+) -> tuple[jax.Array, jax.Array | None]:
+  """Get parameters for linear layer.
+
+  Parameters will be at least float32 or higher precision.
+
+  Arguments:
+    inputs: The input to the Linear layer. Can be either a JAX array or a
+      jax.ShapeDtypeStruct.
+    num_output: The number of output channels. Can be an integer or a sequence
+      of integers.
+    use_bias: Whether to create a bias array.
+    num_input_dims: The number of dimensions to consider as channel dims in the
+      input.
+    initializer: The name of the weight initializer to use.
+    bias_init: A float used to initialize the bias.
+    transpose_weights: If True, will create a transposed version of the weights.
+    name: The Haiku namespace to use for the weight and bias.
+
+  Returns:
+    A tuple[weight, bias] if use_bias otherwise tuple[weight, None].
+  """
+
+  if isinstance(num_output, numbers.Integral):
+    output_shape = (num_output,)
+  else:
+    output_shape = tuple(num_output)
+
+  if num_input_dims > 0:
+    in_shape = inputs.shape[-num_input_dims:]
+  elif num_input_dims == 0:
+    in_shape = ()
+  else:
+    raise ValueError('num_input_dims must be >= 0.')
+
+  weight_init = _get_initializer_scale(initializer, in_shape)
+  with hk.name_scope(name) if name else contextlib.nullcontext():
+
+    if transpose_weights:
+      weight_shape = output_shape + in_shape
+
+      weights = hk.get_parameter(
+          'weights', shape=weight_shape, dtype=inputs.dtype, init=weight_init
+      )
+    else:
+      weight_shape = in_shape + output_shape
+      weights = hk.get_parameter(
+          name='weights',
+          shape=weight_shape,
+          dtype=inputs.dtype,
+          init=weight_init,
+      )
+
+    bias = None
+    if use_bias:
+      bias = hk.get_parameter(
+          name='bias',
+          shape=output_shape,
+          dtype=inputs.dtype,
+          init=hk.initializers.Constant(bias_init),
+      )
+  return weights, bias
+
+
+class Linear(hk.Module):
+  """Custom Linear Module.
+
+  This differs from the standard Linear in a few ways:
+    * It supports inputs of arbitrary rank
+    * It allows to use ntk parametrization
+    * Initializers are specified by strings
+    * It allows to explicitly specify which dimension of the input will map to
+      the tpu sublane/lane dimensions.
+  """
+
+  def __init__(
+      self,
+      num_output: int | Sequence[int],
+      *,
+      initializer: str = 'linear',
+      num_input_dims: int = 1,
+      use_bias: bool = False,
+      bias_init: float = 0.0,
+      precision: PRECISION = None,
+      fast_scalar_mode: bool = True,
+      transpose_weights: bool = False,
+      name: str,
+  ):
+    """Constructs Linear Module.
+
+    Args:
+      num_output: number of output channels. Can be tuple when outputting
+        multiple dimensions.
+      initializer: What initializer to use, should be one of {'linear', 'relu',
+        'zeros'}.
+      num_input_dims: Number of dimensions from the end to project.
+      use_bias: Whether to include trainable bias (False by default).
+      bias_init: Value used to initialize bias.
+      precision: What precision to use for matrix multiplication, defaults to
+        None.
+      fast_scalar_mode: Whether to use optimized path for num_input_dims = 0.
+      transpose_weights: decides whether weights have shape [input, output] or
+        [output, input], True means [output, input], this is helpful to avoid
+        padding on the tensors holding the weights.
+      name: name of module, used for name scopes.
+    """
+    super().__init__(name=name)
+    if isinstance(num_output, numbers.Integral):
+      self.output_shape = (num_output,)
+    else:
+      self.output_shape = tuple(num_output)
+    self.initializer = initializer
+    self.use_bias = use_bias
+    self.bias_init = bias_init
+    self.num_input_dims = num_input_dims
+    self.num_output_dims = len(self.output_shape)
+    self.precision = precision if precision is not None else DEFAULT_PRECISION
+    self.fast_scalar_mode = fast_scalar_mode
+    self.transpose_weights = transpose_weights
+
+  def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
+    """Connects Module.
+
+    Args:
+      inputs: Tensor of shape [..., num_channel]
+
+    Returns:
+      output of shape [..., num_output]
+    """
+
+    num_input_dims = self.num_input_dims
+
+    # Adds specialized path for scalar inputs in Linear layer,
+    # this means the linear Layer does not use the matmul units on the tpu,
+    # which is more efficient and gives compiler more flexibility over layout.
+    if num_input_dims == 0 and self.fast_scalar_mode:
+      weight_shape = self.output_shape
+      if self.initializer == 'zeros':
+        w_init = hk.initializers.Constant(0.0)
+      else:
+        distribution_stddev = jnp.array(1 / TRUNCATED_NORMAL_STDDEV_FACTOR)
+        w_init = hk.initializers.TruncatedNormal(
+            mean=0.0, stddev=distribution_stddev
+        )
+
+      weights = hk.get_parameter('weights', weight_shape, inputs.dtype, w_init)
+
+      inputs = jnp.expand_dims(
+          inputs, tuple(range(-1, -self.num_output_dims - 1, -1))
+      )
+      output = inputs * weights
+    else:
+      if self.num_input_dims > 0:
+        in_shape = inputs.shape[-self.num_input_dims :]
+      else:
+        in_shape = ()
+
+      weight_init = _get_initializer_scale(self.initializer, in_shape)
+
+      in_letters = 'abcde'[: self.num_input_dims]
+      out_letters = 'hijkl'[: self.num_output_dims]
+
+      if self.transpose_weights:
+        weight_shape = self.output_shape + in_shape
+        weights = hk.get_parameter(
+            'weights', weight_shape, inputs.dtype, weight_init
+        )
+        equation = (
+            f'...{in_letters}, {out_letters}{in_letters}->...{out_letters}'
+        )
+      else:
+        weight_shape = in_shape + self.output_shape
+        weights = hk.get_parameter(
+            'weights', weight_shape, inputs.dtype, weight_init
+        )
+
+        equation = (
+            f'...{in_letters}, {in_letters}{out_letters}->...{out_letters}'
+        )
+
+      output = jnp.einsum(equation, inputs, weights, precision=self.precision)
+
+    if self.use_bias:
+      bias = hk.get_parameter(
+          'bias',
+          self.output_shape,
+          inputs.dtype,
+          hk.initializers.Constant(self.bias_init),
+      )
+      output += bias
+
+    return output
+
+
+def _get_initializer_scale(initializer_name, input_shape):
+  """Get initializer for weights."""
+
+  if initializer_name == 'zeros':
+    w_init = hk.initializers.Constant(0.0)
+  else:
+    # fan-in scaling
+    noise_scale = 1.0
+    for channel_dim in input_shape:
+      noise_scale /= channel_dim
+    if initializer_name == 'relu':
+      noise_scale *= 2
+
+    stddev = np.sqrt(noise_scale)
+    # Adjust stddev for truncation.
+    stddev = stddev / TRUNCATED_NORMAL_STDDEV_FACTOR
+    w_init = hk.initializers.TruncatedNormal(mean=0.0, stddev=stddev)
+
+  return w_init

flax_model/alphafold3/model/components/mapping.py

@@ -0,0 +1,240 @@
+
+
+"""Specialized mapping functions."""
+
+from collections.abc import Callable, Sequence
+import functools
+from typing import Any, TypeVar
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+Pytree = Any
+PytreeJaxArray = Any
+
+partial = functools.partial
+PROXY = object()
+
+T = TypeVar("T")
+
+
+def _maybe_slice(array, i, slice_size, axis):
+  if axis is PROXY:
+    return array
+  else:
+    return jax.lax.dynamic_slice_in_dim(
+        array, i, slice_size=slice_size, axis=axis
+    )
+
+
+def _maybe_get_size(array, axis):
+  if axis == PROXY:
+    return -1
+  else:
+    return array.shape[axis]
+
+
+def _expand_axes(axes, values, name="sharded_apply"):
+  values_tree_def = jax.tree_util.tree_structure(values)
+  flat_axes = jax.api_util.flatten_axes(name, values_tree_def, axes)
+  # Replace None's with PROXY.
+  flat_axes = [PROXY if x is None else x for x in flat_axes]
+  return jax.tree_util.tree_unflatten(values_tree_def, flat_axes)
+
+
+def sharded_map(
+    fun: Callable[..., PytreeJaxArray],
+    shard_size: int | None = 1,
+    in_axes: int | Pytree = 0,
+    out_axes: int | Pytree = 0,
+) -> Callable[..., PytreeJaxArray]:
+  """Sharded vmap.
+
+  Maps `fun` over axes, in a way similar to hk.vmap, but does so in shards of
+  `shard_size`. This allows a smooth trade-off between memory usage
+  (as in a plain map) vs higher throughput (as in a vmap).
+
+  Args:
+    fun: Function to apply smap transform to.
+    shard_size: Integer denoting shard size.
+    in_axes: Either integer or pytree describing which axis to map over for each
+      input to `fun`, None denotes broadcasting.
+    out_axes: Integer or pytree denoting to what axis in the output the mapped
+      over axis maps.
+
+  Returns:
+    Function with smap applied.
+  """
+  if hk.running_init():
+    # Guarantees initialisation independent of shard_size. Doesn't incur a high
+    # memory cost, as long as large concrete tensors are not encountered.
+    return hk.vmap(fun, in_axes=in_axes, out_axes=out_axes, split_rng=False)
+  else:
+    vmapped_fun = hk.vmap(fun, in_axes, out_axes, split_rng=True)
+    return sharded_apply(vmapped_fun, shard_size, in_axes, out_axes)
+
+
+def _set_docstring(docstr: str) -> Callable[[T], T]:
+  """Decorator for setting the docstring of a function."""
+
+  def wrapped(fun: T) -> T:
+    fun.__doc__ = docstr.format(fun=getattr(fun, "__name__", repr(fun)))
+    return fun
+
+  return wrapped
+
+
+def sharded_apply(
+    fun: Callable[..., PytreeJaxArray],
+    shard_size: int | None = 1,
+    in_axes: int | Pytree = 0,
+    out_axes: int | Pytree = 0,
+    new_out_axes: bool = False,
+) -> Callable[..., PytreeJaxArray]:
+  """Sharded apply.
+
+  Applies `fun` over shards to axes, in a way similar to vmap,
+  but does so in shards of `shard_size`. Shards are stacked after.
+  This allows a smooth trade-off between
+  memory usage (as in a plain map) vs higher throughput (as in a vmap).
+
+  Args:
+    fun: Function to apply smap transform to.
+    shard_size: Integer denoting shard size. None will return `fun` unchanged.
+    in_axes: Either integer or pytree describing which axis to map over for each
+      input to `fun`, None denotes broadcasting.
+    out_axes: Integer or pytree denoting to what axis in the output the mapped
+      over axis maps.
+    new_out_axes: Whether to stack outputs on new axes. This assumes that the
+      output sizes for each shard (including the possible remainder shard) are
+      the same.
+
+  Returns:
+    Function with smap applied.
+  """
+  docstr = (
+      "Mapped version of {fun}. Takes similar arguments to {fun} "
+      "but with additional array axes over which {fun} is mapped."
+  )
+  if new_out_axes:
+    raise NotImplementedError("New output axes not yet implemented.")
+
+  if shard_size is None:
+    return fun
+
+  @_set_docstring(docstr)
+  @functools.wraps(fun)
+  def mapped_fn(*args, **kwargs):
+    # Expand in axes and determine loop range.
+    in_axes_ = _expand_axes(in_axes, args)
+
+    in_sizes = jax.tree.map(_maybe_get_size, args, in_axes_)
+    in_size = max(jax.tree_util.tree_leaves(in_sizes))
+
+    num_extra_shards = (in_size - 1) // shard_size
+
+    # Fix if necessary.
+    last_shard_size = in_size % shard_size
+    last_shard_size = shard_size if last_shard_size == 0 else last_shard_size
+
+    def apply_fun_to_slice(slice_start, slice_size):
+      input_slice = jax.tree.map(
+          lambda array, axis: _maybe_slice(
+              array, slice_start, slice_size, axis
+          ),
+          args,
+          in_axes_,
+      )
+      return fun(*input_slice, **kwargs)
+
+    remainder_shape_dtype = hk.eval_shape(
+        partial(apply_fun_to_slice, 0, last_shard_size)
+    )
+    out_dtypes = jax.tree.map(lambda x: x.dtype, remainder_shape_dtype)
+    out_shapes = jax.tree.map(lambda x: x.shape, remainder_shape_dtype)
+    out_axes_ = _expand_axes(out_axes, remainder_shape_dtype)
+
+    if num_extra_shards > 0:
+      regular_shard_shape_dtype = hk.eval_shape(
+          partial(apply_fun_to_slice, 0, shard_size)
+      )
+      shard_shapes = jax.tree.map(lambda x: x.shape, regular_shard_shape_dtype)
+
+      def make_output_shape(axis, shard_shape, remainder_shape):
+        return (
+            shard_shape[:axis]
+            + (shard_shape[axis] * num_extra_shards + remainder_shape[axis],)
+            + shard_shape[axis + 1 :]
+        )
+
+      out_shapes = jax.tree.map(
+          make_output_shape, out_axes_, shard_shapes, out_shapes
+      )
+
+    # Calls dynamic Update slice with different argument order.
+    # This is here since tree_map only works with positional arguments.
+    def dynamic_update_slice_in_dim(full_array, update, axis, i):
+      return jax.lax.dynamic_update_slice_in_dim(full_array, update, i, axis)
+
+    def compute_shard(outputs, slice_start, slice_size):
+      slice_out = apply_fun_to_slice(slice_start, slice_size)
+      update_slice = partial(dynamic_update_slice_in_dim, i=slice_start)
+      return jax.tree.map(update_slice, outputs, slice_out, out_axes_)
+
+    def scan_iteration(outputs, i):
+      new_outputs = compute_shard(outputs, i, shard_size)
+      return new_outputs, ()
+
+    slice_starts = jnp.arange(0, in_size - shard_size + 1, shard_size)
+
+    def allocate_buffer(dtype, shape):
+      return jnp.zeros(shape, dtype=dtype)
+
+    outputs = jax.tree.map(allocate_buffer, out_dtypes, out_shapes)
+
+    if slice_starts.shape[0] > 0:
+      outputs, _ = hk.scan(scan_iteration, outputs, slice_starts)
+
+    if last_shard_size != shard_size:
+      remainder_start = in_size - last_shard_size
+      outputs = compute_shard(outputs, remainder_start, last_shard_size)
+
+    return outputs
+
+  return mapped_fn
+
+
+def inference_subbatch(
+    module: Callable[..., PytreeJaxArray],
+    subbatch_size: int,
+    batched_args: Sequence[PytreeJaxArray],
+    nonbatched_args: Sequence[PytreeJaxArray],
+    input_subbatch_dim: int = 0,
+    output_subbatch_dim: int | None = None,
+) -> PytreeJaxArray:
+  """Run through subbatches (like batch apply but with split and concat)."""
+  assert len(batched_args) > 0  # pylint: disable=g-explicit-length-test
+
+  if hk.running_init():
+    args = list(batched_args) + list(nonbatched_args)
+    return module(*args)
+
+  if output_subbatch_dim is None:
+    output_subbatch_dim = input_subbatch_dim
+
+  def run_module(*batched_args):
+    args = list(batched_args) + list(nonbatched_args)
+    res = module(*args)
+    return res
+
+  sharded_module = sharded_apply(
+      run_module,
+      shard_size=subbatch_size,
+      in_axes=input_subbatch_dim,
+      out_axes=output_subbatch_dim,
+  )
+  output = sharded_module(*batched_args)
+
+  return output

flax_model/alphafold3/model/components/utils.py

@@ -0,0 +1,75 @@
+
+
+"""Utility functions for training AlphaFold and similar models."""
+
+from collections import abc
+import contextlib
+import numbers
+
+from flax_model.alphafold3.model import features
+import haiku as hk
+import jax.numpy as jnp
+import numpy as np
+
+
+VALID_DTYPES = [np.float32, np.float64, np.int8, np.int32, np.int64, bool]
+
+
+def remove_invalidly_typed_feats(
+    batch: features.BatchDict,
+) -> features.BatchDict:
+  """Remove features of types we don't want to send to the TPU e.g. strings."""
+  return {
+      k: v
+      for k, v in batch.items()
+      if hasattr(v, 'dtype') and v.dtype in VALID_DTYPES
+  }
+
+
+def bfloat16_getter(next_getter, value, context):
+  """Ensures that a bfloat16 parameter is provided by casting if necessary."""
+  if context.original_dtype == jnp.bfloat16:
+    if value.dtype != jnp.bfloat16:
+      value = value.astype(jnp.bfloat16)
+  return next_getter(value)
+
+
+@contextlib.contextmanager
+def bfloat16_context():
+  with hk.custom_getter(bfloat16_getter):
+    yield
+
+
+def mask_mean(mask, value, axis=None, keepdims=False, eps=1e-10):
+  """Masked mean."""
+
+  mask_shape = mask.shape
+  value_shape = value.shape
+
+  assert len(mask_shape) == len(
+      value_shape
+  ), 'Shapes are not compatible, shapes: {}, {}'.format(mask_shape, value_shape)
+
+  if isinstance(axis, numbers.Integral):
+    axis = [axis]
+  elif axis is None:
+    axis = list(range(len(mask_shape)))
+  assert isinstance(
+      axis, abc.Iterable
+  ), 'axis needs to be either an iterable, integer or "None"'
+
+  broadcast_factor = 1.0
+  for axis_ in axis:
+    value_size = value_shape[axis_]
+    mask_size = mask_shape[axis_]
+    if mask_size == 1:
+      broadcast_factor *= value_size
+    else:
+      error = f'Shapes are not compatible, shapes: {mask_shape}, {value_shape}'
+      assert mask_size == value_size, error
+
+  return jnp.sum(mask * value, keepdims=keepdims, axis=axis) / (
+      jnp.maximum(
+          jnp.sum(mask, keepdims=keepdims, axis=axis) * broadcast_factor, eps
+      )
+  )

flax_model/alphafold3/model/confidence_types.py

@@ -0,0 +1,293 @@
+
+
+"""Confidence categories for predictions."""
+
+import dataclasses
+import enum
+import json
+from typing import Any, Self
+
+from absl import logging
+from flax_model.alphafold3.model import model
+import jax
+import numpy as np
+
+
+class StructureConfidenceFullEncoder(json.JSONEncoder):
+  """JSON encoder for serializing confidence types."""
+
+  def __init__(self, **kwargs):
+    super().__init__(**(kwargs | dict(separators=(',', ':'))))
+
+  def encode(self, o: 'StructureConfidenceFull'):
+    # Cast to np.float64 before rounding, since casting to Python float will
+    # cast to a 64 bit float, potentially undoing np.float32 rounding.
+    atom_plddts = np.round(
+        np.clip(np.asarray(o.atom_plddts, dtype=np.float64), 0.0, 99.99), 2
+    ).astype(float)
+    contact_probs = np.round(
+        np.clip(np.asarray(o.contact_probs, dtype=np.float64), 0.0, 1.0), 2
+    ).astype(float)
+    pae = np.round(
+        np.clip(np.asarray(o.pae, dtype=np.float64), 0.0, 99.9), 1
+    ).astype(float)
+    return """\
+{
+  "atom_chain_ids": %s,
+  "atom_plddts": %s,
+  "contact_probs": %s,
+  "pae": %s,
+  "token_chain_ids": %s,
+  "token_res_ids": %s
+}""" % (
+        super().encode(o.atom_chain_ids),
+        super().encode(list(atom_plddts)).replace('NaN', 'null'),
+        super().encode([list(x) for x in contact_probs]).replace('NaN', 'null'),
+        super().encode([list(x) for x in pae]).replace('NaN', 'null'),
+        super().encode(o.token_chain_ids),
+        super().encode(o.token_res_ids),
+    )
+
+
+def _dump_json(data: Any, indent: int | None = None) -> str:
+  """Dumps a json string with JSON compatible NaN representation."""
+  json_str = json.dumps(
+      data,
+      sort_keys=True,
+      indent=indent,
+      separators=(',', ': '),
+  )
+  return json_str.replace('NaN', 'null')
+
+
+@enum.unique
+class ConfidenceCategory(enum.Enum):
+  """Confidence categories for AlphaFold predictions."""
+
+  HIGH = 0
+  MEDIUM = 1
+  LOW = 2
+  DISORDERED = 3
+
+  @classmethod
+  def from_char(cls, char: str) -> Self:
+    match char:
+      case 'H':
+        return cls.HIGH
+      case 'M':
+        return cls.MEDIUM
+      case 'L':
+        return cls.LOW
+      case 'D':
+        return cls.DISORDERED
+      case _:
+        raise ValueError(
+            f'Unknown character. Expected one of H, M, L or D; got: {char}'
+        )
+
+  def to_char(self) -> str:
+    match self:
+      case self.HIGH:
+        return 'H'
+      case self.MEDIUM:
+        return 'M'
+      case self.LOW:
+        return 'L'
+      case self.DISORDERED:
+        return 'D'
+
+  @classmethod
+  def from_confidence_score(cls, confidence: float) -> Self:
+    if 90 <= confidence <= 100:
+      return cls.HIGH
+    if 70 <= confidence < 90:
+      return cls.MEDIUM
+    if 50 <= confidence < 70:
+      return cls.LOW
+    if 0 <= confidence < 50:
+      return cls.DISORDERED
+    raise ValueError(f'Confidence score out of range [0, 100]: {confidence}')
+
+
+@dataclasses.dataclass()
+class AtomConfidence:
+  """Dataclass for 1D per-atom confidences from AlphaFold."""
+
+  chain_id: list[str]
+  atom_number: list[int]
+  confidence: list[float]
+  confidence_category: list[ConfidenceCategory]
+
+  def __post_init__(self):
+    num_res = len(self.atom_number)
+    if not all(
+        len(v) == num_res
+        for v in [self.chain_id, self.confidence, self.confidence_category]
+    ):
+      raise ValueError('All confidence fields must have the same length.')
+
+  @classmethod
+  def from_inference_result(
+      cls, inference_result: model.InferenceResult
+  ) -> Self:
+    """Instantiates an AtomConfidence from a structure.
+
+    Args:
+      inference_result: Inference result from AlphaFold.
+
+    Returns:
+      Scores in AtomConfidence dataclass.
+    """
+    struc = inference_result.predicted_structure
+    as_dict = {
+        'chain_id': [],
+        'atom_number': [],
+        'confidence': [],
+        'confidence_category': [],
+    }
+    for atom_number, atom in enumerate(struc.iter_atoms()):
+      this_confidence = float(struc.atom_b_factor[atom_number])
+      as_dict['chain_id'].append(atom['chain_id'])
+      as_dict['atom_number'].append(atom_number)
+      as_dict['confidence'].append(round(this_confidence, 2))
+      as_dict['confidence_category'].append(
+          ConfidenceCategory.from_confidence_score(this_confidence)
+      )
+    return cls(**as_dict)
+
+  @classmethod
+  def from_json(cls, json_string: str) -> Self:
+    """Instantiates a AtomConfidence from a json string."""
+    input_dict = json.loads(json_string)
+    input_dict['confidence_category'] = [
+        ConfidenceCategory.from_char(k)
+        for k in input_dict['confidence_category']
+    ]
+    return cls(**input_dict)
+
+  def to_json(self) -> str:
+    output = dataclasses.asdict(self)
+    output['confidence_category'] = [
+        k.to_char() for k in output['confidence_category']
+    ]
+    output['atom_number'] = [int(k) for k in output['atom_number']]
+    return _dump_json(output)
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class StructureConfidenceSummary:
+  """Dataclass for the summary of structure scores from AlphaFold.
+
+  Attributes:
+   ptm: Predicted TM global score.
+   iptm: Interface predicted TM global score.
+   ranking_score: Ranking score extracted from CIF metadata.
+   fraction_disordered: Fraction disordered, measured with RASA.
+   has_clash: Has significant clashing.
+   chain_pair_pae_min: [num_chains, num_chains] Minimum cross chain PAE.
+   chain_pair_iptm: [num_chains, num_chains] Chain pair ipTM.
+   chain_ptm: [num_chains] Chain pTM.
+   chain_iptm: [num_chains] Mean cross chain ipTM for a chain.
+  """
+
+  ptm: float
+  iptm: float
+  ranking_score: float
+  fraction_disordered: float
+  has_clash: float
+  chain_pair_pae_min: np.ndarray
+  chain_pair_iptm: np.ndarray
+  chain_ptm: np.ndarray
+  chain_iptm: np.ndarray
+
+  @classmethod
+  def from_inference_result(
+      cls, inference_result: model.InferenceResult
+  ) -> Self:
+    """Returns a new instance based on a given inference result."""
+    return cls(
+        ptm=float(inference_result.metadata['ptm']),
+        iptm=float(inference_result.metadata['iptm']),
+        ranking_score=float(inference_result.metadata['ranking_score']),
+        fraction_disordered=float(
+            inference_result.metadata['fraction_disordered']
+        ),
+        has_clash=float(inference_result.metadata['has_clash']),
+        chain_pair_pae_min=inference_result.metadata['chain_pair_pae_min'],
+        chain_pair_iptm=inference_result.metadata['chain_pair_iptm'],
+        chain_ptm=inference_result.metadata['iptm_ichain'],
+        chain_iptm=inference_result.metadata['iptm_xchain'],
+    )
+
+  @classmethod
+  def from_json(cls, json_string: str) -> Self:
+    """Returns a new instance from a given json string."""
+    return cls(**json.loads(json_string))
+
+  def to_json(self) -> str:
+    def convert(data):
+      if isinstance(data, np.ndarray):
+        # Cast to np.float64 before rounding, since casting to Python float will
+        # cast to a 64 bit float, potentially undoing np.float32 rounding.
+        rounded_data = np.round(data.astype(np.float64), decimals=2).tolist()
+      else:
+        rounded_data = np.round(data, decimals=2)
+      return rounded_data
+
+    return _dump_json(jax.tree.map(convert, dataclasses.asdict(self)), indent=1)
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class StructureConfidenceFull:
+  """Dataclass for full structure data from AlphaFold."""
+
+  pae: np.ndarray
+  token_chain_ids: list[str]
+  token_res_ids: list[int]
+  atom_plddts: list[float]
+  atom_chain_ids: list[str]
+  contact_probs: np.ndarray  # [num_tokens, num_tokens]
+
+  @classmethod
+  def from_inference_result(
+      cls, inference_result: model.InferenceResult
+  ) -> Self:
+    """Returns a new instance based on a given inference result."""
+
+    pae = inference_result.numerical_data['full_pae']
+    if not isinstance(pae, np.ndarray):
+      logging.info('%s', type(pae))
+      raise TypeError('pae should be a numpy array.')
+
+    contact_probs = inference_result.numerical_data['contact_probs']
+    if not isinstance(contact_probs, np.ndarray):
+      logging.info('%s', type(contact_probs))
+      raise TypeError('contact_probs should be a numpy array.')
+
+    struc = inference_result.predicted_structure
+    chain_ids = struc.chain_id.tolist()
+    atom_plddts = struc.atom_b_factor.tolist()
+    token_chain_ids = [
+        str(token_id)
+        for token_id in inference_result.metadata['token_chain_ids']
+    ]
+    token_res_ids = [
+        int(token_id) for token_id in inference_result.metadata['token_res_ids']
+    ]
+    return cls(
+        pae=pae,
+        token_chain_ids=token_chain_ids,
+        token_res_ids=token_res_ids,
+        atom_plddts=atom_plddts,
+        atom_chain_ids=chain_ids,
+        contact_probs=contact_probs,
+    )
+
+  @classmethod
+  def from_json(cls, json_string: str) -> Self:
+    """Returns a new instance from a given json string."""
+    return cls(**json.loads(json_string))
+
+  def to_json(self) -> str:
+    """Converts StructureConfidenceFull to json string."""
+    return json.dumps(self, cls=StructureConfidenceFullEncoder)

flax_model/alphafold3/model/confidences.py

@@ -0,0 +1,655 @@
+
+
+"""Functions for extracting and processing confidences from model outputs."""
+import warnings
+
+from absl import logging
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import mkdssp
+import jax.numpy as jnp
+import numpy as np
+from scipy import spatial
+
+
+# From Sander & Rost 1994 https://doi.org/10.1002/prot.340200303
+MAX_ACCESSIBLE_SURFACE_AREA = {
+    'ALA': 106.0,
+    'ARG': 248.0,
+    'ASN': 157.0,
+    'ASP': 163.0,
+    'CYS': 135.0,
+    'GLN': 198.0,
+    'GLU': 194.0,
+    'GLY': 84.0,
+    'HIS': 184.0,
+    'ILE': 169.0,
+    'LEU': 164.0,
+    'LYS': 205.0,
+    'MET': 188.0,
+    'PHE': 197.0,
+    'PRO': 136.0,
+    'SER': 130.0,
+    'THR': 142.0,
+    'TRP': 227.0,
+    'TYR': 222.0,
+    'VAL': 142.0,
+}
+
+# Weights for ranking confidence.
+_IPTM_WEIGHT = 0.8
+_FRACTION_DISORDERED_WEIGHT = 0.5
+_CLASH_PENALIZATION_WEIGHT = 100.0
+
+
+def windowed_solvent_accessible_area(cif: str, window: int = 25) -> np.ndarray:
+  """Implementation of AlphaFold-RSA.
+
+  AlphaFold-RSA defined in https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9601767.
+
+  Args:
+    cif: Raw cif string.
+    window: The window over which to average accessible surface area
+
+  Returns:
+    An array of size num_res that predicts disorder by using windowed solvent
+    accessible surface area.
+  """
+  result = mkdssp.get_dssp(cif, calculate_surface_accessibility=True)
+  parse_row = False
+  rasa = []
+  for row in result.splitlines():
+    if parse_row:
+      aa = row[13:14]
+      if aa == '!':
+        continue
+      aa3 = residue_names.PROTEIN_COMMON_ONE_TO_THREE.get(aa, 'ALA')
+      max_acc = MAX_ACCESSIBLE_SURFACE_AREA[aa3]
+      acc = int(row[34:38])
+      norm_acc = acc / max_acc
+      if norm_acc > 1.0:
+        norm_acc = 1.0
+      rasa.append(norm_acc)
+    if row.startswith('  #  RESIDUE'):
+      parse_row = True
+
+  half_w = (window - 1) // 2
+  pad_rasa = np.pad(rasa, (half_w, half_w), 'reflect')
+  rasa = np.convolve(pad_rasa, np.ones(window), 'valid') / window
+  return rasa
+
+
+def fraction_disordered(
+    struc: structure.Structure, rasa_disorder_cutoff: float = 0.581
+) -> float:
+  """Compute fraction of protein residues that are disordered.
+
+  Args:
+    struc: A structure to compute rASA metrics on.
+    rasa_disorder_cutoff: The threshold at which residues are considered
+      disordered. Default value taken from
+      https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9601767.
+
+  Returns:
+    The fraction of protein residues that are disordered
+    (rasa > rasa_disorder_cutoff).
+  """
+  struc = struc.filter_to_entity_type(protein=True)
+  rasa = []
+  seq_rasa = {}
+  for chain_id, chain_seq in struc.chain_single_letter_sequence().items():
+    if chain_seq in seq_rasa:
+      # We assume that identical sequences have approximately similar rasa
+      # values to speed up the computation.
+      rasa.extend(seq_rasa[chain_seq])
+      continue
+    chain_struc = struc.filter(chain_id=chain_id)
+    # Rename the chain to 'A' as MKDSSP supports only single letter chain IDs.
+    chain_struc = chain_struc.rename_chain_ids(new_id_by_old_id={chain_id: 'A'})
+    try:
+      rasa_per_residue = windowed_solvent_accessible_area(
+          chain_struc.to_mmcif()
+      )
+      seq_rasa[chain_seq] = rasa_per_residue
+      rasa.extend(rasa_per_residue)
+    except (ValueError, RuntimeError) as e:
+      logging.warning('%s: rasa calculation failed: %s', struc.name, e)
+
+  if not rasa:
+    return 0.0
+  return np.mean(np.array(rasa) > rasa_disorder_cutoff)
+
+
+def has_clash(
+    struc: structure.Structure,
+    cutoff_radius: float = 1.1,
+    min_clashes_for_overlap: int = 100,
+    min_fraction_for_overlap: float = 0.5,
+) -> bool:
+  """Determine whether the structure has at least one clashing chain.
+
+  A clashing chain is defined as having greater than 100 polymer atoms within
+  1.1A of another polymer atom, or having more than 50% of the chain with
+  clashing atoms.
+
+  Args:
+    struc: A structure to get clash metrics for.
+    cutoff_radius: atom distances under this threshold are considered a clash.
+    min_clashes_for_overlap: The minimum number of atom-atom clashes for a chain
+      to be considered overlapping.
+    min_fraction_for_overlap: The minimum fraction of atoms within a chain that
+      are clashing for the chain to be considered overlapping.
+
+  Returns:
+    True if the structure has at least one clashing chain.
+  """
+  struc = struc.filter_to_entity_type(protein=True, rna=True, dna=True)
+  if not struc.chains:
+    return False
+  coords = struc.coords
+  coord_kdtree = spatial.cKDTree(coords)
+  clashes_per_atom = coord_kdtree.query_ball_point(
+      coords, p=2.0, r=cutoff_radius
+  )
+  per_atom_has_clash = np.zeros(len(coords), dtype=np.int32)
+  for atom_idx, clashing_indices in enumerate(clashes_per_atom):
+    for clashing_idx in clashing_indices:
+      if np.abs(struc.res_id[atom_idx] - struc.res_id[clashing_idx]) > 1 or (
+          struc.chain_id[atom_idx] != struc.chain_id[clashing_idx]
+      ):
+        per_atom_has_clash[atom_idx] = True
+        break
+  for chain_id in struc.chains:
+    mask = struc.chain_id == chain_id
+    num_atoms = np.sum(mask)
+    if num_atoms == 0:
+      continue
+    num_clashes = np.sum(per_atom_has_clash * mask)
+    frac_clashes = num_clashes / num_atoms
+    if (
+        num_clashes > min_clashes_for_overlap
+        or frac_clashes > min_fraction_for_overlap
+    ):
+      return True
+  return False
+
+
+def get_ranking_score(
+    ptm: float, iptm: float, fraction_disordered_: float, has_clash_: bool
+) -> float:
+  # ipTM is NaN for single chain structures. Use pTM for such cases.
+  if np.isnan(iptm):
+    ptm_iptm_average = ptm
+  else:
+    ptm_iptm_average = _IPTM_WEIGHT * iptm + (1.0 - _IPTM_WEIGHT) * ptm
+  return (
+      ptm_iptm_average
+      + _FRACTION_DISORDERED_WEIGHT * fraction_disordered_
+      - _CLASH_PENALIZATION_WEIGHT * has_clash_
+  )
+
+
+def rank_metric(
+    full_pde: jnp.ndarray | np.ndarray, contact_probs: jnp.ndarray | np.ndarray
+) -> jnp.ndarray | np.ndarray:
+  """Compute the metric that will be used to rank predictions, higher is better.
+
+  Args:
+    full_pde: A [num_samples, num_tokens,num_tokens] matrix of predicted
+      distance errors between pairs of tokens.
+    contact_probs: A [num_tokens, num_tokens] matrix consisting of the
+      probability of contact (<8A) that is returned from the distogram head.
+
+  Returns:
+    A scalar that can be used to rank (higher is better).
+  """
+  if not isinstance(full_pde, type(contact_probs)):
+    raise ValueError('full_pde and contact_probs must be of the same type.')
+
+  if isinstance(full_pde, np.ndarray):
+    sum_fn = np.sum
+  elif isinstance(full_pde, jnp.ndarray):
+    sum_fn = jnp.sum
+  else:
+    raise ValueError('full_pde must be a numpy array or a jax array.')
+  # It was found that taking the contact_map weighted average was better than
+  # just the predicted distance error on its own.
+  return -sum_fn(full_pde * contact_probs[None, :, :], axis=(-2, -1)) / (
+      sum_fn(contact_probs) + 1e-6
+  )
+
+
+def weighted_mean(mask, value, axis):
+  return np.mean(mask * value, axis=axis) / (1e-8 + np.mean(mask, axis=axis))
+
+
+def pde_single(
+    num_tokens: int,
+    asym_ids: np.ndarray,
+    full_pde: np.ndarray,
+    contact_probs: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+  """Compute 1D PDE summaries.
+
+  Args:
+    num_tokens: The number of tokens (not including padding).
+    asym_ids: The asym_ids (array of shape num_tokens).
+    full_pde: A [num_samples, num_tokens, num_tokens] matrix of predicted
+      distance errors.
+    contact_probs: A [num_tokens, num_tokens] matrix consisting of the
+      probability of contact (<8A) that is returned from the distogram head.
+
+  Returns:
+    A tuple (ichain, xchain, full_chain) where:
+      `ichain` is a [num_samples, num_chains] matrix where the
+      value assigned to each chain is an average of the full PDE matrix over all
+      its within-chain interactions, weighted by `contact_probs`.
+      `xchain` is a [num_samples, num_chains] matrix where the
+      value assigned to each chain is an average of the full PDE matrix over all
+      its cross-chain interactions, weighted by `contact_probs`.
+      `full_chain` is a [num_samples, num_tokens] matrix where the
+      value assigned to each token is an average of it PDE against all tokens,
+      weighted by `contact_probs`.
+  """
+
+  full_pde = full_pde[:, :num_tokens, :num_tokens]
+  contact_probs = contact_probs[:num_tokens, :num_tokens]
+  asym_ids = asym_ids[:num_tokens]
+  unique_asym_ids = np.unique(asym_ids)
+  num_chains = len(unique_asym_ids)
+  num_samples = full_pde.shape[0]
+
+  asym_ids = asym_ids[None]
+  contact_probs = contact_probs[None]
+
+  ichain = np.zeros((num_samples, num_chains))
+  xchain = np.zeros((num_samples, num_chains))
+
+  for idx, asym_id in enumerate(unique_asym_ids):
+    my_asym_id = asym_ids == asym_id
+    imask = my_asym_id[:, :, None] * my_asym_id[:, None, :]
+    xmask = my_asym_id[:, :, None] * ~my_asym_id[:, None, :]
+    imask = imask * contact_probs
+    xmask = xmask * contact_probs
+    ichain[:, idx] = weighted_mean(mask=imask, value=full_pde, axis=(-2, -1))
+    xchain[:, idx] = weighted_mean(mask=xmask, value=full_pde, axis=(-2, -1))
+
+  full_chain = weighted_mean(mask=contact_probs, value=full_pde, axis=(-1,))
+
+  return ichain, xchain, full_chain
+
+
+def chain_pair_pde(
+    num_tokens: int, asym_ids: np.ndarray, full_pde: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+  """Compute predicted distance errors for all pairs of chains.
+
+  Args:
+    num_tokens: The number of tokens (not including padding).
+    asym_ids: The asym_ids (array of shape num_tokens).
+    full_pde: A [num_samples, num_tokens, num_tokens] matrix of predicted
+      distance errors.
+
+  Returns:
+    chain_pair_pred_err_mean - a [num_chains, num_chains] matrix with average
+      per chain-pair predicted distance error.
+    chain_pair_pred_err_min - a [num_chains, num_chains] matrix with min
+      per chain-pair predicted distance error.
+  """
+  full_pde = full_pde[:, :num_tokens, :num_tokens]
+  asym_ids = asym_ids[:num_tokens]
+  unique_asym_ids = np.unique(asym_ids)
+  num_chains = len(unique_asym_ids)
+  num_samples = full_pde.shape[0]
+  chain_pair_pred_err_mean = np.zeros((num_samples, num_chains, num_chains))
+  chain_pair_pred_err_min = np.zeros((num_samples, num_chains, num_chains))
+
+  for idx1, asym_id_1 in enumerate(unique_asym_ids):
+    subset = full_pde[:, asym_ids == asym_id_1, :]
+    for idx2, asym_id_2 in enumerate(unique_asym_ids):
+      subsubset = subset[:, :, asym_ids == asym_id_2]
+      chain_pair_pred_err_mean[:, idx1, idx2] = np.mean(subsubset, axis=(1, 2))
+      chain_pair_pred_err_min[:, idx1, idx2] = np.min(subsubset, axis=(1, 2))
+  return chain_pair_pred_err_mean, chain_pair_pred_err_min
+
+
+def weighted_nanmean(
+    value: np.ndarray, mask: np.ndarray, axis: int
+) -> np.ndarray:
+  """Nan-mean with weighting -- empty slices return NaN."""
+  assert mask.shape == value.shape
+  assert not np.isnan(mask).all()
+
+  nan_idxs = np.where(np.isnan(value))
+  # Need to NaN the mask to get the correct denominator weighting.
+  mask_with_nan = mask.copy()
+  mask_with_nan[nan_idxs] = np.nan
+  with warnings.catch_warnings():
+    # Mean of empty slice is ok and should return a NaN.
+    warnings.filterwarnings(action='ignore', message='Mean of empty slice')
+    warnings.filterwarnings(
+        action='ignore', message='invalid value encountered in (scalar )?divide'
+    )
+    return np.nanmean(value * mask_with_nan, axis=axis) / np.nanmean(
+        mask_with_nan, axis=axis
+    )
+
+
+def chain_pair_pae(
+    *,
+    num_tokens: int,
+    asym_ids: np.ndarray,
+    full_pae: np.ndarray,
+    mask: np.ndarray | None = None,
+    contact_probs: np.ndarray | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+  """Compute predicted errors for all pairs of chains.
+
+  Args:
+    num_tokens: The number of tokens (not including padding).
+    asym_ids: The asym_ids (array of shape num_tokens).
+    full_pae: A [num_samples, num_tokens, num_tokens] matrix of predicted
+      errors.
+    mask: A [num_tokens, num_tokens] mask matrix.
+    contact_probs: A [num_tokens, num_tokens] matrix consisting of the
+      probability of contact (<8A) that is returned from the distogram head.
+
+  Returns:
+    chain_pair_pred_err_mean - a [num_chains, num_chains] matrix with average
+      per chain-pair predicted error.
+  """
+  if mask is None:
+    mask = np.ones(shape=full_pae.shape[1:], dtype=bool)
+  if contact_probs is None:
+    contact_probs = np.ones(shape=full_pae.shape[1:], dtype=float)
+  assert mask.shape == full_pae.shape[1:]
+
+  full_pae = full_pae[:, :num_tokens, :num_tokens]
+  mask = mask[:num_tokens, :num_tokens]
+  asym_ids = asym_ids[:num_tokens]
+  contact_probs = contact_probs[:num_tokens, :num_tokens]
+  unique_asym_ids = np.unique(asym_ids)
+  num_chains = len(unique_asym_ids)
+  num_samples = full_pae.shape[0]
+  chain_pair_pred_err_mean = np.zeros((num_samples, num_chains, num_chains))
+  chain_pair_pred_err_min = np.zeros((num_samples, num_chains, num_chains))
+
+  for idx1, asym_id_1 in enumerate(unique_asym_ids):
+    subset = full_pae[:, asym_ids == asym_id_1, :]
+    subset_mask = mask[asym_ids == asym_id_1, :]
+    subset_contact_probs = contact_probs[asym_ids == asym_id_1, :]
+    for idx2, asym_id_2 in enumerate(unique_asym_ids):
+      subsubset = subset[:, :, asym_ids == asym_id_2]
+      subsubset_mask = subset_mask[:, asym_ids == asym_id_2]
+      subsubset_contact_probs = subset_contact_probs[:, asym_ids == asym_id_2]
+      (flat_mask_idxs,) = np.where(subsubset_mask.flatten() > 0)
+      flat_subsubset = subsubset.reshape([num_samples, -1])
+      flat_contact_probs = subsubset_contact_probs.flatten()
+      # A ligand chain will have no valid frames if it contains fewer than
+      # three non-colinear atoms (e.g. a sodium ion).
+      if not flat_mask_idxs.size:
+        chain_pair_pred_err_mean[:, idx1, idx2] = np.nan
+        chain_pair_pred_err_min[:, idx1, idx2] = np.nan
+      else:
+        chain_pair_pred_err_min[:, idx1, idx2] = np.min(
+            flat_subsubset[:, flat_mask_idxs], axis=1
+        )
+        chain_pair_pred_err_mean[:, idx1, idx2] = weighted_mean(
+            mask=flat_contact_probs[flat_mask_idxs],
+            value=flat_subsubset[:, flat_mask_idxs],
+            axis=-1,
+        )
+  return chain_pair_pred_err_mean, chain_pair_pred_err_min, unique_asym_ids
+
+
+def reduce_chain_pair(
+    *,
+    chain_pair_met: np.ndarray,
+    num_chain_tokens: np.ndarray,
+    agg_over_col: bool,
+    agg_type: str,
+    weight_method: str,
+) -> tuple[np.ndarray, np.ndarray]:
+  """Compute 1D summaries from a chain-pair summary.
+
+  Args:
+    chain_pair_met: A [num_samples, num_chains, num_chains] aggregate matrix.
+    num_chain_tokens: A [num_chains] array of number of tokens for each chain.
+      Used for 'per_token' weighting.
+    agg_over_col: Whether to aggregate the PAE over rows (i.e. average error
+      when aligned to me) or columns (i.e. my average error when aligned to all
+      others.)
+    agg_type: The type of aggregation to use, 'mean' or 'min'.
+    weight_method: The method to use for weighting the PAE, 'per_token' or
+      'per_chain'.
+
+  Returns:
+    A tuple (ichain, xchain) where:
+      `ichain` is a [num_samples, num_chains] matrix where the
+      value assigned to each chain is an average of the full PAE matrix over all
+      its within-chain interactions, weighted by `contact_probs`.
+      `xchain` is a [num_samples, num_chains] matrix where the
+      value assigned to each chain is an average of the full PAE matrix over all
+      its cross-chain interactions, weighted by `contact_probs`.
+  """
+  num_samples, num_chains, _ = chain_pair_met.shape
+
+  ichain = chain_pair_met.diagonal(axis1=-2, axis2=-1)
+
+  if weight_method == 'per_chain':
+    chain_weight = np.ones((num_chains,), dtype=float)
+  elif weight_method == 'per_token':
+    chain_weight = num_chain_tokens
+  else:
+    raise ValueError(f'Unknown weight method: {weight_method}')
+
+  if agg_over_col:
+    agg_axis = -1
+  else:
+    agg_axis = -2
+
+  if agg_type == 'mean':
+    weight = np.ones((num_samples, num_chains, num_chains), dtype=float)
+    weight -= np.eye(num_chains, dtype=float)
+    weight *= chain_weight[None] * chain_weight[:, None]
+    xchain = weighted_nanmean(chain_pair_met, mask=weight, axis=agg_axis)
+  elif agg_type == 'min':
+    is_self = np.eye(num_chains)
+    with warnings.catch_warnings():
+      # Min over empty slice is ok and should return a NaN.
+      warnings.filterwarnings('ignore', message='All-NaN slice encountered')
+      xchain = np.nanmin(chain_pair_met + 1e8 * is_self, axis=agg_axis)
+  else:
+    raise ValueError(f'Unknown aggregation method: {agg_type}')
+
+  return ichain, xchain
+
+
+def pae_metrics(
+    num_tokens: int,
+    asym_ids: np.ndarray,
+    full_pae: np.ndarray,
+    mask: np.ndarray,
+    contact_probs: np.ndarray,
+    tm_adjusted_pae: np.ndarray,
+):
+  """PAE aggregate metrics."""
+  assert mask.shape == full_pae.shape[1:]
+  assert contact_probs.shape == full_pae.shape[1:]
+
+  chain_pair_contact_weighted, _, unique_asym_ids = chain_pair_pae(
+      num_tokens=num_tokens,
+      asym_ids=asym_ids,
+      full_pae=full_pae,
+      mask=mask,
+      contact_probs=contact_probs,
+  )
+
+  ret = {}
+  ret['chain_pair_pae_mean'], ret['chain_pair_pae_min'], _ = chain_pair_pae(
+      num_tokens=num_tokens,
+      asym_ids=asym_ids,
+      full_pae=full_pae,
+      mask=mask,
+  )
+  chain_pair_iptm = np.stack(
+      [
+          chain_pairwise_predicted_tm_scores(
+              tm_adjusted_pae=sample_tm_adjusted_pae[:num_tokens],
+              asym_id=asym_ids[:num_tokens],
+              pair_mask=mask[:num_tokens, :num_tokens],
+          )
+          for sample_tm_adjusted_pae in tm_adjusted_pae
+      ],
+      axis=0,
+  )
+
+  num_chain_tokens = np.array(
+      [sum(asym_ids == asym_id) for asym_id in unique_asym_ids]
+  )
+
+  def reduce_chain_pair_fn(chain_pair: np.ndarray):
+    def inner(agg_over_col):
+      ichain_pae, xchain_pae = reduce_chain_pair(
+          num_chain_tokens=num_chain_tokens,
+          chain_pair_met=chain_pair,
+          agg_over_col=agg_over_col,
+          agg_type='mean',
+          weight_method='per_chain',
+      )
+      return ichain_pae, xchain_pae
+
+    ichain, xchain_row_agg = inner(False)
+    _, xchain_col_agg = inner(True)
+    with warnings.catch_warnings():
+      # Mean of empty slice is ok and should return a NaN.
+      warnings.filterwarnings(action='ignore', message='Mean of empty slice')
+      xchain = np.nanmean(
+          np.stack([xchain_row_agg, xchain_col_agg], axis=0), axis=0
+      )
+    return ichain, xchain
+
+  pae_ichain, pae_xchain = reduce_chain_pair_fn(chain_pair_contact_weighted)
+  iptm_ichain, iptm_xchain = reduce_chain_pair_fn(chain_pair_iptm)
+
+  ret.update({
+      'chain_pair_iptm': chain_pair_iptm,
+      'iptm_ichain': iptm_ichain,
+      'iptm_xchain': iptm_xchain,
+      'pae_ichain': pae_ichain,
+      'pae_xchain': pae_xchain,
+  })
+
+  return ret
+
+
+def get_iptm_xchain(chain_pair_iptm: np.ndarray) -> np.ndarray:
+  """Cross chain aggregate ipTM."""
+  num_samples, num_chains, _ = chain_pair_iptm.shape
+  weight = np.ones((num_samples, num_chains, num_chains), dtype=float)
+  weight -= np.eye(num_chains, dtype=float)
+  xchain_row_agg = weighted_nanmean(chain_pair_iptm, mask=weight, axis=-2)
+  xchain_col_agg = weighted_nanmean(chain_pair_iptm, mask=weight, axis=-1)
+  with warnings.catch_warnings():
+    # Mean of empty slice is ok and should return a NaN.
+    warnings.filterwarnings(action='ignore', message='Mean of empty slice')
+    iptm_xchain = np.nanmean(
+        np.stack([xchain_row_agg, xchain_col_agg], axis=0), axis=0
+    )
+  return iptm_xchain
+
+
+def predicted_tm_score(
+    tm_adjusted_pae: np.ndarray,
+    pair_mask: np.ndarray,
+    asym_id: np.ndarray,
+    interface: bool = False,
+) -> float:
+  """Computes predicted TM alignment or predicted interface TM alignment score.
+
+  Args:
+    tm_adjusted_pae: [num_res, num_res] Relevant tensor for computing TMScore
+      values.
+    pair_mask: A [num_res, num_res] mask. The TM score will only aggregate over
+      masked-on entries.
+    asym_id: [num_res] asymmetric unit ID (the chain ID). Only needed for ipTM
+      calculation, i.e. when interface=True.
+    interface: If True, the interface predicted TM score is computed. If False,
+      the predicted TM score without any residue pair restrictions is computed.
+
+  Returns:
+   score: pTM or ipTM score.
+  """
+  num_tokens, _ = tm_adjusted_pae.shape
+  if tm_adjusted_pae.shape != (num_tokens, num_tokens):
+    raise ValueError(
+        f'Bad tm_adjusted_pae shape, expected ({num_tokens, num_tokens}), got '
+        f'{tm_adjusted_pae.shape}.'
+    )
+
+  if pair_mask.shape != (num_tokens, num_tokens):
+    raise ValueError(
+        f'Bad pair_mask shape, expected ({num_tokens, num_tokens}), got '
+        f'{pair_mask.shape}.'
+    )
+  if pair_mask.dtype != bool:
+    raise TypeError(f'Bad pair mask type, expected bool, got {pair_mask.dtype}')
+  if asym_id.shape[0] != num_tokens:
+    raise ValueError(
+        f'Bad asym_id shape, expected ({num_tokens},), got {asym_id.shape}.'
+    )
+
+  # Create pair mask.
+  if interface:
+    pair_mask = pair_mask * (asym_id[:, None] != asym_id[None, :])
+
+  # Ions and other ligands with colinear atoms have ill-defined frames.
+  if pair_mask.sum() == 0:
+    return np.nan
+
+  normed_residue_mask = pair_mask / (
+      1e-8 + np.sum(pair_mask, axis=-1, keepdims=True)
+  )
+  per_alignment = np.sum(tm_adjusted_pae * normed_residue_mask, axis=-1)
+  return per_alignment.max()
+
+
+def chain_pairwise_predicted_tm_scores(
+    tm_adjusted_pae: np.ndarray,
+    pair_mask: np.ndarray,
+    asym_id: np.ndarray,
+) -> np.ndarray:
+  """Compute predicted TM (pTM) between each pair of chains independently.
+
+  Args:
+    tm_adjusted_pae: [num_res, num_res] Relevant tensor for computing TMScore
+      values.
+    pair_mask: A [num_res, num_res] mask specifying which frames are valid.
+      Invalid frames can be the result of chains with not enough atoms (e.g.
+      ions).
+    asym_id: [num_res] asymmetric unit ID (the chain ID).
+
+  Returns:
+    A [num_chains, num_chains] matrix, where row i, column j indicates the
+    predicted TM-score for the interface between chain i and chain j.
+  """
+  unique_chains = list(np.unique(asym_id))
+  num_chains = len(unique_chains)
+  all_pairs_iptms = np.zeros((num_chains, num_chains))
+  for i, chain_i in enumerate(unique_chains):
+    chain_i_mask = asym_id == chain_i
+    for j, chain_j in enumerate(unique_chains[i:]):
+      chain_j_mask = asym_id == chain_j
+      mask = chain_i_mask | chain_j_mask
+      (indices,) = np.where(mask)
+      is_interface = chain_i != chain_j
+      indices = np.ix_(indices, indices)
+      iptm = predicted_tm_score(
+          tm_adjusted_pae=tm_adjusted_pae[indices],
+          pair_mask=pair_mask[indices],
+          asym_id=asym_id[mask],
+          interface=is_interface,
+      )
+      all_pairs_iptms[i, i + j] = iptm
+      all_pairs_iptms[i + j, i] = iptm
+  return all_pairs_iptms

flax_model/alphafold3/model/data3.py

@@ -0,0 +1,117 @@
+
+
+"""Protein features that are computed from parsed mmCIF objects."""
+
+from collections.abc import Mapping, MutableMapping
+import datetime
+from typing import TypeAlias
+
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import msa_profile
+from flax_model.alphafold3.model import protein_data_processing
+import numpy as np
+
+
+FeatureDict: TypeAlias = Mapping[str, np.ndarray]
+# MutableFeatureDict: TypeAlias = MutableMapping[str, np.ndarray]
+
+
+# def fix_features(msa_features: MutableFeatureDict) -> MutableFeatureDict:
+#   """Renames the deletion_matrix feature."""
+#   msa_features['deletion_matrix'] = msa_features.pop('deletion_matrix_int')
+#   return msa_features
+
+
+def get_profile_features(
+    msa: np.ndarray, deletion_matrix: np.ndarray
+) -> FeatureDict:
+  """Returns the MSA profile and deletion_mean features."""
+  num_restypes = residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP
+  profile = msa_profile.compute_msa_profile(
+      msa=msa, num_residue_types=num_restypes
+  )
+
+  return {
+      'profile': profile.astype(np.float32),
+      'deletion_mean': np.mean(deletion_matrix, axis=0),
+  }
+
+
+def fix_template_features(
+    template_features: FeatureDict, num_res: int
+) -> FeatureDict:
+  """Convert template features to AlphaFold 3 format.
+
+  Args:
+    template_features: Template features for the protein.
+    num_res: The length of the amino acid sequence of the protein.
+
+  Returns:
+    Updated template_features for the chain.
+  """
+  if not template_features['template_aatype'].shape[0]:
+    template_features = empty_template_features(num_res)
+  else:
+    template_release_timestamp = [
+        _get_timestamp(x.decode('utf-8'))
+        for x in template_features['template_release_date']
+    ]
+
+    # Convert from atom37 to dense atom
+    dense_atom_indices = np.take(
+        protein_data_processing.PROTEIN_AATYPE_DENSE_ATOM_TO_ATOM37,
+        template_features['template_aatype'],
+        axis=0,
+    )
+
+    atom_mask = np.take_along_axis(
+        template_features['template_all_atom_masks'], dense_atom_indices, axis=2
+    )
+    atom_positions = np.take_along_axis(
+        template_features['template_all_atom_positions'],
+        dense_atom_indices[..., None],
+        axis=2,
+    )
+    atom_positions *= atom_mask[..., None]
+
+    template_features = {
+        'template_aatype': template_features['template_aatype'],
+        'template_atom_mask': atom_mask.astype(np.int32),
+        'template_atom_positions': atom_positions.astype(np.float32),
+        'template_domain_names': np.array(
+            template_features['template_domain_names'], dtype=object
+        ),
+        'template_release_timestamp': np.array(
+            template_release_timestamp, dtype=np.float32
+        ),
+    }
+  return template_features
+
+
+def empty_template_features(num_res: int) -> FeatureDict:
+  """Creates a fully masked out template features to allow padding to work.
+
+  Args:
+    num_res: The length of the target chain.
+
+  Returns:
+    Empty template features for the chain.
+  """
+  template_features = {
+      'template_aatype': np.zeros(num_res, dtype=np.int32)[None, ...],
+      'template_atom_mask': np.zeros(
+          (num_res, protein_data_processing.NUM_DENSE), dtype=np.int32
+      )[None, ...],
+      'template_atom_positions': np.zeros(
+          (num_res, protein_data_processing.NUM_DENSE, 3), dtype=np.float32
+      )[None, ...],
+      'template_domain_names': np.array([b''], dtype=object),
+      'template_release_timestamp': np.array([0.0], dtype=np.float32),
+  }
+  return template_features
+
+
+def _get_timestamp(date_str: str):
+  dt = datetime.datetime.fromisoformat(date_str)
+  dt = dt.replace(tzinfo=datetime.timezone.utc)
+  return dt.timestamp()

flax_model/alphafold3/model/data_constants.py

@@ -0,0 +1,19 @@
+
+
+"""Constants shared across modules in the AlphaFold data pipeline."""
+
+from flax_model.alphafold3.constants import residue_names
+
+MSA_GAP_IDX = residue_names.PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP.index(
+    '-'
+)
+
+# Feature groups.
+NUM_SEQ_NUM_RES_MSA_FEATURES = ('msa', 'msa_mask', 'deletion_matrix')
+NUM_SEQ_MSA_FEATURES = ('msa_species_identifiers',)
+TEMPLATE_FEATURES = (
+    'template_aatype',
+    'template_atom_positions',
+    'template_atom_mask',
+)
+MSA_PAD_VALUES = {'msa': MSA_GAP_IDX, 'msa_mask': 1, 'deletion_matrix': 0}

flax_model/alphafold3/model/feat_batch.py

@@ -0,0 +1,77 @@
+
+
+"""Batch dataclass."""
+import dataclasses
+from typing import Self
+
+from flax_model.alphafold3.model import features
+#import chex
+import jax
+
+
+@dataclasses.dataclass(frozen=True)
+class Batch:
+  """Dataclass containing batch."""
+
+  msa: features.MSA
+  templates: features.Templates
+  token_features: features.TokenFeatures
+  ref_structure: features.RefStructure
+  predicted_structure_info: features.PredictedStructureInfo
+  polymer_ligand_bond_info: features.PolymerLigandBondInfo
+  ligand_ligand_bond_info: features.LigandLigandBondInfo
+  pseudo_beta_info: features.PseudoBetaInfo
+  atom_cross_att: features.AtomCrossAtt
+  convert_model_output: features.ConvertModelOutput
+  frames: features.Frames
+
+  @property
+  def num_res(self) -> int:
+    return self.token_features.aatype.shape[-1]
+
+  @classmethod
+  def from_data_dict(cls, batch: features.BatchDict) -> Self:
+    """Construct batch object from dictionary."""
+    return cls(
+        msa=features.MSA.from_data_dict(batch),
+        templates=features.Templates.from_data_dict(batch),
+        token_features=features.TokenFeatures.from_data_dict(batch),
+        ref_structure=features.RefStructure.from_data_dict(batch),
+        predicted_structure_info=features.PredictedStructureInfo.from_data_dict(
+            batch
+        ),
+        polymer_ligand_bond_info=features.PolymerLigandBondInfo.from_data_dict(
+            batch
+        ),
+        ligand_ligand_bond_info=features.LigandLigandBondInfo.from_data_dict(
+            batch
+        ),
+        pseudo_beta_info=features.PseudoBetaInfo.from_data_dict(batch),
+        atom_cross_att=features.AtomCrossAtt.from_data_dict(batch),
+        convert_model_output=features.ConvertModelOutput.from_data_dict(batch),
+        frames=features.Frames.from_data_dict(batch),
+    )
+
+  def as_data_dict(self) -> features.BatchDict:
+    """Converts batch object to dictionary."""
+    output = {
+        **self.msa.as_data_dict(),
+        **self.templates.as_data_dict(),
+        **self.token_features.as_data_dict(),
+        **self.ref_structure.as_data_dict(),
+        **self.predicted_structure_info.as_data_dict(),
+        **self.polymer_ligand_bond_info.as_data_dict(),
+        **self.ligand_ligand_bond_info.as_data_dict(),
+        **self.pseudo_beta_info.as_data_dict(),
+        **self.atom_cross_att.as_data_dict(),
+        **self.convert_model_output.as_data_dict(),
+        **self.frames.as_data_dict(),
+    }
+    return output
+
+
+jax.tree_util.register_dataclass(
+    Batch,
+    data_fields=[f.name for f in dataclasses.fields(Batch)],
+    meta_fields=[],
+)

flax_model/alphafold3/model/features.py

@@ -0,0 +1,2159 @@
+
+
+"""Data-side of the input features processing."""
+
+import dataclasses
+import datetime
+import itertools
+from typing import Any, Self, TypeAlias
+
+from absl import logging
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.common import folding_input
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import periodic_table
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import cif_dict
+from flax_model.alphafold3.data import msa as msa_module
+from flax_model.alphafold3.data import templates
+from flax_model.alphafold3.data.tools import rdkit_utils
+from flax_model.alphafold3.model import data3
+from flax_model.alphafold3.model import data_constants
+from flax_model.alphafold3.model import merging_features
+from flax_model.alphafold3.model import msa_pairing
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.structure import chemical_components as struc_chem_comps
+import jax
+import jax.numpy as jnp
+import numpy as np
+from rdkit import Chem
+
+
+xnp_ndarray: TypeAlias = np.ndarray | jnp.ndarray  # pylint: disable=invalid-name
+BatchDict: TypeAlias = dict[str, xnp_ndarray]
+
+_STANDARD_RESIDUES = frozenset({
+    *residue_names.PROTEIN_TYPES_WITH_UNKNOWN,
+    *residue_names.NUCLEIC_TYPES_WITH_2_UNKS,
+})
+
+
+@dataclasses.dataclass(frozen=True)
+class PaddingShapes:
+  num_tokens: int
+  msa_size: int
+  num_chains: int
+  num_templates: int
+  num_atoms: int
+
+
+def _pad_to(
+    arr: np.ndarray, shape: tuple[int | None, ...], **kwargs
+) -> np.ndarray:
+  """Pads an array to a given shape. Wrapper around np.pad().
+
+  Args:
+    arr: numpy array to pad
+    shape: target shape, use None for axes that should stay the same
+    **kwargs: additional args for np.pad, e.g. constant_values=-1
+
+  Returns:
+    the padded array
+
+  Raises:
+    ValueError if arr and shape have a different number of axes.
+  """
+  if arr.ndim != len(shape):
+    raise ValueError(
+        f'arr and shape have different number of axes. {arr.shape=}, {shape=}'
+    )
+
+  num_pad = []
+  for axis, width in enumerate(shape):
+    if width is None:
+      num_pad.append((0, 0))
+    else:
+      if width >= arr.shape[axis]:
+        num_pad.append((0, width - arr.shape[axis]))
+      else:
+        raise ValueError(
+            f'Can not pad to a smaller shape. {arr.shape=}, {shape=}'
+        )
+  padded_arr = np.pad(arr, pad_width=num_pad, **kwargs)
+  return padded_arr
+
+
+def _unwrap(obj):
+  """Unwrap an object from a zero-dim np.ndarray."""
+  if isinstance(obj, np.ndarray) and obj.ndim == 0:
+    return obj.item()
+  else:
+    return obj
+
+
+@dataclasses.dataclass(frozen=True)
+class Chains:
+  chain_id: np.ndarray
+  asym_id: np.ndarray
+  entity_id: np.ndarray
+  sym_id: np.ndarray
+
+
+jax.tree_util.register_dataclass(
+    Chains,
+    data_fields=[f.name for f in dataclasses.fields(Chains)],
+    meta_fields=[],
+)
+
+
+def _compute_asym_entity_and_sym_id(
+    all_tokens: atom_layout.AtomLayout,
+) -> Chains:
+  """Compute asym_id, entity_id and sym_id.
+
+  Args:
+    all_tokens: atom layout containing a representative atom for each token.
+
+  Returns:
+    A Chains object
+  """
+
+  # Find identical sequences and assign entity_id and sym_id to every chain.
+  seq_to_entity_id_sym_id = {}
+  seen_chain_ids = set()
+  chain_ids = []
+  asym_ids = []
+  entity_ids = []
+  sym_ids = []
+  for chain_id in all_tokens.chain_id:
+    if chain_id not in seen_chain_ids:
+      asym_id = len(seen_chain_ids) + 1
+      seen_chain_ids.add(chain_id)
+      seq = ','.join(all_tokens.res_name[all_tokens.chain_id == chain_id])
+      if seq not in seq_to_entity_id_sym_id:
+        entity_id = len(seq_to_entity_id_sym_id) + 1
+        sym_id = 1
+      else:
+        entity_id, sym_id = seq_to_entity_id_sym_id[seq]
+        sym_id += 1
+      seq_to_entity_id_sym_id[seq] = (entity_id, sym_id)
+
+      chain_ids.append(chain_id)
+      asym_ids.append(asym_id)
+      entity_ids.append(entity_id)
+      sym_ids.append(sym_id)
+
+  return Chains(
+      chain_id=np.array(chain_ids),
+      asym_id=np.array(asym_ids),
+      entity_id=np.array(entity_ids),
+      sym_id=np.array(sym_ids),
+  )
+
+
+def tokenizer(
+    flat_output_layout: atom_layout.AtomLayout,
+    ccd: chemical_components.Ccd,
+    max_atoms_per_token: int,
+    flatten_non_standard_residues: bool,
+    logging_name: str,
+) -> tuple[atom_layout.AtomLayout, atom_layout.AtomLayout, np.ndarray]:
+  """Maps a flat atom layout to tokens for evoformer.
+
+  Creates the evoformer tokens as one token per polymer residue and one token
+  per ligand atom. The tokens are represented as AtomLayouts all_tokens
+  (1 representative atom per token) atoms per residue, and
+  all_token_atoms_layout (num_tokens, max_atoms_per_token). The atoms in a
+  residue token use the layout of the corresponding CCD entry
+
+  Args:
+    flat_output_layout: flat AtomLayout containing all atoms that the model
+      wants to predict.
+    ccd: The chemical components dictionary.
+    max_atoms_per_token: number of slots per token.
+    flatten_non_standard_residues: whether to flatten non-standard residues,
+      i.e. whether to use one token per atom for non-standard residues.
+    logging_name: logging name for debugging (usually the mmcif_id).
+
+  Returns:
+    A tuple (all_tokens, all_tokens_atoms_layout) with
+      all_tokens: AtomLayout shape (num_tokens,) containing one representative
+        atom per token.
+      all_token_atoms_layout: AtomLayout with shape
+        (num_tokens, max_atoms_per_token) containing all atoms per token.
+      standard_token_idxs: The token index that each token would have if not
+        flattening non standard resiudes.
+  """
+  # Select  the representative atom for each token.
+  token_idxs = []
+  single_atom_token = []
+  standard_token_idxs = []
+  current_standard_token_id = 0
+  # Iterate over residues, and provide a group_iter over the atoms of each
+  # residue.
+  for key, group_iter in itertools.groupby(
+      zip(
+          flat_output_layout.chain_type,
+          flat_output_layout.chain_id,
+          flat_output_layout.res_id,
+          flat_output_layout.res_name,
+          flat_output_layout.atom_name,
+          np.arange(flat_output_layout.shape[0]),
+      ),
+      key=lambda x: x[:3],
+  ):
+
+    # Get chain type and chain id of this residue
+    chain_type, chain_id, _ = key
+
+    # Get names and global idxs for all atoms of this residue
+    _, _, _, res_names, atom_names, idxs = zip(*group_iter)
+
+    # As of March 2023, all OTHER CHAINs in pdb are artificial nucleics.
+    is_nucleic_backbone = (
+        chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES
+        or chain_type == mmcif_names.OTHER_CHAIN
+    )
+    if chain_type in mmcif_names.PEPTIDE_CHAIN_TYPES:
+      res_name = res_names[0]
+      if (
+          flatten_non_standard_residues
+          and res_name not in residue_names.PROTEIN_TYPES_WITH_UNKNOWN
+          and res_name != residue_names.MSE
+      ):
+        # For non-standard protein residues take all atoms.
+        # NOTE: This may get very large if we include hydrogens.
+        token_idxs.extend(idxs)
+        single_atom_token += [True] * len(idxs)
+        standard_token_idxs.extend([current_standard_token_id] * len(idxs))
+      else:
+        # For standard protein residues take 'CA' if it exists, else first atom.
+        if 'CA' in atom_names:
+          token_idxs.append(idxs[atom_names.index('CA')])
+        else:
+          token_idxs.append(idxs[0])
+        single_atom_token += [False]
+        standard_token_idxs.append(current_standard_token_id)
+      current_standard_token_id += 1
+    elif is_nucleic_backbone:
+      res_name = res_names[0]
+      if (
+          flatten_non_standard_residues
+          and res_name not in residue_names.NUCLEIC_TYPES_WITH_2_UNKS
+      ):
+        # For non-standard nucleic residues take all atoms.
+        token_idxs.extend(idxs)
+        single_atom_token += [True] * len(idxs)
+        standard_token_idxs.extend([current_standard_token_id] * len(idxs))
+      else:
+        # For standard nucleic residues take C1' if it exists, else first atom.
+        if "C1'" in atom_names:
+          token_idxs.append(idxs[atom_names.index("C1'")])
+        else:
+          token_idxs.append(idxs[0])
+        single_atom_token += [False]
+        standard_token_idxs.append(current_standard_token_id)
+      current_standard_token_id += 1
+    elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
+      # For non-polymers take all atoms
+      token_idxs.extend(idxs)
+      single_atom_token += [True] * len(idxs)
+      standard_token_idxs.extend([current_standard_token_id] * len(idxs))
+      current_standard_token_id += len(idxs)
+    else:
+      # Chain type that we don't handle yet.
+      logging.warning(
+          '%s: ignoring chain %s with chain type %s.',
+          logging_name,
+          chain_id,
+          chain_type,
+      )
+
+  assert len(token_idxs) == len(single_atom_token)
+  assert len(token_idxs) == len(standard_token_idxs)
+  standard_token_idxs = np.array(standard_token_idxs, dtype=np.int32)
+
+  # Create the list of all tokens, represented as a flat AtomLayout with 1
+  # representative atom per token.
+  all_tokens = flat_output_layout[token_idxs]
+
+  # Create the 2D atoms_per_token layout
+  num_tokens = all_tokens.shape[0]
+
+  # Target lists.
+  target_atom_names = []
+  target_atom_elements = []
+  target_res_ids = []
+  target_res_names = []
+  target_chain_ids = []
+  target_chain_types = []
+
+  # uids of all atoms in the flat layout, to check whether the dense atoms
+  # exist -- This is necessary for terminal atoms (e.g. 'OP3' or 'OXT')
+  all_atoms_uids = set(
+      zip(
+          flat_output_layout.chain_id,
+          flat_output_layout.res_id,
+          flat_output_layout.atom_name,
+      )
+  )
+
+  for idx, single_atom in enumerate(single_atom_token):
+    if not single_atom:
+      # Standard protein and nucleic residues have many atoms per token
+      chain_id = all_tokens.chain_id[idx]
+      res_id = all_tokens.res_id[idx]
+      res_name = all_tokens.res_name[idx]
+      atom_names = []
+      atom_elements = []
+
+      res_atoms = struc_chem_comps.get_all_atoms_in_entry(
+          ccd=ccd, res_name=res_name
+      )
+      atom_names_elements = list(
+          zip(
+              res_atoms['_chem_comp_atom.atom_id'],
+              res_atoms['_chem_comp_atom.type_symbol'],
+              strict=True,
+          )
+      )
+
+      for atom_name, atom_element in atom_names_elements:
+        # Remove hydrogens if they are not in flat layout.
+        if atom_element in ['H', 'D'] and (
+            (chain_id, res_id, atom_name) not in all_atoms_uids
+        ):
+          continue
+        elif (chain_id, res_id, atom_name) in all_atoms_uids:
+          atom_names.append(atom_name)
+          atom_elements.append(atom_element)
+        # Leave spaces for OXT etc.
+        else:
+          atom_names.append('')
+          atom_elements.append('')
+
+      if len(atom_names) > max_atoms_per_token:
+        logging.warning(
+            'Atom list for chain %s '
+            'residue %s %s is too long and will be truncated: '
+            '%s to the max atoms limit %s. Dropped atoms: %s',
+            chain_id,
+            res_id,
+            res_name,
+            len(atom_names),
+            max_atoms_per_token,
+            list(
+                zip(
+                    atom_names[max_atoms_per_token:],
+                    atom_elements[max_atoms_per_token:],
+                    strict=True,
+                )
+            ),
+        )
+        atom_names = atom_names[:max_atoms_per_token]
+        atom_elements = atom_elements[:max_atoms_per_token]
+
+      num_pad = max_atoms_per_token - len(atom_names)
+      atom_names.extend([''] * num_pad)
+      atom_elements.extend([''] * num_pad)
+
+    else:
+      # ligands have only 1 atom per token
+      padding = [''] * (max_atoms_per_token - 1)
+      atom_names = [all_tokens.atom_name[idx]] + padding
+      atom_elements = [all_tokens.atom_element[idx]] + padding
+
+    # Append the atoms to the target lists.
+    target_atom_names.append(atom_names)
+    target_atom_elements.append(atom_elements)
+    target_res_names.append([all_tokens.res_name[idx]] * max_atoms_per_token)
+    target_res_ids.append([all_tokens.res_id[idx]] * max_atoms_per_token)
+    target_chain_ids.append([all_tokens.chain_id[idx]] * max_atoms_per_token)
+    target_chain_types.append(
+        [all_tokens.chain_type[idx]] * max_atoms_per_token
+    )
+
+  # Make sure to get the right shape also for 0 tokens
+  trg_shape = (num_tokens, max_atoms_per_token)
+  all_token_atoms_layout = atom_layout.AtomLayout(
+      atom_name=np.array(target_atom_names, dtype=object).reshape(trg_shape),
+      atom_element=np.array(target_atom_elements, dtype=object).reshape(
+          trg_shape
+      ),
+      res_name=np.array(target_res_names, dtype=object).reshape(trg_shape),
+      res_id=np.array(target_res_ids, dtype=int).reshape(trg_shape),
+      chain_id=np.array(target_chain_ids, dtype=object).reshape(trg_shape),
+      chain_type=np.array(target_chain_types, dtype=object).reshape(trg_shape),
+  )
+
+  return all_tokens, all_token_atoms_layout, standard_token_idxs
+
+
+@dataclasses.dataclass(frozen=True)
+class MSA:
+  """Dataclass containing MSA."""
+
+  rows: xnp_ndarray
+  mask: xnp_ndarray
+  deletion_matrix: xnp_ndarray
+  # Occurrence of each residue type along the sequence, averaged over MSA rows.
+  profile: xnp_ndarray
+  # Occurrence of deletions along the sequence, averaged over MSA rows.
+  deletion_mean: xnp_ndarray
+  # Number of MSA alignments.
+  num_alignments: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      *,
+      all_tokens: atom_layout.AtomLayout,
+      standard_token_idxs: np.ndarray,
+      padding_shapes: PaddingShapes,
+      fold_input: folding_input.Input,
+      logging_name: str,
+      max_paired_sequence_per_species: int,
+      resolve_msa_overlaps: bool = True,
+  ) -> Self:
+    """Compute the msa features."""
+    seen_entities = {}
+
+    substruct = atom_layout.make_structure(
+        flat_layout=all_tokens,
+        atom_coords=np.zeros(all_tokens.shape + (3,)),
+        name=logging_name,
+    )
+    prot = substruct.filter_to_entity_type(protein=True)
+    num_unique_chains = len(set(prot.chain_single_letter_sequence().values()))
+    need_msa_pairing = num_unique_chains > 1
+
+    np_chains_list = []
+    input_chains_by_id = {chain.id: chain for chain in fold_input.chains}
+    nonempty_chain_ids = set(all_tokens.chain_id)
+    for asym_id, chain_info in enumerate(substruct.iter_chains(), start=1):
+      b_chain_id = chain_info['chain_id']
+      chain_type = chain_info['chain_type']
+      chain = input_chains_by_id[b_chain_id]
+
+      # Generalised "sequence" for ligands (can't trust residue name)
+      chain_tokens = all_tokens[all_tokens.chain_id == b_chain_id]
+      assert chain_tokens.res_name is not None
+      three_letter_sequence = ','.join(chain_tokens.res_name.tolist())
+      chain_num_tokens = len(chain_tokens.atom_name)
+      if chain_type in mmcif_names.POLYMER_CHAIN_TYPES:
+        sequence = substruct.chain_single_letter_sequence()[b_chain_id]
+        if chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES:
+          # Only allow nucleic residue types for nucleic chains (can have some
+          # protein residues in e.g. tRNA, but that causes MSA search failures).
+          # Replace non nucleic residue types by UNK_NUCLEIC.
+          nucleic_types_one_letter = (
+              residue_names.DNA_TYPES_ONE_LETTER
+              + residue_names.RNA_TYPES_ONE_LETTER_WITH_UNKNOWN
+          )
+          sequence = ''.join([
+              base
+              if base in nucleic_types_one_letter
+              else residue_names.UNK_NUCLEIC_ONE_LETTER
+              for base in sequence
+          ])
+      else:
+        sequence = 'X' * chain_num_tokens
+
+      skip_chain = (
+          chain_type not in mmcif_names.STANDARD_POLYMER_CHAIN_TYPES
+          or len(sequence) <= 4
+          or b_chain_id not in nonempty_chain_ids
+      )
+      if three_letter_sequence in seen_entities:
+        entity_id = seen_entities[three_letter_sequence]
+      else:
+        entity_id = len(seen_entities) + 1
+
+      if chain_type in mmcif_names.STANDARD_POLYMER_CHAIN_TYPES:
+        unpaired_a3m = ''
+        paired_a3m = ''
+        if not skip_chain:
+          if need_msa_pairing and isinstance(chain, folding_input.ProteinChain):
+            paired_a3m = chain.paired_msa
+          if isinstance(
+              chain, folding_input.RnaChain | folding_input.ProteinChain
+          ):
+            unpaired_a3m = chain.unpaired_msa
+        # If we generated the MSA ourselves, it is already deduplicated. If it
+        # is user-provided, keep it as is to prevent destroying desired pairing.
+        unpaired_msa = msa_module.Msa.from_a3m(
+            query_sequence=sequence,
+            chain_poly_type=chain_type,
+            a3m=unpaired_a3m,
+            deduplicate=False,
+        )
+
+        paired_msa = msa_module.Msa.from_a3m(
+            query_sequence=sequence,
+            chain_poly_type=chain_type,
+            a3m=paired_a3m,
+            deduplicate=False,
+        )
+      else:
+        unpaired_msa = msa_module.Msa.from_empty(
+            query_sequence='-' * len(sequence),
+            chain_poly_type=mmcif_names.PROTEIN_CHAIN,
+        )
+        paired_msa = msa_module.Msa.from_empty(
+            query_sequence='-' * len(sequence),
+            chain_poly_type=mmcif_names.PROTEIN_CHAIN,
+        )
+
+      msa_features = unpaired_msa.featurize()
+      all_seqs_msa_features = paired_msa.featurize()
+
+      # msa_features = data3.fix_features(msa_features)
+      # all_seqs_msa_features = data3.fix_features(all_seqs_msa_features)
+
+      msa_features = msa_features | {
+          f'{k}_all_seq': v for k, v in all_seqs_msa_features.items()
+      }
+      feats = msa_features
+      feats['chain_id'] = b_chain_id
+      feats['asym_id'] = np.full(chain_num_tokens, asym_id)
+      feats['entity_id'] = entity_id
+      np_chains_list.append(feats)
+
+    # Add profile features to each chain.
+    for chain in np_chains_list:
+      chain.update(
+          data3.get_profile_features(chain['msa'], chain['deletion_matrix'])
+      )
+
+    # Allow 50% of the MSA to come from MSA pairing.
+    max_paired_sequences = padding_shapes.msa_size // 2
+    if need_msa_pairing:
+      np_chains_list = list(map(dict, np_chains_list))
+      np_chains_list = msa_pairing.create_paired_features(
+          np_chains_list,
+          max_paired_sequences=max_paired_sequences,
+          nonempty_chain_ids=nonempty_chain_ids,
+          max_hits_per_species=max_paired_sequence_per_species,
+      )
+      if resolve_msa_overlaps:
+        np_chains_list = msa_pairing.deduplicate_unpaired_sequences(
+            np_chains_list
+        )
+
+    # Remove all gapped rows from all seqs.
+    nonempty_asym_ids = []
+    for chain in np_chains_list:
+      if chain['chain_id'] in nonempty_chain_ids:
+        nonempty_asym_ids.append(chain['asym_id'][0])
+    if 'msa_all_seq' in np_chains_list[0]:
+      np_chains_list = msa_pairing.remove_all_gapped_rows_from_all_seqs(
+          np_chains_list, asym_ids=nonempty_asym_ids
+      )
+
+    # Crop MSA rows.
+    cropped_chains_list = []
+    for chain in np_chains_list:
+      unpaired_msa_size, paired_msa_size = (
+          msa_pairing.choose_paired_unpaired_msa_crop_sizes(
+              unpaired_msa=chain['msa'],
+              paired_msa=chain.get('msa_all_seq'),
+              total_msa_crop_size=padding_shapes.msa_size,
+              max_paired_sequences=max_paired_sequences,
+          )
+      )
+      cropped_chain = {
+          'asym_id': chain['asym_id'],
+          'chain_id': chain['chain_id'],
+          'profile': chain['profile'],
+          'deletion_mean': chain['deletion_mean'],
+      }
+      for feat in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
+        if feat in chain:
+          cropped_chain[feat] = chain[feat][:unpaired_msa_size]
+        if feat + '_all_seq' in chain:
+          cropped_chain[feat + '_all_seq'] = chain[feat + '_all_seq'][
+              :paired_msa_size
+          ]
+      cropped_chains_list.append(cropped_chain)
+
+    # Merge Chains.
+    # Make sure the chain order is unaltered before slicing with tokens.
+    curr_chain_order = [chain['chain_id'] for chain in cropped_chains_list]
+    orig_chain_order = [chain['chain_id'] for chain in substruct.iter_chains()]
+    assert curr_chain_order == orig_chain_order
+    np_example = {
+        'asym_id': np.concatenate(
+            [c['asym_id'] for c in cropped_chains_list], axis=0
+        ),
+    }
+    for feature in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
+      for feat in [feature, feature + '_all_seq']:
+        if feat in cropped_chains_list[0]:
+          np_example[feat] = merging_features.merge_msa_features(
+              feat, cropped_chains_list
+          )
+    for feature in ['profile', 'deletion_mean']:
+      feature_list = [c[feature] for c in cropped_chains_list]
+      np_example[feature] = np.concatenate(feature_list, axis=0)
+
+    # Crop MSA rows to maximum size given by chains participating in the crop.
+    max_allowed_unpaired = max([
+        len(chain['msa'])
+        for chain in cropped_chains_list
+        if chain['asym_id'][0] in nonempty_asym_ids
+    ])
+    np_example['msa'] = np_example['msa'][:max_allowed_unpaired]
+    if 'msa_all_seq' in np_example:
+      max_allowed_paired = max([
+          len(chain['msa_all_seq'])
+          for chain in cropped_chains_list
+          if chain['asym_id'][0] in nonempty_asym_ids
+      ])
+      np_example['msa_all_seq'] = np_example['msa_all_seq'][:max_allowed_paired]
+
+    np_example = merging_features.merge_paired_and_unpaired_msa(np_example)
+
+    # Crop MSA residues. Need to use the standard token indices, since msa does
+    # not expand non-standard residues. This means that for expanded residues,
+    # we get repeated msa columns.
+    new_cropping_idxs = standard_token_idxs
+    for feature in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
+      if feature in np_example:
+        np_example[feature] = np_example[feature][:, new_cropping_idxs].copy()
+    for feature in ['profile', 'deletion_mean']:
+      np_example[feature] = np_example[feature][new_cropping_idxs]
+
+    # Make MSA mask.
+    np_example['msa_mask'] = np.ones_like(np_example['msa'], dtype=np.float32)
+
+    # Count MSA size before padding.
+    num_alignments = np_example['msa'].shape[0]
+
+    # Pad:
+    msa_size, num_tokens = padding_shapes.msa_size, padding_shapes.num_tokens
+
+    def safe_cast_int8(x):
+      return np.clip(x, np.iinfo(np.int8).min, np.iinfo(np.int8).max).astype(
+          np.int8
+      )
+
+    return MSA(
+        rows=_pad_to(safe_cast_int8(np_example['msa']), (msa_size, num_tokens)),
+        mask=_pad_to(
+            np_example['msa_mask'].astype(bool), (msa_size, num_tokens)
+        ),
+        # deletion_matrix may be out of int8 range, but we mostly care about
+        # small values since we arctan it in the model.
+        deletion_matrix=_pad_to(
+            safe_cast_int8(np_example['deletion_matrix']),
+            (msa_size, num_tokens),
+        ),
+        profile=_pad_to(np_example['profile'], (num_tokens, None)),
+        deletion_mean=_pad_to(np_example['deletion_mean'], (num_tokens,)),
+        num_alignments=np.array(num_alignments, dtype=np.int32),
+    )
+
+  def index_msa_rows(self, indices: xnp_ndarray) -> Self:
+    assert indices.ndim == 1
+
+    return MSA(
+        rows=self.rows[indices, :],
+        mask=self.mask[indices, :],
+        deletion_matrix=self.deletion_matrix[indices, :],
+        profile=self.profile,
+        deletion_mean=self.deletion_mean,
+        num_alignments=self.num_alignments,
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    output = cls(
+        rows=batch['msa'],
+        mask=batch['msa_mask'],
+        deletion_matrix=batch['deletion_matrix'],
+        profile=batch['profile'],
+        deletion_mean=batch['deletion_mean'],
+        num_alignments=batch['num_alignments'],
+    )
+    return output
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'msa': self.rows,
+        'msa_mask': self.mask,
+        'deletion_matrix': self.deletion_matrix,
+        'profile': self.profile,
+        'deletion_mean': self.deletion_mean,
+        'num_alignments': self.num_alignments,
+    }
+
+
+jax.tree_util.register_dataclass(
+    MSA,
+    data_fields=[f.name for f in dataclasses.fields(MSA)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class Templates:
+  """Dataclass containing templates."""
+
+  # aatype of templates, int32 w shape [num_templates, num_res]
+  aatype: xnp_ndarray
+  # atom positions of templates, float32 w shape [num_templates, num_res, 24, 3]
+  atom_positions: xnp_ndarray
+  # atom mask of templates, bool w shape [num_templates, num_res, 24]
+  atom_mask: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      standard_token_idxs: np.ndarray,
+      padding_shapes: PaddingShapes,
+      fold_input: folding_input.Input,
+      max_templates: int,
+      logging_name: str,
+  ) -> Self:
+    """Compute the template features."""
+
+    seen_entities = {}
+    polymer_entity_features = {True: {}, False: {}}
+
+    substruct = atom_layout.make_structure(
+        flat_layout=all_tokens,
+        atom_coords=np.zeros(all_tokens.shape + (3,)),
+        name=logging_name,
+    )
+    np_chains_list = []
+
+    input_chains_by_id = {chain.id: chain for chain in fold_input.chains}
+
+    nonempty_chain_ids = set(all_tokens.chain_id)
+    for chain_info in substruct.iter_chains():
+      chain_id = chain_info['chain_id']
+      chain_type = chain_info['chain_type']
+      chain = input_chains_by_id[chain_id]
+
+      # Generalised "sequence" for ligands (can't trust residue name)
+      chain_tokens = all_tokens[all_tokens.chain_id == chain_id]
+      assert chain_tokens.res_name is not None
+      three_letter_sequence = ','.join(chain_tokens.res_name.tolist())
+      chain_num_tokens = len(chain_tokens.atom_name)
+
+      # Don't compute features for chains not included in the crop, or ligands.
+      skip_chain = (
+          chain_type != mmcif_names.PROTEIN_CHAIN
+          or chain_num_tokens <= 4  # not cache filled
+          or chain_id not in nonempty_chain_ids
+      )
+
+      if three_letter_sequence in seen_entities:
+        entity_id = seen_entities[three_letter_sequence]
+      else:
+        entity_id = len(seen_entities) + 1
+
+      if entity_id not in polymer_entity_features[skip_chain]:
+        if skip_chain:
+          template_features = data3.empty_template_features(chain_num_tokens)
+        else:
+          assert isinstance(chain, folding_input.ProteinChain)
+
+          sorted_features = []
+          for template in chain.templates:
+            struc = structure.from_mmcif(
+                template.mmcif,
+                fix_mse_residues=True,
+                fix_arginines=True,
+                include_bonds=False,
+                include_water=False,
+                include_other=True,  # For non-standard polymer chains.
+            )
+            hit_features = templates.get_polymer_features(
+                chain=struc,
+                chain_poly_type=mmcif_names.PROTEIN_CHAIN,
+                query_sequence_length=len(chain.sequence),
+                query_to_hit_mapping=dict(template.query_to_template_map),
+            )
+            sorted_features.append(hit_features)
+
+          template_features = templates.package_template_features(
+              hit_features=sorted_features,
+              include_ligand_features=False,
+          )
+
+          template_features = data3.fix_template_features(
+              template_features=template_features, num_res=len(chain.sequence)
+          )
+
+        template_features = _reduce_template_features(
+            template_features, max_templates
+        )
+        polymer_entity_features[skip_chain][entity_id] = template_features
+
+      seen_entities[three_letter_sequence] = entity_id
+      feats = polymer_entity_features[skip_chain][entity_id].copy()
+      feats['chain_id'] = chain_id
+      np_chains_list.append(feats)
+
+    # We pad the num_templates dimension before merging, so that different
+    # chains can be concatenated on the num_res dimension.  Masking will be
+    # applied so that each chains templates can't see each other.
+    for chain in np_chains_list:
+      chain['template_aatype'] = _pad_to(
+          chain['template_aatype'], (max_templates, None)
+      )
+      chain['template_atom_positions'] = _pad_to(
+          chain['template_atom_positions'], (max_templates, None, None, None)
+      )
+      chain['template_atom_mask'] = _pad_to(
+          chain['template_atom_mask'], (max_templates, None, None)
+      )
+
+    # Merge on token dimension.
+    np_example = {
+        ft: np.concatenate([c[ft] for c in np_chains_list], axis=1)
+        for ft in np_chains_list[0]
+        if ft in data_constants.TEMPLATE_FEATURES
+    }
+
+    # Crop template data. Need to use the standard token indices, since msa does
+    # not expand non-standard residues. This means that for expanded residues,
+    # we get repeated template information.
+    for feature_name, v in np_example.items():
+      np_example[feature_name] = v[:max_templates, standard_token_idxs, ...]
+
+    # Pad along the token dimension.
+    templates_features = Templates(
+        aatype=_pad_to(
+            np_example['template_aatype'], (None, padding_shapes.num_tokens)
+        ),
+        atom_positions=_pad_to(
+            np_example['template_atom_positions'],
+            (None, padding_shapes.num_tokens, None, None),
+        ),
+        atom_mask=_pad_to(
+            np_example['template_atom_mask'].astype(bool),
+            (None, padding_shapes.num_tokens, None),
+        ),
+    )
+    return templates_features
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    """Make Template from batch dictionary."""
+    return cls(
+        aatype=batch['template_aatype'],
+        atom_positions=batch['template_atom_positions'],
+        atom_mask=batch['template_atom_mask'],
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'template_aatype': self.aatype,
+        'template_atom_positions': self.atom_positions,
+        'template_atom_mask': self.atom_mask,
+    }
+
+
+jax.tree_util.register_dataclass(
+    Templates,
+    data_fields=[f.name for f in dataclasses.fields(Templates)],
+    meta_fields=[],
+)
+
+
+def _reduce_template_features(
+    template_features: data3.FeatureDict,
+    max_templates: int,
+) -> data3.FeatureDict:
+  """Reduces template features to max num templates and defined feature set."""
+  num_templates = template_features['template_aatype'].shape[0]
+  template_keep_mask = np.arange(num_templates) < max_templates
+  template_fields = data_constants.TEMPLATE_FEATURES + (
+      'template_release_timestamp',
+  )
+  template_features = {
+      k: v[template_keep_mask]
+      for k, v in template_features.items()
+      if k in template_fields
+  }
+  return template_features
+
+
+@dataclasses.dataclass(frozen=True)
+class TokenFeatures:
+  """Dataclass containing features for tokens."""
+
+  residue_index: xnp_ndarray
+  token_index: xnp_ndarray
+  aatype: xnp_ndarray
+  mask: xnp_ndarray
+  seq_length: xnp_ndarray
+
+  # Chain symmetry identifiers
+  # for an A3B2 stoichiometry the meaning of these features is as follows:
+  # asym_id:    1 2 3 4 5
+  # entity_id:  1 1 1 2 2
+  # sym_id:     1 2 3 1 2
+  asym_id: xnp_ndarray
+  entity_id: xnp_ndarray
+  sym_id: xnp_ndarray
+
+  # token type features
+  is_protein: xnp_ndarray
+  is_rna: xnp_ndarray
+  is_dna: xnp_ndarray
+  is_ligand: xnp_ndarray
+  is_nonstandard_polymer_chain: xnp_ndarray
+  is_water: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Compute the per-token features."""
+
+    residue_index = all_tokens.res_id.astype(np.int32)
+
+    token_index = np.arange(1, len(all_tokens.atom_name) + 1).astype(np.int32)
+
+    aatype = []
+    for res_name, chain_type in zip(all_tokens.res_name, all_tokens.chain_type):
+      if chain_type in mmcif_names.POLYMER_CHAIN_TYPES:
+        res_name = mmcif_names.fix_non_standard_polymer_res(
+            res_name=res_name, chain_type=chain_type
+        )
+        if (
+            chain_type == mmcif_names.DNA_CHAIN
+            and res_name == residue_names.UNK_DNA
+        ):
+          res_name = residue_names.UNK_NUCLEIC_ONE_LETTER
+      elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
+        res_name = residue_names.UNK
+      else:
+        raise ValueError(f'Chain type {chain_type} not polymer or ligand.')
+      aa = residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP[res_name]
+      aatype.append(aa)
+    aatype = np.array(aatype, dtype=np.int32)
+
+    mask = np.ones(all_tokens.shape[0], dtype=bool)
+    chains = _compute_asym_entity_and_sym_id(all_tokens)
+    m = dict(zip(chains.chain_id, chains.asym_id))
+    asym_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)
+
+    m = dict(zip(chains.chain_id, chains.entity_id))
+    entity_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)
+
+    m = dict(zip(chains.chain_id, chains.sym_id))
+    sym_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)
+
+    seq_length = np.array(all_tokens.shape[0], dtype=np.int32)
+
+    is_protein = all_tokens.chain_type == mmcif_names.PROTEIN_CHAIN
+    is_rna = all_tokens.chain_type == mmcif_names.RNA_CHAIN
+    is_dna = all_tokens.chain_type == mmcif_names.DNA_CHAIN
+    is_ligand = np.isin(
+        all_tokens.chain_type, list(mmcif_names.LIGAND_CHAIN_TYPES)
+    )
+    standard_polymer_chain = list(mmcif_names.NON_POLYMER_CHAIN_TYPES) + list(
+        mmcif_names.STANDARD_POLYMER_CHAIN_TYPES
+    )
+    is_nonstandard_polymer_chain = np.isin(
+        all_tokens.chain_type, standard_polymer_chain, invert=True
+    )
+    is_water = all_tokens.chain_type == mmcif_names.WATER
+
+    return TokenFeatures(
+        residue_index=_pad_to(residue_index, (padding_shapes.num_tokens,)),
+        token_index=_pad_to(token_index, (padding_shapes.num_tokens,)),
+        aatype=_pad_to(aatype, (padding_shapes.num_tokens,)),
+        mask=_pad_to(mask, (padding_shapes.num_tokens,)),
+        asym_id=_pad_to(asym_id, (padding_shapes.num_tokens,)),
+        entity_id=_pad_to(entity_id, (padding_shapes.num_tokens,)),
+        sym_id=_pad_to(sym_id, (padding_shapes.num_tokens,)),
+        seq_length=seq_length,
+        is_protein=_pad_to(is_protein, (padding_shapes.num_tokens,)),
+        is_rna=_pad_to(is_rna, (padding_shapes.num_tokens,)),
+        is_dna=_pad_to(is_dna, (padding_shapes.num_tokens,)),
+        is_ligand=_pad_to(is_ligand, (padding_shapes.num_tokens,)),
+        is_nonstandard_polymer_chain=_pad_to(
+            is_nonstandard_polymer_chain, (padding_shapes.num_tokens,)
+        ),
+        is_water=_pad_to(is_water, (padding_shapes.num_tokens,)),
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        residue_index=batch['residue_index'],
+        token_index=batch['token_index'],
+        aatype=batch['aatype'],
+        mask=batch['seq_mask'],
+        entity_id=batch['entity_id'],
+        asym_id=batch['asym_id'],
+        sym_id=batch['sym_id'],
+        seq_length=batch['seq_length'],
+        is_protein=batch['is_protein'],
+        is_rna=batch['is_rna'],
+        is_dna=batch['is_dna'],
+        is_ligand=batch['is_ligand'],
+        is_nonstandard_polymer_chain=batch['is_nonstandard_polymer_chain'],
+        is_water=batch['is_water'],
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'residue_index': self.residue_index,
+        'token_index': self.token_index,
+        'aatype': self.aatype,
+        'seq_mask': self.mask,
+        'entity_id': self.entity_id,
+        'asym_id': self.asym_id,
+        'sym_id': self.sym_id,
+        'seq_length': self.seq_length,
+        'is_protein': self.is_protein,
+        'is_rna': self.is_rna,
+        'is_dna': self.is_dna,
+        'is_ligand': self.is_ligand,
+        'is_nonstandard_polymer_chain': self.is_nonstandard_polymer_chain,
+        'is_water': self.is_water,
+    }
+
+
+jax.tree_util.register_dataclass(
+    TokenFeatures,
+    data_fields=[f.name for f in dataclasses.fields(TokenFeatures)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class PredictedStructureInfo:
+  """Contains information necessary to work with predicted structure."""
+
+  atom_mask: xnp_ndarray
+  residue_center_index: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Compute the PredictedStructureInfo features.
+
+    Args:
+      all_tokens: flat AtomLayout with 1 representative atom per token, shape
+        (num_tokens,)
+      all_token_atoms_layout: AtomLayout for all atoms per token, shape
+        (num_tokens, max_atoms_per_token)
+      padding_shapes: padding shapes.
+
+    Returns:
+      A PredictedStructureInfo object.
+    """
+    atom_mask = _pad_to(
+        all_token_atoms_layout.atom_name.astype(bool),
+        (padding_shapes.num_tokens, None),
+    )
+    residue_center_index = np.zeros(padding_shapes.num_tokens, dtype=np.int32)
+    for idx in range(all_tokens.shape[0]):
+      repr_atom = all_tokens.atom_name[idx]
+      atoms = list(all_token_atoms_layout.atom_name[idx, :])
+      if repr_atom in atoms:
+        residue_center_index[idx] = atoms.index(repr_atom)
+      else:
+        # Representative atoms can be missing if cropping the number of atoms
+        # per residue.
+        logging.warning(
+            'The representative atom in all_tokens (%s) is not in '
+            'all_token_atoms_layout (%s)',
+            all_tokens[idx : idx + 1],
+            all_token_atoms_layout[idx, :],
+        )
+        residue_center_index[idx] = 0
+    return cls(atom_mask=atom_mask, residue_center_index=residue_center_index)
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        atom_mask=batch['pred_dense_atom_mask'],
+        residue_center_index=batch['residue_center_index'],
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'pred_dense_atom_mask': self.atom_mask,
+        'residue_center_index': self.residue_center_index,
+    }
+
+
+jax.tree_util.register_dataclass(
+    PredictedStructureInfo,
+    data_fields=[f.name for f in dataclasses.fields(PredictedStructureInfo)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class PolymerLigandBondInfo:
+  """Contains information about polymer-ligand bonds."""
+
+  tokens_to_polymer_ligand_bonds: atom_layout.GatherInfo
+  # Gather indices to convert from cropped dense atom layout to bonds layout
+  # (num_tokens, 2)
+  token_atoms_to_bonds: atom_layout.GatherInfo
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      bond_layout: atom_layout.AtomLayout | None,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Computes the InterChainBondInfo features.
+
+    Args:
+      all_tokens: AtomLayout for tokens; shape (num_tokens,).
+      all_token_atoms_layout: Atom Layout for all atoms (num_tokens,
+        max_atoms_per_token)
+      bond_layout: Bond layout for polymer-ligand bonds.
+      padding_shapes: Padding shapes.
+
+    Returns:
+      A PolymerLigandBondInfo object.
+    """
+
+    if bond_layout is not None:
+      # Must convert to list before calling np.isin, will not work raw.
+      peptide_types = list(mmcif_names.PEPTIDE_CHAIN_TYPES)
+      nucleic_types = list(mmcif_names.NUCLEIC_ACID_CHAIN_TYPES) + [
+          mmcif_names.OTHER_CHAIN
+      ]
+      # These atom renames are so that we can use the atom layout code with
+      # all_tokens, which only has a single atom per token.
+      atom_names = bond_layout.atom_name.copy()
+      atom_names[np.isin(bond_layout.chain_type, peptide_types)] = 'CA'
+      atom_names[np.isin(bond_layout.chain_type, nucleic_types)] = "C1'"
+      adjusted_bond_layout = atom_layout.AtomLayout(
+          atom_name=atom_names,
+          res_id=bond_layout.res_id,
+          chain_id=bond_layout.chain_id,
+          chain_type=bond_layout.chain_type,
+      )
+      # Remove bonds that are not in the crop.
+      cropped_tokens_to_bonds = atom_layout.compute_gather_idxs(
+          source_layout=all_tokens, target_layout=adjusted_bond_layout
+      )
+      bond_is_in_crop = np.all(
+          cropped_tokens_to_bonds.gather_mask, axis=1
+      ).astype(bool)
+      adjusted_bond_layout = adjusted_bond_layout[bond_is_in_crop, :]
+    else:
+      # Create layout with correct shape when bond_layout is None.
+      s = (0, 2)
+      adjusted_bond_layout = atom_layout.AtomLayout(
+          atom_name=np.array([], dtype=object).reshape(s),
+          res_id=np.array([], dtype=int).reshape(s),
+          chain_id=np.array([], dtype=object).reshape(s),
+      )
+    adjusted_bond_layout = adjusted_bond_layout.copy_and_pad_to(
+        (padding_shapes.num_tokens, 2)
+    )
+    tokens_to_polymer_ligand_bonds = atom_layout.compute_gather_idxs(
+        source_layout=all_tokens, target_layout=adjusted_bond_layout
+    )
+
+    # Stuff for computing the bond loss.
+    if bond_layout is not None:
+      # Pad to num_tokens (hoping that there are never more bonds than tokens).
+      padded_bond_layout = bond_layout.copy_and_pad_to(
+          (padding_shapes.num_tokens, 2)
+      )
+      token_atoms_to_bonds = atom_layout.compute_gather_idxs(
+          source_layout=all_token_atoms_layout, target_layout=padded_bond_layout
+      )
+    else:
+      token_atoms_to_bonds = atom_layout.GatherInfo(
+          gather_idxs=np.zeros((padding_shapes.num_tokens, 2), dtype=int),
+          gather_mask=np.zeros((padding_shapes.num_tokens, 2), dtype=bool),
+          input_shape=np.array((
+              padding_shapes.num_tokens,
+              all_token_atoms_layout.shape[1],
+          )),
+      )
+
+    return cls(
+        tokens_to_polymer_ligand_bonds=tokens_to_polymer_ligand_bonds,
+        token_atoms_to_bonds=token_atoms_to_bonds,
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        tokens_to_polymer_ligand_bonds=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='tokens_to_polymer_ligand_bonds'
+        ),
+        token_atoms_to_bonds=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='token_atoms_to_polymer_ligand_bonds'
+        ),
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        **self.tokens_to_polymer_ligand_bonds.as_dict(
+            key_prefix='tokens_to_polymer_ligand_bonds'
+        ),
+        **self.token_atoms_to_bonds.as_dict(
+            key_prefix='token_atoms_to_polymer_ligand_bonds'
+        ),
+    }
+
+
+jax.tree_util.register_dataclass(
+    PolymerLigandBondInfo,
+    data_fields=[f.name for f in dataclasses.fields(PolymerLigandBondInfo)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class LigandLigandBondInfo:
+  """Contains information about the location of ligand-ligand bonds."""
+
+  tokens_to_ligand_ligand_bonds: atom_layout.GatherInfo
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      bond_layout: atom_layout.AtomLayout | None,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Computes the InterChainBondInfo features.
+
+    Args:
+      all_tokens: AtomLayout for tokens; shape (num_tokens,).
+      bond_layout: Bond layout for ligand-ligand bonds.
+      padding_shapes: Padding shapes.
+
+    Returns:
+      A LigandLigandBondInfo object.
+    """
+
+    if bond_layout is not None:
+      # Discard any bonds that do not join to an existing atom.
+      keep_mask = []
+      all_atom_ids = {
+          uid
+          for uid in zip(
+              all_tokens.chain_id,
+              all_tokens.res_id,
+              all_tokens.atom_name,
+              strict=True,
+          )
+      }
+      for chain_id, res_id, atom_name in zip(
+          bond_layout.chain_id,
+          bond_layout.res_id,
+          bond_layout.atom_name,
+          strict=True,
+      ):
+        atom_a = (chain_id[0], res_id[0], atom_name[0])
+        atom_b = (chain_id[1], res_id[1], atom_name[1])
+        if atom_a in all_atom_ids and atom_b in all_atom_ids:
+          keep_mask.append(True)
+        else:
+          keep_mask.append(False)
+      keep_mask = np.array(keep_mask).astype(bool)
+      bond_layout = bond_layout[keep_mask]
+      # Remove any bonds to Hydrogen atoms.
+      bond_layout = bond_layout[
+          ~np.char.startswith(bond_layout.atom_name.astype(str), 'H').any(
+              axis=1
+          )
+      ]
+      atom_names = bond_layout.atom_name
+      adjusted_bond_layout = atom_layout.AtomLayout(
+          atom_name=atom_names,
+          res_id=bond_layout.res_id,
+          chain_id=bond_layout.chain_id,
+          chain_type=bond_layout.chain_type,
+      )
+    else:
+      # Create layout with correct shape when bond_layout is None.
+      s = (0, 2)
+      adjusted_bond_layout = atom_layout.AtomLayout(
+          atom_name=np.array([], dtype=object).reshape(s),
+          res_id=np.array([], dtype=int).reshape(s),
+          chain_id=np.array([], dtype=object).reshape(s),
+      )
+    # 10 x num_tokens as max_inter_bonds_ratio + max_intra_bonds_ration = 2.061.
+    adjusted_bond_layout = adjusted_bond_layout.copy_and_pad_to(
+        (padding_shapes.num_tokens * 10, 2)
+    )
+    gather_idx = atom_layout.compute_gather_idxs(
+        source_layout=all_tokens, target_layout=adjusted_bond_layout
+    )
+    return cls(tokens_to_ligand_ligand_bonds=gather_idx)
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        tokens_to_ligand_ligand_bonds=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='tokens_to_ligand_ligand_bonds'
+        )
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        **self.tokens_to_ligand_ligand_bonds.as_dict(
+            key_prefix='tokens_to_ligand_ligand_bonds'
+        )
+    }
+
+
+jax.tree_util.register_dataclass(
+    LigandLigandBondInfo,
+    data_fields=[f.name for f in dataclasses.fields(LigandLigandBondInfo)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class PseudoBetaInfo:
+  """Contains information for extracting pseudo-beta and equivalent atoms."""
+
+  token_atoms_to_pseudo_beta: atom_layout.GatherInfo
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      ccd: chemical_components.Ccd,
+      padding_shapes: PaddingShapes,
+      logging_name: str,
+  ) -> Self:
+    """Compute the PseudoBetaInfo features.
+
+    Args:
+      all_token_atoms_layout:  AtomLayout for all atoms per token, shape
+        (num_tokens, max_atoms_per_token)
+      ccd: The chemical components dictionary.
+      padding_shapes: padding shapes.
+      logging_name: logging name for debugging (usually the mmcif_id)
+
+    Returns:
+      A PseudoBetaInfo object.
+    """
+    token_idxs = []
+    atom_idxs = []
+    for token_idx in range(all_token_atoms_layout.shape[0]):
+      chain_type = all_token_atoms_layout.chain_type[token_idx, 0]
+      atom_names = list(all_token_atoms_layout.atom_name[token_idx, :])
+      atom_idx = None
+      is_nucleic_backbone = (
+          chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES
+          or chain_type == mmcif_names.OTHER_CHAIN
+      )
+      if chain_type == mmcif_names.PROTEIN_CHAIN:
+        # Protein chains
+        if 'CB' in atom_names:
+          atom_idx = atom_names.index('CB')
+        elif 'CA' in atom_names:
+          atom_idx = atom_names.index('CA')
+      elif is_nucleic_backbone:
+        # RNA / DNA chains
+        res_name = all_token_atoms_layout.res_name[token_idx, 0]
+        cifdict = ccd.get(res_name)
+        if cifdict:
+          parent = cifdict['_chem_comp.mon_nstd_parent_comp_id'][0]
+          if parent != '?':
+            res_name = parent
+        if res_name in {'A', 'G', 'DA', 'DG'}:
+          if 'C4' in atom_names:
+            atom_idx = atom_names.index('C4')
+        else:
+          if 'C2' in atom_names:
+            atom_idx = atom_names.index('C2')
+      elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
+        # Ligands: there is only one atom per token
+        atom_idx = 0
+      else:
+        logging.warning(
+            '%s: Unknown chain type for token %i. (%s)',
+            logging_name,
+            token_idx,
+            all_token_atoms_layout[token_idx : token_idx + 1],
+        )
+        atom_idx = 0
+      if atom_idx is None:
+        (valid_atom_idxs,) = np.nonzero(
+            all_token_atoms_layout.atom_name[token_idx, :]
+        )
+        if valid_atom_idxs.shape[0] > 0:
+          atom_idx = valid_atom_idxs[0]
+        else:
+          atom_idx = 0
+        logging.warning(
+            '%s token %i (%s), does not contain a pseudo-beta atom.'
+            'Using first valid atom (%s) instead.',
+            logging_name,
+            token_idx,
+            all_token_atoms_layout[token_idx : token_idx + 1],
+            all_token_atoms_layout.atom_name[token_idx, atom_idx],
+        )
+
+      token_idxs.append(token_idx)
+      atom_idxs.append(atom_idx)
+
+    pseudo_beta_layout = all_token_atoms_layout[token_idxs, atom_idxs]
+    pseudo_beta_layout = pseudo_beta_layout.copy_and_pad_to((
+        padding_shapes.num_tokens,
+    ))
+    token_atoms_to_pseudo_beta = atom_layout.compute_gather_idxs(
+        source_layout=all_token_atoms_layout, target_layout=pseudo_beta_layout
+    )
+
+    return cls(
+        token_atoms_to_pseudo_beta=token_atoms_to_pseudo_beta,
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        token_atoms_to_pseudo_beta=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='token_atoms_to_pseudo_beta'
+        ),
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        **self.token_atoms_to_pseudo_beta.as_dict(
+            key_prefix='token_atoms_to_pseudo_beta'
+        ),
+    }
+
+
+jax.tree_util.register_dataclass(
+    PseudoBetaInfo,
+    data_fields=[f.name for f in dataclasses.fields(PseudoBetaInfo)],
+    meta_fields=[],
+)
+
+
+_DEFAULT_BLANK_REF = {
+    'positions': np.zeros(3),
+    'mask': 0,
+    'element': 0,
+    'charge': 0,
+    'atom_name_chars': np.zeros(4),
+}
+
+
+def random_rotation(random_state: np.random.RandomState) -> np.ndarray:
+  # Create a random rotation (Gram-Schmidt orthogonalization of two
+  # random normal vectors)
+  v0, v1 = random_state.normal(size=(2, 3))
+  e0 = v0 / np.maximum(1e-10, np.linalg.norm(v0))
+  v1 = v1 - e0 * np.dot(v1, e0)
+  e1 = v1 / np.maximum(1e-10, np.linalg.norm(v1))
+  e2 = np.cross(e0, e1)
+  return np.stack([e0, e1, e2])
+
+
+def random_augmentation(
+    positions: np.ndarray,
+    random_state: np.random.RandomState,
+) -> np.ndarray:
+  """Center then apply random translation and rotation."""
+
+  center = np.mean(positions, axis=0)
+  rot = random_rotation(random_state)
+  positions_target = np.einsum('ij,kj->ki', rot, positions - center)
+
+  translation = random_state.normal(size=(3,))
+  positions_target = positions_target + translation
+  return positions_target
+
+
+def _get_reference_positions_from_ccd_cif(
+    ccd_cif: cif_dict.CifDict,
+    ref_max_modified_date: datetime.date,
+    logging_name: str,
+) -> np.ndarray:
+  """Creates reference positions from a CCD mmcif data block."""
+  num_atoms = len(ccd_cif['_chem_comp_atom.atom_id'])
+  if '_chem_comp_atom.pdbx_model_Cartn_x_ideal' in ccd_cif:
+    atom_x = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_x_ideal']
+    atom_y = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_y_ideal']
+    atom_z = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_z_ideal']
+  else:
+    atom_x = np.array(['?'] * num_atoms)
+    atom_y = np.array(['?'] * num_atoms)
+    atom_z = np.array(['?'] * num_atoms)
+  pos = np.array([[x, y, z] for x, y, z in zip(atom_x, atom_y, atom_z)])
+  # Unknown reference coordinates are specified by '?' in chem comp dict.
+  # Replace unknown reference coords with 0.
+  if '?' in pos and '_chem_comp.pdbx_modified_date' in ccd_cif:
+    # Use reference coordinates if modifed date is before cutoff.
+    modified_dates = [
+        datetime.date.fromisoformat(date)
+        for date in ccd_cif['_chem_comp.pdbx_modified_date']
+    ]
+    max_modified_date = max(modified_dates)
+    if max_modified_date < ref_max_modified_date:
+      atom_x = ccd_cif['_chem_comp_atom.model_Cartn_x']
+      atom_y = ccd_cif['_chem_comp_atom.model_Cartn_y']
+      atom_z = ccd_cif['_chem_comp_atom.model_Cartn_z']
+      pos = np.array([[x, y, z] for x, y, z in zip(atom_x, atom_y, atom_z)])
+  if '?' in pos:
+    if np.all(pos == '?'):
+      logging.warning('All ref positions unknown for: %s', logging_name)
+    else:
+      logging.warning('Some ref positions unknown for: %s', logging_name)
+    pos[pos == '?'] = 0
+  return np.array(pos, dtype=np.float32)
+
+
+def get_reference(
+    res_name: str,
+    chemical_components_data: struc_chem_comps.ChemicalComponentsData,
+    ccd: chemical_components.Ccd,
+    random_state: np.random.RandomState,
+    ref_max_modified_date: datetime.date,
+    conformer_max_iterations: int | None,
+) -> tuple[dict[str, Any], Any, Any]:
+  """Reference structure for residue from CCD or SMILES.
+
+  Uses CCD entry if available, otherwise uses SMILES from chemical components
+  data. Conformer generation is done using RDKit, with a fallback to CCD ideal
+  or reference coordinates if RDKit fails and those coordinates are supplied.
+
+  Args:
+    res_name: ccd code of the residue.
+    chemical_components_data: ChemicalComponentsData for making ref structure.
+    ccd: The chemical components dictionary.
+    random_state: Numpy RandomState
+    ref_max_modified_date: date beyond which reference structures must not be
+      modified to be allowed to use reference coordinates.
+    conformer_max_iterations: Optional override for maximum number of iterations
+      to run for RDKit conformer search.
+
+  Returns:
+    Mapping from atom names to features, from_atoms, dest_atoms.
+  """
+
+  ccd_cif = ccd.get(res_name)
+
+  mol = None
+  if ccd_cif:
+    try:
+      mol = rdkit_utils.mol_from_ccd_cif(ccd_cif, remove_hydrogens=False)
+    except rdkit_utils.MolFromMmcifError:
+      logging.warning('Failed to construct mol from ccd_cif for: %s', res_name)
+  else:  # No CCD entry, use SMILES from chemical components data.
+    if not (
+        chemical_components_data.chem_comp
+        and res_name in chemical_components_data.chem_comp
+        and chemical_components_data.chem_comp[res_name].pdbx_smiles
+    ):
+      raise ValueError(f'No CCD entry or SMILES for {res_name}.')
+    smiles_string = chemical_components_data.chem_comp[res_name].pdbx_smiles
+    logging.info('Using SMILES for: %s - %s', res_name, smiles_string)
+
+    mol = Chem.MolFromSmiles(smiles_string)
+    if mol is None:
+      # In this case the model will not have any information about this molecule
+      # and will not be able to predict anything about it.
+      raise ValueError(
+          f'Failed to construct RDKit Mol for {res_name} from SMILES string: '
+          f'{smiles_string} . This is likely due to an issue with the SMILES '
+          'string. Note that the userCCD input format provides an alternative '
+          'way to define custom molecules directly without RDKit or SMILES.'
+      )
+    mol = Chem.AddHs(mol)
+    # No existing names, we assign them from the graph.
+    mol = rdkit_utils.assign_atom_names_from_graph(mol)
+    # Temporary CCD cif with just atom and bond information, no coordinates.
+    ccd_cif = rdkit_utils.mol_to_ccd_cif(mol, component_id='fake_cif')
+
+  conformer = None
+  atom_names = []
+  elements = []
+  charges = []
+  pos = []
+
+  # If mol is not None (must be True for SMILES case), then we try and generate
+  # an RDKit conformer.
+  if mol is not None:
+    conformer_random_seed = int(random_state.randint(1, 1 << 31))
+    conformer = rdkit_utils.get_random_conformer(
+        mol=mol,
+        random_seed=conformer_random_seed,
+        max_iterations=conformer_max_iterations,
+        logging_name=res_name,
+    )
+    if conformer:
+      for idx, atom in enumerate(mol.GetAtoms()):
+        atom_names.append(atom.GetProp('atom_name'))
+        elements.append(atom.GetAtomicNum())
+        charges.append(atom.GetFormalCharge())
+        coords = conformer.GetAtomPosition(idx)
+        pos.append([coords.x, coords.y, coords.z])
+      pos = np.array(pos, dtype=np.float32)
+
+  # If no mol could be generated (can only happen when using CCD), or no
+  # conformer could be generated from the mol (can happen in either case), then
+  # use CCD cif instead (which will have zero coordinates for SMILES case).
+  if conformer is None:
+    atom_names = ccd_cif['_chem_comp_atom.atom_id']
+    charges = ccd_cif['_chem_comp_atom.charge']
+    type_symbols = ccd_cif['_chem_comp_atom.type_symbol']
+    elements = [
+        periodic_table.ATOMIC_NUMBER.get(elem_type.capitalize(), 0)
+        for elem_type in type_symbols
+    ]
+    pos = _get_reference_positions_from_ccd_cif(
+        ccd_cif=ccd_cif,
+        ref_max_modified_date=ref_max_modified_date,
+        logging_name=res_name,
+    )
+
+  # Augment reference positions.
+  pos = random_augmentation(pos, random_state)
+
+  # Extract atom and bond information from CCD cif.
+  from_atom = ccd_cif.get('_chem_comp_bond.atom_id_1', None)
+  dest_atom = ccd_cif.get('_chem_comp_bond.atom_id_2', None)
+
+  features = {}
+  for atom_name in atom_names:
+    features[atom_name] = {}
+    idx = atom_names.index(atom_name)
+    charge = 0 if charges[idx] == '?' else int(charges[idx])
+    atom_name_chars = np.array([ord(c) - 32 for c in atom_name], dtype=int)
+    atom_name_chars = _pad_to(atom_name_chars, (4,))
+    features[atom_name]['positions'] = pos[idx]
+    features[atom_name]['mask'] = 1
+    features[atom_name]['element'] = elements[idx]
+    features[atom_name]['charge'] = charge
+    features[atom_name]['atom_name_chars'] = atom_name_chars
+  return features, from_atom, dest_atom
+
+
+@dataclasses.dataclass(frozen=True)
+class RefStructure:
+  """Contains ref structure information."""
+
+  # Array with positions, float32, shape [num_res, max_atoms_per_token, 3]
+  positions: xnp_ndarray
+  # Array with masks, bool, shape [num_res, max_atoms_per_token]
+  mask: xnp_ndarray
+  # Array with elements, int32, shape [num_res, max_atoms_per_token]
+  element: xnp_ndarray
+  # Array with charges, float32, shape [num_res, max_atoms_per_token]
+  charge: xnp_ndarray
+  # Array with atom name characters, int32, [num_res, max_atoms_per_token, 4]
+  atom_name_chars: xnp_ndarray
+  # Array with reference space uids, int32, [num_res, max_atoms_per_token]
+  ref_space_uid: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      ccd: chemical_components.Ccd,
+      padding_shapes: PaddingShapes,
+      chemical_components_data: struc_chem_comps.ChemicalComponentsData,
+      random_state: np.random.RandomState,
+      ref_max_modified_date: datetime.date,
+      conformer_max_iterations: int | None,
+      ligand_ligand_bonds: atom_layout.AtomLayout | None = None,
+  ) -> tuple[Self, Any]:
+    """Reference structure information for each residue."""
+
+    # Get features per atom
+    padded_shape = (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
+    result = {
+        'positions': np.zeros((*padded_shape, 3), 'float32'),
+        'mask': np.zeros(padded_shape, 'bool'),
+        'element': np.zeros(padded_shape, 'int32'),
+        'charge': np.zeros(padded_shape, 'float32'),
+        'atom_name_chars': np.zeros((*padded_shape, 4), 'int32'),
+        'ref_space_uid': np.zeros((*padded_shape,), 'int32'),
+    }
+
+    atom_names_all = []
+    chain_ids_all = []
+    res_ids_all = []
+
+    # Cache reference conformations for each residue.
+    conformations = {}
+    ref_space_uids = {}
+    for idx in np.ndindex(all_token_atoms_layout.shape):
+      chain_id = all_token_atoms_layout.chain_id[idx]
+      res_id = all_token_atoms_layout.res_id[idx]
+      res_name = all_token_atoms_layout.res_name[idx]
+      is_non_standard = res_name not in _STANDARD_RESIDUES
+      atom_name = all_token_atoms_layout.atom_name[idx]
+      if not atom_name:
+        ref = _DEFAULT_BLANK_REF
+      else:
+        if (chain_id, res_id) not in conformations:
+          conf, from_atom, dest_atom = get_reference(
+              res_name=res_name,
+              chemical_components_data=chemical_components_data,
+              ccd=ccd,
+              random_state=random_state,
+              ref_max_modified_date=ref_max_modified_date,
+              conformer_max_iterations=conformer_max_iterations,
+          )
+          conformations[(chain_id, res_id)] = conf
+
+          if (
+              is_non_standard
+              and (from_atom is not None)
+              and (dest_atom is not None)
+          ):
+            # Add intra-ligand bond graph
+            atom_names_ligand = np.stack(
+                [from_atom, dest_atom], axis=1, dtype=object
+            )
+            atom_names_all.append(atom_names_ligand)
+            res_ids_all.append(
+                np.full_like(atom_names_ligand, res_id, dtype=int)
+            )
+            chain_ids_all.append(
+                np.full_like(atom_names_ligand, chain_id, dtype=object)
+            )
+
+        conformation = conformations.get(
+            (chain_id, res_id), {atom_name: _DEFAULT_BLANK_REF}
+        )
+        if atom_name not in conformation:
+          logging.warning(
+              'Missing atom "%s" for CCD "%s"',
+              atom_name,
+              all_token_atoms_layout.res_name[idx],
+          )
+        ref = conformation.get(atom_name, _DEFAULT_BLANK_REF)
+      for k in ref:
+        result[k][idx] = ref[k]
+
+      # Assign a unique reference space id to each component, to determine which
+      # reference positions live in the same reference space.
+      space_str_id = (
+          all_token_atoms_layout.chain_id[idx],
+          all_token_atoms_layout.res_id[idx],
+      )
+      if space_str_id not in ref_space_uids:
+        ref_space_uids[space_str_id] = len(ref_space_uids)
+      result['ref_space_uid'][idx] = ref_space_uids[space_str_id]
+
+    if atom_names_all:
+      atom_names_all = np.concatenate(atom_names_all, axis=0)
+      res_ids_all = np.concatenate(res_ids_all, axis=0)
+      chain_ids_all = np.concatenate(chain_ids_all, axis=0)
+      if ligand_ligand_bonds is not None:
+        adjusted_ligand_ligand_bonds = atom_layout.AtomLayout(
+            atom_name=np.concatenate(
+                [ligand_ligand_bonds.atom_name, atom_names_all], axis=0
+            ),
+            chain_id=np.concatenate(
+                [ligand_ligand_bonds.chain_id, chain_ids_all], axis=0
+            ),
+            res_id=np.concatenate(
+                [ligand_ligand_bonds.res_id, res_ids_all], axis=0
+            ),
+        )
+      else:
+        adjusted_ligand_ligand_bonds = atom_layout.AtomLayout(
+            atom_name=atom_names_all,
+            chain_id=chain_ids_all,
+            res_id=res_ids_all,
+        )
+    else:
+      adjusted_ligand_ligand_bonds = ligand_ligand_bonds
+
+    return cls(**result), adjusted_ligand_ligand_bonds
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        positions=batch['ref_pos'],
+        mask=batch['ref_mask'],
+        element=batch['ref_element'],
+        charge=batch['ref_charge'],
+        atom_name_chars=batch['ref_atom_name_chars'],
+        ref_space_uid=batch['ref_space_uid'],
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'ref_pos': self.positions,
+        'ref_mask': self.mask,
+        'ref_element': self.element,
+        'ref_charge': self.charge,
+        'ref_atom_name_chars': self.atom_name_chars,
+        'ref_space_uid': self.ref_space_uid,
+    }
+
+
+jax.tree_util.register_dataclass(
+    RefStructure,
+    data_fields=[f.name for f in dataclasses.fields(RefStructure)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class ConvertModelOutput:
+  """Contains atom layout info."""
+
+  cleaned_struc: structure.Structure
+  token_atoms_layout: atom_layout.AtomLayout
+  flat_output_layout: atom_layout.AtomLayout
+  empty_output_struc: structure.Structure
+  polymer_ligand_bonds: atom_layout.AtomLayout
+  ligand_ligand_bonds: atom_layout.AtomLayout
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      padding_shapes: PaddingShapes,
+      cleaned_struc: structure.Structure,
+      flat_output_layout: atom_layout.AtomLayout,
+      empty_output_struc: structure.Structure,
+      polymer_ligand_bonds: atom_layout.AtomLayout,
+      ligand_ligand_bonds: atom_layout.AtomLayout,
+  ) -> Self:
+    """Pads the all_token_atoms_layout and stores other data."""
+    # Crop and pad the all_token_atoms_layout.
+    token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
+        (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
+    )
+
+    return cls(
+        cleaned_struc=cleaned_struc,
+        token_atoms_layout=token_atoms_layout,
+        flat_output_layout=flat_output_layout,
+        empty_output_struc=empty_output_struc,
+        polymer_ligand_bonds=polymer_ligand_bonds,
+        ligand_ligand_bonds=ligand_ligand_bonds,
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    """Construct atom layout object from dictionary."""
+
+    return cls(
+        cleaned_struc=_unwrap(batch.get('cleaned_struc', None)),
+        token_atoms_layout=_unwrap(batch.get('token_atoms_layout', None)),
+        flat_output_layout=_unwrap(batch.get('flat_output_layout', None)),
+        empty_output_struc=_unwrap(batch.get('empty_output_struc', None)),
+        polymer_ligand_bonds=_unwrap(batch.get('polymer_ligand_bonds', None)),
+        ligand_ligand_bonds=_unwrap(batch.get('ligand_ligand_bonds', None)),
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        'cleaned_struc': np.array(self.cleaned_struc, object),
+        'token_atoms_layout': np.array(self.token_atoms_layout, object),
+        'flat_output_layout': np.array(self.flat_output_layout, object),
+        'empty_output_struc': np.array(self.empty_output_struc, object),
+        'polymer_ligand_bonds': np.array(self.polymer_ligand_bonds, object),
+        'ligand_ligand_bonds': np.array(self.ligand_ligand_bonds, object),
+    }
+
+
+jax.tree_util.register_dataclass(
+    ConvertModelOutput,
+    data_fields=[f.name for f in dataclasses.fields(ConvertModelOutput)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class AtomCrossAtt:
+  """Operate on flat atoms."""
+
+  token_atoms_to_queries: atom_layout.GatherInfo
+  tokens_to_queries: atom_layout.GatherInfo
+  tokens_to_keys: atom_layout.GatherInfo
+  queries_to_keys: atom_layout.GatherInfo
+  queries_to_token_atoms: atom_layout.GatherInfo
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_token_atoms_layout: atom_layout.AtomLayout,  # (num_tokens, num_dense)
+      queries_subset_size: int,
+      keys_subset_size: int,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Computes gather indices and meta data to work with a flat atom list."""
+
+    token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
+        (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
+    )
+    token_atoms_mask = token_atoms_layout.atom_name.astype(bool)
+    flat_layout = token_atoms_layout[token_atoms_mask]
+    num_atoms = flat_layout.shape[0]
+
+    padded_flat_layout = flat_layout.copy_and_pad_to((
+        padding_shapes.num_atoms,
+    ))
+
+    # Create the layout for queries
+    num_subsets = padding_shapes.num_atoms // queries_subset_size
+    lay_arr = padded_flat_layout.to_array()
+    queries_layout = atom_layout.AtomLayout.from_array(
+        lay_arr.reshape((6, num_subsets, queries_subset_size))
+    )
+
+    # Create the layout for the keys (the key subsets are centered around the
+    # query subsets)
+    # Create initial gather indices (contain out-of-bound indices)
+    subset_centers = np.arange(
+        queries_subset_size / 2, padding_shapes.num_atoms, queries_subset_size
+    )
+    flat_to_key_gathers = (
+        subset_centers[:, None]
+        + np.arange(-keys_subset_size / 2, keys_subset_size / 2)[None, :]
+    )
+    flat_to_key_gathers = flat_to_key_gathers.astype(int)
+    # Shift subsets with out-of-bound indices, such that they are fully within
+    # the bounds.
+    for row in range(flat_to_key_gathers.shape[0]):
+      if flat_to_key_gathers[row, 0] < 0:
+        flat_to_key_gathers[row, :] -= flat_to_key_gathers[row, 0]
+      elif flat_to_key_gathers[row, -1] > num_atoms - 1:
+        overflow = flat_to_key_gathers[row, -1] - (num_atoms - 1)
+        flat_to_key_gathers[row, :] -= overflow
+    # Create the keys layout.
+    keys_layout = padded_flat_layout[flat_to_key_gathers]
+
+    # Create gather indices for conversion between token atoms layout,
+    # queries layout and keys layout.
+    token_atoms_to_queries = atom_layout.compute_gather_idxs(
+        source_layout=token_atoms_layout, target_layout=queries_layout
+    )
+
+    token_atoms_to_keys = atom_layout.compute_gather_idxs(
+        source_layout=token_atoms_layout, target_layout=keys_layout
+    )
+
+    queries_to_keys = atom_layout.compute_gather_idxs(
+        source_layout=queries_layout, target_layout=keys_layout
+    )
+
+    queries_to_token_atoms = atom_layout.compute_gather_idxs(
+        source_layout=queries_layout, target_layout=token_atoms_layout
+    )
+
+    # Create gather indices for conversion of tokens layout to
+    # queries and keys layout
+    token_idxs = np.arange(padding_shapes.num_tokens).astype(np.int64)
+    token_idxs = np.broadcast_to(token_idxs[:, None], token_atoms_layout.shape)
+    tokens_to_queries = atom_layout.GatherInfo(
+        gather_idxs=atom_layout.convert(
+            token_atoms_to_queries, token_idxs, layout_axes=(0, 1)
+        ),
+        gather_mask=atom_layout.convert(
+            token_atoms_to_queries, token_atoms_mask, layout_axes=(0, 1)
+        ),
+        input_shape=np.array((padding_shapes.num_tokens,)),
+    )
+
+    tokens_to_keys = atom_layout.GatherInfo(
+        gather_idxs=atom_layout.convert(
+            token_atoms_to_keys, token_idxs, layout_axes=(0, 1)
+        ),
+        gather_mask=atom_layout.convert(
+            token_atoms_to_keys, token_atoms_mask, layout_axes=(0, 1)
+        ),
+        input_shape=np.array((padding_shapes.num_tokens,)),
+    )
+
+    return cls(
+        token_atoms_to_queries=token_atoms_to_queries,
+        tokens_to_queries=tokens_to_queries,
+        tokens_to_keys=tokens_to_keys,
+        queries_to_keys=queries_to_keys,
+        queries_to_token_atoms=queries_to_token_atoms,
+    )
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(
+        token_atoms_to_queries=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='token_atoms_to_queries'
+        ),
+        tokens_to_queries=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='tokens_to_queries'
+        ),
+        tokens_to_keys=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='tokens_to_keys'
+        ),
+        queries_to_keys=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='queries_to_keys'
+        ),
+        queries_to_token_atoms=atom_layout.GatherInfo.from_dict(
+            batch, key_prefix='queries_to_token_atoms'
+        ),
+    )
+
+  def as_data_dict(self) -> BatchDict:
+    return {
+        **self.token_atoms_to_queries.as_dict(
+            key_prefix='token_atoms_to_queries'
+        ),
+        **self.tokens_to_queries.as_dict(key_prefix='tokens_to_queries'),
+        **self.tokens_to_keys.as_dict(key_prefix='tokens_to_keys'),
+        **self.queries_to_keys.as_dict(key_prefix='queries_to_keys'),
+        **self.queries_to_token_atoms.as_dict(
+            key_prefix='queries_to_token_atoms'
+        ),
+    }
+
+
+jax.tree_util.register_dataclass(
+    AtomCrossAtt,
+    data_fields=[f.name for f in dataclasses.fields(AtomCrossAtt)],
+    meta_fields=[],
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class Frames:
+  """Features for backbone frames."""
+
+  mask: xnp_ndarray
+
+  @classmethod
+  def compute_features(
+      cls,
+      all_tokens: atom_layout.AtomLayout,
+      all_token_atoms_layout: atom_layout.AtomLayout,
+      ref_structure: RefStructure,
+      padding_shapes: PaddingShapes,
+  ) -> Self:
+    """Computes features for backbone frames."""
+    num_tokens = padding_shapes.num_tokens
+    all_token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
+        (num_tokens, all_token_atoms_layout.shape[1])
+    )
+
+    all_token_atoms_to_all_tokens = atom_layout.compute_gather_idxs(
+        source_layout=all_token_atoms_layout, target_layout=all_tokens
+    )
+    ref_coordinates = atom_layout.convert(
+        all_token_atoms_to_all_tokens,
+        ref_structure.positions.astype(np.float32),
+        layout_axes=(0, 1),
+    )
+    ref_mask = atom_layout.convert(
+        all_token_atoms_to_all_tokens,
+        ref_structure.mask.astype(bool),
+        layout_axes=(0, 1),
+    )
+    ref_mask = ref_mask & all_token_atoms_to_all_tokens.gather_mask.astype(bool)
+
+    all_frame_mask = []
+
+    # Iterate over tokens
+    for idx, args in enumerate(
+        zip(all_tokens.chain_type, all_tokens.chain_id, all_tokens.res_id)
+    ):
+
+      chain_type, chain_id, res_id = args
+
+      if chain_type in list(mmcif_names.PEPTIDE_CHAIN_TYPES):
+        frame_mask = True
+      elif chain_type in list(mmcif_names.NUCLEIC_ACID_CHAIN_TYPES):
+        frame_mask = True
+      elif chain_type in list(mmcif_names.NON_POLYMER_CHAIN_TYPES):
+        # For ligands, build frames from closest atoms from the same molecule.
+        (local_token_idxs,) = np.where(
+            (all_tokens.chain_type == chain_type)
+            & (all_tokens.chain_id == chain_id)
+            & (all_tokens.res_id == res_id)
+        )
+
+        if len(local_token_idxs) < 3:
+          frame_mask = False
+
+        else:
+          # [local_tokens]
+          local_dist = np.linalg.norm(
+              ref_coordinates[idx] - ref_coordinates[local_token_idxs], axis=-1
+          )
+          local_mask = ref_mask[local_token_idxs]
+          cost = local_dist + 1e8 * ~local_mask
+          cost = cost + 1e8 * (idx == local_token_idxs)
+          # [local_tokens]
+          closest_idxs = np.argsort(cost, axis=0)
+
+          # The closest indices index an array of local tokens. Convert this
+          # to indices of the full (num_tokens,) array.
+          global_closest_idxs = local_token_idxs[closest_idxs]
+
+          # Construct frame by placing the current token at the origin and two
+          # nearest atoms on either side.
+          global_frame_idxs = np.array(
+              (global_closest_idxs[0], idx, global_closest_idxs[1])
+          )
+
+          # Check that the frame atoms are not colinear.
+          a, b, c = ref_coordinates[global_frame_idxs]
+          vec1 = a - b
+          vec2 = c - b
+          # Reference coordinates can be all zeros, in which case we have
+          # to explicitly set colinearity.
+          if np.isclose(np.linalg.norm(vec1, axis=-1), 0) or np.isclose(
+              np.linalg.norm(vec2, axis=-1), 0
+          ):
+            is_colinear = True
+            logging.info('Found identical coordinates: Assigning as colinear.')
+          else:
+            vec1 = vec1 / np.linalg.norm(vec1, axis=-1)
+            vec2 = vec2 / np.linalg.norm(vec2, axis=-1)
+            cos_angle = np.einsum('...k,...k->...', vec1, vec2)
+            # <25 degree deviation is considered colinear.
+            is_colinear = 1 - np.abs(cos_angle) < 0.0937
+
+          frame_mask = not is_colinear
+      else:
+        # No frame for other chain types.
+        frame_mask = False
+
+      all_frame_mask.append(frame_mask)
+
+    all_frame_mask = np.array(all_frame_mask, dtype=bool)
+
+    mask = _pad_to(all_frame_mask, (padding_shapes.num_tokens,))
+
+    return cls(mask=mask)
+
+  @classmethod
+  def from_data_dict(cls, batch: BatchDict) -> Self:
+    return cls(mask=batch['frames_mask'])
+
+  def as_data_dict(self) -> BatchDict:
+    return {'frames_mask': self.mask}
+
+
+jax.tree_util.register_dataclass(
+    Frames,
+    data_fields=[f.name for f in dataclasses.fields(Frames)],
+    meta_fields=[],
+)

flax_model/alphafold3/model/mkdssp_pybind.cc

@@ -0,0 +1,63 @@
+// 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 "alphafold3/model/mkdssp_pybind.h"
+
+#include <filesystem>
+
+#include <cif++/file.hpp>
+#include <cif++/pdb.hpp>
+#include <dssp.hpp>
+#include <sstream>
+
+#include "absl/strings/string_view.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+
+namespace alphafold3 {
+namespace py = pybind11;
+
+void RegisterModuleMkdssp(pybind11::module m) {
+  py::module site = py::module::import("site");
+  py::list paths = py::cast<py::list>(site.attr("getsitepackages")());
+  // Find the first path that contains the libcifpp components.cif file.
+  bool found = false;
+  for (const auto& py_path : paths) {
+    auto path_str =
+        std::filesystem::path(py::cast<absl::string_view>(py_path)) /
+        "share/libcifpp/components.cif";
+    if (std::filesystem::exists(path_str)) {
+      setenv("LIBCIFPP_DATA_DIR", path_str.parent_path().c_str(), 0);
+      found = true;
+      break;
+    }
+  }
+  if (!found) {
+    throw py::type_error("Could not find the libcifpp components.cif file.");
+  }
+  m.def(
+      "get_dssp",
+      [](absl::string_view mmcif, int model_no,
+         int min_poly_proline_stretch_length,
+         bool calculate_surface_accessibility) {
+        cif::file cif_file(mmcif.data(), mmcif.size());
+        dssp result(cif_file.front(), model_no, min_poly_proline_stretch_length,
+                    calculate_surface_accessibility);
+        std::stringstream sstream;
+        result.write_legacy_output(sstream);
+        return sstream.str();
+      },
+      py::arg("mmcif"), py::arg("model_no") = 1,
+      py::arg("min_poly_proline_stretch_length") = 3,
+      py::arg("calculate_surface_accessibility") = false,
+      py::doc("Gets secondary structure from an mmCIF file."));
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/model/mkdssp_pybind.h

@@ -0,0 +1,26 @@
+/*
+ * 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
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_MODEL_MKDSSP_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_MODEL_MKDSSP_PYBIND_H_
+
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleMkdssp(pybind11::module m);
+
+}
+
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_MODEL_MKDSSP_PYBIND_H_

flax_model/alphafold3/model/mmcif_metadata.py

@@ -0,0 +1,222 @@
+
+
+"""Adds mmCIF metadata (to be ModelCIF-conformant) and author and legal info."""
+
+from typing import Final
+
+from flax_model.alphafold3.structure import mmcif
+import numpy as np
+
+_LICENSE_URL: Final[str] = (
+    'https://github.com/google-deepmind/alphafold3/blob/main/OUTPUT_TERMS_OF_USE.md'
+)
+
+_LICENSE: Final[str] = f"""
+Non-commercial use only, by using this file you agree to the terms of use found
+at {_LICENSE_URL}.
+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.
+""".strip()
+
+_DISCLAIMER: Final[str] = """\
+AlphaFold 3 and its output are not intended for, have not been validated for,
+and are not approved for clinical use. They are provided "as-is" without any
+warranty of any kind, whether expressed or implied. No warranty is given that
+use shall not infringe the rights of any third party.
+""".strip()
+
+_MMCIF_PAPER_AUTHORS: Final[tuple[str, ...]] = (
+    'Google DeepMind',
+    'Isomorphic Labs',
+)
+
+# Authors of the mmCIF - we set them to be equal to the authors of the paper.
+_MMCIF_AUTHORS: Final[tuple[str, ...]] = _MMCIF_PAPER_AUTHORS
+
+
+def add_metadata_to_mmcif(
+    old_cif: mmcif.Mmcif, version: str, model_id: bytes
+) -> mmcif.Mmcif:
+  """Adds metadata to a mmCIF to make it ModelCIF-conformant."""
+  cif = {}
+
+  # ModelCIF conformation dictionary.
+  cif['_audit_conform.dict_name'] = ['mmcif_ma.dic']
+  cif['_audit_conform.dict_version'] = ['1.4.5']
+  cif['_audit_conform.dict_location'] = [
+      'https://raw.githubusercontent.com/ihmwg/ModelCIF/master/dist/mmcif_ma.dic'
+  ]
+
+  cif['_pdbx_data_usage.id'] = ['1', '2']
+  cif['_pdbx_data_usage.type'] = ['license', 'disclaimer']
+  cif['_pdbx_data_usage.details'] = [_LICENSE, _DISCLAIMER]
+  cif['_pdbx_data_usage.url'] = [_LICENSE_URL, '?']
+
+  # Structure author details.
+  cif['_audit_author.name'] = []
+  cif['_audit_author.pdbx_ordinal'] = []
+  for author_index, author_name in enumerate(_MMCIF_AUTHORS, start=1):
+    cif['_audit_author.name'].append(author_name)
+    cif['_audit_author.pdbx_ordinal'].append(str(author_index))
+
+  # Paper author details.
+  cif['_citation_author.citation_id'] = []
+  cif['_citation_author.name'] = []
+  cif['_citation_author.ordinal'] = []
+  for author_index, author_name in enumerate(_MMCIF_PAPER_AUTHORS, start=1):
+    cif['_citation_author.citation_id'].append('primary')
+    cif['_citation_author.name'].append(author_name)
+    cif['_citation_author.ordinal'].append(str(author_index))
+
+  # Paper citation details.
+  cif['_citation.id'] = ['primary']
+  cif['_citation.title'] = [
+      'Accurate structure prediction of biomolecular interactions with'
+      ' AlphaFold 3'
+  ]
+  cif['_citation.journal_full'] = ['Nature']
+  cif['_citation.journal_volume'] = ['630']
+  cif['_citation.page_first'] = ['493']
+  cif['_citation.page_last'] = ['500']
+  cif['_citation.year'] = ['2024']
+  cif['_citation.journal_id_ASTM'] = ['NATUAS']
+  cif['_citation.country'] = ['UK']
+  cif['_citation.journal_id_ISSN'] = ['0028-0836']
+  cif['_citation.journal_id_CSD'] = ['0006']
+  cif['_citation.book_publisher'] = ['?']
+  cif['_citation.pdbx_database_id_PubMed'] = ['38718835']
+  cif['_citation.pdbx_database_id_DOI'] = ['10.1038/s41586-024-07487-w']
+
+  # Type of data in the dataset including data used in the model generation.
+  cif['_ma_data.id'] = ['1']
+  cif['_ma_data.name'] = ['Model']
+  cif['_ma_data.content_type'] = ['model coordinates']
+
+  # Description of number of instances for each entity.
+  cif['_ma_target_entity_instance.asym_id'] = old_cif['_struct_asym.id']
+  cif['_ma_target_entity_instance.entity_id'] = old_cif[
+      '_struct_asym.entity_id'
+  ]
+  cif['_ma_target_entity_instance.details'] = ['.'] * len(
+      cif['_ma_target_entity_instance.entity_id']
+  )
+
+  # Details about the target entities.
+  cif['_ma_target_entity.entity_id'] = cif[
+      '_ma_target_entity_instance.entity_id'
+  ]
+  cif['_ma_target_entity.data_id'] = ['1'] * len(
+      cif['_ma_target_entity.entity_id']
+  )
+  cif['_ma_target_entity.origin'] = ['.'] * len(
+      cif['_ma_target_entity.entity_id']
+  )
+
+  # Details of the models being deposited.
+  cif['_ma_model_list.ordinal_id'] = ['1']
+  cif['_ma_model_list.model_id'] = ['1']
+  cif['_ma_model_list.model_group_id'] = ['1']
+  cif['_ma_model_list.model_name'] = ['Top ranked model']
+
+  cif['_ma_model_list.model_group_name'] = [
+      f'AlphaFold-beta-20231127 ({version})'
+  ]
+  cif['_ma_model_list.data_id'] = ['1']
+  cif['_ma_model_list.model_type'] = ['Ab initio model']
+
+  # Software used.
+  cif['_software.pdbx_ordinal'] = ['1']
+  cif['_software.name'] = ['AlphaFold']
+  cif['_software.version'] = [
+      f'AlphaFold-beta-20231127 ({model_id.decode("ascii")})'
+  ]
+  cif['_software.type'] = ['package']
+  cif['_software.description'] = ['Structure prediction']
+  cif['_software.classification'] = ['other']
+  cif['_software.date'] = ['?']
+
+  # Collection of software into groups.
+  cif['_ma_software_group.ordinal_id'] = ['1']
+  cif['_ma_software_group.group_id'] = ['1']
+  cif['_ma_software_group.software_id'] = ['1']
+
+  # Method description to conform with ModelCIF.
+  cif['_ma_protocol_step.ordinal_id'] = ['1', '2', '3']
+  cif['_ma_protocol_step.protocol_id'] = ['1', '1', '1']
+  cif['_ma_protocol_step.step_id'] = ['1', '2', '3']
+  cif['_ma_protocol_step.method_type'] = [
+      'coevolution MSA',
+      'template search',
+      'modeling',
+  ]
+
+  # Details of the metrics use to assess model confidence.
+  cif['_ma_qa_metric.id'] = ['1', '2']
+  cif['_ma_qa_metric.name'] = ['pLDDT', 'pLDDT']
+  # Accepted values are distance, energy, normalised score, other, zscore.
+  cif['_ma_qa_metric.type'] = ['pLDDT', 'pLDDT']
+  cif['_ma_qa_metric.mode'] = ['global', 'local']
+  cif['_ma_qa_metric.software_group_id'] = ['1', '1']
+
+  # Global model confidence pLDDT value.
+  cif['_ma_qa_metric_global.ordinal_id'] = ['1']
+  cif['_ma_qa_metric_global.model_id'] = ['1']
+  cif['_ma_qa_metric_global.metric_id'] = ['1']
+  # Mean over all atoms, since AlphaFold 3 outputs pLDDT per-atom.
+  global_plddt = np.mean(
+      [float(v) for v in old_cif['_atom_site.B_iso_or_equiv']]
+  )
+  cif['_ma_qa_metric_global.metric_value'] = [f'{global_plddt:.2f}']
+
+  # Local (per residue) model confidence pLDDT value.
+  cif['_ma_qa_metric_local.ordinal_id'] = []
+  cif['_ma_qa_metric_local.model_id'] = []
+  cif['_ma_qa_metric_local.label_asym_id'] = []
+  cif['_ma_qa_metric_local.label_seq_id'] = []
+  cif['_ma_qa_metric_local.label_comp_id'] = []
+  cif['_ma_qa_metric_local.metric_id'] = []
+  cif['_ma_qa_metric_local.metric_value'] = []
+
+  plddt_grouped_by_res = {}
+  for *res, atom_plddt in zip(
+      old_cif['_atom_site.label_asym_id'],
+      old_cif['_atom_site.label_seq_id'],
+      old_cif['_atom_site.label_comp_id'],
+      old_cif['_atom_site.B_iso_or_equiv'],
+  ):
+    plddt_grouped_by_res.setdefault(tuple(res), []).append(float(atom_plddt))
+
+  for ordinal_id, ((chain_id, res_id, res_name), res_plddts) in enumerate(
+      plddt_grouped_by_res.items(), start=1
+  ):
+    res_plddt = np.mean(res_plddts)
+    cif['_ma_qa_metric_local.ordinal_id'].append(str(ordinal_id))
+    cif['_ma_qa_metric_local.model_id'].append('1')
+    cif['_ma_qa_metric_local.label_asym_id'].append(chain_id)
+    cif['_ma_qa_metric_local.label_seq_id'].append(res_id)
+    cif['_ma_qa_metric_local.label_comp_id'].append(res_name)
+    cif['_ma_qa_metric_local.metric_id'].append('2')  # See _ma_qa_metric.id.
+    cif['_ma_qa_metric_local.metric_value'].append(f'{res_plddt:.2f}')
+
+  cif['_atom_type.symbol'] = sorted(set(old_cif['_atom_site.type_symbol']))
+
+  return old_cif.copy_and_update(cif)
+
+
+def add_legal_comment(cif: str) -> str:
+  """Adds legal comment at the top of the mmCIF."""
+  # fmt: off
+  # pylint: disable=line-too-long
+  comment = (
+      '# By using this file you agree to the legally binding terms of use found at\n'
+      f'# {_LICENSE_URL}.\n'
+      '# To request access to the AlphaFold 3 model parameters, follow the process set\n'
+      '# out at https://github.com/google-deepmind/alphafold3. You may only use these if\n'
+      '# received directly from Google. Use is subject to terms of use available at\n'
+      '# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md.'
+  )
+  # pylint: enable=line-too-long
+  # fmt: on
+  return f'{comment}\n{cif}'

flax_model/alphafold3/model/model.py

@@ -0,0 +1,533 @@
+
+
+"""AlphaFold3 model."""
+
+from collections.abc import Iterable, Mapping
+import concurrent
+import dataclasses
+import functools
+from typing import Any, TypeAlias
+
+from absl import logging
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.model import confidences
+from flax_model.alphafold3.model import feat_batch
+from flax_model.alphafold3.model import features
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.model.components import mapping
+from flax_model.alphafold3.model.components import utils
+from flax_model.alphafold3.model.network import atom_cross_attention
+from flax_model.alphafold3.model.network import confidence_head
+from flax_model.alphafold3.model.network import diffusion_head
+from flax_model.alphafold3.model.network import distogram_head
+from flax_model.alphafold3.model.network import evoformer as evoformer_network
+from flax_model.alphafold3.model.network import featurization
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+ModelResult: TypeAlias = Mapping[str, Any]
+_ScalarNumberOrArray: TypeAlias = Mapping[str, float | int | np.ndarray]
+
+
+@dataclasses.dataclass(frozen=True)
+class InferenceResult:
+  """Postprocessed model result.
+
+  Attributes:
+    predicted_structure: Predicted protein structure.
+    numerical_data: Useful numerical data (scalars or arrays) to be saved at
+      inference time.
+    metadata: Smaller numerical data (usually scalar) to be saved as inference
+      metadata.
+    debug_outputs: Additional dict for debugging, e.g. raw outputs of a model
+      forward pass.
+    model_id: Model identifier.
+  """
+
+  predicted_structure: structure.Structure = dataclasses.field()
+  numerical_data: _ScalarNumberOrArray = dataclasses.field(default_factory=dict)
+  metadata: _ScalarNumberOrArray = dataclasses.field(default_factory=dict)
+  debug_outputs: Mapping[str, Any] = dataclasses.field(default_factory=dict)
+  model_id: bytes = b''
+
+
+def get_predicted_structure(
+    result: ModelResult, batch: feat_batch.Batch
+) -> structure.Structure:
+  """Creates the predicted structure and ion preditions.
+
+  Args:
+    result: model output in a model specific layout
+    batch: model input batch
+
+  Returns:
+    Predicted structure.
+  """
+  model_output_coords = result['diffusion_samples']['atom_positions']
+
+  # Rearrange model output coordinates to the flat output layout.
+  model_output_to_flat = atom_layout.compute_gather_idxs(
+      source_layout=batch.convert_model_output.token_atoms_layout,
+      target_layout=batch.convert_model_output.flat_output_layout,
+  )
+  pred_flat_atom_coords = atom_layout.convert(
+      gather_info=model_output_to_flat,
+      arr=model_output_coords,
+      layout_axes=(-3, -2),
+  )
+
+  predicted_lddt = result.get('predicted_lddt')
+
+  if predicted_lddt is not None:
+    pred_flat_b_factors = atom_layout.convert(
+        gather_info=model_output_to_flat,
+        arr=predicted_lddt,
+        layout_axes=(-2, -1),
+    )
+  else:
+    # Handle models which don't have predicted_lddt outputs.
+    pred_flat_b_factors = np.zeros(pred_flat_atom_coords.shape[:-1])
+
+  (missing_atoms_indices,) = np.nonzero(model_output_to_flat.gather_mask == 0)
+  if missing_atoms_indices.shape[0] > 0:
+    missing_atoms_flat_layout = batch.convert_model_output.flat_output_layout[
+        missing_atoms_indices
+    ]
+    missing_atoms_uids = list(
+        zip(
+            missing_atoms_flat_layout.chain_id,
+            missing_atoms_flat_layout.res_id,
+            missing_atoms_flat_layout.res_name,
+            missing_atoms_flat_layout.atom_name,
+        )
+    )
+    logging.warning(
+        'Target %s: warning: %s atoms were not predicted by the '
+        'model, setting their coordinates to (0, 0, 0). '
+        'Missing atoms: %s',
+        batch.convert_model_output.empty_output_struc.name,
+        missing_atoms_indices.shape[0],
+        missing_atoms_uids,
+    )
+
+  # Put them into a structure
+  pred_struc = batch.convert_model_output.empty_output_struc
+  pred_struc = pred_struc.copy_and_update_atoms(
+      atom_x=pred_flat_atom_coords[..., 0],
+      atom_y=pred_flat_atom_coords[..., 1],
+      atom_z=pred_flat_atom_coords[..., 2],
+      atom_b_factor=pred_flat_b_factors,
+      atom_occupancy=np.ones(pred_flat_atom_coords.shape[:-1]),  # Always 1.0.
+  )
+  # Set manually/differently when adding metadata.
+  pred_struc = pred_struc.copy_and_update_globals(release_date=None)
+  return pred_struc
+
+
+def create_target_feat_embedding(
+    batch: feat_batch.Batch,
+    config: evoformer_network.Evoformer.Config,
+    global_config: model_config.GlobalConfig,
+) -> jnp.ndarray:
+  """Create target feature embedding."""
+
+  dtype = jnp.bfloat16 if global_config.bfloat16 == 'all' else jnp.float32
+
+  with utils.bfloat16_context():
+    target_feat = featurization.create_target_feat(
+        batch,
+        append_per_atom_features=False,
+    ).astype(dtype)
+
+    enc = atom_cross_attention.atom_cross_att_encoder(
+        token_atoms_act=None,
+        trunk_single_cond=None,
+        trunk_pair_cond=None,
+        config=config.per_atom_conditioning,
+        global_config=global_config,
+        batch=batch,
+        name='evoformer_conditioning',
+    )
+    target_feat = jnp.concatenate([target_feat, enc.token_act], axis=-1).astype(
+        dtype
+    )
+
+  return target_feat
+
+
+def _compute_ptm(
+    result: ModelResult,
+    num_tokens: int,
+    asym_id: np.ndarray,
+    pae_single_mask: np.ndarray,
+    interface: bool,
+) -> np.ndarray:
+  """Computes the pTM metrics from PAE."""
+  return np.stack(
+      [
+          confidences.predicted_tm_score(
+              tm_adjusted_pae=tm_adjusted_pae[:num_tokens, :num_tokens],
+              asym_id=asym_id,
+              pair_mask=pae_single_mask[:num_tokens, :num_tokens],
+              interface=interface,
+          )
+          for tm_adjusted_pae in result['tmscore_adjusted_pae_global']
+      ],
+      axis=0,
+  )
+
+
+def _compute_chain_pair_iptm(
+    num_tokens: int,
+    asym_ids: np.ndarray,
+    mask: np.ndarray,
+    tm_adjusted_pae: np.ndarray,
+) -> np.ndarray:
+  """Computes the chain pair ipTM metrics from PAE."""
+  return np.stack(
+      [
+          confidences.chain_pairwise_predicted_tm_scores(
+              tm_adjusted_pae=sample_tm_adjusted_pae[:num_tokens],
+              asym_id=asym_ids[:num_tokens],
+              pair_mask=mask[:num_tokens, :num_tokens],
+          )
+          for sample_tm_adjusted_pae in tm_adjusted_pae
+      ],
+      axis=0,
+  )
+
+
+class Model(hk.Module):
+  """Full model. Takes in data batch and returns model outputs."""
+
+  class HeadsConfig(base_config.BaseConfig):
+    diffusion: diffusion_head.DiffusionHead.Config = base_config.autocreate()
+    confidence: confidence_head.ConfidenceHead.Config = base_config.autocreate()
+    distogram: distogram_head.DistogramHead.Config = base_config.autocreate()
+
+  class Config(base_config.BaseConfig):
+    evoformer: evoformer_network.Evoformer.Config = base_config.autocreate()
+    global_config: model_config.GlobalConfig = base_config.autocreate()
+    heads: 'Model.HeadsConfig' = base_config.autocreate()
+    num_recycles: int = 10
+    return_embeddings: bool = False
+    return_distogram: bool = False
+
+  def __init__(self, config: Config, name: str = 'diffuser'):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = config.global_config
+    self.diffusion_module = diffusion_head.DiffusionHead(
+        self.config.heads.diffusion, self.global_config
+    )
+
+  @hk.transparent
+  def _sample_diffusion(
+      self,
+      batch: feat_batch.Batch,
+      embeddings: dict[str, jnp.ndarray],
+      *,
+      sample_config: diffusion_head.SampleConfig,
+  ) -> dict[str, jnp.ndarray]:
+    denoising_step = functools.partial(
+        self.diffusion_module,
+        batch=batch,
+        embeddings=embeddings,
+        use_conditioning=True,
+    )
+
+    sample = diffusion_head.sample(
+        denoising_step=denoising_step,
+        batch=batch,
+        key=hk.next_rng_key(),
+        config=sample_config,
+    )
+    return sample
+
+  def __call__(
+      self, batch: features.BatchDict, key: jax.Array | None = None
+  ) -> ModelResult:
+    if key is None:
+      key = hk.next_rng_key()
+
+    batch = feat_batch.Batch.from_data_dict(batch)
+
+    embedding_module = evoformer_network.Evoformer(
+        self.config.evoformer, self.global_config
+    )
+    target_feat = create_target_feat_embedding(
+        batch=batch,
+        config=embedding_module.config,
+        global_config=self.global_config,
+    )
+
+    num_res = batch.num_res
+    if num_res>=4000:
+      def recycle_body(_, args):
+        prev, key = args
+        key, subkey = jax.random.split(key)
+        embeddings = embedding_module(
+            batch=batch,
+            prev=prev,
+            target_feat=target_feat,
+            key=subkey,
+        )
+        embeddings['pair'] = embeddings['pair'].astype(jnp.bfloat16)
+        embeddings['single'] = embeddings['single'].astype(jnp.bfloat16)
+        return embeddings, key
+      
+      embeddings = {
+          'pair': jnp.zeros(
+              [num_res, num_res, self.config.evoformer.pair_channel],
+              dtype=jnp.bfloat16,
+          ),
+          'single': jnp.zeros(
+              [num_res, self.config.evoformer.seq_channel], dtype=jnp.bfloat16
+          ),
+          'target_feat': target_feat,
+      }
+    else:
+      def recycle_body(_, args):
+        prev, key = args
+        key, subkey = jax.random.split(key)
+        embeddings = embedding_module(
+            batch=batch,
+            prev=prev,
+            target_feat=target_feat,
+            key=subkey,
+        )
+        embeddings['pair'] = embeddings['pair'].astype(jnp.float32)
+        embeddings['single'] = embeddings['single'].astype(jnp.float32)
+        return embeddings, key
+      
+      embeddings = {
+          'pair': jnp.zeros(
+              [num_res, num_res, self.config.evoformer.pair_channel],
+              dtype=jnp.float32,
+          ),
+          'single': jnp.zeros(
+              [num_res, self.config.evoformer.seq_channel], dtype=jnp.float32
+          ),
+          'target_feat': target_feat,
+      }
+    if hk.running_init():
+      embeddings, _ = recycle_body(None, (embeddings, key))
+    else:
+      # Number of recycles is number of additional forward trunk passes.
+      num_iter = self.config.num_recycles + 1
+      embeddings, _ = hk.fori_loop(0, num_iter, recycle_body, (embeddings, key))
+
+    samples = self._sample_diffusion(
+        batch,
+        embeddings,
+        sample_config=self.config.heads.diffusion.eval,
+    )
+
+    # Compute dist_error_fn over all samples for distance error logging.
+    confidence_output = mapping.sharded_map(
+        lambda dense_atom_positions: confidence_head.ConfidenceHead(
+            self.config.heads.confidence, self.global_config
+        )(
+            dense_atom_positions=dense_atom_positions,
+            embeddings=embeddings,
+            seq_mask=batch.token_features.mask,
+            token_atoms_to_pseudo_beta=batch.pseudo_beta_info.token_atoms_to_pseudo_beta,
+            asym_id=batch.token_features.asym_id,
+        ),
+        in_axes=0,
+    )(samples['atom_positions'])
+
+    distogram = distogram_head.DistogramHead(
+        self.config.heads.distogram, self.global_config
+    )(batch, embeddings, return_distogram=self.config.return_distogram)
+
+    output = {
+        'diffusion_samples': samples,
+        'distogram': distogram,
+        **confidence_output,
+    }
+    if self.config.return_embeddings:
+      output['single_embeddings'] = embeddings['single']
+      output['pair_embeddings'] = embeddings['pair']
+    return output
+
+  @classmethod
+  def get_inference_result(
+      cls,
+      batch: features.BatchDict,
+      result: ModelResult,
+      target_name: str = '',
+  ) -> Iterable[InferenceResult]:
+    """Get the predicted structure, scalars, and arrays for inference.
+
+    This function also computes any inference-time quantities, which are not a
+    part of the forward-pass, e.g. additional confidence scores. Note that this
+    function is not serialized, so it should be slim if possible.
+
+    Args:
+      batch: data batch used for model inference, incl. TPU invalid types.
+      result: output dict from the model's forward pass.
+      target_name: target name to be saved within structure.
+
+    Yields:
+      inference_result: dataclass object that contains a predicted structure,
+      important inference-time scalars and arrays, as well as a slightly trimmed
+      dictionary of raw model result from the forward pass (for debugging).
+    """
+    del target_name
+    batch = feat_batch.Batch.from_data_dict(batch)
+
+    # Retrieve structure and construct a predicted structure.
+    pred_structure = get_predicted_structure(result=result, batch=batch)
+
+    num_tokens = batch.token_features.seq_length.item()
+
+    pae_single_mask = np.tile(
+        batch.frames.mask[:, None],
+        [1, batch.frames.mask.shape[0]],
+    )
+    ptm = _compute_ptm(
+        result=result,
+        num_tokens=num_tokens,
+        asym_id=batch.token_features.asym_id[:num_tokens],
+        pae_single_mask=pae_single_mask,
+        interface=False,
+    )
+    iptm = _compute_ptm(
+        result=result,
+        num_tokens=num_tokens,
+        asym_id=batch.token_features.asym_id[:num_tokens],
+        pae_single_mask=pae_single_mask,
+        interface=True,
+    )
+    ptm_iptm_average = 0.8 * iptm + 0.2 * ptm
+
+    asym_ids = batch.token_features.asym_id[:num_tokens]
+    # Map asym IDs back to chain IDs. Asym IDs are constructed from chain IDs by
+    # iterating over the chain IDs, and for each unique chain ID incrementing
+    # the asym ID by 1 and mapping it to the particular chain ID. Asym IDs are
+    # 1-indexed, so subtract 1 to get back to the chain ID.
+    chain_ids = [pred_structure.chains[asym_id - 1] for asym_id in asym_ids]
+    res_ids = batch.token_features.residue_index[:num_tokens]
+
+    if len(np.unique(asym_ids[:num_tokens])) > 1:
+      # There is more than one chain, hence interface pTM (i.e. ipTM) defined,
+      # so use it.
+      ranking_confidence = ptm_iptm_average
+    else:
+      # There is only one chain, hence ipTM=NaN, so use just pTM.
+      ranking_confidence = ptm
+
+    contact_probs = result['distogram']['contact_probs']
+    # Compute PAE related summaries.
+    _, chain_pair_pae_min, _ = confidences.chain_pair_pae(
+        num_tokens=num_tokens,
+        asym_ids=batch.token_features.asym_id,
+        full_pae=result['full_pae'],
+        mask=pae_single_mask,
+    )
+    chain_pair_pde_mean, chain_pair_pde_min = confidences.chain_pair_pde(
+        num_tokens=num_tokens,
+        asym_ids=batch.token_features.asym_id,
+        full_pde=result['full_pde'],
+    )
+    intra_chain_single_pde, cross_chain_single_pde, _ = confidences.pde_single(
+        num_tokens,
+        batch.token_features.asym_id,
+        result['full_pde'],
+        contact_probs,
+    )
+    pae_metrics = confidences.pae_metrics(
+        num_tokens=num_tokens,
+        asym_ids=batch.token_features.asym_id,
+        full_pae=result['full_pae'],
+        mask=pae_single_mask,
+        contact_probs=contact_probs,
+        tm_adjusted_pae=result['tmscore_adjusted_pae_interface'],
+    )
+    ranking_confidence_pae = confidences.rank_metric(
+        result['full_pae'],
+        contact_probs * batch.frames.mask[:, None].astype(float),
+    )
+    chain_pair_iptm = _compute_chain_pair_iptm(
+        num_tokens=num_tokens,
+        asym_ids=batch.token_features.asym_id,
+        mask=pae_single_mask,
+        tm_adjusted_pae=result['tmscore_adjusted_pae_interface'],
+    )
+    # iptm_ichain is a vector of per-chain ptm values. iptm_ichain[0],
+    # for example, is just the zeroth diagonal entry of the chain pair iptm
+    # matrix:
+    # [[x, , ],
+    #  [ , , ],
+    #  [ , , ]]]
+    iptm_ichain = chain_pair_iptm.diagonal(axis1=-2, axis2=-1)
+    # iptm_xchain is a vector of cross-chain interactions for each chain.
+    # iptm_xchain[0], for example, is an average of chain 0's interactions with
+    # other chains:
+    # [[ ,x,x],
+    #  [x, , ],
+    #  [x, , ]]]
+    iptm_xchain = confidences.get_iptm_xchain(chain_pair_iptm)
+
+    predicted_distance_errors = result['average_pde']
+
+    # Computing solvent accessible area with dssp can be slow for large
+    # structures with lots of chains, so we parallelize the call.
+    pred_structures = pred_structure.unstack()
+    num_workers = len(pred_structures)
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=num_workers
+    ) as executor:
+      has_clash = list(executor.map(confidences.has_clash, pred_structures))
+      fraction_disordered = list(
+          executor.map(confidences.fraction_disordered, pred_structures)
+      )
+
+    for idx, pred_structure in enumerate(pred_structures):
+      ranking_score = confidences.get_ranking_score(
+          ptm=ptm[idx],
+          iptm=iptm[idx],
+          fraction_disordered_=fraction_disordered[idx],
+          has_clash_=has_clash[idx],
+      )
+      yield InferenceResult(
+          predicted_structure=pred_structure,
+          numerical_data={
+              'full_pde': result['full_pde'][idx, :num_tokens, :num_tokens],
+              'full_pae': result['full_pae'][idx, :num_tokens, :num_tokens],
+              'contact_probs': contact_probs[:num_tokens, :num_tokens],
+          },
+          metadata={
+              'predicted_distance_error': predicted_distance_errors[idx],
+              'ranking_score': ranking_score,
+              'fraction_disordered': fraction_disordered[idx],
+              'has_clash': has_clash[idx],
+              'predicted_tm_score': ptm[idx],
+              'interface_predicted_tm_score': iptm[idx],
+              'chain_pair_pde_mean': chain_pair_pde_mean[idx],
+              'chain_pair_pde_min': chain_pair_pde_min[idx],
+              'chain_pair_pae_min': chain_pair_pae_min[idx],
+              'ptm': ptm[idx],
+              'iptm': iptm[idx],
+              'ptm_iptm_average': ptm_iptm_average[idx],
+              'intra_chain_single_pde': intra_chain_single_pde[idx],
+              'cross_chain_single_pde': cross_chain_single_pde[idx],
+              'pae_ichain': pae_metrics['pae_ichain'][idx],
+              'pae_xchain': pae_metrics['pae_xchain'][idx],
+              'ranking_confidence': ranking_confidence[idx],
+              'ranking_confidence_pae': ranking_confidence_pae[idx],
+              'chain_pair_iptm': chain_pair_iptm[idx],
+              'iptm_ichain': iptm_ichain[idx],
+              'iptm_xchain': iptm_xchain[idx],
+              'token_chain_ids': chain_ids,
+              'token_res_ids': res_ids,
+          },
+          model_id=result['__identifier__'],
+          debug_outputs={},
+      )

flax_model/alphafold3/model/model_config.py

@@ -0,0 +1,24 @@
+
+
+"""Global config for the model."""
+
+from collections.abc import Sequence
+from typing import Literal, TypeAlias
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.jax.attention import attention
+
+
+_Shape2DType: TypeAlias = tuple[int | None, int | None]
+
+
+class GlobalConfig(base_config.BaseConfig):
+  bfloat16: Literal['all', 'none', 'intermediate'] = 'all'
+  final_init: Literal['zeros', 'linear'] = 'zeros'
+  pair_attention_chunk_size: Sequence[_Shape2DType] = ((1536, 128), (None, 32))
+  pair_transition_shard_spec: Sequence[_Shape2DType] = (
+      (2048, None),
+      (None, 1024),
+  )
+  # Note: flash_attention_implementation = 'xla' means no flash attention.
+  flash_attention_implementation: attention.Implementation = 'triton'

flax_model/alphafold3/model/msa_pairing.py

@@ -0,0 +1,304 @@
+
+
+"""Functions for producing "paired" and "unpaired" MSA features for each chain.
+
+The paired MSA:
+- Is made from the result of the all_seqs MSA query.
+- Is ordered such that you can concatenate features across chains and related
+  sequences will end up on the same row. Related here means "from the same
+  species". Gaps are added to facilitate this whenever a sequence has no
+  suitable pair.
+
+The unpaired MSA:
+- Is made from the results of the remaining MSA queries.
+- Has no special ordering properties.
+- Is deduplicated such that it doesn't contain any sequences in the paired MSA.
+"""
+
+from typing import Mapping, MutableMapping, Sequence
+from flax_model.alphafold3.model import data_constants
+import numpy as np
+
+
+def _align_species(
+    all_species: Sequence[bytes],
+    chains_species_to_rows: Sequence[Mapping[bytes, np.ndarray]],
+    min_hits_per_species: Mapping[bytes, int],
+) -> np.ndarray:
+  """Aligns MSA row indices based on species.
+
+  Within a species, MSAs are aligned based on their original order (the first
+  sequence for a species in the first chain's MSA is aligned to the first
+  sequence for the same species in the second chain's MSA).
+
+  Args:
+    all_species: A list of all unique species identifiers.
+    chains_species_to_rows: A dictionary for each chain, that maps species to
+      the set of MSA row indices from that species in that chain.
+    min_hits_per_species: A mapping from species id, to the minimum MSA size
+      across chains for that species (ignoring chains with zero hits).
+
+  Returns:
+    A matrix of size [num_msa_rows, num_chains], where the i,j element is an
+    index into the jth chains MSA.  Each row consists of sequences from each
+    chain for the same species (or -1 if that chain has no sequences for that
+    species).
+  """
+  # Each species block is of size [num_seqs x num_chains] and consists of
+  # indices into the respective MSAs that have been aligned and are all for the
+  # same species.
+  species_blocks = []
+  for species in all_species:
+    chain_row_indices = []
+    for species_to_rows in chains_species_to_rows:
+      min_msa_size = min_hits_per_species[species]
+      if species not in species_to_rows:
+        # If a given chain has no hits for a species then we pad it with -1's,
+        # later on these values are used to make sure each feature is padded
+        # with its appropriate pad value.
+        row_indices = np.full(min_msa_size, fill_value=-1, dtype=np.int32)
+      else:
+        # We crop down to the smallest MSA for a given species across chains.
+        row_indices = species_to_rows[species][:min_msa_size]
+      chain_row_indices.append(row_indices)
+    species_block = np.stack(chain_row_indices, axis=1)
+    species_blocks.append(species_block)
+  aligned_matrix = np.concatenate(species_blocks, axis=0)
+  return aligned_matrix
+
+
+def create_paired_features(
+    chains: Sequence[MutableMapping[str, np.ndarray]],
+    max_paired_sequences: int,
+    nonempty_chain_ids: set[str],
+    max_hits_per_species: int,
+) -> Sequence[MutableMapping[str, np.ndarray]]:
+  """Creates per-chain MSA features where the MSAs have been aligned.
+
+  Args:
+    chains: A list of feature dicts, one for each chain.
+    max_paired_sequences: No more than this many paired sequences will be
+      returned from this function.
+    nonempty_chain_ids: A set of chain ids (str) that are included in the crop
+      there is no reason to process chains not in this list.
+    max_hits_per_species: No more than this number of sequences will be returned
+      for a given species.
+
+  Returns:
+    An updated feature dictionary for each chain, where the {}_all_seq features
+    have been aligned so that the nth row in chain 1 is aligned to the nth row
+    in chain 2's features.
+  """
+  # The number of chains that the given species appears in - we rank hits
+  # across more chains higher.
+  species_num_chains = {}
+
+  # For each chain we keep a mapping from species to the row indices in the
+  # original MSA for that chain.
+  chains_species_to_rows = []
+
+  # Keep track of the minimum number of hits across chains for a given species.
+  min_hits_per_species = {}
+
+  for chain in chains:
+    species_ids = chain['msa_species_identifiers_all_seq']
+
+    # The query gets an empty species_id, so no pairing happens for this row.
+    if (
+        species_ids.size == 0
+        or (species_ids.size == 1 and not species_ids[0])
+        or chain['chain_id'] not in nonempty_chain_ids
+    ):
+      chains_species_to_rows.append({})
+      continue
+
+    # For each species keep track of which row indices in the original MSA are
+    # from this species.
+    row_indices = np.arange(len(species_ids))
+    # The grouping np.split code requires that the input is already clustered
+    # by species id.
+    sort_idxs = species_ids.argsort()
+    species_ids = species_ids[sort_idxs]
+    row_indices = row_indices[sort_idxs]
+
+    species, unique_row_indices = np.unique(species_ids, return_index=True)
+    grouped_row_indices = np.split(row_indices, unique_row_indices[1:])
+    species_to_rows = dict(zip(species, grouped_row_indices, strict=True))
+    chains_species_to_rows.append(species_to_rows)
+
+    for s in species:
+      species_num_chains[s] = species_num_chains.get(s, 0) + 1
+
+    for species, row_indices in species_to_rows.items():
+      min_hits_per_species[species] = min(
+          min_hits_per_species.get(species, max_hits_per_species),
+          len(row_indices),
+      )
+
+  # Construct a mapping from the number of chains a species appears in to
+  # the list of species with that count.
+  num_chains_to_species = {}
+  for species, num_chains in species_num_chains.items():
+    if not species or num_chains <= 1:
+      continue
+    if num_chains not in num_chains_to_species:
+      num_chains_to_species[num_chains] = []
+    num_chains_to_species[num_chains].append(species)
+
+  num_rows_seen = 0
+  # We always keep the first row as it is the query sequence.
+  all_rows = [np.array([[0] * len(chains)], dtype=np.int32)]
+
+  # We prioritize species that have hits across more chains.
+  for num_chains in sorted(num_chains_to_species, reverse=True):
+    all_species = num_chains_to_species[num_chains]
+
+    # Align all the per-chain row indices by species, so every paired row is
+    # for a single species.
+    rows = _align_species(
+        all_species, chains_species_to_rows, min_hits_per_species
+    )
+    # Sort rows by the product of the original indices in the respective chain
+    # MSAS, so as to rank hits that appear higher in the original MSAs higher.
+    rank_metric = np.abs(np.prod(rows.astype(np.float32), axis=1))
+    sorted_rows = rows[np.argsort(rank_metric), :]
+    all_rows.append(sorted_rows)
+    num_rows_seen += rows.shape[0]
+    if num_rows_seen >= max_paired_sequences:
+      break
+
+  all_rows = np.concatenate(all_rows, axis=0)
+  all_rows = all_rows[:max_paired_sequences, :]
+
+  # Now we just have to select the relevant rows from the original msa and
+  # deletion matrix features
+  paired_chains = []
+  for chain_idx, chain in enumerate(chains):
+    out_chain = {k: v for k, v in chain.items() if 'all_seq' not in k}
+    selected_row_indices = all_rows[:, chain_idx]
+    for feat_name in {'msa', 'deletion_matrix'}:
+      all_seq_name = f'{feat_name}_all_seq'
+      feat_value = chain[all_seq_name]
+
+      # The selected row indices are padded to be the same shape for each chain,
+      # they are padded with -1's, so we add a single row onto the feature with
+      # the appropriate pad value.  This has the effect that we correctly pad
+      # each feature since all padded indices will select this padding row.
+      pad_value = data_constants.MSA_PAD_VALUES[feat_name]
+      feat_value = np.concatenate([
+          feat_value,
+          np.full((1, feat_value.shape[1]), pad_value, feat_value.dtype),
+      ])
+
+      feat_value = feat_value[selected_row_indices, :]
+      out_chain[all_seq_name] = feat_value
+    out_chain['num_alignments_all_seq'] = np.array(
+        out_chain['msa_all_seq'].shape[0]
+    )
+    paired_chains.append(out_chain)
+  return paired_chains
+
+
+def deduplicate_unpaired_sequences(
+    np_chains: Sequence[MutableMapping[str, np.ndarray]],
+) -> Sequence[MutableMapping[str, np.ndarray]]:
+  """Deduplicates unpaired sequences based on paired sequences."""
+
+  feature_names = np_chains[0].keys()
+  msa_features = (
+      data_constants.NUM_SEQ_MSA_FEATURES
+      + data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES
+  )
+
+  for chain in np_chains:
+    sequence_set = set(
+        hash(s.data.tobytes()) for s in chain['msa_all_seq'].astype(np.int8)
+    )
+    keep_rows = []
+    # Go through unpaired MSA seqs and remove any rows that correspond to the
+    # sequences that are already present in the paired MSA.
+    for row_num, seq in enumerate(chain['msa'].astype(np.int8)):
+      if hash(seq.data.tobytes()) not in sequence_set:
+        keep_rows.append(row_num)
+    for feature_name in feature_names:
+      if feature_name in msa_features:
+        chain[feature_name] = chain[feature_name][keep_rows]
+    chain['num_alignments'] = np.array(chain['msa'].shape[0], dtype=np.int32)
+  return np_chains
+
+
+def choose_paired_unpaired_msa_crop_sizes(
+    unpaired_msa: np.ndarray,
+    paired_msa: np.ndarray | None,
+    total_msa_crop_size: int,
+    max_paired_sequences: int,
+) -> tuple[int, int | None]:
+  """Returns the sizes of the MSA crop and MSA_all_seq crop.
+
+  NOTE: Unpaired + paired MSA sizes can exceed total_msa_size when
+    there are lots of gapped rows. Through the pairing logic another chain(s)
+    will have fewer than total_msa_size.
+
+  Args:
+     unpaired_msa: The unpaired MSA array (not all_seq).
+     paired_msa: The paired MSA array (all_seq).
+     total_msa_crop_size: The maximum total number of sequences to crop to.
+     max_paired_sequences: The maximum number of sequences that can come from
+       MSA pairing.
+
+  Returns:
+    A tuple of:
+      The size of the reduced MSA crop (not all_seq features).
+      The size of the unreduced MSA crop (for all_seq features) or None, if
+        paired_msa is None.
+  """
+  if paired_msa is not None:
+    paired_crop_size = np.minimum(paired_msa.shape[0], max_paired_sequences)
+
+    # We reduce the number of un-paired sequences, by the number of times a
+    # sequence from this chains MSA is included in the paired MSA.  This keeps
+    # the MSA size for each chain roughly constant.
+    cropped_all_seq_msa = paired_msa[:max_paired_sequences]
+    num_non_gapped_pairs = cropped_all_seq_msa.shape[0]
+
+    assert num_non_gapped_pairs <= max_paired_sequences
+    unpaired_crop_size = np.minimum(
+        unpaired_msa.shape[0], total_msa_crop_size - num_non_gapped_pairs
+    )
+    assert unpaired_crop_size >= 0
+  else:
+    unpaired_crop_size = np.minimum(unpaired_msa.shape[0], total_msa_crop_size)
+    paired_crop_size = None
+  return unpaired_crop_size, paired_crop_size
+
+
+def remove_all_gapped_rows_from_all_seqs(
+    chains_list: Sequence[dict[str, np.ndarray]], asym_ids: Sequence[float]
+) -> Sequence[dict[str, np.ndarray]]:
+  """Removes all gapped rows from all_seq feat based on selected asym_ids."""
+
+  merged_msa_all_seq = np.concatenate(
+      [
+          chain['msa_all_seq']
+          for chain in chains_list
+          if chain['asym_id'][0] in asym_ids
+      ],
+      axis=1,
+  )
+
+  non_gapped_keep_rows = np.any(
+      merged_msa_all_seq != data_constants.MSA_GAP_IDX, axis=1
+  )
+  for chain in chains_list:
+    for feat_name in list(chains_list)[0]:
+      if '_all_seq' in feat_name:
+        feat_name_split = feat_name.split('_all_seq')[0]
+        if feat_name_split in (
+            data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES
+            + data_constants.NUM_SEQ_MSA_FEATURES
+        ):
+          # For consistency we do this for all chains even though the
+          # gapped rows are based on a selected set asym_ids.
+          chain[feat_name] = chain[feat_name][non_gapped_keep_rows]
+    chain['num_alignments_all_seq'] = np.sum(non_gapped_keep_rows)
+  return chains_list

flax_model/alphafold3/model/network/atom_cross_attention.py

@@ -0,0 +1,413 @@
+
+
+"""Per-atom cross attention."""
+import dataclasses
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.model import feat_batch
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.model.components import haiku_modules as hm
+from flax_model.alphafold3.model.components import utils
+from flax_model.alphafold3.model.network import diffusion_transformer
+# import chex
+import jax
+import jax.numpy as jnp
+
+
+class AtomCrossAttEncoderConfig(base_config.BaseConfig):
+  per_token_channels: int = 768
+  per_atom_channels: int = 128
+  atom_transformer: diffusion_transformer.CrossAttTransformer.Config = (
+      base_config.autocreate(num_intermediate_factor=2, num_blocks=3)
+  )
+  per_atom_pair_channels: int = 16
+
+
+def _per_atom_conditioning(
+    config: AtomCrossAttEncoderConfig, batch: feat_batch.Batch, name: str
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+  """computes single and pair conditioning for all atoms in each token."""
+
+  c = config
+  # Compute per-atom single conditioning
+  # Shape (num_tokens, num_dense, channels)
+  act = hm.Linear(
+      c.per_atom_channels, precision='highest', name=f'{name}_embed_ref_pos'
+  )(batch.ref_structure.positions)
+  act += hm.Linear(c.per_atom_channels, name=f'{name}_embed_ref_mask')(
+      batch.ref_structure.mask.astype(jnp.float32)[:, :, None]
+  )
+  # Element is encoded as atomic number if the periodic table, so
+  # 128 should be fine.
+  act += hm.Linear(c.per_atom_channels, name=f'{name}_embed_ref_element')(
+      jax.nn.one_hot(batch.ref_structure.element, 128)
+  )
+  act += hm.Linear(c.per_atom_channels, name=f'{name}_embed_ref_charge')(
+      jnp.arcsinh(batch.ref_structure.charge)[:, :, None]
+  )
+  # Characters are encoded as ASCII code minus 32, so we need 64 classes,
+  # to encode all standard ASCII characters between 32 and 96.
+  atom_name_chars_1hot = jax.nn.one_hot(batch.ref_structure.atom_name_chars, 64)
+  num_token, num_dense, _ = act.shape
+  act += hm.Linear(c.per_atom_channels, name=f'{name}_embed_ref_atom_name')(
+      atom_name_chars_1hot.reshape(num_token, num_dense, -1)
+  )
+  act *= batch.ref_structure.mask.astype(jnp.float32)[:, :, None]
+
+  # Compute pair conditioning
+  # shape (num_tokens, num_dense, num_dense, channels)
+  # Embed single features
+  row_act = hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_single_to_pair_cond_row'
+  )(jax.nn.relu(act))
+  col_act = hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_single_to_pair_cond_col'
+  )(jax.nn.relu(act))
+  pair_act = row_act[:, :, None, :] + col_act[:, None, :, :]
+  # Embed pairwise offsets
+  pair_act += hm.Linear(
+      c.per_atom_pair_channels,
+      precision='highest',
+      name=f'{name}_embed_pair_offsets',
+  )(
+      batch.ref_structure.positions[:, :, None, :]
+      - batch.ref_structure.positions[:, None, :, :]
+  )
+  # Embed pairwise inverse squared distances
+  sq_dists = jnp.sum(
+      jnp.square(
+          batch.ref_structure.positions[:, :, None, :]
+          - batch.ref_structure.positions[:, None, :, :]
+      ),
+      axis=-1,
+  )
+  pair_act += hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_embed_pair_distances'
+  )(1.0 / (1 + sq_dists[:, :, :, None]))
+
+  return act, pair_act
+
+
+@dataclasses.dataclass(frozen=True)
+class AtomCrossAttEncoderOutput:
+  token_act: jnp.ndarray  # (num_tokens, ch)
+  skip_connection: jnp.ndarray  # (num_subsets, num_queries, ch)
+  queries_mask: jnp.ndarray  # (num_subsets, num_queries)
+  queries_single_cond: jnp.ndarray  # (num_subsets, num_queries, ch)
+  keys_mask: jnp.ndarray  # (num_subsets, num_keys)
+  keys_single_cond: jnp.ndarray  # (num_subsets, num_keys, ch)
+  pair_cond: jnp.ndarray  # (num_subsets, num_queries, num_keys, ch)
+
+
+jax.tree_util.register_dataclass(
+    AtomCrossAttEncoderOutput,
+    data_fields=[f.name for f in dataclasses.fields(AtomCrossAttEncoderOutput)],
+    meta_fields=[],
+)
+
+
+def atom_cross_att_encoder(
+    token_atoms_act: jnp.ndarray | None,  # (num_tokens, max_atoms_per_token, 3)
+    trunk_single_cond: jnp.ndarray | None,  # (num_tokens, ch)
+    trunk_pair_cond: jnp.ndarray | None,  # (num_tokens, num_tokens, ch)
+    config: AtomCrossAttEncoderConfig,
+    global_config: model_config.GlobalConfig,
+    batch: feat_batch.Batch,
+    name: str,
+) -> AtomCrossAttEncoderOutput:
+  """Cross-attention on flat atom subsets and mapping to per-token features."""
+  c = config
+
+  # Compute single conditioning from atom meta data and convert to queries
+  # layout.
+  # (num_subsets, num_queries, channels)
+  token_atoms_single_cond, _ = _per_atom_conditioning(config, batch, name)
+  token_atoms_mask = batch.predicted_structure_info.atom_mask
+  queries_single_cond = atom_layout.convert(
+      batch.atom_cross_att.token_atoms_to_queries,
+      token_atoms_single_cond,
+      layout_axes=(-3, -2),
+  )
+  queries_mask = atom_layout.convert(
+      batch.atom_cross_att.token_atoms_to_queries,
+      token_atoms_mask,
+      layout_axes=(-2, -1),
+  )
+
+  # If provided, broadcast single conditioning from trunk to all queries
+  if trunk_single_cond is not None:
+    trunk_single_cond = hm.Linear(
+        c.per_atom_channels,
+        precision='highest',
+        initializer=global_config.final_init,
+        name=f'{name}_embed_trunk_single_cond',
+    )(
+        hm.LayerNorm(
+            use_fast_variance=False,
+            create_offset=False,
+            name=f'{name}_lnorm_trunk_single_cond',
+        )(trunk_single_cond)
+    )
+    queries_single_cond += atom_layout.convert(
+        batch.atom_cross_att.tokens_to_queries,
+        trunk_single_cond,
+        layout_axes=(-2,),
+    )
+
+  if token_atoms_act is None:
+    # if no token_atoms_act is given (e.g. begin of evoformer), we use the
+    # static conditioning only
+    queries_act = queries_single_cond
+  else:
+    # Convert token_atoms_act to queries layout and map to per_atom_channels
+    # (num_subsets, num_queries, channels)
+    queries_act = atom_layout.convert(
+        batch.atom_cross_att.token_atoms_to_queries,
+        token_atoms_act,
+        layout_axes=(-3, -2),
+    )
+    queries_act = hm.Linear(
+        c.per_atom_channels,
+        precision='highest',
+        name=f'{name}_atom_positions_to_features',
+    )(queries_act)
+    queries_act *= queries_mask[..., None]
+    queries_act += queries_single_cond
+
+  # Gather the keys from the queries.
+  keys_single_cond = atom_layout.convert(
+      batch.atom_cross_att.queries_to_keys,
+      queries_single_cond,
+      layout_axes=(-3, -2),
+  )
+  keys_mask = atom_layout.convert(
+      batch.atom_cross_att.queries_to_keys, queries_mask, layout_axes=(-2, -1)
+  )
+
+  # Embed single features into the pair conditioning.
+  # shape (num_subsets, num_queries, num_keys, ch)
+  row_act = hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_single_to_pair_cond_row'
+  )(jax.nn.relu(queries_single_cond))
+  pair_cond_keys_input = atom_layout.convert(
+      batch.atom_cross_att.queries_to_keys,
+      queries_single_cond,
+      layout_axes=(-3, -2),
+  )
+  col_act = hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_single_to_pair_cond_col'
+  )(jax.nn.relu(pair_cond_keys_input))
+  pair_act = row_act[:, :, None, :] + col_act[:, None, :, :]
+
+  if trunk_pair_cond is not None:
+    # If provided, broadcast the pair conditioning for the trunk (evoformer
+    # pairs) to the atom pair activations. This should boost ligands, but also
+    # help for cross attention within proteins, because we always have atoms
+    # from multiple residues in a subset.
+    # Map trunk pair conditioning to per_atom_pair_channels
+    # (num_tokens, num_tokens, per_atom_pair_channels)
+    trunk_pair_cond = hm.Linear(
+        c.per_atom_pair_channels,
+        precision='highest',
+        initializer=global_config.final_init,
+        name=f'{name}_embed_trunk_pair_cond',
+    )(
+        hm.LayerNorm(
+            use_fast_variance=False,
+            create_offset=False,
+            name=f'{name}_lnorm_trunk_pair_cond',
+        )(trunk_pair_cond)
+    )
+
+    # Create the GatherInfo into a flattened trunk_pair_cond from the
+    # queries and keys gather infos.
+    num_tokens = trunk_pair_cond.shape[0]
+    # (num_subsets, num_queries)
+    tokens_to_queries = batch.atom_cross_att.tokens_to_queries
+    # (num_subsets, num_keys)
+    tokens_to_keys = batch.atom_cross_att.tokens_to_keys
+    # (num_subsets, num_queries, num_keys)
+    trunk_pair_to_atom_pair = atom_layout.GatherInfo(
+        gather_idxs=(
+            num_tokens * tokens_to_queries.gather_idxs[:, :, None]
+            + tokens_to_keys.gather_idxs[:, None, :]
+        ),
+        gather_mask=(
+            tokens_to_queries.gather_mask[:, :, None]
+            & tokens_to_keys.gather_mask[:, None, :]
+        ),
+        input_shape=jnp.array((num_tokens, num_tokens)),
+    )
+    # Gather the conditioning and add it to the atom-pair activations.
+    pair_act += atom_layout.convert(
+        trunk_pair_to_atom_pair, trunk_pair_cond, layout_axes=(-3, -2)
+    )
+
+  # Embed pairwise offsets
+  queries_ref_pos = atom_layout.convert(
+      batch.atom_cross_att.token_atoms_to_queries,
+      batch.ref_structure.positions,
+      layout_axes=(-3, -2),
+  )
+  queries_ref_space_uid = atom_layout.convert(
+      batch.atom_cross_att.token_atoms_to_queries,
+      batch.ref_structure.ref_space_uid,
+      layout_axes=(-2, -1),
+  )
+  keys_ref_pos = atom_layout.convert(
+      batch.atom_cross_att.queries_to_keys,
+      queries_ref_pos,
+      layout_axes=(-3, -2),
+  )
+  keys_ref_space_uid = atom_layout.convert(
+      batch.atom_cross_att.queries_to_keys,
+      batch.ref_structure.ref_space_uid,
+      layout_axes=(-2, -1),
+  )
+
+  offsets_valid = (
+      queries_ref_space_uid[:, :, None] == keys_ref_space_uid[:, None, :]
+  )
+  offsets = queries_ref_pos[:, :, None, :] - keys_ref_pos[:, None, :, :]
+  pair_act += (
+      hm.Linear(
+          c.per_atom_pair_channels,
+          precision='highest',
+          name=f'{name}_embed_pair_offsets',
+      )(offsets)
+      * offsets_valid[:, :, :, None]
+  )
+
+  # Embed pairwise inverse squared distances
+  sq_dists = jnp.sum(jnp.square(offsets), axis=-1)
+  pair_act += (
+      hm.Linear(c.per_atom_pair_channels, name=f'{name}_embed_pair_distances')(
+          1.0 / (1 + sq_dists[:, :, :, None])
+      )
+      * offsets_valid[:, :, :, None]
+  )
+  # Embed offsets valid mask
+  pair_act += hm.Linear(
+      c.per_atom_pair_channels, name=f'{name}_embed_pair_offsets_valid'
+  )(offsets_valid[:, :, :, None].astype(jnp.float32))
+
+  # Run a small MLP on the pair acitvations
+  pair_act2 = hm.Linear(
+      c.per_atom_pair_channels, initializer='relu', name=f'{name}_pair_mlp_1'
+  )(jax.nn.relu(pair_act))
+  pair_act2 = hm.Linear(
+      c.per_atom_pair_channels, initializer='relu', name=f'{name}_pair_mlp_2'
+  )(jax.nn.relu(pair_act2))
+  pair_act += hm.Linear(
+      c.per_atom_pair_channels,
+      initializer=global_config.final_init,
+      name=f'{name}_pair_mlp_3',
+  )(jax.nn.relu(pair_act2))
+
+  # Run the atom cross attention transformer.
+  queries_act = diffusion_transformer.CrossAttTransformer(
+      c.atom_transformer, global_config, name=f'{name}_atom_transformer_encoder'
+  )(
+      queries_act=queries_act,
+      queries_mask=queries_mask,
+      queries_to_keys=batch.atom_cross_att.queries_to_keys,
+      keys_mask=keys_mask,
+      queries_single_cond=queries_single_cond,
+      keys_single_cond=keys_single_cond,
+      pair_cond=pair_act,
+  )
+  queries_act *= queries_mask[..., None]
+  skip_connection = queries_act
+
+  # Convert back to token-atom layout and aggregate to tokens
+  queries_act = hm.Linear(
+      c.per_token_channels, name=f'{name}_project_atom_features_for_aggr'
+  )(queries_act)
+  token_atoms_act = atom_layout.convert(
+      batch.atom_cross_att.queries_to_token_atoms,
+      queries_act,
+      layout_axes=(-3, -2),
+  )
+  token_act = utils.mask_mean(
+      token_atoms_mask[..., None], jax.nn.relu(token_atoms_act), axis=-2
+  )
+
+  return AtomCrossAttEncoderOutput(
+      token_act=token_act,
+      skip_connection=skip_connection,
+      queries_mask=queries_mask,
+      queries_single_cond=queries_single_cond,
+      keys_mask=keys_mask,
+      keys_single_cond=keys_single_cond,
+      pair_cond=pair_act,
+  )
+
+
+class AtomCrossAttDecoderConfig(base_config.BaseConfig):
+  per_atom_channels: int = 128
+  atom_transformer: diffusion_transformer.CrossAttTransformer.Config = (
+      base_config.autocreate(num_intermediate_factor=2, num_blocks=3)
+  )
+
+
+def atom_cross_att_decoder(
+    token_act: jnp.ndarray,  # (num_tokens, ch)
+    enc: AtomCrossAttEncoderOutput,
+    config: AtomCrossAttDecoderConfig,
+    global_config: model_config.GlobalConfig,
+    batch: feat_batch.Batch,
+    name: str,
+):  # (num_tokens, max_atoms_per_token, 3)
+  """Mapping to per-atom features and self-attention on subsets."""
+  c = config
+  # map per-token act down to per_atom channels
+  token_act = hm.Linear(
+      c.per_atom_channels, name=f'{name}_project_token_features_for_broadcast'
+  )(token_act)
+  # Broadcast to token-atoms layout and convert to queries layout.
+  num_token, max_atoms_per_token = (
+      batch.atom_cross_att.queries_to_token_atoms.shape
+  )
+  token_atom_act = jnp.broadcast_to(
+      token_act[:, None, :],
+      (num_token, max_atoms_per_token, c.per_atom_channels),
+  )
+  queries_act = atom_layout.convert(
+      batch.atom_cross_att.token_atoms_to_queries,
+      token_atom_act,
+      layout_axes=(-3, -2),
+  )
+  queries_act += enc.skip_connection
+  queries_act *= enc.queries_mask[..., None]
+
+  # Run the atom cross attention transformer.
+  queries_act = diffusion_transformer.CrossAttTransformer(
+      c.atom_transformer, global_config, name=f'{name}_atom_transformer_decoder'
+  )(
+      queries_act=queries_act,
+      queries_mask=enc.queries_mask,
+      queries_to_keys=batch.atom_cross_att.queries_to_keys,
+      keys_mask=enc.keys_mask,
+      queries_single_cond=enc.queries_single_cond,
+      keys_single_cond=enc.keys_single_cond,
+      pair_cond=enc.pair_cond,
+  )
+  queries_act *= enc.queries_mask[..., None]
+  queries_act = hm.LayerNorm(
+      use_fast_variance=False,
+      create_offset=False,
+      name=f'{name}_atom_features_layer_norm',
+  )(queries_act)
+  queries_position_update = hm.Linear(
+      3,
+      initializer=global_config.final_init,
+      precision='highest',
+      name=f'{name}_atom_features_to_position_update',
+  )(queries_act)
+  position_update = atom_layout.convert(
+      batch.atom_cross_att.queries_to_token_atoms,
+      queries_position_update,
+      layout_axes=(-3, -2),
+  )
+  return position_update

flax_model/alphafold3/model/network/confidence_head.py

@@ -0,0 +1,312 @@
+
+
+"""Confidence Head."""
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.model.components import haiku_modules as hm
+from flax_model.alphafold3.model.components import utils
+from flax_model.alphafold3.model.network import modules
+from flax_model.alphafold3.model.network import template_modules
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+def _safe_norm(x, keepdims, axis, eps=1e-8):
+  return jnp.sqrt(eps + jnp.sum(jnp.square(x), axis=axis, keepdims=keepdims))
+
+
+class ConfidenceHead(hk.Module):
+  """Head to predict the distance errors in a prediction."""
+
+  class PAEConfig(base_config.BaseConfig):
+    max_error_bin: float = 31.0
+    num_bins: int = 64
+
+  class Config(base_config.BaseConfig):
+    """Configuration for ConfidenceHead."""
+
+    pairformer: modules.PairFormerIteration.Config = base_config.autocreate(
+        single_attention=base_config.autocreate(),
+        single_transition=base_config.autocreate(),
+        num_layer=4,
+    )
+    max_error_bin: float = 31.0
+    num_plddt_bins: int = 50
+    num_bins: int = 64
+    no_embedding_prob: float = 0.2
+    pae: 'ConfidenceHead.PAEConfig' = base_config.autocreate()
+    dgram_features: template_modules.DistogramFeaturesConfig = (
+        base_config.autocreate()
+    )
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name='confidence_head',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def _embed_features(
+      self,
+      dense_atom_positions,
+      token_atoms_to_pseudo_beta,
+      pair_mask,
+      pair_act,
+      target_feat,
+  ):
+    out = hm.Linear(pair_act.shape[-1], name='left_target_feat_project')(
+        target_feat
+    ).astype(pair_act.dtype)
+    out += hm.Linear(pair_act.shape[-1], name='right_target_feat_project')(
+        target_feat
+    ).astype(pair_act.dtype)[:, None]
+    positions = atom_layout.convert(
+        token_atoms_to_pseudo_beta,
+        dense_atom_positions,
+        layout_axes=(-3, -2),
+    )
+    dgram = template_modules.dgram_from_positions(
+        positions, self.config.dgram_features
+    )
+    dgram *= pair_mask[..., None]
+
+    out += hm.Linear(pair_act.shape[-1], name='distogram_feat_project')(
+        dgram.astype(pair_act.dtype)
+    )
+    return out
+
+  def __call__(
+      self,
+      dense_atom_positions: jnp.ndarray,
+      embeddings: dict[str, jnp.ndarray],
+      seq_mask: jnp.ndarray,
+      token_atoms_to_pseudo_beta: atom_layout.GatherInfo,
+      asym_id: jnp.ndarray,
+  ) -> dict[str, jnp.ndarray]:
+    """Builds ConfidenceHead module.
+
+    Arguments:
+      dense_atom_positions: [N_res, N_atom, 3] array of positions.
+      embeddings: Dictionary of representations.
+      seq_mask: Sequence mask.
+      token_atoms_to_pseudo_beta: Pseudo beta info for atom tokens.
+      asym_id: Asym ID token features.
+
+    Returns:
+      Dictionary of results.
+    """
+    dtype = (
+        jnp.bfloat16 if self.global_config.bfloat16 == 'all' else jnp.float32
+    )
+    with utils.bfloat16_context():
+      seq_mask_cast = seq_mask.astype(dtype)
+      pair_mask = seq_mask_cast[:, None] * seq_mask_cast[None, :]
+      pair_mask = pair_mask.astype(dtype)
+
+      pair_act = embeddings['pair'].astype(dtype)
+      single_act = embeddings['single'].astype(dtype)
+      target_feat = embeddings['target_feat'].astype(dtype)
+
+      num_residues = seq_mask.shape[0]
+      num_pair_channels = pair_act.shape[2]
+
+      pair_act += self._embed_features(
+          dense_atom_positions,
+          token_atoms_to_pseudo_beta,
+          pair_mask,
+          pair_act,
+          target_feat,
+      )
+
+      def pairformer_fn(act):
+        pair_act, single_act = act
+        return modules.PairFormerIteration(
+            self.config.pairformer,
+            self.global_config,
+            with_single=True,
+            name='confidence_pairformer',
+        )(
+            act=pair_act,
+            single_act=single_act,
+            pair_mask=pair_mask,
+            seq_mask=seq_mask,
+        )
+
+      pairformer_stack = hk.experimental.layer_stack(
+          self.config.pairformer.num_layer
+      )(pairformer_fn)
+
+      pair_act, single_act = pairformer_stack((pair_act, single_act))
+      pair_act = pair_act.astype(jnp.float32)
+      assert pair_act.shape == (num_residues, num_residues, num_pair_channels)
+
+      # Produce logits to predict a distogram of pairwise distance errors
+      # between the input prediction and the ground truth.
+
+      # Shape (num_res, num_res, num_bins)
+      left_distance_logits = hm.Linear(
+          self.config.num_bins,
+          initializer=self.global_config.final_init,
+          name='left_half_distance_logits',
+      )(hm.LayerNorm(name='logits_ln')(pair_act))
+      right_distance_logits = left_distance_logits
+      distance_logits = left_distance_logits + jnp.swapaxes(  # Symmetrize.
+          right_distance_logits, -2, -3
+      )
+      # Shape (num_bins,)
+      distance_breaks = jnp.linspace(
+          0.0, self.config.max_error_bin, self.config.num_bins - 1
+      )
+
+      step = distance_breaks[1] - distance_breaks[0]
+
+      # Add half-step to get the center
+      bin_centers = distance_breaks + step / 2
+      # Add a catch-all bin at the end.
+      bin_centers = jnp.concatenate(
+          [bin_centers, bin_centers[-1:] + step], axis=0
+      )
+
+      distance_probs = jax.nn.softmax(distance_logits, axis=-1)
+
+      pred_distance_error = (
+          jnp.sum(distance_probs * bin_centers, axis=-1) * pair_mask
+      )
+      average_pred_distance_error = jnp.sum(
+          pred_distance_error, axis=[-2, -1]
+      ) / jnp.sum(pair_mask, axis=[-2, -1])
+
+      # Predicted aligned error
+      pae_outputs = {}
+      # Shape (num_res, num_res, num_bins)
+      pae_logits = hm.Linear(
+          self.config.pae.num_bins,
+          initializer=self.global_config.final_init,
+          name='pae_logits',
+      )(hm.LayerNorm(name='pae_logits_ln')(pair_act))
+      # Shape (num_bins,)
+      pae_breaks = jnp.linspace(
+          0.0, self.config.pae.max_error_bin, self.config.pae.num_bins - 1
+      )
+      step = pae_breaks[1] - pae_breaks[0]
+      # Add half-step to get the center
+      bin_centers = pae_breaks + step / 2
+      # Add a catch-all bin at the end.
+      bin_centers = jnp.concatenate(
+          [bin_centers, bin_centers[-1:] + step], axis=0
+      )
+      pae_probs = jax.nn.softmax(pae_logits, axis=-1)
+
+      seq_mask_bool = seq_mask.astype(bool)
+      pair_mask_bool = seq_mask_bool[:, None] * seq_mask_bool[None, :]
+      pae = jnp.sum(pae_probs * bin_centers, axis=-1) * pair_mask_bool
+      pae_outputs.update({
+          'full_pae': pae,
+      })
+
+    # The pTM is computed outside of bfloat16 context.
+    tmscore_adjusted_pae_global, tmscore_adjusted_pae_interface = (
+        self._get_tmscore_adjusted_pae(
+            asym_id=asym_id,
+            seq_mask=seq_mask,
+            pair_mask=pair_mask_bool,
+            bin_centers=bin_centers,
+            pae_probs=pae_probs,
+        )
+    )
+    pae_outputs.update({
+        'tmscore_adjusted_pae_global': tmscore_adjusted_pae_global,
+        'tmscore_adjusted_pae_interface': tmscore_adjusted_pae_interface,
+    })
+    single_act = single_act.astype('float32')
+
+    # pLDDT
+    # Shape (num_res, num_atom, num_bins)
+    plddt_logits = hm.Linear(
+        (dense_atom_positions.shape[-2], self.config.num_plddt_bins),
+        initializer=self.global_config.final_init,
+        name='plddt_logits',
+    )(hm.LayerNorm(name='plddt_logits_ln')(single_act))
+
+    bin_width = 1.0 / self.config.num_plddt_bins
+    bin_centers = jnp.arange(0.5 * bin_width, 1.0, bin_width)
+    predicted_lddt = jnp.sum(
+        jax.nn.softmax(plddt_logits, axis=-1) * bin_centers, axis=-1
+    )
+    predicted_lddt = predicted_lddt * 100.0
+
+    # Experimentally resolved
+    # Shape (num_res, num_atom, 2)
+    experimentally_resolved_logits = hm.Linear(
+        (dense_atom_positions.shape[-2], 2),
+        initializer=self.global_config.final_init,
+        name='experimentally_resolved_logits',
+    )(hm.LayerNorm(name='experimentally_resolved_ln')(single_act))
+
+    predicted_experimentally_resolved = jax.nn.softmax(
+        experimentally_resolved_logits, axis=-1
+    )[..., 1]
+
+    return {
+        'predicted_lddt': predicted_lddt,
+        'predicted_experimentally_resolved': predicted_experimentally_resolved,
+        'full_pde': pred_distance_error,
+        'average_pde': average_pred_distance_error,
+        **pae_outputs,
+    }
+
+  def _get_tmscore_adjusted_pae(
+      self,
+      asym_id: jnp.ndarray,
+      seq_mask: jnp.ndarray,
+      pair_mask: jnp.ndarray,
+      bin_centers: jnp.ndarray,
+      pae_probs: jnp.ndarray,
+  ):
+    def get_tmscore_adjusted_pae(num_interface_tokens, bin_centers, pae_probs):
+      # Clip to avoid negative/undefined d0.
+      clipped_num_res = jnp.maximum(num_interface_tokens, 19)
+
+      # Compute d_0(num_res) as defined by TM-score, eqn. (5) in
+      # http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf
+      # Yang & Skolnick "Scoring function for automated
+      # assessment of protein structure template quality" 2004.
+      d0 = 1.24 * (clipped_num_res - 15) ** (1.0 / 3) - 1.8
+
+      # Make compatible with [num_tokens, num_tokens, num_bins]
+      d0 = d0[:, :, None]
+      bin_centers = bin_centers[None, None, :]
+
+      # TM-Score term for every bin.
+      tm_per_bin = 1.0 / (1 + jnp.square(bin_centers) / jnp.square(d0))
+      # E_distances tm(distance).
+      predicted_tm_term = jnp.sum(pae_probs * tm_per_bin, axis=-1)
+      return predicted_tm_term
+
+    # Interface version
+    x = asym_id[None, :] == asym_id[:, None]
+    num_chain_tokens = jnp.sum(x * pair_mask, axis=-1)
+    num_interface_tokens = num_chain_tokens[None, :] + num_chain_tokens[:, None]
+    # Don't double-count within a single chain
+    num_interface_tokens -= x * (num_interface_tokens // 2)
+    num_interface_tokens = num_interface_tokens * pair_mask
+
+    num_global_tokens = jnp.full(
+        shape=pair_mask.shape, fill_value=seq_mask.sum()
+    )
+
+    assert num_global_tokens.dtype == 'int32'
+    assert num_interface_tokens.dtype == 'int32'
+    global_apae = get_tmscore_adjusted_pae(
+        num_global_tokens, bin_centers, pae_probs
+    )
+    interface_apae = get_tmscore_adjusted_pae(
+        num_interface_tokens, bin_centers, pae_probs
+    )
+    return global_apae, interface_apae

flax_model/alphafold3/model/network/diffusion_head.py

@@ -0,0 +1,363 @@
+
+
+"""Diffusion Head."""
+
+from collections.abc import Callable
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.model import feat_batch
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.components import haiku_modules as hm
+from flax_model.alphafold3.model.components import utils
+from flax_model.alphafold3.model.network import atom_cross_attention
+from flax_model.alphafold3.model.network import diffusion_transformer
+from flax_model.alphafold3.model.network import featurization
+from flax_model.alphafold3.model.network import noise_level_embeddings
+#import chex
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+# Carefully measured by averaging multimer training set.
+SIGMA_DATA = 16.0
+
+
+def random_rotation(key):
+  # Create a random rotation (Gram-Schmidt orthogonalization of two
+  # random normal vectors)
+  v0, v1 = jax.random.normal(key, shape=(2, 3))
+  e0 = v0 / jnp.maximum(1e-10, jnp.linalg.norm(v0))
+  v1 = v1 - e0 * jnp.dot(v1, e0, precision=jax.lax.Precision.HIGHEST)
+  e1 = v1 / jnp.maximum(1e-10, jnp.linalg.norm(v1))
+  e2 = jnp.cross(e0, e1)
+  return jnp.stack([e0, e1, e2])
+
+
+def random_augmentation(
+    rng_key: jnp.ndarray,
+    positions: jnp.ndarray,
+    mask: jnp.ndarray,
+) -> jnp.ndarray:
+  """Apply random rigid augmentation.
+
+  Args:
+    rng_key: random key
+    positions: atom positions of shape (<common_axes>, 3)
+    mask: per-atom mask of shape (<common_axes>,)
+
+  Returns:
+    Transformed positions with the same shape as input positions.
+  """
+  rotation_key, translation_key = jax.random.split(rng_key)
+
+  center = utils.mask_mean(
+      mask[..., None], positions, axis=(-2, -3), keepdims=True, eps=1e-6
+  )
+  rot = random_rotation(rotation_key)
+  translation = jax.random.normal(translation_key, shape=(3,))
+
+  augmented_positions = (
+      jnp.einsum(
+          '...i,ij->...j',
+          positions - center,
+          rot,
+          precision=jax.lax.Precision.HIGHEST,
+      )
+      + translation
+  )
+  return augmented_positions * mask[..., None]
+
+
+def noise_schedule(t, smin=0.0004, smax=160.0, p=7):
+  return (
+      SIGMA_DATA
+      * (smax ** (1 / p) + t * (smin ** (1 / p) - smax ** (1 / p))) ** p
+  )
+
+
+class ConditioningConfig(base_config.BaseConfig):
+  pair_channel: int
+  seq_channel: int
+  prob: float
+
+
+class SampleConfig(base_config.BaseConfig):
+  steps: int
+  gamma_0: float = 0.8
+  gamma_min: float = 1.0
+  noise_scale: float = 1.003
+  step_scale: float = 1.5
+  num_samples: int = 1
+
+
+class DiffusionHead(hk.Module):
+  """Denoising Diffusion Head."""
+
+  class Config(
+      atom_cross_attention.AtomCrossAttEncoderConfig,
+      atom_cross_attention.AtomCrossAttDecoderConfig,
+  ):
+    """Configuration for DiffusionHead."""
+
+    eval_batch_size: int = 5
+    eval_batch_dim_shard_size: int = 5
+    conditioning: ConditioningConfig = base_config.autocreate(
+        prob=0.8, pair_channel=128, seq_channel=384
+    )
+    eval: SampleConfig = base_config.autocreate(
+        num_samples=5,
+        steps=200,
+    )
+    transformer: diffusion_transformer.Transformer.Config = (
+        base_config.autocreate()
+    )
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name='diffusion_head',
+  ):
+    self.config = config
+    self.global_config = global_config
+    super().__init__(name=name)
+
+  @hk.transparent
+  def _conditioning(
+      self,
+      batch: feat_batch.Batch,
+      embeddings: dict[str, jnp.ndarray],
+      noise_level: jnp.ndarray,
+      use_conditioning: bool,
+  ) -> tuple[jnp.ndarray, jnp.ndarray]:
+    single_embedding = use_conditioning * embeddings['single']
+    pair_embedding = use_conditioning * embeddings['pair']
+
+    rel_features = featurization.create_relative_encoding(
+        seq_features=batch.token_features,
+        max_relative_idx=32,
+        max_relative_chain=2,
+    ).astype(pair_embedding.dtype)
+    features_2d = jnp.concatenate([pair_embedding, rel_features], axis=-1)
+    pair_cond = hm.Linear(
+        self.config.conditioning.pair_channel,
+        precision='highest',
+        name='pair_cond_initial_projection',
+    )(
+        hm.LayerNorm(
+            use_fast_variance=False,
+            create_offset=False,
+            name='pair_cond_initial_norm',
+        )(features_2d)
+    )
+
+    for idx in range(2):
+      pair_cond += diffusion_transformer.transition_block(
+          pair_cond, 2, self.global_config, name=f'pair_transition_{idx}'
+      )
+
+    target_feat = embeddings['target_feat']
+    features_1d = jnp.concatenate([single_embedding, target_feat], axis=-1)
+    single_cond = hm.LayerNorm(
+        use_fast_variance=False,
+        create_offset=False,
+        name='single_cond_initial_norm',
+    )(features_1d)
+    single_cond = hm.Linear(
+        self.config.conditioning.seq_channel,
+        precision='highest',
+        name='single_cond_initial_projection',
+    )(single_cond)
+
+    noise_embedding = noise_level_embeddings.noise_embeddings(
+        sigma_scaled_noise_level=noise_level / SIGMA_DATA
+    )
+    single_cond += hm.Linear(
+        self.config.conditioning.seq_channel,
+        precision='highest',
+        name='noise_embedding_initial_projection',
+    )(
+        hm.LayerNorm(
+            use_fast_variance=False,
+            create_offset=False,
+            name='noise_embedding_initial_norm',
+        )(noise_embedding)
+    )
+
+    for idx in range(2):
+      single_cond += diffusion_transformer.transition_block(
+          single_cond, 2, self.global_config, name=f'single_transition_{idx}'
+      )
+
+    return single_cond, pair_cond
+
+  def __call__(
+      self,
+      # positions_noisy.shape: (num_token, max_atoms_per_token, 3)
+      positions_noisy: jnp.ndarray,
+      noise_level: jnp.ndarray,
+      batch: feat_batch.Batch,
+      embeddings: dict[str, jnp.ndarray],
+      use_conditioning: bool,
+  ) -> jnp.ndarray:
+
+    with utils.bfloat16_context():
+      # Get conditioning
+      trunk_single_cond, trunk_pair_cond = self._conditioning(
+          batch=batch,
+          embeddings=embeddings,
+          noise_level=noise_level,
+          use_conditioning=use_conditioning,
+      )
+
+      # Extract features
+      sequence_mask = batch.token_features.mask
+      atom_mask = batch.predicted_structure_info.atom_mask
+
+      # Position features
+      act = positions_noisy * atom_mask[..., None]
+      act = act / jnp.sqrt(noise_level**2 + SIGMA_DATA**2)
+
+      enc = atom_cross_attention.atom_cross_att_encoder(
+          token_atoms_act=act,
+          trunk_single_cond=embeddings['single'],
+          trunk_pair_cond=trunk_pair_cond,
+          config=self.config,
+          global_config=self.global_config,
+          batch=batch,
+          name='diffusion',
+      )
+      act = enc.token_act
+
+      # Token-token attention
+      # chex.assert_shape(act, (None, self.config.per_token_channels))
+      act = jnp.asarray(act, dtype=jnp.float32)
+
+      act += hm.Linear(
+          act.shape[-1],
+          precision='highest',
+          initializer=self.global_config.final_init,
+          name='single_cond_embedding_projection',
+      )(
+          hm.LayerNorm(
+              use_fast_variance=False,
+              create_offset=False,
+              name='single_cond_embedding_norm',
+          )(trunk_single_cond)
+      )
+
+      act = jnp.asarray(act, dtype=jnp.float32)
+      trunk_single_cond = jnp.asarray(trunk_single_cond, dtype=jnp.float32)
+      trunk_pair_cond = jnp.asarray(trunk_pair_cond, dtype=jnp.float32)
+      sequence_mask = jnp.asarray(sequence_mask, dtype=jnp.float32)
+
+      transformer = diffusion_transformer.Transformer(
+          self.config.transformer, self.global_config
+      )
+      act = transformer(
+          act=act,
+          single_cond=trunk_single_cond,
+          mask=sequence_mask,
+          pair_cond=trunk_pair_cond,
+      )
+      act = hm.LayerNorm(
+          use_fast_variance=False, create_offset=False, name='output_norm'
+      )(act)
+      # (n_tokens, per_token_channels)
+
+      # (Possibly) atom-granularity decoder
+      assert isinstance(enc, atom_cross_attention.AtomCrossAttEncoderOutput)
+      position_update = atom_cross_attention.atom_cross_att_decoder(
+          token_act=act,
+          enc=enc,
+          config=self.config,
+          global_config=self.global_config,
+          batch=batch,
+          name='diffusion',
+      )
+
+      skip_scaling = SIGMA_DATA**2 / (noise_level**2 + SIGMA_DATA**2)
+      out_scaling = (
+          noise_level * SIGMA_DATA / jnp.sqrt(noise_level**2 + SIGMA_DATA**2)
+      )
+    # End `with utils.bfloat16_context()`.
+
+    return (
+        skip_scaling * positions_noisy + out_scaling * position_update
+    ) * atom_mask[..., None]
+
+
+def sample(
+    denoising_step: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
+    batch: feat_batch.Batch,
+    key: jnp.ndarray,
+    config: SampleConfig,
+) -> dict[str, jnp.ndarray]:
+  """Sample using denoiser on batch.
+
+  Args:
+    denoising_step: the denoising function.
+    batch: the batch
+    key: random key
+    config: config for the sampling process (e.g. number of denoising steps,
+      etc.)
+
+  Returns:
+    a dict
+      {
+         'atom_positions': jnp.array(...)       # shape (<common_axes>, 3)
+         'mask': jnp.array(...)                 # shape (<common_axes>,)
+      }
+    where the <common_axes> are
+    (num_samples, num_tokens, max_atoms_per_token)
+  """
+
+  mask = batch.predicted_structure_info.atom_mask
+
+  def apply_denoising_step(carry, noise_level):
+    key, positions, noise_level_prev = carry
+    key, key_noise, key_aug = jax.random.split(key, 3)
+
+    positions = random_augmentation(
+        rng_key=key_aug, positions=positions, mask=mask
+    )
+
+    gamma = config.gamma_0 * (noise_level > config.gamma_min)
+    t_hat = noise_level_prev * (1 + gamma)
+
+    noise_scale = config.noise_scale * jnp.sqrt(t_hat**2 - noise_level_prev**2)
+    noise = noise_scale * jax.random.normal(key_noise, positions.shape)
+    positions_noisy = positions + noise
+
+    positions_denoised = denoising_step(positions_noisy, t_hat)
+    grad = (positions_noisy - positions_denoised) / t_hat
+
+    d_t = noise_level - t_hat
+    positions_out = positions_noisy + config.step_scale * d_t * grad
+
+    return (key, positions_out, noise_level), positions_out
+
+  num_samples = config.num_samples
+
+  noise_levels = noise_schedule(jnp.linspace(0, 1, config.steps + 1))
+
+  key, noise_key = jax.random.split(key)
+  positions = jax.random.normal(noise_key, (num_samples,) + mask.shape + (3,))
+  positions *= noise_levels[0]
+
+  init = (
+      jax.random.split(key, num_samples),
+      positions,
+      jnp.tile(noise_levels[None, 0], (num_samples,)),
+  )
+
+  apply_denoising_step = hk.vmap(
+      apply_denoising_step, in_axes=(0, None), split_rng=(not hk.running_init())
+  )
+  result, _ = hk.scan(apply_denoising_step, init, noise_levels[1:], unroll=4)
+  _, positions_out, _ = result
+
+  final_dense_atom_mask = jnp.tile(mask[None], (num_samples, 1, 1))
+
+  return {'atom_positions': positions_out, 'mask': final_dense_atom_mask}

flax_model/alphafold3/model/network/diffusion_transformer.py

@@ -0,0 +1,396 @@
+
+
+"""Diffusion transformer model."""
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.jax.gated_linear_unit import gated_linear_unit
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.model.components import haiku_modules as hm
+import haiku as hk
+import jax
+from jax import numpy as jnp
+
+
+def adaptive_layernorm(x, single_cond, name):
+  """Adaptive LayerNorm."""
+  # Adopted from Scalable Diffusion Models with Transformers
+  # https://arxiv.org/abs/2212.09748
+  if single_cond is None:
+    x = hm.LayerNorm(name=f'{name}layer_norm', use_fast_variance=False)(x)
+  else:
+    x = hm.LayerNorm(
+        name=f'{name}layer_norm',
+        use_fast_variance=False,
+        create_scale=False,
+        create_offset=False,
+    )(x)
+    single_cond = hm.LayerNorm(
+        name=f'{name}single_cond_layer_norm',
+        use_fast_variance=False,
+        create_offset=False,
+    )(single_cond)
+    single_scale = hm.Linear(
+        x.shape[-1],
+        initializer='zeros',
+        use_bias=True,
+        name=f'{name}single_cond_scale',
+    )(single_cond)
+    single_bias = hm.Linear(
+        x.shape[-1], initializer='zeros', name=f'{name}single_cond_bias'
+    )(single_cond)
+    x = jax.nn.sigmoid(single_scale) * x + single_bias
+  return x
+
+
+def adaptive_zero_init(
+    x, num_channels, single_cond, global_config: model_config.GlobalConfig, name
+):
+  """Adaptive zero init, from AdaLN-zero."""
+  if single_cond is None:
+    output = hm.Linear(
+        num_channels,
+        initializer=global_config.final_init,
+        name=f'{name}transition2',
+    )(x)
+  else:
+    output = hm.Linear(num_channels, name=f'{name}transition2')(x)
+    # Init to a small gain, sigmoid(-2) ~ 0.1
+    cond = hm.Linear(
+        output.shape[-1],
+        initializer='zeros',
+        use_bias=True,
+        bias_init=-2.0,
+        name=f'{name}adaptive_zero_cond',
+    )(single_cond)
+    output = jax.nn.sigmoid(cond) * output
+  return output
+
+
+def transition_block(
+    x: jnp.ndarray,
+    num_intermediate_factor: int,
+    global_config: model_config.GlobalConfig,
+    single_cond: jnp.ndarray | None = None,
+    use_glu_kernel: bool = True,
+    name: str = '',
+) -> jnp.ndarray:
+  """Transition Block."""
+  num_channels = x.shape[-1]
+  num_intermediates = num_intermediate_factor * num_channels
+
+  x = adaptive_layernorm(x, single_cond, name=f'{name}ffw_')
+
+  if use_glu_kernel:
+    weights, _ = hm.haiku_linear_get_params(
+        x,
+        num_output=num_intermediates * 2,
+        initializer='relu',
+        name=f'{name}ffw_transition1',
+    )
+    weights = jnp.reshape(weights, (len(weights), 2, num_intermediates))
+    c = gated_linear_unit.gated_linear_unit(
+        x=x, weight=weights, implementation=None, activation=jax.nn.swish
+    )
+  else:
+    x = hm.Linear(
+        num_intermediates * 2, initializer='relu', name=f'{name}ffw_transition1'
+    )(x)
+    a, b = jnp.split(x, 2, axis=-1)
+    c = jax.nn.swish(a) * b
+
+  output = adaptive_zero_init(
+      c, num_channels, single_cond, global_config, f'{name}ffw_'
+  )
+  return output
+
+
+class SelfAttentionConfig(base_config.BaseConfig):
+  num_head: int = 16
+  key_dim: int | None = None
+  value_dim: int | None = None
+
+
+def self_attention(
+    x: jnp.ndarray,  # (num_tokens, ch)
+    mask: jnp.ndarray,  # (num_tokens,)
+    pair_logits: jnp.ndarray | None,  # (num_heads, num_tokens, num_tokens)
+    config: SelfAttentionConfig,
+    global_config: model_config.GlobalConfig,
+    single_cond: jnp.ndarray | None = None,  # (num_tokens, ch)
+    name: str = '',
+) -> jnp.ndarray:
+  """Multihead self-attention."""
+  assert len(mask.shape) == len(x.shape) - 1, f'{mask.shape}, {x.shape}'
+  # bias: ... x heads (1) x query (1) x key
+  bias = (1e9 * (mask - 1.0))[..., None, None, :]
+
+  x = adaptive_layernorm(x, single_cond, name=name)
+
+  num_channels = x.shape[-1]
+  # Sensible default for when the config keys are missing
+  key_dim = config.key_dim if config.key_dim is not None else num_channels
+  value_dim = config.value_dim if config.value_dim is not None else num_channels
+  num_head = config.num_head
+  assert key_dim % num_head == 0, f'{key_dim=} % {num_head=} != 0'
+  assert value_dim % num_head == 0, f'{value_dim=} % {num_head=} != 0'
+  key_dim = key_dim // num_head
+  value_dim = value_dim // num_head
+
+  qk_shape = (num_head, key_dim)
+  q = hm.Linear(qk_shape, use_bias=True, name=f'{name}q_projection')(x)
+  k = hm.Linear(qk_shape, use_bias=False, name=f'{name}k_projection')(x)
+
+  # In some situations the gradient norms can blow up without running this
+  # einsum in float32.
+  q = q.astype(jnp.float32)
+  k = k.astype(jnp.float32)
+  bias = bias.astype(jnp.float32)
+  logits = jnp.einsum('...qhc,...khc->...hqk', q * key_dim ** (-0.5), k) + bias
+  if pair_logits is not None:
+    logits += pair_logits  # (num_heads, seq_len, seq_len)
+  weights = jax.nn.softmax(logits, axis=-1)
+  weights = jnp.asarray(weights, dtype=x.dtype)
+
+  v_shape = (num_head, value_dim)
+  v = hm.Linear(v_shape, use_bias=False, name=f'{name}v_projection')(x)
+  weighted_avg = jnp.einsum('...hqk,...khc->...qhc', weights, v)
+  weighted_avg = jnp.reshape(weighted_avg, weighted_avg.shape[:-2] + (-1,))
+
+  gate_logits = hm.Linear(
+      num_head * value_dim,
+      bias_init=1.0,
+      initializer='zeros',
+      name=f'{name}gating_query',
+  )(x)
+  weighted_avg *= jax.nn.sigmoid(gate_logits)
+
+  output = adaptive_zero_init(
+      weighted_avg, num_channels, single_cond, global_config, name
+  )
+  return output
+
+
+class Transformer(hk.Module):
+  """Simple transformer stack."""
+
+  class Config(base_config.BaseConfig):
+    attention: SelfAttentionConfig = base_config.autocreate()
+    num_blocks: int = 24
+    block_remat: bool = False
+    super_block_size: int = 4
+    num_intermediate_factor: int = 2
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name: str = 'transformer',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def __call__(
+      self,
+      act: jnp.ndarray,
+      mask: jnp.ndarray,
+      single_cond: jnp.ndarray,
+      pair_cond: jnp.ndarray | None,
+  ) -> jnp.ndarray:
+    def block(act, pair_logits):
+      act += self_attention(
+          act,
+          mask,
+          pair_logits,
+          self.config.attention,
+          self.global_config,
+          single_cond,
+          name=self.name,
+      )
+      act += transition_block(
+          act,
+          self.config.num_intermediate_factor,
+          self.global_config,
+          single_cond,
+          name=self.name,
+      )
+      return act, None
+
+    # Precompute pair logits for performance
+    if pair_cond is None:
+      pair_act = None
+    else:
+      pair_act = hm.LayerNorm(
+          name='pair_input_layer_norm',
+          use_fast_variance=False,
+          create_offset=False,
+      )(pair_cond)
+
+    assert self.config.num_blocks % self.config.super_block_size == 0
+    num_super_blocks = self.config.num_blocks // self.config.super_block_size
+
+    def super_block(act):
+      if pair_act is None:
+        pair_logits = None
+      else:
+        pair_logits = hm.Linear(
+            (self.config.super_block_size, self.config.attention.num_head),
+            name='pair_logits_projection',
+        )(pair_act)
+        pair_logits = jnp.transpose(pair_logits, [2, 3, 0, 1])
+      return hk.experimental.layer_stack(
+          self.config.super_block_size, with_per_layer_inputs=True
+      )(block)(act, pair_logits)
+
+    return hk.experimental.layer_stack(
+        num_super_blocks, with_per_layer_inputs=True
+    )(super_block)(act)[0]
+
+
+class CrossAttentionConfig(base_config.BaseConfig):
+  num_head: int = 4
+  key_dim: int = 128
+  value_dim: int = 128
+
+
+def cross_attention(
+    x_q: jnp.ndarray,  # (..., Q, C)
+    x_k: jnp.ndarray,  # (..., K, C)
+    mask_q: jnp.ndarray,  # (..., Q)
+    mask_k: jnp.ndarray,  # (..., K)
+    config: CrossAttentionConfig,
+    global_config: model_config.GlobalConfig,
+    pair_logits: jnp.ndarray | None = None,  # (..., Q, K)
+    single_cond_q: jnp.ndarray | None = None,  # (..., Q, C)
+    single_cond_k: jnp.ndarray | None = None,  # (..., K, C)
+    name: str = '',
+) -> jnp.ndarray:
+  """Multihead self-attention."""
+  assert len(mask_q.shape) == len(x_q.shape) - 1, f'{mask_q.shape}, {x_q.shape}'
+  assert len(mask_k.shape) == len(x_k.shape) - 1, f'{mask_k.shape}, {x_k.shape}'
+  # bias: ... x heads (1) x query x key
+  bias = (
+      1e9
+      * (mask_q - 1.0)[..., None, :, None]
+      * (mask_k - 1.0)[..., None, None, :]
+  )
+
+  x_q = adaptive_layernorm(x_q, single_cond_q, name=f'{name}q')
+  x_k = adaptive_layernorm(x_k, single_cond_k, name=f'{name}k')
+
+  assert config.key_dim % config.num_head == 0
+  assert config.value_dim % config.num_head == 0
+  key_dim = config.key_dim // config.num_head
+  value_dim = config.value_dim // config.num_head
+
+  q = hm.Linear(
+      (config.num_head, key_dim), use_bias=True, name=f'{name}q_projection'
+  )(x_q)
+  k = hm.Linear(
+      (config.num_head, key_dim), use_bias=False, name=f'{name}k_projection'
+  )(x_k)
+
+  # In some situations the gradient norms can blow up without running this
+  # einsum in float32.
+  q = q.astype(jnp.float32)
+  k = k.astype(jnp.float32)
+  bias = bias.astype(jnp.float32)
+  logits = jnp.einsum('...qhc,...khc->...hqk', q * key_dim ** (-0.5), k) + bias
+  if pair_logits is not None:
+    logits += pair_logits
+  weights = jax.nn.softmax(logits, axis=-1)
+  weights = jnp.asarray(weights, dtype=x_q.dtype)
+
+  v = hm.Linear(
+      (config.num_head, value_dim), use_bias=False, name=f'{name}v_projection'
+  )(x_k)
+  weighted_avg = jnp.einsum('...hqk,...khc->...qhc', weights, v)
+  weighted_avg = jnp.reshape(weighted_avg, weighted_avg.shape[:-2] + (-1,))
+
+  gate_logits = hm.Linear(
+      config.num_head * value_dim,
+      bias_init=1.0,
+      initializer='zeros',
+      name=f'{name}gating_query',
+  )(x_q)
+  weighted_avg *= jax.nn.sigmoid(gate_logits)
+
+  output = adaptive_zero_init(
+      weighted_avg, x_q.shape[-1], single_cond_q, global_config, name
+  )
+  return output
+
+
+class CrossAttTransformer(hk.Module):
+  """Transformer that applies cross attention between two sets of subsets."""
+
+  class Config(base_config.BaseConfig):
+    num_intermediate_factor: int
+    num_blocks: int
+    attention: CrossAttentionConfig = base_config.autocreate()
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name: str = 'transformer',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def __call__(
+      self,
+      queries_act: jnp.ndarray,  # (num_subsets, num_queries, ch)
+      queries_mask: jnp.ndarray,  # (num_subsets, num_queries)
+      queries_to_keys: atom_layout.GatherInfo,  # (num_subsets, num_keys)
+      keys_mask: jnp.ndarray,  # (num_subsets, num_keys)
+      queries_single_cond: jnp.ndarray,  # (num_subsets, num_queries, ch)
+      keys_single_cond: jnp.ndarray,  # (num_subsets, num_keys, ch)
+      pair_cond: jnp.ndarray,  # (num_subsets, num_queries, num_keys, ch)
+  ) -> jnp.ndarray:
+    def block(queries_act, pair_logits):
+      # copy the queries activations to the keys layout
+      keys_act = atom_layout.convert(
+          queries_to_keys, queries_act, layout_axes=(-3, -2)
+      )
+      # cross attention
+      queries_act += cross_attention(
+          x_q=queries_act,
+          x_k=keys_act,
+          mask_q=queries_mask,
+          mask_k=keys_mask,
+          config=self.config.attention,
+          global_config=self.global_config,
+          pair_logits=pair_logits,
+          single_cond_q=queries_single_cond,
+          single_cond_k=keys_single_cond,
+          name=self.name,
+      )
+      queries_act += transition_block(
+          queries_act,
+          self.config.num_intermediate_factor,
+          self.global_config,
+          queries_single_cond,
+          name=self.name,
+      )
+      return queries_act, None
+
+    # Precompute pair logits for performance
+    pair_act = hm.LayerNorm(
+        name='pair_input_layer_norm',
+        use_fast_variance=False,
+        create_offset=False,
+    )(pair_cond)
+    # (num_subsets, num_queries, num_keys, num_blocks, num_heads)
+    pair_logits = hm.Linear(
+        (self.config.num_blocks, self.config.attention.num_head),
+        name='pair_logits_projection',
+    )(pair_act)
+    # (num_block, num_subsets, num_heads, num_queries, num_keys)
+    pair_logits = jnp.transpose(pair_logits, [3, 0, 4, 1, 2])
+
+    return hk.experimental.layer_stack(
+        self.config.num_blocks, with_per_layer_inputs=True
+    )(block)(queries_act, pair_logits)[0]

flax_model/alphafold3/model/network/distogram_head.py

@@ -0,0 +1,75 @@
+
+
+"""Distogram head."""
+
+from typing import Final
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.model import feat_batch
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model.components import haiku_modules as hm
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+_CONTACT_THRESHOLD: Final[float] = 8.0
+_CONTACT_EPSILON: Final[float] = 1e-3
+
+
+class DistogramHead(hk.Module):
+  """Distogram head."""
+
+  class Config(base_config.BaseConfig):
+    first_break: float = 2.3125
+    last_break: float = 21.6875
+    num_bins: int = 64
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name='distogram_head',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def __call__(
+      self,
+      batch: feat_batch.Batch,
+      embeddings: dict[str, jnp.ndarray],
+      return_distogram: bool = False,
+  ) -> dict[str, jnp.ndarray]:
+    pair_act = embeddings['pair']
+    seq_mask = batch.token_features.mask.astype(bool)
+    pair_mask = seq_mask[:, None] * seq_mask[None, :]
+
+    left_half_logits = hm.Linear(
+        self.config.num_bins,
+        initializer=self.global_config.final_init,
+        name='half_logits',
+    )(pair_act)
+
+    right_half_logits = left_half_logits
+    logits = left_half_logits + jnp.swapaxes(right_half_logits, -2, -3)
+    probs = jax.nn.softmax(logits, axis=-1)
+    breaks = jnp.linspace(
+        self.config.first_break,
+        self.config.last_break,
+        self.config.num_bins - 1,
+    )
+
+    bin_tops = jnp.append(breaks, breaks[-1] + (breaks[-1] - breaks[-2]))
+    threshold = _CONTACT_THRESHOLD + _CONTACT_EPSILON
+    is_contact_bin = 1.0 * (bin_tops <= threshold)
+    contact_probs = jnp.einsum(
+        'ijk,k->ij', probs, is_contact_bin, precision=jax.lax.Precision.HIGHEST
+    )
+    contact_probs = pair_mask * contact_probs
+
+    return_dict = {'bin_edges': breaks, 'contact_probs': contact_probs}
+    if return_distogram:
+      return_dict['distogram'] = logits
+
+    return return_dict

flax_model/alphafold3/model/network/featurization.py

@@ -0,0 +1,261 @@
+
+
+"""Model-side of the input features processing."""
+
+import functools
+
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.model import feat_batch
+from flax_model.alphafold3.model import features
+from flax_model.alphafold3.model.components import utils
+#import chex
+import jax
+import jax.numpy as jnp
+
+
+def _grid_keys(key, shape):
+  """Generate a grid of rng keys that is consistent with different padding.
+
+  Generate random keys such that the keys will be identical, regardless of
+  how much padding is added to any dimension.
+
+  Args:
+    key: A PRNG key.
+    shape: The shape of the output array of keys that will be generated.
+
+  Returns:
+    An array of shape `shape` consisting of random keys.
+  """
+  if not shape:
+    return key
+  new_keys = jax.vmap(functools.partial(jax.random.fold_in, key))(
+      jnp.arange(shape[0])
+  )
+  return jax.vmap(functools.partial(_grid_keys, shape=shape[1:]))(new_keys)
+
+
+def _padding_consistent_rng(f):
+  """Modify any element-wise random function to be consistent with padding.
+
+  Normally if you take a function like jax.random.normal and generate an array,
+  say of size (10,10), you will get a different set of random numbers to if you
+  add padding and take the first (10,10) sub-array.
+
+  This function makes a random function that is consistent regardless of the
+  amount of padding added.
+
+  Note: The padding-consistent function is likely to be slower to compile and
+  run than the function it is wrapping, but these slowdowns are likely to be
+  negligible in a large network.
+
+  Args:
+    f: Any element-wise function that takes (PRNG key, shape) as the first 2
+      arguments.
+
+  Returns:
+    An equivalent function to f, that is now consistent for different amounts of
+    padding.
+  """
+
+  def inner(key, shape, **kwargs):
+    keys = _grid_keys(key, shape)
+    signature = (
+        '()->()'
+        if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
+        else '(2)->()'
+    )
+    return jnp.vectorize(
+        functools.partial(f, shape=(), **kwargs), signature=signature
+    )(keys)
+
+  return inner
+
+
+def gumbel_argsort_sample_idx(
+    key: jnp.ndarray, logits: jnp.ndarray
+) -> jnp.ndarray:
+  """Samples with replacement from a distribution given by 'logits'.
+
+  This uses Gumbel trick to implement the sampling an efficient manner. For a
+  distribution over k items this samples k times without replacement, so this
+  is effectively sampling a random permutation with probabilities over the
+  permutations derived from the logprobs.
+
+  Args:
+    key: prng key
+    logits: logarithm of probabilities to sample from, probabilities can be
+      unnormalized.
+
+  Returns:
+    Sample from logprobs in one-hot form.
+  """
+  gumbel = _padding_consistent_rng(jax.random.gumbel)
+  z = gumbel(key, logits.shape)
+  # This construction is equivalent to jnp.argsort, but using a non stable sort,
+  # since stable sort's aren't supported by jax2tf
+  axis = len(logits.shape) - 1
+  iota = jax.lax.broadcasted_iota(jnp.int64, logits.shape, axis)
+  _, perm = jax.lax.sort_key_val(
+      logits + z, iota, dimension=-1, is_stable=False
+  )
+  return perm[::-1]
+
+
+def create_msa_feat(msa: features.MSA) -> jax.Array:
+  """Create and concatenate MSA features."""
+  msa_1hot = jax.nn.one_hot(
+      msa.rows, residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP + 1
+  )
+  deletion_matrix = msa.deletion_matrix
+  has_deletion = jnp.clip(deletion_matrix, 0.0, 1.0)[..., None]
+  deletion_value = (jnp.arctan(deletion_matrix / 3.0) * (2.0 / jnp.pi))[
+      ..., None
+  ]
+
+  msa_feat = [
+      msa_1hot,
+      has_deletion,
+      deletion_value,
+  ]
+
+  return jnp.concatenate(msa_feat, axis=-1)
+
+
+def truncate_msa_batch(msa: features.MSA, num_msa: int) -> features.MSA:
+  indices = jnp.arange(num_msa)
+  return msa.index_msa_rows(indices)
+
+
+def create_target_feat(
+    batch: feat_batch.Batch,
+    append_per_atom_features: bool,
+) -> jax.Array:
+  """Make target feat."""
+  token_features = batch.token_features
+  target_features = []
+  target_features.append(
+      jax.nn.one_hot(
+          token_features.aatype,
+          residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP,
+      )
+  )
+  target_features.append(batch.msa.profile)
+  target_features.append(batch.msa.deletion_mean[..., None])
+
+  # Reference structure features
+  if append_per_atom_features:
+    ref_mask = batch.ref_structure.mask
+    element_feat = jax.nn.one_hot(batch.ref_structure.element, 128)
+    element_feat = utils.mask_mean(
+        mask=ref_mask[..., None], value=element_feat, axis=-2, eps=1e-6
+    )
+    target_features.append(element_feat)
+    pos_feat = batch.ref_structure.positions
+    pos_feat = pos_feat.reshape([pos_feat.shape[0], -1])
+    target_features.append(pos_feat)
+    target_features.append(ref_mask)
+
+  return jnp.concatenate(target_features, axis=-1)
+
+
+def create_relative_encoding(
+    seq_features: features.TokenFeatures,
+    max_relative_idx: int,
+    max_relative_chain: int,
+) -> jax.Array:
+  """Add relative position encodings."""
+  rel_feats = []
+  token_index = seq_features.token_index
+  residue_index = seq_features.residue_index
+  asym_id = seq_features.asym_id
+  entity_id = seq_features.entity_id
+  sym_id = seq_features.sym_id
+
+  left_asym_id = asym_id[:, None]
+  right_asym_id = asym_id[None, :]
+
+  left_residue_index = residue_index[:, None]
+  right_residue_index = residue_index[None, :]
+
+  left_token_index = token_index[:, None]
+  right_token_index = token_index[None, :]
+
+  left_entity_id = entity_id[:, None]
+  right_entity_id = entity_id[None, :]
+
+  left_sym_id = sym_id[:, None]
+  right_sym_id = sym_id[None, :]
+
+  # Embed relative positions using a one-hot embedding of distance along chain
+  offset = left_residue_index - right_residue_index
+  clipped_offset = jnp.clip(
+      offset + max_relative_idx, min=0, max=2 * max_relative_idx
+  )
+  asym_id_same = left_asym_id == right_asym_id
+  final_offset = jnp.where(
+      asym_id_same,
+      clipped_offset,
+      (2 * max_relative_idx + 1) * jnp.ones_like(clipped_offset),
+  )
+  rel_pos = jax.nn.one_hot(final_offset, 2 * max_relative_idx + 2)
+  rel_feats.append(rel_pos)
+
+  # Embed relative token index as a one-hot embedding of distance along residue
+  token_offset = left_token_index - right_token_index
+  clipped_token_offset = jnp.clip(
+      token_offset + max_relative_idx, min=0, max=2 * max_relative_idx
+  )
+  residue_same = (left_asym_id == right_asym_id) & (
+      left_residue_index == right_residue_index
+  )
+  final_token_offset = jnp.where(
+      residue_same,
+      clipped_token_offset,
+      (2 * max_relative_idx + 1) * jnp.ones_like(clipped_token_offset),
+  )
+  rel_token = jax.nn.one_hot(final_token_offset, 2 * max_relative_idx + 2)
+  rel_feats.append(rel_token)
+
+  # Embed same entity ID
+  entity_id_same = left_entity_id == right_entity_id
+  rel_feats.append(entity_id_same.astype(rel_pos.dtype)[..., None])
+
+  # Embed relative chain ID inside each symmetry class
+  rel_sym_id = left_sym_id - right_sym_id
+
+  max_rel_chain = max_relative_chain
+
+  clipped_rel_chain = jnp.clip(
+      rel_sym_id + max_rel_chain, min=0, max=2 * max_rel_chain
+  )
+
+  final_rel_chain = jnp.where(
+      entity_id_same,
+      clipped_rel_chain,
+      (2 * max_rel_chain + 1) * jnp.ones_like(clipped_rel_chain),
+  )
+  rel_chain = jax.nn.one_hot(final_rel_chain, 2 * max_relative_chain + 2)
+
+  rel_feats.append(rel_chain)
+
+  return jnp.concatenate(rel_feats, axis=-1)
+
+
+def shuffle_msa(
+    key: jax.Array, msa: features.MSA
+) -> tuple[features.MSA, jax.Array]:
+  """Shuffle MSA randomly, return batch with shuffled MSA.
+
+  Args:
+    key: rng key for random number generation.
+    msa: MSA object to sample msa from.
+
+  Returns:
+    Protein with sampled msa.
+  """
+  key, sample_key = jax.random.split(key)
+  # Sample uniformly among sequences with at least one non-masked position.
+  logits = (jnp.clip(jnp.sum(msa.mask, axis=-1), 0.0, 1.0) - 1.0) * 1e6
+  index_order = gumbel_argsort_sample_idx(sample_key, logits)
+
+  return msa.index_msa_rows(index_order), key

flax_model/alphafold3/model/network/noise_level_embeddings.py

@@ -0,0 +1,133 @@
+
+
+"""Fourier embeddings for given noise levels.
+
+We supply fixed weights and biases for the Fourier embeddings. These were
+initially generated by the following code, but we make them into constants
+to future proof against changes in jax rng generation:
+
+```
+dim = 256
+w_key, b_key = jax.random.split(jax.random.PRNGKey(42))
+weight = jax.random.normal(w_key, shape=[dim])
+bias = jax.random.uniform(b_key, shape=[dim])
+```
+"""
+
+import jax.numpy as jnp
+
+# pyformat: disable
+# pylint: disable=bad-whitespace
+# pylint: disable=bad-continuation
+_WEIGHT = [
+     0.45873642,  0.06516238, -0.07278306, -0.26992258,  0.64292115,
+    -0.40763968,  3.60116863,  0.54461384, -0.32644904,  2.10888267,
+     1.30805349,  1.19838560, -1.37745857,  1.99475312, -1.64120293,
+     1.07823789, -0.02288206,  0.88305283,  0.48099944,  0.17655374,
+     0.30281949,  0.80646873,  0.62605333, -0.23965347, -1.02609432,
+     0.75006109, -0.19913037,  0.07466396,  0.66431236, -0.60990530,
+    -0.69709194, -0.44453633, -1.77656078,  0.02299878,  0.04095552,
+     0.35485864, -0.47602659, -0.98820388, -0.24106771, -1.07254291,
+    -0.99741757,  0.22697604,  1.41390419,  1.54984057, -0.12237291,
+     0.20156337,  0.61767143,  0.23959029,  0.92454034,  1.84082258,
+     0.89030224,  0.39598912, -1.52224910,  0.29669049,  1.52356744,
+    -0.33968377,  0.24155144, -0.52308381, -0.23622665,  0.92825454,
+    -0.63864607, -0.62169307,  0.78623551, -0.80352145, -0.45496067,
+     1.30877995, -0.06686528,  1.00248849, -0.63593471,  0.16372502,
+    -1.46133232,  1.10562658, -0.01693927,  0.28684548, -0.72843230,
+     0.66133535, -1.92225552,  0.70241231, -0.96868867, -0.47309339,
+    -1.66894221,  0.46018723, -0.56806105,  0.32694784, -0.46529883,
+     1.02299964,  0.84688205,  1.19581807, -1.82454145,  0.05999713,
+    -0.59530073,  1.44862521, -0.34933713, -0.46564487, -0.55005538,
+    -1.61170268,  0.17502306,  0.38670063, -1.12133658, -0.29343036,
+    -0.52527446, -1.26285112,  1.07982683,  0.51215219,  1.48963666,
+     1.09847653, -0.01563358,  0.32574457,  1.94779706, -1.29198587,
+     1.06249654, -0.86965990,  0.22975266, -0.27182648, -0.21130897,
+    -0.41773933, -0.02329035,  1.31049252,  0.05579265, -1.23127055,
+    -0.99691105,  0.27058721, -0.72509319, -0.14421797, -1.48605061,
+     1.35041201,  1.29619241, -1.01022530, -0.79787987, -0.16166858,
+     0.87210685,  1.69248152,  1.42469788, -0.72325104, -1.24823737,
+     0.07051118,  0.71332991, -0.07360429, -0.91955227, -2.68856549,
+    -0.44033936,  0.35482934, -0.57933813,  0.97468042, -0.31050494,
+    -0.88454425, -2.08785224,  0.47322822, -0.02400172,  0.26644820,
+    -0.19147627, -2.10538960, -1.27962470, -1.35999286,  2.09867334,
+     0.65099514,  0.21604492, -0.45951018,  0.15994427, -0.31420693,
+    -0.65202618, -0.61077976, -1.06100249, -1.47254968,  1.18165290,
+    -0.78656220,  1.28182006,  1.80323684,  1.09196901,  0.26118696,
+    -0.30168581,  0.39749333,  0.26812574, -1.51995814, -0.46909946,
+     0.03874255, -1.36774313,  2.30143976,  2.06959820, -0.41647521,
+     1.85624206,  0.49019700, -0.06726539,  0.00457313,  0.23915423,
+    -1.84971249, -0.20482327, -0.34097880, -0.57933033, -1.10541213,
+    -0.30269983, -0.16430426, -0.82371718,  0.10345812,  1.78753936,
+     0.04786763,  1.86778629, -0.65214992,  0.81544143, -0.28214937,
+     0.31187257,  0.57661986,  1.21938801, -1.56046617,  0.38046429,
+    -0.18235965,  0.81794524, -0.40474343,  0.46538028, -1.15558851,
+     0.59625793, -1.07801270,  0.07310858,  0.61526084,  0.55518496,
+    -0.49787554,  0.92703879, -1.27780271, -0.83373469, -0.43015575,
+     0.41877759, -1.03987372, -1.46055734,  0.61282396,  0.15590595,
+    -0.34269521,  0.56509072, -1.17904210,  0.11374855, -1.83310866,
+     0.38734794, -0.58623004,  0.77931106,  1.53930688, -0.70299625,
+    -0.11389336, -1.14818096, -0.44400632,  1.21887410,  0.64066756,
+    -0.70249403, -0.27244881,  0.38586098, -1.07925785,  0.12448707,
+    -1.28286278,  0.37827531,  0.68812364,  1.65695465,  0.12440517,
+    -0.03689830,  1.10224664, -0.28323629, -0.47939169,  0.70120829,
+    -0.67204583
+]
+
+_BIAS = [
+    0.00465965, 0.21738243, 0.22277749, 0.68463874, 0.84596848, 0.17337036,
+    0.39573753, 0.78153563, 0.86311185, 0.21782327, 0.24377882, 0.42310703,
+    0.19887352, 0.10486019, 0.48707581, 0.22205460, 0.97263455, 0.29714966,
+    0.11244559, 0.53020525, 0.36796236, 0.37294638, 0.80261672, 0.04669094,
+    0.86319661, 0.75907171, 0.77297020, 0.01114798, 0.55850804, 0.91799915,
+    0.23032320, 0.12154722, 0.26701927, 0.42934716, 0.47951782, 0.96782577,
+    0.86785042, 0.61985648, 0.05743814, 0.41800117, 0.68881893, 0.60575199,
+    0.21058667, 0.64412105, 0.63958526, 0.89390790, 0.69755554, 0.89345169,
+    0.53330755, 0.56985939, 0.30724049, 0.00984561, 0.91407037, 0.92118979,
+    0.94153070, 0.81097460, 0.70537627, 0.32810748, 0.47227263, 0.11821401,
+    0.44983089, 0.30767226, 0.31756389, 0.62969446, 0.69892538, 0.16949117,
+    0.06207097, 0.46717727, 0.95348179, 0.62363589, 0.49018729, 0.06920040,
+    0.39333904, 0.41299903, 0.52514863, 0.61197245, 0.56871891, 0.65053988,
+    0.22203422, 0.46748531, 0.86931503, 0.87050021, 0.40208721, 0.32084906,
+    0.55084610, 0.94584596, 0.76279902, 0.36250532, 0.74272907, 0.66682065,
+    0.96452832, 0.64768302, 0.88070846, 0.56995463, 0.06395614, 0.69499350,
+    0.44494808, 0.39775658, 0.20280898, 0.33363521, 0.05999005, 0.44414878,
+    0.65227020, 0.01199079, 0.71995056, 0.19045687, 0.48342144, 0.25127733,
+    0.66515994, 0.22465158, 0.22313106, 0.06302810, 0.55783665, 0.93625581,
+    0.58800840, 0.72525370, 0.52879298, 0.77195418, 0.15548682, 0.01028740,
+    0.39325142, 0.45401239, 0.71494079, 0.33011997, 0.05050695, 0.26381660,
+    0.63064706, 0.47604024, 0.08593416, 0.00383425, 0.06352687, 0.05510247,
+    0.03552997, 0.35810637, 0.56094289, 0.60922170, 0.88599777, 0.45419788,
+    0.40486634, 0.71297824, 0.34976673, 0.97825217, 0.12915993, 0.09566259,
+    0.64318919, 0.16717327, 0.82308614, 0.32672071, 0.81688786, 0.84857118,
+    0.99922776, 0.07551706, 0.18766022, 0.13051236, 0.39136350, 0.08768725,
+    0.92048228, 0.87185788, 0.39158428, 0.79224777, 0.17492688, 0.68902445,
+    0.81980729, 0.70458186, 0.59489477, 0.93324888, 0.49986637, 0.40705478,
+    0.89202917, 0.20673239, 0.39339757, 0.20996964, 0.02923799, 0.53992438,
+    0.40119815, 0.10366607, 0.08044600, 0.95551598, 0.20518017, 0.68826210,
+    0.90159297, 0.69008791, 0.86880815, 0.16246438, 0.89628279, 0.11481643,
+    0.61353648, 0.41545081, 0.92478311, 0.78212476, 0.48292696, 0.79621077,
+    0.11947489, 0.01747024, 0.22928023, 0.87387264, 0.86349785, 0.89526737,
+    0.58904779, 0.13896775, 0.68194926, 0.55824125, 0.44428205, 0.55422378,
+    0.28189969, 0.27923775, 0.09979951, 0.66994715, 0.45943546, 0.71207762,
+    0.17300689, 0.83434916, 0.02573085, 0.45858085, 0.55934799, 0.30676675,
+    0.52219367, 0.34544575, 0.19280875, 0.26937950, 0.07147646, 0.06295013,
+    0.76382887, 0.38737607, 0.58825982, 0.17423475, 0.05509448, 0.97228825,
+    0.94380617, 0.91664016, 0.18800116, 0.41771865, 0.59420645, 0.77371931,
+    0.64687788, 0.27284670, 0.22310913, 0.15663862, 0.45573199, 0.50386798,
+    0.66712272, 0.71649647, 0.28475654, 0.83415413, 0.75261366, 0.61517799,
+    0.93544555, 0.76141870, 0.85474241, 0.74766934, 0.33459592, 0.78477907,
+    0.07250881, 0.10174239, 0.95332730, 0.80793905
+]
+# pyformat: enable
+# pylint: enable=bad-whitespace
+# pylint: enable=bad-continuation
+
+
+def noise_embeddings(sigma_scaled_noise_level: jnp.ndarray) -> jnp.ndarray:
+  """Returns Fourier noise level embeddings for diffusion model."""
+  transformed_noise_level = (1 / 4) * jnp.log(sigma_scaled_noise_level)
+  weight = jnp.array(_WEIGHT, dtype=jnp.float32)
+  bias = jnp.array(_BIAS, dtype=jnp.float32)
+  embeddings = transformed_noise_level[..., None] * weight + bias
+  return jnp.cos(2 * jnp.pi * embeddings)

flax_model/alphafold3/model/network/template_modules.py

@@ -0,0 +1,343 @@
+
+
+"""Modules for embedding templates."""
+
+from flax_model.alphafold3.common import base_config
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.jax import geometry
+from flax_model.alphafold3.model import features
+from flax_model.alphafold3.model import model_config
+from flax_model.alphafold3.model import protein_data_processing
+from flax_model.alphafold3.model.components import haiku_modules as hm
+from flax_model.alphafold3.model.network import modules
+from flax_model.alphafold3.model.scoring import scoring
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+class DistogramFeaturesConfig(base_config.BaseConfig):
+  # The left edge of the first bin.
+  min_bin: float = 3.25
+  # The left edge of the final bin. The final bin catches everything larger than
+  # `max_bin`.
+  max_bin: float = 50.75
+  # The number of bins in the distogram.
+  num_bins: int = 39
+
+
+def dgram_from_positions(positions, config: DistogramFeaturesConfig):
+  """Compute distogram from amino acid positions.
+
+  Args:
+    positions: (num_res, 3) Position coordinates.
+    config: Distogram bin configuration.
+
+  Returns:
+    Distogram with the specified number of bins.
+  """
+  lower_breaks = jnp.linspace(config.min_bin, config.max_bin, config.num_bins)
+  lower_breaks = jnp.square(lower_breaks)
+  upper_breaks = jnp.concatenate(
+      [lower_breaks[1:], jnp.array([1e8], dtype=jnp.float32)], axis=-1
+  )
+  dist2 = jnp.sum(
+      jnp.square(
+          jnp.expand_dims(positions, axis=-2)
+          - jnp.expand_dims(positions, axis=-3)
+      ),
+      axis=-1,
+      keepdims=True,
+  )
+
+  dgram = (dist2 > lower_breaks).astype(jnp.float32) * (
+      dist2 < upper_breaks
+  ).astype(jnp.float32)
+  return dgram
+
+
+def make_backbone_rigid(
+    positions: geometry.Vec3Array,
+    mask: jnp.ndarray,
+    group_indices: jnp.ndarray,
+) -> tuple[geometry.Rigid3Array, jnp.ndarray]:
+  """Make backbone Rigid3Array and mask.
+
+  Args:
+    positions: (num_res, num_atoms) of atom positions as Vec3Array.
+    mask: (num_res, num_atoms) for atom mask.
+    group_indices: (num_res, num_group, 3) for atom indices forming groups.
+
+  Returns:
+    tuple of backbone Rigid3Array and mask (num_res,).
+  """
+  backbone_indices = group_indices[:, 0]
+
+  # main backbone frames differ in sidechain frame convention.
+  # for sidechain it's (C, CA, N), for backbone it's (N, CA, C)
+  # Hence using c, b, a, each of shape (num_res,).
+  c, b, a = [backbone_indices[..., i] for i in range(3)]
+
+  slice_index = jax.vmap(lambda x, i: x[i])
+  rigid_mask = (
+      slice_index(mask, a) * slice_index(mask, b) * slice_index(mask, c)
+  ).astype(jnp.float32)
+
+  frame_positions = []
+  for indices in [a, b, c]:
+    frame_positions.append(
+        jax.tree.map(lambda x, idx=indices: slice_index(x, idx), positions)
+    )
+
+  rotation = geometry.Rot3Array.from_two_vectors(
+      frame_positions[2] - frame_positions[1],
+      frame_positions[0] - frame_positions[1],
+  )
+  rigid = geometry.Rigid3Array(rotation, frame_positions[1])
+
+  return rigid, rigid_mask
+
+
+class TemplateEmbedding(hk.Module):
+  """Embed a set of templates."""
+
+  class Config(base_config.BaseConfig):
+    num_channels: int = 64
+    template_stack: modules.PairFormerIteration.Config = base_config.autocreate(
+        num_layer=2,
+        pair_transition=base_config.autocreate(num_intermediate_factor=2),
+    )
+    dgram_features: DistogramFeaturesConfig = base_config.autocreate()
+
+  def __init__(
+      self,
+      config: Config,
+      global_config: model_config.GlobalConfig,
+      name='template_embedding',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def __call__(
+      self,
+      query_embedding: jnp.ndarray,
+      templates: features.Templates,
+      padding_mask_2d: jnp.ndarray,
+      multichain_mask_2d: jnp.ndarray,
+      key: jnp.ndarray,
+  ) -> jnp.ndarray:
+    """Generate an embedding for a set of templates.
+
+    Args:
+      query_embedding: [num_res, num_res, num_channel] a query tensor that will
+        be used to attend over the templates to remove the num_templates
+        dimension.
+      templates: A 'Templates' object.
+      padding_mask_2d: [num_res, num_res] Pair mask for attention operations.
+      multichain_mask_2d: [num_res, num_res] Pair mask for multichain.
+      key: random key generator.
+
+    Returns:
+      An embedding of size [num_res, num_res, num_channels]
+    """
+    c = self.config
+    num_residues = query_embedding.shape[0]
+    num_templates = templates.aatype.shape[0]
+    query_num_channels = query_embedding.shape[2]
+    num_atoms = 24
+    assert query_embedding.shape == (
+        num_residues,
+        num_residues,
+        query_num_channels,
+    )
+    assert templates.aatype.shape == (num_templates, num_residues)
+    assert templates.atom_positions.shape == (
+        num_templates,
+        num_residues,
+        num_atoms,
+        3,
+    )
+    assert templates.atom_mask.shape == (num_templates, num_residues, num_atoms)
+    assert padding_mask_2d.shape == (num_residues, num_residues)
+
+    num_templates = templates.aatype.shape[0]
+    num_res, _, query_num_channels = query_embedding.shape
+
+    # Embed each template separately.
+    template_embedder = SingleTemplateEmbedding(self.config, self.global_config)
+
+    subkeys = jnp.array(jax.random.split(key, num_templates))
+
+    def scan_fn(carry, x):
+      templates, key = x
+      embedding = template_embedder(
+          query_embedding,
+          templates,
+          padding_mask_2d,
+          multichain_mask_2d,
+          key,
+      )
+      return carry + embedding, None
+
+    scan_init = jnp.zeros(
+        (num_res, num_res, c.num_channels), dtype=query_embedding.dtype
+    )
+    summed_template_embeddings, _ = hk.scan(
+        scan_fn, scan_init, (templates, subkeys)
+    )
+
+    embedding = summed_template_embeddings / (1e-7 + num_templates)
+    embedding = jax.nn.relu(embedding)
+    embedding = hm.Linear(
+        query_num_channels, initializer='relu', name='output_linear'
+    )(embedding)
+
+    assert embedding.shape == (num_residues, num_residues, query_num_channels)
+    return embedding
+
+
+class SingleTemplateEmbedding(hk.Module):
+  """Embed a single template."""
+
+  def __init__(
+      self,
+      config: TemplateEmbedding.Config,
+      global_config: model_config.GlobalConfig,
+      name='single_template_embedding',
+  ):
+    super().__init__(name=name)
+    self.config = config
+    self.global_config = global_config
+
+  def __call__(
+      self,
+      query_embedding: jnp.ndarray,
+      templates: features.Templates,
+      padding_mask_2d: jnp.ndarray,
+      multichain_mask_2d: jnp.ndarray,
+      key: jnp.ndarray,
+  ) -> jnp.ndarray:
+    """Build the single template embedding graph.
+
+    Args:
+      query_embedding: (num_res, num_res, num_channels) - embedding of the query
+        sequence/msa.
+      templates: 'Templates' object containing single Template.
+      padding_mask_2d: Padding mask (Note: this doesn't care if a template
+        exists, unlike the template_pseudo_beta_mask).
+      multichain_mask_2d: A mask indicating intra-chain residue pairs, used to
+        mask out between chain distances/features when templates are for single
+        chains.
+      key: Random key generator.
+
+    Returns:
+      A template embedding (num_res, num_res, num_channels).
+    """
+    gc = self.global_config
+    c = self.config
+    assert padding_mask_2d.dtype == query_embedding.dtype
+    dtype = query_embedding.dtype
+    num_channels = self.config.num_channels
+
+    def construct_input(
+        query_embedding, templates: features.Templates, multichain_mask_2d
+    ):
+
+      # Compute distogram feature for the template.
+      aatype = templates.aatype
+      dense_atom_mask = templates.atom_mask
+
+      dense_atom_positions = templates.atom_positions
+      dense_atom_positions *= dense_atom_mask[..., None]
+
+      pseudo_beta_positions, pseudo_beta_mask = scoring.pseudo_beta_fn(
+          templates.aatype, dense_atom_positions, dense_atom_mask
+      )
+      pseudo_beta_mask_2d = (
+          pseudo_beta_mask[:, None] * pseudo_beta_mask[None, :]
+      )
+      pseudo_beta_mask_2d *= multichain_mask_2d
+      dgram = dgram_from_positions(
+          pseudo_beta_positions, self.config.dgram_features
+      )
+      dgram *= pseudo_beta_mask_2d[..., None]
+      dgram = dgram.astype(dtype)
+      pseudo_beta_mask_2d = pseudo_beta_mask_2d.astype(dtype)
+      to_concat = [(dgram, 1), (pseudo_beta_mask_2d, 0)]
+
+      aatype = jax.nn.one_hot(
+          aatype,
+          residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP,
+          axis=-1,
+          dtype=dtype,
+      )
+      to_concat.append((aatype[None, :, :], 1))
+      to_concat.append((aatype[:, None, :], 1))
+
+      # Compute a feature representing the normalized vector between each
+      # backbone affine - i.e. in each residues local frame, what direction are
+      # each of the other residues.
+
+      template_group_indices = jnp.take(
+          protein_data_processing.RESTYPE_RIGIDGROUP_DENSE_ATOM_IDX,
+          templates.aatype,
+          axis=0,
+      )
+      rigid, backbone_mask = make_backbone_rigid(
+          geometry.Vec3Array.from_array(dense_atom_positions),
+          dense_atom_mask,
+          template_group_indices.astype(jnp.int32),
+      )
+      points = rigid.translation
+      rigid_vec = rigid[:, None].inverse().apply_to_point(points)
+      unit_vector = rigid_vec.normalized()
+      unit_vector = [unit_vector.x, unit_vector.y, unit_vector.z]
+
+      unit_vector = [x.astype(dtype) for x in unit_vector]
+      backbone_mask = backbone_mask.astype(dtype)
+
+      backbone_mask_2d = backbone_mask[:, None] * backbone_mask[None, :]
+      backbone_mask_2d *= multichain_mask_2d
+      unit_vector = [x * backbone_mask_2d for x in unit_vector]
+
+      # Note that the backbone_mask takes into account C, CA and N (unlike
+      # pseudo beta mask which just needs CB) so we add both masks as features.
+      to_concat.extend([(x, 0) for x in unit_vector])
+      to_concat.append((backbone_mask_2d, 0))
+
+      query_embedding = hm.LayerNorm(name='query_embedding_norm')(
+          query_embedding
+      )
+      # Allow the template embedder to see the query embedding.  Note this
+      # contains the position relative feature, so this is how the network knows
+      # which residues are next to each other.
+      to_concat.append((query_embedding, 1))
+
+      act = 0
+
+      for i, (x, n_input_dims) in enumerate(to_concat):
+        act += hm.Linear(
+            num_channels,
+            num_input_dims=n_input_dims,
+            initializer='relu',
+            name=f'template_pair_embedding_{i}',
+        )(x)
+      return act
+
+    act = construct_input(query_embedding, templates, multichain_mask_2d)
+
+    if c.template_stack.num_layer:
+
+      def template_iteration_fn(x):
+        return modules.PairFormerIteration(
+            c.template_stack, gc, name='template_embedding_iteration'
+        )(act=x, pair_mask=padding_mask_2d)
+
+      template_stack = hk.experimental.layer_stack(c.template_stack.num_layer)(
+          template_iteration_fn
+      )
+      act = template_stack(act)
+
+    act = hm.LayerNorm(name='output_layer_norm')(act)
+    return act

flax_model/alphafold3/model/params.py

@@ -0,0 +1,202 @@
+
+
+"""Model param loading."""
+
+import bisect
+import collections
+from collections.abc import Iterator
+import contextlib
+import io
+import os
+import pathlib
+import re
+import struct
+import sys
+from typing import IO
+
+import haiku as hk
+import jax.numpy as jnp
+import numpy as np
+import zstandard
+
+
+class RecordError(Exception):
+  """Error reading a record."""
+
+
+def encode_record(scope: str, name: str, arr: np.ndarray) -> bytes:
+  """Encodes a single haiku param as bytes, preserving non-numpy dtypes."""
+  scope = scope.encode('utf-8')
+  name = name.encode('utf-8')
+  shape = arr.shape
+  dtype = str(arr.dtype).encode('utf-8')
+  arr = np.ascontiguousarray(arr)
+  if sys.byteorder == 'big':
+    arr = arr.byteswap()
+  arr_buffer = arr.tobytes('C')
+  header = struct.pack(
+      '<5i', len(scope), len(name), len(dtype), len(shape), len(arr_buffer)
+  )
+  return header + b''.join(
+      (scope, name, dtype, struct.pack(f'{len(shape)}i', *shape), arr_buffer)
+  )
+
+
+def _read_record(stream: IO[bytes]) -> tuple[str, str, np.ndarray] | None:
+  """Reads a record encoded by `_encode_record` from a byte stream."""
+  header_size = struct.calcsize('<5i')
+  header = stream.read(header_size)
+  if not header:
+    return None
+  if len(header) < header_size:
+    raise RecordError(f'Incomplete header: {len(header)=} < {header_size=}')
+  (scope_len, name_len, dtype_len, shape_len, arr_buffer_len) = struct.unpack(
+      '<5i', header
+  )
+  fmt = f'<{scope_len}s{name_len}s{dtype_len}s{shape_len}i'
+  payload_size = struct.calcsize(fmt) + arr_buffer_len
+  payload = stream.read(payload_size)
+  if len(payload) < payload_size:
+    raise RecordError(f'Incomplete payload: {len(payload)=} < {payload_size=}')
+  scope, name, dtype, *shape = struct.unpack_from(fmt, payload)
+  scope = scope.decode('utf-8')
+  name = name.decode('utf-8')
+  dtype = dtype.decode('utf-8')
+  arr = np.frombuffer(payload[-arr_buffer_len:], dtype=dtype)
+  arr = np.reshape(arr, shape)
+  if sys.byteorder == 'big':
+    arr = arr.byteswap()
+  return scope, name, arr
+
+
+def read_records(stream: IO[bytes]) -> Iterator[tuple[str, str, np.ndarray]]:
+  """Fully reads the contents of a byte stream."""
+  while record := _read_record(stream):
+    yield record
+
+
+class _MultiFileIO(io.RawIOBase):
+  """A file-like object that presents a concatenated view of multiple files."""
+
+  def __init__(self, files: list[pathlib.Path]):
+    self._files = files
+    self._stack = contextlib.ExitStack()
+    self._handles = [
+        self._stack.enter_context(file.open('rb')) for file in files
+    ]
+    self._sizes = []
+    for handle in self._handles:
+      handle.seek(0, os.SEEK_END)
+      self._sizes.append(handle.tell())
+    self._length = sum(self._sizes)
+    self._offsets = [0]
+    for s in self._sizes[:-1]:
+      self._offsets.append(self._offsets[-1] + s)
+    self._abspos = 0
+    self._relpos = (0, 0)
+
+  def _abs_to_rel(self, pos: int) -> tuple[int, int]:
+    idx = bisect.bisect_right(self._offsets, pos) - 1
+    return idx, pos - self._offsets[idx]
+
+  def close(self):
+    self._stack.close()
+
+  def closed(self) -> bool:
+    return all(handle.closed for handle in self._handles)
+
+  def fileno(self) -> int:
+    return -1
+
+  def readable(self) -> bool:
+    return True
+
+  def tell(self) -> int:
+    return self._abspos
+
+  def seek(self, pos: int, whence: int = os.SEEK_SET, /):
+    match whence:
+      case os.SEEK_SET:
+        pass
+      case os.SEEK_CUR:
+        pos += self._abspos
+      case os.SEEK_END:
+        pos = self._length - pos
+      case _:
+        raise ValueError(f'Invalid whence: {whence}')
+    self._abspos = pos
+    self._relpos = self._abs_to_rel(pos)
+
+  def readinto(self, b: bytearray | memoryview) -> int:
+    result = 0
+    mem = memoryview(b)
+    while mem:
+      self._handles[self._relpos[0]].seek(self._relpos[1])
+      count = self._handles[self._relpos[0]].readinto(mem)
+      result += count
+      self._abspos += count
+      self._relpos = self._abs_to_rel(self._abspos)
+      mem = mem[count:]
+      if self._abspos == self._length:
+        break
+    return result
+
+
+@contextlib.contextmanager
+def open_for_reading(model_files: list[pathlib.Path], is_compressed: bool):
+  with contextlib.closing(_MultiFileIO(model_files)) as f:
+    if is_compressed:
+      yield zstandard.ZstdDecompressor().stream_reader(f)
+    else:
+      yield f
+
+
+def _match_model(
+    paths: list[pathlib.Path], pattern: re.Pattern[str]
+) -> dict[str, list[pathlib.Path]]:
+  """Match files in a directory with a pattern, and group by model name."""
+  models = collections.defaultdict(list)
+  for path in paths:
+    match = pattern.fullmatch(path.name)
+    if match:
+      models[match.group('model_name')].append(path)
+  return {k: sorted(v) for k, v in models.items()}
+
+
+def select_model_files(
+    model_dir: pathlib.Path, model_name: str | None = None
+) -> tuple[list[pathlib.Path], bool]:
+  """Select the model files from a model directory."""
+  files = [file for file in model_dir.iterdir() if file.is_file()]
+
+  for pattern, is_compressed in (
+      (r'(?P<model_name>.*)\.[0-9]+\.bin\.zst$', True),
+      (r'(?P<model_name>.*)\.bin\.zst\.[0-9]+$', True),
+      (r'(?P<model_name>.*)\.[0-9]+\.bin$', False),
+      (r'(?P<model_name>.*)\.bin]\.[0-9]+$', False),
+      (r'(?P<model_name>.*)\.bin\.zst$', True),
+      (r'(?P<model_name>.*)\.bin$', False),
+  ):
+    models = _match_model(files, re.compile(pattern))
+    if model_name is not None:
+      if model_name in models:
+        return models[model_name], is_compressed
+    else:
+      if models:
+        if len(models) > 1:
+          raise RuntimeError(f'Multiple models matched in {model_dir}')
+        _, model_files = models.popitem()
+        return model_files, is_compressed
+  raise FileNotFoundError(f'No models matched in {model_dir}')
+
+
+def get_model_haiku_params(model_dir: pathlib.Path) -> hk.Params:
+  """Get the Haiku parameters from a model name."""
+  params: dict[str, dict[str, jnp.Array]] = {}
+  model_files, is_compressed = select_model_files(model_dir)
+  with open_for_reading(model_files, is_compressed) as stream:
+    for scope, name, arr in read_records(stream):
+      params.setdefault(scope, {})[name] = jnp.array(arr)
+  if not params:
+    raise FileNotFoundError(f'Model missing from "{model_dir}"')
+  return params

flax_model/alphafold3/model/pipeline/structure_cleaning.py

@@ -0,0 +1,360 @@
+
+
+"""Prepare PDB structure for training or inference."""
+
+from typing import Any
+
+from absl import logging
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.constants import chemical_component_sets
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.model.atom_layout import atom_layout
+from flax_model.alphafold3.model.pipeline import inter_chain_bonds
+from flax_model.alphafold3.model.scoring import covalent_bond_cleaning
+from flax_model.alphafold3.structure import sterics
+import numpy as np
+
+
+def _get_leaving_atom_mask(
+    struc: structure.Structure,
+    polymer_ligand_bonds: atom_layout.AtomLayout | None,
+    ligand_ligand_bonds: atom_layout.AtomLayout | None,
+    chain_id: str,
+    chain_type: str,
+    res_id: int,
+    res_name: str,
+) -> np.ndarray:
+  """Updates a drop_leaving_atoms mask with new leaving atom locations."""
+  bonded_atoms = atom_layout.get_bonded_atoms(
+      polymer_ligand_bonds,
+      ligand_ligand_bonds,
+      res_id,
+      chain_id,
+  )
+  # Connect the amino-acids, i.e. remove OXT, HXT and H2.
+  drop_atoms = atom_layout.get_link_drop_atoms(
+      res_name=res_name,
+      chain_type=chain_type,
+      is_start_terminus=False,
+      is_end_terminus=False,
+      bonded_atoms=bonded_atoms,
+      drop_ligand_leaving_atoms=True,
+  )
+  # Default mask where everything is false, which equates to being kept.
+  drop_atom_filter_atoms = struc.chain_id != struc.chain_id
+  for drop_atom in drop_atoms:
+    drop_atom_filter_atom = np.logical_and(
+        np.logical_and(
+            struc.atom_name == drop_atom,
+            struc.chain_id == chain_id,
+        ),
+        struc.res_id == res_id,
+    )
+    drop_atom_filter_atoms = np.logical_or(
+        drop_atom_filter_atoms, drop_atom_filter_atom
+    )
+  return drop_atom_filter_atoms
+
+
+def clean_structure(
+    struc: structure.Structure,
+    ccd: chemical_components.Ccd,
+    *,
+    drop_missing_sequence: bool,
+    filter_clashes: bool,
+    drop_non_standard_atoms: bool,
+    filter_crystal_aids: bool,
+    filter_waters: bool,
+    filter_hydrogens: bool,
+    filter_leaving_atoms: bool,
+    only_glycan_ligands_for_leaving_atoms: bool,
+    covalent_bonds_only: bool,
+    remove_polymer_polymer_bonds: bool,
+    remove_bad_bonds: bool,
+    remove_nonsymmetric_bonds: bool,
+) -> tuple[structure.Structure, dict[str, Any]]:
+  """Cleans structure.
+
+  Args:
+    struc: Structure to clean.
+    ccd: The chemical components dictionary.
+    drop_missing_sequence: Whether to drop chains without specified sequences.
+    filter_clashes: Whether to drop clashing chains.
+    drop_non_standard_atoms: Whether to drop non CCD standard atoms.
+    filter_crystal_aids: Whether to drop ligands in the crystal aid set.
+    filter_waters: Whether to drop water chains.
+    filter_hydrogens: Whether to drop hyrdogen atoms.
+    filter_leaving_atoms: Whether to drop leaving atoms based on heuristics.
+    only_glycan_ligands_for_leaving_atoms: Whether to only include glycan
+      ligands when filtering leaving atoms.
+    covalent_bonds_only: Only include covalent bonds.
+    remove_polymer_polymer_bonds: Remove polymer-polymer bonds.
+    remove_bad_bonds: Whether to remove badly bonded ligands.
+    remove_nonsymmetric_bonds: Whether to remove nonsymmetric polymer-ligand
+      bonds from symmetric polymer chains.
+
+  Returns:
+    Tuple of structure and metadata dict. The metadata dict has
+    information about what was cleaned from the original.
+  """
+
+  metadata = {}
+  # Crop crystallization aids.
+  if (
+      filter_crystal_aids
+      and struc.structure_method in mmcif_names.CRYSTALLIZATION_METHODS
+  ):
+    struc = struc.filter_out(
+        res_name=chemical_component_sets.COMMON_CRYSTALLIZATION_AIDS
+    )
+
+  # Drop chains without specified sequences.
+  if drop_missing_sequence:
+    chains_with_unk_sequence = struc.find_chains_with_unknown_sequence()
+    num_with_unk_sequence = len(chains_with_unk_sequence)
+    if chains_with_unk_sequence:
+      struc = struc.filter_out(chain_id=chains_with_unk_sequence)
+  else:
+    num_with_unk_sequence = 0
+  metadata['num_with_unk_sequence'] = num_with_unk_sequence
+
+  # Remove intersecting chains.
+  if filter_clashes and struc.num_chains > 1:
+    clashing_chains = sterics.find_clashing_chains(struc)
+    if clashing_chains:
+      struc = struc.filter_out(chain_id=clashing_chains)
+  else:
+    clashing_chains = []
+  metadata['num_clashing_chains_removed'] = len(clashing_chains)
+  metadata['chains_removed'] = clashing_chains
+
+  # Drop non-standard atoms
+  if drop_non_standard_atoms:
+    struc = struc.drop_non_standard_atoms(
+        ccd=ccd, drop_unk=False, drop_non_ccd=False
+    )
+
+  # Sort chains in "reverse-spreadsheet" order.
+  struc = struc.with_sorted_chains
+
+  if filter_hydrogens:
+    struc = struc.without_hydrogen()
+
+  if filter_waters:
+    struc = struc.filter_out(chain_type=mmcif_names.WATER)
+
+  if filter_leaving_atoms:
+    drop_leaving_atoms_all = struc.chain_id != struc.chain_id
+    polymer_ligand_bonds = inter_chain_bonds.get_polymer_ligand_bonds(
+        struc,
+        only_glycan_ligands=only_glycan_ligands_for_leaving_atoms,
+    )
+    ligand_ligand_bonds = inter_chain_bonds.get_ligand_ligand_bonds(
+        struc,
+        only_glycan_ligands=only_glycan_ligands_for_leaving_atoms,
+    )
+    all_glycans = {
+        *chemical_component_sets.GLYCAN_OTHER_LIGANDS,
+        *chemical_component_sets.GLYCAN_LINKING_LIGANDS,
+    }
+    # If only glycan ligands and no O1 atoms, we can do parallel drop.
+    if (
+        only_glycan_ligands_for_leaving_atoms
+        and (not (ligand_ligand_bonds.atom_name == 'O1').any())
+        and (not (polymer_ligand_bonds.atom_name == 'O1').any())
+    ):
+      drop_leaving_atoms_all = np.logical_and(
+          np.isin(struc.atom_name, 'O1'),
+          np.isin(struc.res_name, list(all_glycans)),
+      )
+    else:
+      substruct = struc.group_by_residue
+      glycan_mask = np.isin(substruct.res_name, list(all_glycans))
+      substruct = substruct.filter(glycan_mask)
+      # We need to iterate over all glycan residues for this.
+      for res in substruct.iter_residues():
+        # Only need to do drop leaving atoms for glycans depending on bonds.
+        if (res_name := res['res_name']) in all_glycans:
+          drop_atom_filter = _get_leaving_atom_mask(
+              struc=struc,
+              polymer_ligand_bonds=polymer_ligand_bonds,
+              ligand_ligand_bonds=ligand_ligand_bonds,
+              chain_id=res['chain_id'],
+              chain_type=res['chain_type'],
+              res_id=res['res_id'],
+              res_name=res_name,
+          )
+          drop_leaving_atoms_all = np.logical_or(
+              drop_leaving_atoms_all, drop_atom_filter
+          )
+
+    num_atoms_before = struc.num_atoms
+    struc = struc.filter_out(drop_leaving_atoms_all)
+    num_atoms_after = struc.num_atoms
+
+    if num_atoms_before > num_atoms_after:
+      logging.error(
+          'Dropped %s atoms from GT struc: chain_id %s res_id %s res_name %s',
+          num_atoms_before - num_atoms_after,
+          struc.chain_id,
+          struc.res_id,
+          struc.res_name,
+      )
+
+  # Can filter by bond type without having to iterate over bonds.
+  if struc.bonds and covalent_bonds_only:
+    is_covalent = np.isin(struc.bonds.type, ['covale'])
+    if sum(is_covalent) > 0:
+      new_bonds = struc.bonds[is_covalent]
+    else:
+      new_bonds = structure.Bonds.make_empty()
+    struc = struc.copy_and_update(bonds=new_bonds)
+
+  # Other bond filters require iterating over individual bonds.
+  if struc.bonds and (remove_bad_bonds or remove_polymer_polymer_bonds):
+    include_bond = []
+    num_pp_bonds = 0
+    num_bad_bonds = 0
+    for bond in struc.iter_bonds():
+      dest_atom = bond.dest_atom
+      from_atom = bond.from_atom
+      if remove_polymer_polymer_bonds:
+        if (
+            from_atom['chain_type'] in mmcif_names.POLYMER_CHAIN_TYPES
+            and dest_atom['chain_type'] in mmcif_names.POLYMER_CHAIN_TYPES
+        ):
+          num_pp_bonds += 1
+          include_bond.append(False)
+          continue
+      if remove_bad_bonds:
+        dest_coords = np.array(
+            [dest_atom['atom_x'], dest_atom['atom_y'], dest_atom['atom_z']]
+        )
+        from_coords = np.array(
+            [from_atom['atom_x'], from_atom['atom_y'], from_atom['atom_z']]
+        )
+        squared_dist = np.sum(np.square(dest_coords - from_coords))
+        squared_threshold = 2.4 * 2.4
+        if squared_dist > squared_threshold:
+          num_bad_bonds += 1
+          include_bond.append(False)
+          continue
+      include_bond.append(True)
+    if sum(include_bond) < len(struc.bonds):
+      logging.info(
+          'Reducing number of bonds for %s from %s to %s, of which %s are'
+          ' polymer-polymer bonds and %s are bad bonds.',
+          struc.name,
+          len(struc.bonds),
+          sum(include_bond),
+          num_pp_bonds,
+          num_bad_bonds,
+      )
+      if sum(include_bond) > 0:
+        # Need to index bonds with bond keys or arrays of bools with same length
+        # as num bonds. In this case, we use array of bools (as elsewhere in the
+        # cleaning code).
+        new_bonds = struc.bonds[np.array(include_bond, dtype=bool)]
+      else:
+        new_bonds = structure.Bonds.make_empty()
+      struc = struc.copy_and_update(bonds=new_bonds)
+
+  if struc.bonds and remove_nonsymmetric_bonds:
+    # Check for asymmetric polymer-ligand bonds and remove if these exist.
+    polymer_ligand_bonds = inter_chain_bonds.get_polymer_ligand_bonds(
+        struc,
+        only_glycan_ligands=False,
+    )
+    if polymer_ligand_bonds:
+      if covalent_bond_cleaning.has_nonsymmetric_bonds_on_symmetric_polymer_chains(
+          struc, polymer_ligand_bonds
+      ):
+        from_atom_idxs, dest_atom_idxs = struc.bonds.get_atom_indices(
+            struc.atom_key
+        )
+        poly_chain_types = list(mmcif_names.POLYMER_CHAIN_TYPES)
+        is_polymer_bond = np.logical_or(
+            np.isin(struc.chain_type[from_atom_idxs], poly_chain_types),
+            np.isin(struc.chain_type[dest_atom_idxs], poly_chain_types),
+        )
+        struc = struc.copy_and_update(bonds=struc.bonds[~is_polymer_bond])
+
+  return struc, metadata
+
+
+def create_empty_output_struc_and_layout(
+    struc: structure.Structure,
+    ccd: chemical_components.Ccd,
+    *,
+    with_hydrogens: bool = False,
+    skip_unk: bool = False,
+    polymer_ligand_bonds: atom_layout.AtomLayout | None = None,
+    ligand_ligand_bonds: atom_layout.AtomLayout | None = None,
+    drop_ligand_leaving_atoms: bool = False,
+) -> tuple[structure.Structure, atom_layout.AtomLayout]:
+  """Make zero-coordinate structure from all physical residues.
+
+  Args:
+    struc: Structure object.
+    ccd: The chemical components dictionary.
+    with_hydrogens: Whether to keep hydrogen atoms in structure.
+    skip_unk: Whether to remove unknown residues from structure.
+    polymer_ligand_bonds: Bond information for polymer-ligand pairs.
+    ligand_ligand_bonds: Bond information for ligand-ligand pairs.
+    drop_ligand_leaving_atoms: Flag for handling leaving atoms for ligands.
+
+  Returns:
+    Tuple of structure with all bonds, physical residues and coordinates set to
+    0 and a flat atom layout of empty structure.
+  """
+  bonded_atom_pairs = []
+  if polymer_ligand_bonds:
+    for chain_ids, res_ids, atom_names in zip(
+        polymer_ligand_bonds.chain_id,
+        polymer_ligand_bonds.res_id,
+        polymer_ligand_bonds.atom_name,
+        strict=True,
+    ):
+      bonded_atom_pairs.append((
+          (chain_ids[0], res_ids[0], atom_names[0]),
+          (chain_ids[1], res_ids[1], atom_names[1]),
+      ))
+  if ligand_ligand_bonds:
+    for chain_ids, res_ids, atom_names in zip(
+        ligand_ligand_bonds.chain_id,
+        ligand_ligand_bonds.res_id,
+        ligand_ligand_bonds.atom_name,
+        strict=True,
+    ):
+      bonded_atom_pairs.append((
+          (chain_ids[0], res_ids[0], atom_names[0]),
+          (chain_ids[1], res_ids[1], atom_names[1]),
+      ))
+  residues = atom_layout.residues_from_structure(
+      struc, include_missing_residues=True
+  )
+
+  flat_output_layout = atom_layout.make_flat_atom_layout(
+      residues,
+      ccd=ccd,
+      with_hydrogens=with_hydrogens,
+      skip_unk_residues=skip_unk,
+      polymer_ligand_bonds=polymer_ligand_bonds,
+      ligand_ligand_bonds=ligand_ligand_bonds,
+      drop_ligand_leaving_atoms=drop_ligand_leaving_atoms,
+  )
+
+  empty_output_struc = atom_layout.make_structure(
+      flat_layout=flat_output_layout,
+      atom_coords=np.zeros((flat_output_layout.shape[0], 3)),
+      name=struc.name,
+      atom_b_factors=None,
+      all_physical_residues=residues,
+  )
+  if bonded_atom_pairs:
+    empty_output_struc = empty_output_struc.add_bonds(
+        bonded_atom_pairs, bond_type=mmcif_names.COVALENT_BOND
+    )
+
+  return empty_output_struc, flat_output_layout

flax_model/alphafold3/model/protein_data_processing.py

@@ -0,0 +1,118 @@
+
+
+"""Process Structure Data."""
+
+from flax_model.alphafold3.constants import atom_types
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.constants import side_chains
+import numpy as np
+
+
+NUM_DENSE = atom_types.DENSE_ATOM_NUM
+NUM_AA = len(residue_names.PROTEIN_TYPES)
+NUM_AA_WITH_UNK_AND_GAP = len(
+    residue_names.PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP
+)
+NUM_RESTYPES_WITH_UNK_AND_GAP = (
+    residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP
+)
+
+
+def _make_restype_rigidgroup_dense_atom_idx():
+  """Create Mapping from rigid_groups to dense_atom indices."""
+  # Create an array with the atom names.
+  # shape (num_restypes, num_rigidgroups, 3_atoms):
+  # (31, 8, 3)
+  base_atom_indices = np.zeros(
+      (NUM_RESTYPES_WITH_UNK_AND_GAP, 8, 3), dtype=np.int32
+  )
+
+  # 4,5,6,7: 'chi1,2,3,4-group'
+  for restype, restype_letter in enumerate(
+      residue_names.PROTEIN_TYPES_ONE_LETTER
+  ):
+    resname = residue_names.PROTEIN_COMMON_ONE_TO_THREE[restype_letter]
+
+    dense_atom_names = atom_types.ATOM14[resname]
+    # 0: backbone frame
+    base_atom_indices[restype, 0, :] = [
+        dense_atom_names.index(atom) for atom in ['C', 'CA', 'N']
+    ]
+
+    # 3: 'psi-group'
+    base_atom_indices[restype, 3, :] = [
+        dense_atom_names.index(atom) for atom in ['CA', 'C', 'O']
+    ]
+    for chi_idx in range(4):
+      if side_chains.CHI_ANGLES_MASK[restype][chi_idx]:
+        atom_names = side_chains.CHI_ANGLES_ATOMS[resname][chi_idx]
+        base_atom_indices[restype, chi_idx + 4, :] = [
+            dense_atom_names.index(atom) for atom in atom_names[1:]
+        ]
+  dense_atom_names = atom_types.DENSE_ATOM['A']
+  nucleic_rigid_atoms = [
+      dense_atom_names.index(atom) for atom in ["C1'", "C3'", "C4'"]
+  ]
+  for nanum, _ in enumerate(residue_names.NUCLEIC_TYPES):
+    # 0: backbone frame only.
+    # we have aa + unk + gap, so we want to start after those
+    resnum = nanum + NUM_AA_WITH_UNK_AND_GAP
+    base_atom_indices[resnum, 0, :] = nucleic_rigid_atoms
+
+  return base_atom_indices
+
+
+RESTYPE_RIGIDGROUP_DENSE_ATOM_IDX = _make_restype_rigidgroup_dense_atom_idx()
+
+
+def _make_restype_pseudobeta_idx():
+  """Returns indices of residue's pseudo-beta."""
+  restype_pseudobeta_index = np.zeros(
+      (NUM_RESTYPES_WITH_UNK_AND_GAP,), dtype=np.int32
+  )
+  for restype, restype_letter in enumerate(
+      residue_names.PROTEIN_TYPES_ONE_LETTER
+  ):
+    restype_name = residue_names.PROTEIN_COMMON_ONE_TO_THREE[restype_letter]
+    atom_names = list(atom_types.ATOM14[restype_name])
+    if restype_name in {'GLY'}:
+      restype_pseudobeta_index[restype] = atom_names.index('CA')
+    else:
+      restype_pseudobeta_index[restype] = atom_names.index('CB')
+  for nanum, resname in enumerate(residue_names.NUCLEIC_TYPES):
+    atom_names = list(atom_types.DENSE_ATOM[resname])
+    # 0: backbone frame only.
+    # we have aa + unk , so we want to start after those
+    restype = nanum + NUM_AA_WITH_UNK_AND_GAP
+    if resname in {'A', 'G', 'DA', 'DG'}:
+      restype_pseudobeta_index[restype] = atom_names.index('C4')
+    else:
+      restype_pseudobeta_index[restype] = atom_names.index('C2')
+  return restype_pseudobeta_index
+
+
+RESTYPE_PSEUDOBETA_INDEX = _make_restype_pseudobeta_idx()
+
+
+def _make_aatype_dense_atom_to_atom37():
+  """Map from dense_atom to atom37 per residue type."""
+  restype_dense_atom_to_atom37 = []  # mapping (restype, dense_atom) --> atom37
+  for rt in residue_names.PROTEIN_TYPES_ONE_LETTER:
+    atom_names = list(
+        atom_types.ATOM14_PADDED[residue_names.PROTEIN_COMMON_ONE_TO_THREE[rt]]
+    )
+    atom_names.extend([''] * (NUM_DENSE - len(atom_names)))
+    restype_dense_atom_to_atom37.append(
+        [(atom_types.ATOM37_ORDER[name] if name else 0) for name in atom_names]
+    )
+  # Add dummy mapping for restype 'UNK', '-' (gap), and nucleics [but not DN].
+  for _ in range(2 + len(residue_names.NUCLEIC_TYPES_WITH_UNKNOWN)):
+    restype_dense_atom_to_atom37.append([0] * NUM_DENSE)
+
+  restype_dense_atom_to_atom37 = np.array(
+      restype_dense_atom_to_atom37, dtype=np.int32
+  )
+  return restype_dense_atom_to_atom37
+
+
+PROTEIN_AATYPE_DENSE_ATOM_TO_ATOM37 = _make_aatype_dense_atom_to_atom37()