File size: 765 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
"""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'
|