diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 28912b6..f62f829 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -1,7 +1,6 @@ name: Continuous Integration on: - # Trigger the workflow on push or pull request, only on the main branch. push: branches: - main @@ -18,21 +17,28 @@ jobs: matrix: include: - name-prefix: "all tests" - python-version: '3.11' + python-version: '3.12' os: ubuntu-latest + steps: - - uses: actions/checkout@v4 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - name: Set up uv + uses: astral-sh/setup-uv@v7 with: - python-version: ${{ matrix.python-version }} + enable-cache: true + cache-dependency-glob: "uv.lock" + + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + - name: Install dependencies - run: | - sudo apt-get install -y hmmer - pip install --upgrade pip - pip install -r dev-requirements.txt - pip install --no-deps . - build_data + run: sudo apt-get install -y hmmer + + - name: Install Python dependencies + run: uv sync --frozen --all-groups + + - name: Build data + run: uv run build_data + - name: Run CPU-only tests - run: | - python run_alphafold_data_test.py + run: uv run python run_alphafold_data_test.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 8116272..6d07902 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -9,6 +18,11 @@ # https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md cmake_minimum_required(VERSION 3.28) + +# This forces Git to use the 'files' backend for all FetchContent operations. +# This fixes libcifpp and dssp incompatibility with newer git versions. +set(ENV{GIT_CONFIG_PARAMETERS} "'init.defaultRefFormat=files'") + project( "${SKBUILD_PROJECT_NAME}" LANGUAGES CXX diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..ccc2b8d --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,36 @@ +# How to Contribute + +We welcome small patches related to bug fixes and documentation, but we do not +plan to make any major changes to this repository. + +## AI Generated Code + +We welcome the use of AI tools for the generation of code, documentation and/or +Pull Request (PR) description as long as: + +1. It has been transparently labelled as such. Make sure to declare it in the + PR message. +2. You have manually reviewed the code before sending the PR. +3. The change has been manually tested. We might ask you to fold a certain + input to check correctness of the PR. + +Please do not submit AI generated PRs where test results have been hallucinated. + +## Contributor License Agreement + +Contributions to this project must be accompanied by a Contributor License +Agreement. You (or your employer) retain the copyright to your contribution, +this simply gives us permission to use and redistribute your contributions as +part of the project. Head over to to see +your current agreements on file or to sign a new one. + +You generally only need to submit a CLA once, so if you've already submitted one +(even if it was for a different project), you probably don't need to do it +again. + +## Code reviews + +All submissions, including submissions by project members, require review. We +use GitHub pull requests for this purpose. Consult +[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more +information on using pull requests. diff --git a/LICENSE b/LICENSE index cbe5ad1..d645695 100644 --- a/LICENSE +++ b/LICENSE @@ -1,437 +1,202 @@ -Attribution-NonCommercial-ShareAlike 4.0 International - -======================================================================= - -Creative Commons Corporation ("Creative Commons") is not a law firm and -does not provide legal services or legal advice. Distribution of -Creative Commons public licenses does not create a lawyer-client or -other relationship. Creative Commons makes its licenses and related -information available on an "as-is" basis. Creative Commons gives no -warranties regarding its licenses, any material licensed under their -terms and conditions, or any related information. Creative Commons -disclaims all liability for damages resulting from their use to the -fullest extent possible. - -Using Creative Commons Public Licenses - -Creative Commons public licenses provide a standard set of terms and -conditions that creators and other rights holders may use to share -original works of authorship and other material subject to copyright -and certain other rights specified in the public license below. The -following considerations are for informational purposes only, are not -exhaustive, and do not form part of our licenses. - - Considerations for licensors: Our public licenses are - intended for use by those authorized to give the public - permission to use material in ways otherwise restricted by - copyright and certain other rights. Our licenses are - irrevocable. Licensors should read and understand the terms - and conditions of the license they choose before applying it. - Licensors should also secure all rights necessary before - applying our licenses so that the public can reuse the - material as expected. Licensors should clearly mark any - material not subject to the license. This includes other CC- - licensed material, or material used under an exception or - limitation to copyright. More considerations for licensors: - wiki.creativecommons.org/Considerations_for_licensors - - Considerations for the public: By using one of our public - licenses, a licensor grants the public permission to use the - licensed material under specified terms and conditions. If - the licensor's permission is not necessary for any reason--for - example, because of any applicable exception or limitation to - copyright--then that use is not regulated by the license. Our - licenses grant only permissions under copyright and certain - other rights that a licensor has authority to grant. Use of - the licensed material may still be restricted for other - reasons, including because others have copyright or other - rights in the material. A licensor may make special requests, - such as asking that all changes be marked or described. - Although not required by our licenses, you are encouraged to - respect those requests where reasonable. More considerations - for the public: - wiki.creativecommons.org/Considerations_for_licensees - -======================================================================= - -Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International -Public License - -By exercising the Licensed Rights (defined below), You accept and agree -to be bound by the terms and conditions of this Creative Commons -Attribution-NonCommercial-ShareAlike 4.0 International Public License -("Public License"). To the extent this Public License may be -interpreted as a contract, You are granted the Licensed Rights in -consideration of Your acceptance of these terms and conditions, and the -Licensor grants You such rights in consideration of benefits the -Licensor receives from making the Licensed Material available under -these terms and conditions. - - -Section 1 -- Definitions. - - a. Adapted Material means material subject to Copyright and Similar - Rights that is derived from or based upon the Licensed Material - and in which the Licensed Material is translated, altered, - arranged, transformed, or otherwise modified in a manner requiring - permission under the Copyright and Similar Rights held by the - Licensor. For purposes of this Public License, where the Licensed - Material is a musical work, performance, or sound recording, - Adapted Material is always produced where the Licensed Material is - synched in timed relation with a moving image. - - b. Adapter's License means the license You apply to Your Copyright - and Similar Rights in Your contributions to Adapted Material in - accordance with the terms and conditions of this Public License. - - c. BY-NC-SA Compatible License means a license listed at - creativecommons.org/compatiblelicenses, approved by Creative - Commons as essentially the equivalent of this Public License. - - d. Copyright and Similar Rights means copyright and/or similar rights - closely related to copyright including, without limitation, - performance, broadcast, sound recording, and Sui Generis Database - Rights, without regard to how the rights are labeled or - categorized. For purposes of this Public License, the rights - specified in Section 2(b)(1)-(2) are not Copyright and Similar - Rights. - - e. Effective Technological Measures means those measures that, in the - absence of proper authority, may not be circumvented under laws - fulfilling obligations under Article 11 of the WIPO Copyright - Treaty adopted on December 20, 1996, and/or similar international - agreements. - - f. Exceptions and Limitations means fair use, fair dealing, and/or - any other exception or limitation to Copyright and Similar Rights - that applies to Your use of the Licensed Material. - - g. License Elements means the license attributes listed in the name - of a Creative Commons Public License. The License Elements of this - Public License are Attribution, NonCommercial, and ShareAlike. - - h. Licensed Material means the artistic or literary work, database, - or other material to which the Licensor applied this Public - License. - - i. Licensed Rights means the rights granted to You subject to the - terms and conditions of this Public License, which are limited to - all Copyright and Similar Rights that apply to Your use of the - Licensed Material and that the Licensor has authority to license. - - j. Licensor means the individual(s) or entity(ies) granting rights - under this Public License. - - k. NonCommercial means not primarily intended for or directed towards - commercial advantage or monetary compensation. For purposes of - this Public License, the exchange of the Licensed Material for - other material subject to Copyright and Similar Rights by digital - file-sharing or similar means is NonCommercial provided there is - no payment of monetary compensation in connection with the - exchange. - - l. Share means to provide material to the public by any means or - process that requires permission under the Licensed Rights, such - as reproduction, public display, public performance, distribution, - dissemination, communication, or importation, and to make material - available to the public including in ways that members of the - public may access the material from a place and at a time - individually chosen by them. - - m. Sui Generis Database Rights means rights other than copyright - resulting from Directive 96/9/EC of the European Parliament and of - the Council of 11 March 1996 on the legal protection of databases, - as amended and/or succeeded, as well as other essentially - equivalent rights anywhere in the world. - - n. You means the individual or entity exercising the Licensed Rights - under this Public License. Your has a corresponding meaning. - - -Section 2 -- Scope. - - a. License grant. - - 1. Subject to the terms and conditions of this Public License, - the Licensor hereby grants You a worldwide, royalty-free, - non-sublicensable, non-exclusive, irrevocable license to - exercise the Licensed Rights in the Licensed Material to: - - a. reproduce and Share the Licensed Material, in whole or - in part, for NonCommercial purposes only; and - - b. produce, reproduce, and Share Adapted Material for - NonCommercial purposes only. - - 2. Exceptions and Limitations. For the avoidance of doubt, where - Exceptions and Limitations apply to Your use, this Public - License does not apply, and You do not need to comply with - its terms and conditions. - - 3. Term. The term of this Public License is specified in Section - 6(a). - - 4. Media and formats; technical modifications allowed. The - Licensor authorizes You to exercise the Licensed Rights in - all media and formats whether now known or hereafter created, - and to make technical modifications necessary to do so. The - Licensor waives and/or agrees not to assert any right or - authority to forbid You from making technical modifications - necessary to exercise the Licensed Rights, including - technical modifications necessary to circumvent Effective - Technological Measures. For purposes of this Public License, - simply making modifications authorized by this Section 2(a) - (4) never produces Adapted Material. - - 5. Downstream recipients. - - a. Offer from the Licensor -- Licensed Material. Every - recipient of the Licensed Material automatically - receives an offer from the Licensor to exercise the - Licensed Rights under the terms and conditions of this - Public License. - - b. Additional offer from the Licensor -- Adapted Material. - Every recipient of Adapted Material from You - automatically receives an offer from the Licensor to - exercise the Licensed Rights in the Adapted Material - under the conditions of the Adapter's License You apply. - - c. No downstream restrictions. You may not offer or impose - any additional or different terms or conditions on, or - apply any Effective Technological Measures to, the - Licensed Material if doing so restricts exercise of the - Licensed Rights by any recipient of the Licensed - Material. - - 6. No endorsement. Nothing in this Public License constitutes or - may be construed as permission to assert or imply that You - are, or that Your use of the Licensed Material is, connected - with, or sponsored, endorsed, or granted official status by, - the Licensor or others designated to receive attribution as - provided in Section 3(a)(1)(A)(i). - - b. Other rights. - - 1. Moral rights, such as the right of integrity, are not - licensed under this Public License, nor are publicity, - privacy, and/or other similar personality rights; however, to - the extent possible, the Licensor waives and/or agrees not to - assert any such rights held by the Licensor to the limited - extent necessary to allow You to exercise the Licensed - Rights, but not otherwise. - - 2. Patent and trademark rights are not licensed under this - Public License. - - 3. To the extent possible, the Licensor waives any right to - collect royalties from You for the exercise of the Licensed - Rights, whether directly or through a collecting society - under any voluntary or waivable statutory or compulsory - licensing scheme. In all other cases the Licensor expressly - reserves any right to collect such royalties, including when - the Licensed Material is used other than for NonCommercial - purposes. - - -Section 3 -- License Conditions. - -Your exercise of the Licensed Rights is expressly made subject to the -following conditions. - - a. Attribution. - - 1. If You Share the Licensed Material (including in modified - form), You must: - - a. retain the following if it is supplied by the Licensor - with the Licensed Material: - - i. identification of the creator(s) of the Licensed - Material and any others designated to receive - attribution, in any reasonable manner requested by - the Licensor (including by pseudonym if - designated); - - ii. a copyright notice; - - iii. a notice that refers to this Public License; - - iv. a notice that refers to the disclaimer of - warranties; - - v. a URI or hyperlink to the Licensed Material to the - extent reasonably practicable; - - b. indicate if You modified the Licensed Material and - retain an indication of any previous modifications; and - - c. indicate the Licensed Material is licensed under this - Public License, and include the text of, or the URI or - hyperlink to, this Public License. - - 2. You may satisfy the conditions in Section 3(a)(1) in any - reasonable manner based on the medium, means, and context in - which You Share the Licensed Material. For example, it may be - reasonable to satisfy the conditions by providing a URI or - hyperlink to a resource that includes the required - information. - 3. If requested by the Licensor, You must remove any of the - information required by Section 3(a)(1)(A) to the extent - reasonably practicable. - - b. ShareAlike. - In addition to the conditions in Section 3(a), if You Share - Adapted Material You produce, the following conditions also apply. - - 1. The Adapter's License You apply must be a Creative Commons - license with the same License Elements, this version or - later, or a BY-NC-SA Compatible License. - - 2. You must include the text of, or the URI or hyperlink to, the - Adapter's License You apply. You may satisfy this condition - in any reasonable manner based on the medium, means, and - context in which You Share Adapted Material. - - 3. You may not offer or impose any additional or different terms - or conditions on, or apply any Effective Technological - Measures to, Adapted Material that restrict exercise of the - rights granted under the Adapter's License You apply. - - -Section 4 -- Sui Generis Database Rights. - -Where the Licensed Rights include Sui Generis Database Rights that -apply to Your use of the Licensed Material: - - a. for the avoidance of doubt, Section 2(a)(1) grants You the right - to extract, reuse, reproduce, and Share all or a substantial - portion of the contents of the database for NonCommercial purposes - only; - - b. if You include all or a substantial portion of the database - contents in a database in which You have Sui Generis Database - Rights, then the database in which You have Sui Generis Database - Rights (but not its individual contents) is Adapted Material, - including for purposes of Section 3(b); and - - c. You must comply with the conditions in Section 3(a) if You Share - all or a substantial portion of the contents of the database. - -For the avoidance of doubt, this Section 4 supplements and does not -replace Your obligations under this Public License where the Licensed -Rights include other Copyright and Similar Rights. - - -Section 5 -- Disclaimer of Warranties and Limitation of Liability. - - a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE - EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS - AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF - ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, - IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, - WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR - PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, - ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT - KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT - ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. - - b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE - TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, - NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, - INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, - COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR - USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN - ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR - DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR - IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. - - c. The disclaimer of warranties and limitation of liability provided - above shall be interpreted in a manner that, to the extent - possible, most closely approximates an absolute disclaimer and - waiver of all liability. - - -Section 6 -- Term and Termination. - - a. This Public License applies for the term of the Copyright and - Similar Rights licensed here. However, if You fail to comply with - this Public License, then Your rights under this Public License - terminate automatically. - - b. Where Your right to use the Licensed Material has terminated under - Section 6(a), it reinstates: - - 1. automatically as of the date the violation is cured, provided - it is cured within 30 days of Your discovery of the - violation; or - - 2. upon express reinstatement by the Licensor. - - For the avoidance of doubt, this Section 6(b) does not affect any - right the Licensor may have to seek remedies for Your violations - of this Public License. - - c. For the avoidance of doubt, the Licensor may also offer the - Licensed Material under separate terms or conditions or stop - distributing the Licensed Material at any time; however, doing so - will not terminate this Public License. - - d. Sections 1, 5, 6, 7, and 8 survive termination of this Public - License. - - -Section 7 -- Other Terms and Conditions. - - a. The Licensor shall not be bound by any additional or different - terms or conditions communicated by You unless expressly agreed. - - b. Any arrangements, understandings, or agreements regarding the - Licensed Material not stated herein are separate from and - independent of the terms and conditions of this Public License. - - -Section 8 -- Interpretation. - - a. For the avoidance of doubt, this Public License does not, and - shall not be interpreted to, reduce, limit, restrict, or impose - conditions on any use of the Licensed Material that could lawfully - be made without permission under this Public License. - - b. To the extent possible, if any provision of this Public License is - deemed unenforceable, it shall be automatically reformed to the - minimum extent necessary to make it enforceable. If the provision - cannot be reformed, it shall be severed from this Public License - without affecting the enforceability of the remaining terms and - conditions. - - c. No term or condition of this Public License will be waived and no - failure to comply consented to unless expressly agreed to by the - Licensor. - - d. Nothing in this Public License constitutes or may be interpreted - as a limitation upon, or waiver of, any privileges and immunities - that apply to the Licensor or You, including from the legal - processes of any jurisdiction or authority. - -======================================================================= - -Creative Commons is not a party to its public -licenses. Notwithstanding, Creative Commons may elect to apply one of -its public licenses to material it publishes and in those instances -will be considered the “Licensor.” The text of the Creative Commons -public licenses is dedicated to the public domain under the CC0 Public -Domain Dedication. Except for the limited purpose of indicating that -material is shared under a Creative Commons public license or as -otherwise permitted by the Creative Commons policies published at -creativecommons.org/policies, Creative Commons does not authorize the -use of the trademark "Creative Commons" or any other trademark or logo -of Creative Commons without its prior written consent including, -without limitation, in connection with any unauthorized modifications -to any of its public licenses or any other arrangements, -understandings, or agreements concerning use of licensed material. For -the avoidance of doubt, this paragraph does not form part of the -public licenses. - -Creative Commons may be contacted at creativecommons.org. + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md index af7b8d8..7d455d0 100644 --- a/README.md +++ b/README.md @@ -3,9 +3,10 @@ # AlphaFold 3 This package provides an implementation of the inference pipeline of AlphaFold -3. See below for how to access the model parameters. You may only use AlphaFold -3 model parameters if received directly from Google. Use is subject to these -[terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md). + +1. See below for how to access the model parameters. You may only use AlphaFold + 3 model parameters if received directly from Google. Use is subject to these + [terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md). Any publication that discloses findings arising from using this source code, the model parameters or outputs produced by those should [cite](#citing-this-work) @@ -145,15 +146,12 @@ AlphaFold 3 uses the following separate libraries and packages: * [HMMER Suite](https://github.com/EddyRivasLab/hmmer) * [Haiku](https://github.com/deepmind/dm-haiku) * [JAX](https://github.com/jax-ml/jax/) -* [jax-triton](https://github.com/jax-ml/jax-triton) -* [jaxtyping](https://github.com/patrick-kidger/jaxtyping) * [libcifpp](https://github.com/pdb-redo/libcifpp) * [NumPy](https://github.com/numpy/numpy) * [pybind11](https://github.com/pybind/pybind11) and [pybind11_abseil](https://github.com/pybind/pybind11_abseil) * [RDKit](https://github.com/rdkit/rdkit) -* [Tree](https://github.com/deepmind/tree) -* [Triton](https://github.com/triton-lang/triton) +* [Tokamax](https://github.com/openxla/tokamax) * [tqdm](https://github.com/tqdm/tqdm) We thank all their contributors and maintainers! @@ -175,11 +173,10 @@ Copyright 2024 DeepMind Technologies Limited. ### AlphaFold 3 Source Code and Model Parameters -The AlphaFold 3 source code is licensed under the Creative Commons -Attribution-Non-Commercial ShareAlike International License, Version 4.0 -(CC-BY-NC-SA 4.0) (the "License"); you may not use this file except in -compliance with the License. You may obtain a copy of the License at -[https://github.com/google-deepmind/alphafold3/blob/main/LICENSE](https://github.com/google-deepmind/alphafold3/blob/main/LICENSE). +AlphaFold 3 source code is licensed under the Apache License, Version 2.0 (the +"License"); you may not use its source code except in compliance with the +License. You may obtain a copy of the License at +http://www.apache.org/licenses/LICENSE-2.0. The AlphaFold 3 model parameters are made available under the [AlphaFold 3 Model Parameters Terms of Use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) diff --git a/WEIGHTS_TERMS_OF_USE.md b/WEIGHTS_TERMS_OF_USE.md index cfcc596..9f7db84 100644 --- a/WEIGHTS_TERMS_OF_USE.md +++ b/WEIGHTS_TERMS_OF_USE.md @@ -62,9 +62,8 @@ As used in these Terms: "**AlphaFold 3**" means: (a) the AlphaFold 3 source code made available [here](https://github.com/google-deepmind/alphafold3/) and licensed under the -terms of the Creative Commons Attribution-NonCommercial-Sharealike 4.0 -International (CC-BY-NC-SA 4.0) license and any derivative source code, and (b) -Model Parameters. +terms of the Apache 2.0 license and any derivative source code, and (b) Model +Parameters. "**AlphaFold 3 Assets**" means the Model Parameters and Output. diff --git a/dev-requirements.txt b/dev-requirements.txt deleted file mode 100644 index b9db0e4..0000000 --- a/dev-requirements.txt +++ /dev/null @@ -1,538 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.12 -# by the following command: -# -# pip-compile --extra=dev --generate-hashes --output-file=dev-requirements.txt pyproject.toml -# -absl-py==2.1.0 \ - --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \ - --hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff - # via - # alphafold3 (pyproject.toml) - # dm-haiku - # jax-triton -dm-haiku==0.0.13 \ - --hash=sha256:029bb91b5b1edb0d3fe23304d3bf12a545ea6e485041f7f5d8c8d85ebcf6e17d \ - --hash=sha256:ee9562c68a059f146ad07f555ca591cb8c11ef751afecc38353863562bd23f43 - # via alphafold3 (pyproject.toml) -dm-tree==0.1.8 \ - --hash=sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6 \ - --hash=sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760 \ - --hash=sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c \ - --hash=sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf \ - --hash=sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430 \ - --hash=sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de \ - --hash=sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317 \ - --hash=sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca \ - --hash=sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913 \ - --hash=sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf \ - --hash=sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef \ - --hash=sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426 \ - --hash=sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1 \ - --hash=sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e \ - --hash=sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60 \ - --hash=sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5 \ - --hash=sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f \ - --hash=sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b \ - --hash=sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410 \ - --hash=sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134 \ - --hash=sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb \ - --hash=sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b \ - --hash=sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7 \ - --hash=sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393 \ - --hash=sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571 \ - --hash=sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368 \ - --hash=sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22 \ - --hash=sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715 \ - --hash=sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80 \ - --hash=sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7 \ - --hash=sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d \ - --hash=sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a \ - --hash=sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d \ - --hash=sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6 \ - --hash=sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5 \ - --hash=sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e \ - --hash=sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68 \ - --hash=sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8 \ - --hash=sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f \ - --hash=sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436 \ - --hash=sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee \ - --hash=sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb \ - --hash=sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144 \ - --hash=sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8 \ - --hash=sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb \ - --hash=sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d - # via alphafold3 (pyproject.toml) -filelock==3.16.1 \ - --hash=sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0 \ - --hash=sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435 - # via triton -jax[cuda12]==0.4.34 \ - --hash=sha256:44196854f40c5f9cea3142824b9f1051f85afc3fcf7593ec5479fc8db01c58db \ - --hash=sha256:b957ca1fc91f7343f91a186af9f19c7f342c946f95a8c11c7f1e5cdfe2e58d9e - # via - # alphafold3 (pyproject.toml) - # jax-triton -jax-cuda12-pjrt==0.4.34 \ - --hash=sha256:0c7cc98f962cc7fc8e0a5ea6331b42a0cee516f202f1c3019f6aa5cd9530cca0 \ - --hash=sha256:6247459827550e7d11480f5707d7d2843cbf57bcda5d833565493456075cd143 - # via jax-cuda12-plugin -jax-cuda12-plugin[with-cuda]==0.4.34 \ - --hash=sha256:15fa8ff678fad34df424e01f1a9151abc966700670d3b22f52b87a2740028602 \ - --hash=sha256:1b037a408b9d9de59367c9513a9445b705da526657a0c4d593e84234312261e7 \ - --hash=sha256:b2099a4407225122ff76f6dcdc8dbdae47e6f29343bdfd21460ad337dc34a209 \ - --hash=sha256:d035ea72bd9b8a65a6ea621bca1affdd33127fa3a52e7bded7692670d360adab \ - --hash=sha256:dadb507949b70ba148567219bf0e7752eb63baec551f7ec860f8fcbb4ee48c48 \ - --hash=sha256:db988b7ba5063483a936ddbf162f04d1b4412e0d64340f11788c7bbc877e8a43 \ - --hash=sha256:e23721d1654b311b47cd6b35768520284bd036f8c7e6b11600143258b4a0409a \ - --hash=sha256:fb72bda44ae9f8a1e5ee9c5f82e2d273abe019874b0110b5c2d8395f1fe966b4 - # via jax -jax-triton==0.2.0 \ - --hash=sha256:6db5f6147327f2462c35becab612d415cf9ee70c96e43020c3c7ea873eb8576e \ - --hash=sha256:ed564a5ffb9e404557dc8d296e7eb30e501da5a4d7b03408acd0837c1c618c21 - # via alphafold3 (pyproject.toml) -jaxlib==0.4.34 \ - --hash=sha256:096f0ca309d41fa692a9d1f2f9baab1c5c8ca0749876ebb3f748e738a27c7ff4 \ - --hash=sha256:133070d4fec5525ffea4dc72956398c1cf647a04dcb37f8a935ee82af78d9965 \ - --hash=sha256:1a30771d85fa77f9ab8f18e63240f455ab3a3f87660ed7b8d5eea6ceecbe5c1e \ - --hash=sha256:3bcfa639ca3cfaf86c8ceebd5fc0d47300fd98a078014a1d0cc03133e1523d5f \ - --hash=sha256:3e60bc826933082e99b19b87c21818a8d26fcdb01f418d47cedff554746fd6cc \ - --hash=sha256:45d719a2ce0ebf21255a277b71d756f3609b7b5be70cddc5d88fd58c35219de0 \ - --hash=sha256:48272e9034ff868d4328cf0055a07882fd2be93f59dfb6283af7de491f9d1290 \ - --hash=sha256:571ef03259835458111596a71a2f4a6fabf4ec34595df4cea555035362ac5bf0 \ - --hash=sha256:6b43a974c5d91a19912d138f2658dd8dbb7d30dcdff5c961d896c673e872b611 \ - --hash=sha256:72e22e99a5dc890a64443c3fc12f13f20091f578c405a76de077ba42b4c62cd7 \ - --hash=sha256:7be673a876ebd1aef440fb7e3ebaf99a91abeb550c9728c644b7d7c7b5d7c108 \ - --hash=sha256:87f25a477cd279840e53718403f97092eba0e8a945fcab47bcf435b6f9119dda \ - --hash=sha256:8ee3f93836e53c86556ccd9449a4ea43516ee05184d031a71dd692e81259f7d9 \ - --hash=sha256:901cb4040ed24eae40071d8114ea8d10dff436277fa74a1a5b9e7206f641151c \ - --hash=sha256:b0001c8f0e2b1c7bc99e4f314b524a340d25653505c1a1484d4041a9d3617f6f \ - --hash=sha256:b7a212a3cb5c6acc201c32ae4f4b5f5a9ac09457fbb77ba8db5ce7e7d4adc214 \ - --hash=sha256:c303f5acaf6c56ce5ff133a923c9b6247bdebedde15bd2c893c24be4d8f71306 \ - --hash=sha256:c7b3e724a30426a856070aba0192b5d199e95b4411070e7ad96ad8b196877b10 \ - --hash=sha256:c9d3adcae43a33aad4332be9c2aedc5ef751d1e755f917a5afb30c7872eacaa8 \ - --hash=sha256:d840e64b85f8865404d6d225b9bb340e158df1457152a361b05680e24792b232 - # via jax -jaxtyping==0.2.34 \ - --hash=sha256:2f81fb6d1586e497a6ea2d28c06dcab37b108a096cbb36ea3fe4fa2e1c1f32e5 \ - --hash=sha256:eed9a3458ec8726c84ea5457ebde53c964f65d2c22c0ec40d0555ae3fed5bbaf - # via alphafold3 (pyproject.toml) -jmp==0.0.4 \ - --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \ - --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d - # via dm-haiku -ml-dtypes==0.5.0 \ - --hash=sha256:099e09edd54e676903b4538f3815b5ab96f5b119690514602d96bfdb67172cbe \ - --hash=sha256:2e7534392682c3098bc7341648c650864207169c654aed83143d7a19c67ae06f \ - --hash=sha256:3e7d3a380fe73a63c884f06136f8baa7a5249cc8e9fdec677997dd78549f8128 \ - --hash=sha256:54415257f00eb44fbcc807454efac3356f75644f1cbfc2d4e5522a72ae1dacab \ - --hash=sha256:5f2b59233a0dbb6a560b3137ed6125433289ccba2f8d9c3695a52423a369ed15 \ - --hash=sha256:60275f2b51b56834e840c4809fca840565f9bf8e9a73f6d8c94f5b5935701215 \ - --hash=sha256:76942f6aeb5c40766d5ea62386daa4148e6a54322aaf5b53eae9e7553240222f \ - --hash=sha256:7ee9c320bb0f9ffdf9f6fa6a696ef2e005d1f66438d6f1c1457338e00a02e8cf \ - --hash=sha256:8c32138975797e681eb175996d64356bcfa124bdbb6a70460b9768c2b35a6fa4 \ - --hash=sha256:968fede07d1f9b926a63df97d25ac656cac1a57ebd33701734eaf704bc55d8d8 \ - --hash=sha256:a03fc861b86cc586728e3d093ba37f0cc05e65330c3ebd7688e7bae8290f8859 \ - --hash=sha256:a38df8df61194aeaae1ab7579075779b4ad32cd1cffd012c28be227fa7f2a70a \ - --hash=sha256:a988bac6572630e1e9c2edd9b1277b4eefd1c86209e52b0d061b775ac33902ff \ - --hash=sha256:ab046f2ff789b1f11b2491909682c5d089934835f9a760fafc180e47dcb676b8 \ - --hash=sha256:afa08343069874a30812871d639f9c02b4158ace065601406a493a8511180c02 \ - --hash=sha256:c7a9152f5876fef565516aa5dd1dccd6fc298a5891b2467973905103eb5c7856 \ - --hash=sha256:cb5cc7b25acabd384f75bbd78892d0c724943f3e2e1986254665a1aa10982e07 \ - --hash=sha256:d3b3db9990c3840986a0e70524e122cfa32b91139c3653df76121ba7776e015f \ - --hash=sha256:d4b1a70a3e5219790d6b55b9507606fc4e02911d1497d16c18dd721eb7efe7d0 \ - --hash=sha256:dc74fd9995513d33eac63d64e436240f5494ec74d522a9f0920194942fc3d2d7 \ - --hash=sha256:e04fde367b2fe901b1d47234426fe8819909bd1dd862a5adb630f27789c20599 - # via - # jax - # jaxlib -numpy==2.1.3 \ - --hash=sha256:016d0f6f5e77b0f0d45d77387ffa4bb89816b57c835580c3ce8e099ef830befe \ - --hash=sha256:02135ade8b8a84011cbb67dc44e07c58f28575cf9ecf8ab304e51c05528c19f0 \ - --hash=sha256:08788d27a5fd867a663f6fc753fd7c3ad7e92747efc73c53bca2f19f8bc06f48 \ - --hash=sha256:0d30c543f02e84e92c4b1f415b7c6b5326cbe45ee7882b6b77db7195fb971e3a \ - --hash=sha256:0fa14563cc46422e99daef53d725d0c326e99e468a9320a240affffe87852564 \ - --hash=sha256:13138eadd4f4da03074851a698ffa7e405f41a0845a6b1ad135b81596e4e9958 \ - --hash=sha256:14e253bd43fc6b37af4921b10f6add6925878a42a0c5fe83daee390bca80bc17 \ - --hash=sha256:15cb89f39fa6d0bdfb600ea24b250e5f1a3df23f901f51c8debaa6a5d122b2f0 \ - --hash=sha256:17ee83a1f4fef3c94d16dc1802b998668b5419362c8a4f4e8a491de1b41cc3ee \ - --hash=sha256:2312b2aa89e1f43ecea6da6ea9a810d06aae08321609d8dc0d0eda6d946a541b \ - --hash=sha256:2564fbdf2b99b3f815f2107c1bbc93e2de8ee655a69c261363a1172a79a257d4 \ - --hash=sha256:3522b0dfe983a575e6a9ab3a4a4dfe156c3e428468ff08ce582b9bb6bd1d71d4 \ - --hash=sha256:4394bc0dbd074b7f9b52024832d16e019decebf86caf909d94f6b3f77a8ee3b6 \ - --hash=sha256:45966d859916ad02b779706bb43b954281db43e185015df6eb3323120188f9e4 \ - --hash=sha256:4d1167c53b93f1f5d8a139a742b3c6f4d429b54e74e6b57d0eff40045187b15d \ - --hash=sha256:4f2015dfe437dfebbfce7c85c7b53d81ba49e71ba7eadbf1df40c915af75979f \ - --hash=sha256:50ca6aba6e163363f132b5c101ba078b8cbd3fa92c7865fd7d4d62d9779ac29f \ - --hash=sha256:50d18c4358a0a8a53f12a8ba9d772ab2d460321e6a93d6064fc22443d189853f \ - --hash=sha256:5641516794ca9e5f8a4d17bb45446998c6554704d888f86df9b200e66bdcce56 \ - --hash=sha256:576a1c1d25e9e02ed7fa5477f30a127fe56debd53b8d2c89d5578f9857d03ca9 \ - --hash=sha256:6a4825252fcc430a182ac4dee5a505053d262c807f8a924603d411f6718b88fd \ - --hash=sha256:72dcc4a35a8515d83e76b58fdf8113a5c969ccd505c8a946759b24e3182d1f23 \ - --hash=sha256:747641635d3d44bcb380d950679462fae44f54b131be347d5ec2bce47d3df9ed \ - --hash=sha256:762479be47a4863e261a840e8e01608d124ee1361e48b96916f38b119cfda04a \ - --hash=sha256:78574ac2d1a4a02421f25da9559850d59457bac82f2b8d7a44fe83a64f770098 \ - --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ - --hash=sha256:8637dcd2caa676e475503d1f8fdb327bc495554e10838019651b76d17b98e512 \ - --hash=sha256:96fe52fcdb9345b7cd82ecd34547fca4321f7656d500eca497eb7ea5a926692f \ - --hash=sha256:973faafebaae4c0aaa1a1ca1ce02434554d67e628b8d805e61f874b84e136b09 \ - --hash=sha256:996bb9399059c5b82f76b53ff8bb686069c05acc94656bb259b1d63d04a9506f \ - --hash=sha256:a38c19106902bb19351b83802531fea19dee18e5b37b36454f27f11ff956f7fc \ - --hash=sha256:a6b46587b14b888e95e4a24d7b13ae91fa22386c199ee7b418f449032b2fa3b8 \ - --hash=sha256:a9f7f672a3388133335589cfca93ed468509cb7b93ba3105fce780d04a6576a0 \ - --hash=sha256:aa08e04e08aaf974d4458def539dece0d28146d866a39da5639596f4921fd761 \ - --hash=sha256:b0df3635b9c8ef48bd3be5f862cf71b0a4716fa0e702155c45067c6b711ddcef \ - --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ - --hash=sha256:baed7e8d7481bfe0874b566850cb0b85243e982388b7b23348c6db2ee2b2ae8e \ - --hash=sha256:bc6f24b3d1ecc1eebfbf5d6051faa49af40b03be1aaa781ebdadcbc090b4539b \ - --hash=sha256:c006b607a865b07cd981ccb218a04fc86b600411d83d6fc261357f1c0966755d \ - --hash=sha256:c181ba05ce8299c7aa3125c27b9c2167bca4a4445b7ce73d5febc411ca692e43 \ - --hash=sha256:c7662f0e3673fe4e832fe07b65c50342ea27d989f92c80355658c7f888fcc83c \ - --hash=sha256:c80e4a09b3d95b4e1cac08643f1152fa71a0a821a2d4277334c88d54b2219a41 \ - --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff \ - --hash=sha256:d7aac50327da5d208db2eec22eb11e491e3fe13d22653dce51b0f4109101b408 \ - --hash=sha256:d89dd2b6da69c4fff5e39c28a382199ddedc3a5be5390115608345dec660b9e2 \ - --hash=sha256:d9beb777a78c331580705326d2367488d5bc473b49a9bc3036c154832520aca9 \ - --hash=sha256:dc258a761a16daa791081d026f0ed4399b582712e6fc887a95af09df10c5ca57 \ - --hash=sha256:e14e26956e6f1696070788252dcdff11b4aca4c3e8bd166e0df1bb8f315a67cb \ - --hash=sha256:e6988e90fcf617da2b5c78902fe8e668361b43b4fe26dbf2d7b0f8034d4cafb9 \ - --hash=sha256:e711e02f49e176a01d0349d82cb5f05ba4db7d5e7e0defd026328e5cfb3226d3 \ - --hash=sha256:ea4dedd6e394a9c180b33c2c872b92f7ce0f8e7ad93e9585312b0c5a04777a4a \ - --hash=sha256:ecc76a9ba2911d8d37ac01de72834d8849e55473457558e12995f4cd53e778e0 \ - --hash=sha256:f55ba01150f52b1027829b50d70ef1dafd9821ea82905b63936668403c3b471e \ - --hash=sha256:f653490b33e9c3a4c1c01d41bc2aef08f9475af51146e4a7710c450cf9761598 \ - --hash=sha256:fa2d1337dc61c8dc417fbccf20f6d1e139896a30721b7f1e832b2bb6ef4eb6c4 - # via - # alphafold3 (pyproject.toml) - # dm-haiku - # jax - # jaxlib - # jmp - # ml-dtypes - # rdkit - # scipy -nvidia-cublas-cu12==12.6.3.3 \ - --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \ - --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \ - --hash=sha256:f33fb68e101d99470c82d17f92a0dd9f74de2a21685c217f4716cdd63b1316eb - # via - # jax-cuda12-plugin - # nvidia-cudnn-cu12 - # nvidia-cusolver-cu12 -nvidia-cuda-cupti-cu12==12.6.80 \ - --hash=sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4 \ - --hash=sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73 \ - --hash=sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a - # via jax-cuda12-plugin -nvidia-cuda-nvcc-cu12==12.6.77 \ - --hash=sha256:26e118b4dae4ebd316ab6f5ee3f3c4cb929f8265f717a8e016ba57e254c0b4f1 \ - --hash=sha256:9a2aad43b8f4e4be9c237ece04f64450df77871eeebffad5a5a872714a97ca77 \ - --hash=sha256:b66bb5dd6b8ae62262586691977d3b4a425e91db61b4b9bc2f6b42bcd4154b96 - # via jax-cuda12-plugin -nvidia-cuda-runtime-cu12==12.6.77 \ - --hash=sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f \ - --hash=sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8 \ - --hash=sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e - # via jax-cuda12-plugin -nvidia-cudnn-cu12==9.5.1.17 \ - --hash=sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2 \ - --hash=sha256:9fd4584468533c61873e5fda8ca41bac3a38bcb2d12350830c69b0a96a7e4def \ - --hash=sha256:d7af0f8a4f3b4b9dbb3122f2ef553b45694ed9c384d5a75bab197b8eefb79ab8 - # via jax-cuda12-plugin -nvidia-cufft-cu12==11.3.0.4 \ - --hash=sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464 \ - --hash=sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca \ - --hash=sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb - # via jax-cuda12-plugin -nvidia-cusolver-cu12==11.7.1.2 \ - --hash=sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0 \ - --hash=sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7 \ - --hash=sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6 - # via jax-cuda12-plugin -nvidia-cusparse-cu12==12.5.4.2 \ - --hash=sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f \ - --hash=sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20 \ - --hash=sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1 - # via - # jax-cuda12-plugin - # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.23.4 \ - --hash=sha256:aa946c8327e22ced28e7cef508a334673abc42064ec85f02d005ba1785ea4cec \ - --hash=sha256:b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1 - # via jax-cuda12-plugin -nvidia-nvjitlink-cu12==12.6.77 \ - --hash=sha256:3bf10d85bb1801e9c894c6e197e44dd137d2a0a9e43f8450e9ad13f2df0dd52d \ - --hash=sha256:410718cd44962bed862a31dd0318620f6f9a8b28a6291967bcfcb446a6516771 \ - --hash=sha256:9ae346d16203ae4ea513be416495167a0101d33d2d14935aa9c1829a3fb45142 - # via - # jax-cuda12-plugin - # nvidia-cufft-cu12 - # nvidia-cusolver-cu12 - # nvidia-cusparse-cu12 -opt-einsum==3.4.0 \ - --hash=sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd \ - --hash=sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac - # via jax -pillow==11.0.0 \ - --hash=sha256:00177a63030d612148e659b55ba99527803288cea7c75fb05766ab7981a8c1b7 \ - --hash=sha256:006bcdd307cc47ba43e924099a038cbf9591062e6c50e570819743f5607404f5 \ - --hash=sha256:084a07ef0821cfe4858fe86652fffac8e187b6ae677e9906e192aafcc1b69903 \ - --hash=sha256:0ae08bd8ffc41aebf578c2af2f9d8749d91f448b3bfd41d7d9ff573d74f2a6b2 \ - --hash=sha256:0e038b0745997c7dcaae350d35859c9715c71e92ffb7e0f4a8e8a16732150f38 \ - --hash=sha256:1187739620f2b365de756ce086fdb3604573337cc28a0d3ac4a01ab6b2d2a6d2 \ - --hash=sha256:16095692a253047fe3ec028e951fa4221a1f3ed3d80c397e83541a3037ff67c9 \ - --hash=sha256:1a61b54f87ab5786b8479f81c4b11f4d61702830354520837f8cc791ebba0f5f \ - --hash=sha256:1c1d72714f429a521d8d2d018badc42414c3077eb187a59579f28e4270b4b0fc \ - --hash=sha256:1e2688958a840c822279fda0086fec1fdab2f95bf2b717b66871c4ad9859d7e8 \ - --hash=sha256:20ec184af98a121fb2da42642dea8a29ec80fc3efbaefb86d8fdd2606619045d \ - --hash=sha256:21a0d3b115009ebb8ac3d2ebec5c2982cc693da935f4ab7bb5c8ebe2f47d36f2 \ - --hash=sha256:224aaa38177597bb179f3ec87eeefcce8e4f85e608025e9cfac60de237ba6316 \ - --hash=sha256:2679d2258b7f1192b378e2893a8a0a0ca472234d4c2c0e6bdd3380e8dfa21b6a \ - --hash=sha256:27a7860107500d813fcd203b4ea19b04babe79448268403172782754870dac25 \ - --hash=sha256:290f2cc809f9da7d6d622550bbf4c1e57518212da51b6a30fe8e0a270a5b78bd \ - --hash=sha256:2e46773dc9f35a1dd28bd6981332fd7f27bec001a918a72a79b4133cf5291dba \ - --hash=sha256:3107c66e43bda25359d5ef446f59c497de2b5ed4c7fdba0894f8d6cf3822dafc \ - --hash=sha256:375b8dd15a1f5d2feafff536d47e22f69625c1aa92f12b339ec0b2ca40263273 \ - --hash=sha256:45c566eb10b8967d71bf1ab8e4a525e5a93519e29ea071459ce517f6b903d7fa \ - --hash=sha256:499c3a1b0d6fc8213519e193796eb1a86a1be4b1877d678b30f83fd979811d1a \ - --hash=sha256:4ad70c4214f67d7466bea6a08061eba35c01b1b89eaa098040a35272a8efb22b \ - --hash=sha256:4b60c9520f7207aaf2e1d94de026682fc227806c6e1f55bba7606d1c94dd623a \ - --hash=sha256:5178952973e588b3f1360868847334e9e3bf49d19e169bbbdfaf8398002419ae \ - --hash=sha256:52a2d8323a465f84faaba5236567d212c3668f2ab53e1c74c15583cf507a0291 \ - --hash=sha256:598b4e238f13276e0008299bd2482003f48158e2b11826862b1eb2ad7c768b97 \ - --hash=sha256:5bd2d3bdb846d757055910f0a59792d33b555800813c3b39ada1829c372ccb06 \ - --hash=sha256:5c39ed17edea3bc69c743a8dd3e9853b7509625c2462532e62baa0732163a904 \ - --hash=sha256:5d203af30149ae339ad1b4f710d9844ed8796e97fda23ffbc4cc472968a47d0b \ - --hash=sha256:5ddbfd761ee00c12ee1be86c9c0683ecf5bb14c9772ddbd782085779a63dd55b \ - --hash=sha256:607bbe123c74e272e381a8d1957083a9463401f7bd01287f50521ecb05a313f8 \ - --hash=sha256:61b887f9ddba63ddf62fd02a3ba7add935d053b6dd7d58998c630e6dbade8527 \ - --hash=sha256:6619654954dc4936fcff82db8eb6401d3159ec6be81e33c6000dfd76ae189947 \ - --hash=sha256:674629ff60030d144b7bca2b8330225a9b11c482ed408813924619c6f302fdbb \ - --hash=sha256:6ec0d5af64f2e3d64a165f490d96368bb5dea8b8f9ad04487f9ab60dc4bb6003 \ - --hash=sha256:6f4dba50cfa56f910241eb7f883c20f1e7b1d8f7d91c750cd0b318bad443f4d5 \ - --hash=sha256:70fbbdacd1d271b77b7721fe3cdd2d537bbbd75d29e6300c672ec6bb38d9672f \ - --hash=sha256:72bacbaf24ac003fea9bff9837d1eedb6088758d41e100c1552930151f677739 \ - --hash=sha256:7326a1787e3c7b0429659e0a944725e1b03eeaa10edd945a86dead1913383944 \ - --hash=sha256:73853108f56df97baf2bb8b522f3578221e56f646ba345a372c78326710d3830 \ - --hash=sha256:73e3a0200cdda995c7e43dd47436c1548f87a30bb27fb871f352a22ab8dcf45f \ - --hash=sha256:75acbbeb05b86bc53cbe7b7e6fe00fbcf82ad7c684b3ad82e3d711da9ba287d3 \ - --hash=sha256:8069c5179902dcdce0be9bfc8235347fdbac249d23bd90514b7a47a72d9fecf4 \ - --hash=sha256:846e193e103b41e984ac921b335df59195356ce3f71dcfd155aa79c603873b84 \ - --hash=sha256:8594f42df584e5b4bb9281799698403f7af489fba84c34d53d1c4bfb71b7c4e7 \ - --hash=sha256:86510e3f5eca0ab87429dd77fafc04693195eec7fd6a137c389c3eeb4cfb77c6 \ - --hash=sha256:8853a3bf12afddfdf15f57c4b02d7ded92c7a75a5d7331d19f4f9572a89c17e6 \ - --hash=sha256:88a58d8ac0cc0e7f3a014509f0455248a76629ca9b604eca7dc5927cc593c5e9 \ - --hash=sha256:8ba470552b48e5835f1d23ecb936bb7f71d206f9dfeee64245f30c3270b994de \ - --hash=sha256:8c676b587da5673d3c75bd67dd2a8cdfeb282ca38a30f37950511766b26858c4 \ - --hash=sha256:8ec4a89295cd6cd4d1058a5e6aec6bf51e0eaaf9714774e1bfac7cfc9051db47 \ - --hash=sha256:94f3e1780abb45062287b4614a5bc0874519c86a777d4a7ad34978e86428b8dd \ - --hash=sha256:9a0f748eaa434a41fccf8e1ee7a3eed68af1b690e75328fd7a60af123c193b50 \ - --hash=sha256:a5629742881bcbc1f42e840af185fd4d83a5edeb96475a575f4da50d6ede337c \ - --hash=sha256:a65149d8ada1055029fcb665452b2814fe7d7082fcb0c5bed6db851cb69b2086 \ - --hash=sha256:b3c5ac4bed7519088103d9450a1107f76308ecf91d6dabc8a33a2fcfb18d0fba \ - --hash=sha256:b4fd7bd29610a83a8c9b564d457cf5bd92b4e11e79a4ee4716a63c959699b306 \ - --hash=sha256:bcd1fb5bb7b07f64c15618c89efcc2cfa3e95f0e3bcdbaf4642509de1942a699 \ - --hash=sha256:c12b5ae868897c7338519c03049a806af85b9b8c237b7d675b8c5e089e4a618e \ - --hash=sha256:c26845094b1af3c91852745ae78e3ea47abf3dbcd1cf962f16b9a5fbe3ee8488 \ - --hash=sha256:c6a660307ca9d4867caa8d9ca2c2658ab685de83792d1876274991adec7b93fa \ - --hash=sha256:c809a70e43c7977c4a42aefd62f0131823ebf7dd73556fa5d5950f5b354087e2 \ - --hash=sha256:c8b2351c85d855293a299038e1f89db92a2f35e8d2f783489c6f0b2b5f3fe8a3 \ - --hash=sha256:cb929ca942d0ec4fac404cbf520ee6cac37bf35be479b970c4ffadf2b6a1cad9 \ - --hash=sha256:d2c0a187a92a1cb5ef2c8ed5412dd8d4334272617f532d4ad4de31e0495bd923 \ - --hash=sha256:d69bfd8ec3219ae71bcde1f942b728903cad25fafe3100ba2258b973bd2bc1b2 \ - --hash=sha256:daffdf51ee5db69a82dd127eabecce20729e21f7a3680cf7cbb23f0829189790 \ - --hash=sha256:e58876c91f97b0952eb766123bfef372792ab3f4e3e1f1a2267834c2ab131734 \ - --hash=sha256:eda2616eb2313cbb3eebbe51f19362eb434b18e3bb599466a1ffa76a033fb916 \ - --hash=sha256:ee217c198f2e41f184f3869f3e485557296d505b5195c513b2bfe0062dc537f1 \ - --hash=sha256:f02541ef64077f22bf4924f225c0fd1248c168f86e4b7abdedd87d6ebaceab0f \ - --hash=sha256:f1b82c27e89fffc6da125d5eb0ca6e68017faf5efc078128cfaa42cf5cb38798 \ - --hash=sha256:fba162b8872d30fea8c52b258a542c5dfd7b235fb5cb352240c8d63b414013eb \ - --hash=sha256:fbbcb7b57dc9c794843e3d1258c0fbf0f48656d46ffe9e09b63bbd6e8cd5d0a2 \ - --hash=sha256:fcb4621042ac4b7865c179bb972ed0da0218a076dc1820ffc48b1d74c1e37fe9 - # via rdkit -rdkit==2024.3.5 \ - --hash=sha256:01fd323f5902a720925293c0ce08a33e630422374cf263d04ef173c106c63c36 \ - --hash=sha256:04079aef5f91965bd2a48bf432b32c36ea1bf088e26efb2c9a43b625b6a75d15 \ - --hash=sha256:1cb7cdb29956649d4c875790b94142868c08c0735dee4d90b2d70dfd378e9d2e \ - --hash=sha256:292141f0ba2359be07d8566d9a39ce697c1d6e7eb406031a2bf0387026503a47 \ - --hash=sha256:2affa1b4cf601d27b157bbc44255c640baac9bc04cfdf5fa88f43c34c3f360d5 \ - --hash=sha256:316bb5273bff5be76d4e5b0def998455e9b1a20898df44548b0e299f90786cb2 \ - --hash=sha256:455d510beff8806e62e48b977d7acbfbc351474fa132124738a07223440c1b9a \ - --hash=sha256:53638d0d94f7c511273ff5765abc6a1fe60fe6c8633db13cad03bc527550a7a9 \ - --hash=sha256:6af38c62d4af9d669f8996551d3ae809de9ab245570b861be3f7a2b1a536f79f \ - --hash=sha256:73fa2fe3a9a2334edf8448a5a3093a6f949ee715de77d3baa685751e2cf9317b \ - --hash=sha256:7c91c5ccbf8d51cd3dd894bf770ca99da5fe37714f6e2a66016e8af4f57d2b56 \ - --hash=sha256:7d5576bad6672959d0c1939e8d7e2fcd8656f8daf9788ce58e5c313db719b5ec \ - --hash=sha256:89f7a72466ffbd45320d845ee1bcde1bddb1a322352836835d26ebb712228a41 \ - --hash=sha256:8d4ebafb7e845bdc55eb6e2e188123ec84fb2bf724d3a7eb2da9bcfade7acc65 \ - --hash=sha256:90376467f4988ee44f3f47d1e558ffc27646254a99fe65d19a7d9a123e832ef2 \ - --hash=sha256:913494a03fcf6121c103e5db43fa33bca0fcc65ebd5a6c46661baa6851f79d91 \ - --hash=sha256:98dac94a5e414ed1880b215eacec63e5fd77b0ea351e0a18fc5229acb8e886f5 \ - --hash=sha256:a51cc403792ba37d1ee8ad84ad829177a4453547a4e085e1115ae2b237084ddb \ - --hash=sha256:a98099cef42c1eb52512b343beb68a8b7cc4da79924a67181f58692a6376d91b \ - --hash=sha256:c971eba2a67f8ae301b9b8c66b3431cd62472f648cdec9694e5413d8ebc59720 \ - --hash=sha256:cf2747e7c82d6c90a8a631fe063d24d77b31f535e0f8a17ac9dbfbe4ab456886 \ - --hash=sha256:d8714c03558086d022649617197c0c45c75c301f6a9fb67205695493a7e5b826 \ - --hash=sha256:e5b0dfb28aa0949152ff816fba0a2f6267154f5c25159a9ec84c27cf18f72322 \ - --hash=sha256:e69cec14cbb42c5ca912346b594d3f6a7aac572463949c55c7749da257e68123 \ - --hash=sha256:f662ed10e14c7874d0e011c1f8260a6cf4a5547100806f9f4540f8550a3f1d36 - # via alphafold3 (pyproject.toml) -scipy==1.14.1 \ - --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \ - --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \ - --hash=sha256:278266012eb69f4a720827bdd2dc54b2271c97d84255b2faaa8f161a158c3b37 \ - --hash=sha256:2843f2d527d9eebec9a43e6b406fb7266f3af25a751aa91d62ff416f54170bc5 \ - --hash=sha256:2da0469a4ef0ecd3693761acbdc20f2fdeafb69e6819cc081308cc978153c675 \ - --hash=sha256:2ff0a7e01e422c15739ecd64432743cf7aae2b03f3084288f399affcefe5222d \ - --hash=sha256:2ff38e22128e6c03ff73b6bb0f85f897d2362f8c052e3b8ad00532198fbdae3f \ - --hash=sha256:30ac8812c1d2aab7131a79ba62933a2a76f582d5dbbc695192453dae67ad6310 \ - --hash=sha256:3a1b111fac6baec1c1d92f27e76511c9e7218f1695d61b59e05e0fe04dc59617 \ - --hash=sha256:4079b90df244709e675cdc8b93bfd8a395d59af40b72e339c2287c91860deb8e \ - --hash=sha256:5149e3fd2d686e42144a093b206aef01932a0059c2a33ddfa67f5f035bdfe13e \ - --hash=sha256:5a275584e726026a5699459aa72f828a610821006228e841b94275c4a7c08417 \ - --hash=sha256:631f07b3734d34aced009aaf6fedfd0eb3498a97e581c3b1e5f14a04164a456d \ - --hash=sha256:716e389b694c4bb564b4fc0c51bc84d381735e0d39d3f26ec1af2556ec6aad94 \ - --hash=sha256:8426251ad1e4ad903a4514712d2fa8fdd5382c978010d1c6f5f37ef286a713ad \ - --hash=sha256:8475230e55549ab3f207bff11ebfc91c805dc3463ef62eda3ccf593254524ce8 \ - --hash=sha256:8bddf15838ba768bb5f5083c1ea012d64c9a444e16192762bd858f1e126196d0 \ - --hash=sha256:8e32dced201274bf96899e6491d9ba3e9a5f6b336708656466ad0522d8528f69 \ - --hash=sha256:8f9ea80f2e65bdaa0b7627fb00cbeb2daf163caa015e59b7516395fe3bd1e066 \ - --hash=sha256:97c5dddd5932bd2a1a31c927ba5e1463a53b87ca96b5c9bdf5dfd6096e27efc3 \ - --hash=sha256:a49f6ed96f83966f576b33a44257d869756df6cf1ef4934f59dd58b25e0327e5 \ - --hash=sha256:af29a935803cc707ab2ed7791c44288a682f9c8107bc00f0eccc4f92c08d6e07 \ - --hash=sha256:b05d43735bb2f07d689f56f7b474788a13ed8adc484a85aa65c0fd931cf9ccd2 \ - --hash=sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389 \ - --hash=sha256:b99722ea48b7ea25e8e015e8341ae74624f72e5f21fc2abd45f3a93266de4c5d \ - --hash=sha256:baff393942b550823bfce952bb62270ee17504d02a1801d7fd0719534dfb9c84 \ - --hash=sha256:c0ee987efa6737242745f347835da2cc5bb9f1b42996a4d97d5c7ff7928cb6f2 \ - --hash=sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3 \ - --hash=sha256:e0cf28db0f24a38b2a0ca33a85a54852586e43cf6fd876365c86e0657cfe7d73 \ - --hash=sha256:e4f5a7c49323533f9103d4dacf4e4f07078f360743dec7f7596949149efeec06 \ - --hash=sha256:eb58ca0abd96911932f688528977858681a59d61a7ce908ffd355957f7025cfc \ - --hash=sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1 \ - --hash=sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2 - # via - # jax - # jaxlib -tabulate==0.9.0 \ - --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \ - --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f - # via dm-haiku -tqdm==4.67.0 \ - --hash=sha256:0cd8af9d56911acab92182e88d763100d4788bdf421d251616040cc4d44863be \ - --hash=sha256:fe5a6f95e6fe0b9755e9469b77b9c3cf850048224ecaa8293d7d2d31f97d869a - # via alphafold3 (pyproject.toml) -triton==3.1.0 \ - --hash=sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c \ - --hash=sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8 \ - --hash=sha256:6dadaca7fc24de34e180271b5cf864c16755702e9f63a16f62df714a8099126a \ - --hash=sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11 \ - --hash=sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc - # via - # alphafold3 (pyproject.toml) - # jax-triton -typeguard==2.13.3 \ - --hash=sha256:00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4 \ - --hash=sha256:5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1 - # via - # alphafold3 (pyproject.toml) - # jaxtyping -zstandard==0.23.0 \ - --hash=sha256:034b88913ecc1b097f528e42b539453fa82c3557e414b3de9d5632c80439a473 \ - --hash=sha256:0a7f0804bb3799414af278e9ad51be25edf67f78f916e08afdb983e74161b916 \ - --hash=sha256:11e3bf3c924853a2d5835b24f03eeba7fc9b07d8ca499e247e06ff5676461a15 \ - --hash=sha256:12a289832e520c6bd4dcaad68e944b86da3bad0d339ef7989fb7e88f92e96072 \ - --hash=sha256:1516c8c37d3a053b01c1c15b182f3b5f5eef19ced9b930b684a73bad121addf4 \ - --hash=sha256:157e89ceb4054029a289fb504c98c6a9fe8010f1680de0201b3eb5dc20aa6d9e \ - --hash=sha256:1bfe8de1da6d104f15a60d4a8a768288f66aa953bbe00d027398b93fb9680b26 \ - --hash=sha256:1e172f57cd78c20f13a3415cc8dfe24bf388614324d25539146594c16d78fcc8 \ - --hash=sha256:1fd7e0f1cfb70eb2f95a19b472ee7ad6d9a0a992ec0ae53286870c104ca939e5 \ - --hash=sha256:203d236f4c94cd8379d1ea61db2fce20730b4c38d7f1c34506a31b34edc87bdd \ - --hash=sha256:27d3ef2252d2e62476389ca8f9b0cf2bbafb082a3b6bfe9d90cbcbb5529ecf7c \ - --hash=sha256:29a2bc7c1b09b0af938b7a8343174b987ae021705acabcbae560166567f5a8db \ - --hash=sha256:2ef230a8fd217a2015bc91b74f6b3b7d6522ba48be29ad4ea0ca3a3775bf7dd5 \ - --hash=sha256:2ef3775758346d9ac6214123887d25c7061c92afe1f2b354f9388e9e4d48acfc \ - --hash=sha256:2f146f50723defec2975fb7e388ae3a024eb7151542d1599527ec2aa9cacb152 \ - --hash=sha256:2fb4535137de7e244c230e24f9d1ec194f61721c86ebea04e1581d9d06ea1269 \ - --hash=sha256:32ba3b5ccde2d581b1e6aa952c836a6291e8435d788f656fe5976445865ae045 \ - --hash=sha256:34895a41273ad33347b2fc70e1bff4240556de3c46c6ea430a7ed91f9042aa4e \ - --hash=sha256:379b378ae694ba78cef921581ebd420c938936a153ded602c4fea612b7eaa90d \ - --hash=sha256:38302b78a850ff82656beaddeb0bb989a0322a8bbb1bf1ab10c17506681d772a \ - --hash=sha256:3aa014d55c3af933c1315eb4bb06dd0459661cc0b15cd61077afa6489bec63bb \ - --hash=sha256:4051e406288b8cdbb993798b9a45c59a4896b6ecee2f875424ec10276a895740 \ - --hash=sha256:40b33d93c6eddf02d2c19f5773196068d875c41ca25730e8288e9b672897c105 \ - --hash=sha256:43da0f0092281bf501f9c5f6f3b4c975a8a0ea82de49ba3f7100e64d422a1274 \ - --hash=sha256:445e4cb5048b04e90ce96a79b4b63140e3f4ab5f662321975679b5f6360b90e2 \ - --hash=sha256:48ef6a43b1846f6025dde6ed9fee0c24e1149c1c25f7fb0a0585572b2f3adc58 \ - --hash=sha256:50a80baba0285386f97ea36239855f6020ce452456605f262b2d33ac35c7770b \ - --hash=sha256:519fbf169dfac1222a76ba8861ef4ac7f0530c35dd79ba5727014613f91613d4 \ - --hash=sha256:53dd9d5e3d29f95acd5de6802e909ada8d8d8cfa37a3ac64836f3bc4bc5512db \ - --hash=sha256:53ea7cdc96c6eb56e76bb06894bcfb5dfa93b7adcf59d61c6b92674e24e2dd5e \ - --hash=sha256:576856e8594e6649aee06ddbfc738fec6a834f7c85bf7cadd1c53d4a58186ef9 \ - --hash=sha256:59556bf80a7094d0cfb9f5e50bb2db27fefb75d5138bb16fb052b61b0e0eeeb0 \ - --hash=sha256:5d41d5e025f1e0bccae4928981e71b2334c60f580bdc8345f824e7c0a4c2a813 \ - --hash=sha256:61062387ad820c654b6a6b5f0b94484fa19515e0c5116faf29f41a6bc91ded6e \ - --hash=sha256:61f89436cbfede4bc4e91b4397eaa3e2108ebe96d05e93d6ccc95ab5714be512 \ - --hash=sha256:62136da96a973bd2557f06ddd4e8e807f9e13cbb0bfb9cc06cfe6d98ea90dfe0 \ - --hash=sha256:64585e1dba664dc67c7cdabd56c1e5685233fbb1fc1966cfba2a340ec0dfff7b \ - --hash=sha256:65308f4b4890aa12d9b6ad9f2844b7ee42c7f7a4fd3390425b242ffc57498f48 \ - --hash=sha256:66b689c107857eceabf2cf3d3fc699c3c0fe8ccd18df2219d978c0283e4c508a \ - --hash=sha256:6a41c120c3dbc0d81a8e8adc73312d668cd34acd7725f036992b1b72d22c1772 \ - --hash=sha256:6f77fa49079891a4aab203d0b1744acc85577ed16d767b52fc089d83faf8d8ed \ - --hash=sha256:72c68dda124a1a138340fb62fa21b9bf4848437d9ca60bd35db36f2d3345f373 \ - --hash=sha256:752bf8a74412b9892f4e5b58f2f890a039f57037f52c89a740757ebd807f33ea \ - --hash=sha256:76e79bc28a65f467e0409098fa2c4376931fd3207fbeb6b956c7c476d53746dd \ - --hash=sha256:774d45b1fac1461f48698a9d4b5fa19a69d47ece02fa469825b442263f04021f \ - --hash=sha256:77da4c6bfa20dd5ea25cbf12c76f181a8e8cd7ea231c673828d0386b1740b8dc \ - --hash=sha256:77ea385f7dd5b5676d7fd943292ffa18fbf5c72ba98f7d09fc1fb9e819b34c23 \ - --hash=sha256:80080816b4f52a9d886e67f1f96912891074903238fe54f2de8b786f86baded2 \ - --hash=sha256:80a539906390591dd39ebb8d773771dc4db82ace6372c4d41e2d293f8e32b8db \ - --hash=sha256:82d17e94d735c99621bf8ebf9995f870a6b3e6d14543b99e201ae046dfe7de70 \ - --hash=sha256:837bb6764be6919963ef41235fd56a6486b132ea64afe5fafb4cb279ac44f259 \ - --hash=sha256:84433dddea68571a6d6bd4fbf8ff398236031149116a7fff6f777ff95cad3df9 \ - --hash=sha256:8c24f21fa2af4bb9f2c492a86fe0c34e6d2c63812a839590edaf177b7398f700 \ - --hash=sha256:8ed7d27cb56b3e058d3cf684d7200703bcae623e1dcc06ed1e18ecda39fee003 \ - --hash=sha256:9206649ec587e6b02bd124fb7799b86cddec350f6f6c14bc82a2b70183e708ba \ - --hash=sha256:983b6efd649723474f29ed42e1467f90a35a74793437d0bc64a5bf482bedfa0a \ - --hash=sha256:98da17ce9cbf3bfe4617e836d561e433f871129e3a7ac16d6ef4c680f13a839c \ - --hash=sha256:9c236e635582742fee16603042553d276cca506e824fa2e6489db04039521e90 \ - --hash=sha256:9da6bc32faac9a293ddfdcb9108d4b20416219461e4ec64dfea8383cac186690 \ - --hash=sha256:a05e6d6218461eb1b4771d973728f0133b2a4613a6779995df557f70794fd60f \ - --hash=sha256:a0817825b900fcd43ac5d05b8b3079937073d2b1ff9cf89427590718b70dd840 \ - --hash=sha256:a4ae99c57668ca1e78597d8b06d5af837f377f340f4cce993b551b2d7731778d \ - --hash=sha256:a8c86881813a78a6f4508ef9daf9d4995b8ac2d147dcb1a450448941398091c9 \ - --hash=sha256:a8fffdbd9d1408006baaf02f1068d7dd1f016c6bcb7538682622c556e7b68e35 \ - --hash=sha256:a9b07268d0c3ca5c170a385a0ab9fb7fdd9f5fd866be004c4ea39e44edce47dd \ - --hash=sha256:ab19a2d91963ed9e42b4e8d77cd847ae8381576585bad79dbd0a8837a9f6620a \ - --hash=sha256:ac184f87ff521f4840e6ea0b10c0ec90c6b1dcd0bad2f1e4a9a1b4fa177982ea \ - --hash=sha256:b0e166f698c5a3e914947388c162be2583e0c638a4703fc6a543e23a88dea3c1 \ - --hash=sha256:b2170c7e0367dde86a2647ed5b6f57394ea7f53545746104c6b09fc1f4223573 \ - --hash=sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09 \ - --hash=sha256:b4567955a6bc1b20e9c31612e615af6b53733491aeaa19a6b3b37f3b65477094 \ - --hash=sha256:b69bb4f51daf461b15e7b3db033160937d3ff88303a7bc808c67bbc1eaf98c78 \ - --hash=sha256:b8c0bd73aeac689beacd4e7667d48c299f61b959475cdbb91e7d3d88d27c56b9 \ - --hash=sha256:be9b5b8659dff1f913039c2feee1aca499cfbc19e98fa12bc85e037c17ec6ca5 \ - --hash=sha256:bf0a05b6059c0528477fba9054d09179beb63744355cab9f38059548fedd46a9 \ - --hash=sha256:c16842b846a8d2a145223f520b7e18b57c8f476924bda92aeee3a88d11cfc391 \ - --hash=sha256:c363b53e257246a954ebc7c488304b5592b9c53fbe74d03bc1c64dda153fb847 \ - --hash=sha256:c7c517d74bea1a6afd39aa612fa025e6b8011982a0897768a2f7c8ab4ebb78a2 \ - --hash=sha256:d20fd853fbb5807c8e84c136c278827b6167ded66c72ec6f9a14b863d809211c \ - --hash=sha256:d2240ddc86b74966c34554c49d00eaafa8200a18d3a5b6ffbf7da63b11d74ee2 \ - --hash=sha256:d477ed829077cd945b01fc3115edd132c47e6540ddcd96ca169facff28173057 \ - --hash=sha256:d50d31bfedd53a928fed6707b15a8dbeef011bb6366297cc435accc888b27c20 \ - --hash=sha256:dc1d33abb8a0d754ea4763bad944fd965d3d95b5baef6b121c0c9013eaf1907d \ - --hash=sha256:dc5d1a49d3f8262be192589a4b72f0d03b72dcf46c51ad5852a4fdc67be7b9e4 \ - --hash=sha256:e2d1a054f8f0a191004675755448d12be47fa9bebbcffa3cdf01db19f2d30a54 \ - --hash=sha256:e7792606d606c8df5277c32ccb58f29b9b8603bf83b48639b7aedf6df4fe8171 \ - --hash=sha256:ed1708dbf4d2e3a1c5c69110ba2b4eb6678262028afd6c6fbcc5a8dac9cda68e \ - --hash=sha256:f2d4380bf5f62daabd7b751ea2339c1a21d1c9463f1feb7fc2bdcea2c29c3160 \ - --hash=sha256:f3513916e8c645d0610815c257cbfd3242adfd5c4cfa78be514e5a3ebb42a41b \ - --hash=sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58 \ - --hash=sha256:f83fa6cae3fff8e98691248c9320356971b59678a17f20656a9e59cd32cee6d8 \ - --hash=sha256:fa6ce8b52c5987b3e34d5674b0ab529a4602b632ebab0a93b07bfb4dfc8f8a33 \ - --hash=sha256:fb2b1ecfef1e67897d336de3a0e3f52478182d6a47eda86cbd42504c5cbd009a \ - --hash=sha256:fc9ca1c9718cb3b06634c7c8dec57d24e9438b2aa9a0f02b8bb36bf478538880 \ - --hash=sha256:fd30d9c67d13d891f2360b2a120186729c111238ac63b43dbd37a5a40670b8ca \ - --hash=sha256:fd7699e8fd9969f455ef2926221e0233f81a2542921471382e77a9e2f2b57f4b \ - --hash=sha256:fe3b385d996ee0822fd46528d9f0443b880d4d05528fd26a9119a54ec3f91c69 - # via alphafold3 (pyproject.toml) diff --git a/docker/Dockerfile b/docker/Dockerfile index cf3eb68..bbeba22 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -16,17 +25,21 @@ FROM nvidia/cuda:12.6.3-base-ubuntu24.04 # * gcc, g++, make are required for compiling HMMER and AlphaFold 3 libaries. # * zlib is a required dependency of AlphaFold 3. RUN DEBIAN_FRONTEND=noninteractive \ - apt-get update --quiet \ - && apt-get install --yes --quiet python3.12 python3-pip python3.12-venv python3.12-dev \ - && apt-get install --yes --quiet git wget gcc g++ make zlib1g-dev zstd +apt-get update --quiet \ +&& apt-get install --yes --quiet python3.12 python3.12-dev \ +&& apt-get install --yes --quiet git wget gcc g++ make zlib1g-dev zstd + +# Install uv from the official repository. The version is pinned for +# reproducibility. +COPY --from=ghcr.io/astral-sh/uv:0.9.24 /uv /uvx /bin/ + +# UV_COMPILE_BYTECODE=1 speeds up future container starts. +# UV_PROJECT_ENVIRONMENT explicitly sets the virtual environment location. +ENV UV_COMPILE_BYTECODE=1 +ENV UV_PROJECT_ENVIRONMENT=/alphafold3_venv +RUN uv venv $UV_PROJECT_ENVIRONMENT -# The base image uses Python 3.12. For local installation use 3.11 or greater. -RUN python3 -m venv /alphafold3_venv ENV PATH="/hmmer/bin:/alphafold3_venv/bin:$PATH" -# Update pip to the latest version. Not necessary in Docker, but good to do when -# this is used as a recipe for local installation since we rely on new pip -# features for secure installs. -RUN pip3 install --no-cache-dir --upgrade pip # Install HMMER. Do so before copying the source code, so that docker can cache # the image layer containing HMMER. Alternatively, you could also install it @@ -47,10 +60,10 @@ COPY docker/jackhmmer_seq_limit.patch /hmmer_build/ RUN (cd /hmmer_build && patch -p0 < jackhmmer_seq_limit.patch) # Build HMMER. -RUN (cd /hmmer_build/hmmer-3.4 && ./configure --prefix /hmmer) ; \ - (cd /hmmer_build/hmmer-3.4 && make -j) ; \ - (cd /hmmer_build/hmmer-3.4 && make install) ; \ - (cd /hmmer_build/hmmer-3.4/easel && make install) ; \ +RUN (cd /hmmer_build/hmmer-3.4 && ./configure --prefix /hmmer) && \ + (cd /hmmer_build/hmmer-3.4 && make -j) && \ + (cd /hmmer_build/hmmer-3.4 && make install) && \ + (cd /hmmer_build/hmmer-3.4/easel && make install) && \ rm -R /hmmer_build # Copy the AlphaFold 3 source code from the local machine to the container and @@ -58,13 +71,17 @@ RUN (cd /hmmer_build/hmmer-3.4 && ./configure --prefix /hmmer) ; \ COPY . /app/alphafold WORKDIR /app/alphafold -# Install setuptools (removed in Python 3.12) which some libraries still need. -# Then install the Python dependencies of AlphaFold 3. -RUN pip3 install --no-cache-dir setuptools \ - && pip3 install --no-cache-dir -r dev-requirements.txt \ - && pip3 install --no-cache-dir --no-deps . -# Build chemical components database (this binary was installed by pip). -RUN build_data +# Install the exact dependency tree using uv and cache the build artifacts. +# --frozen: do not update the lockfile during build. +# --all-groups: install development/test dependencies defined in pyproject.toml. +# --no-editable: install as a static package. +# If using this as a recipe for local installation, we recommend removing the +# --frozen and --no-editable flags. +RUN --mount=type=cache,target=/root/.cache/uv \ + UV_LINK_MODE=copy uv sync --frozen --all-groups --no-editable + +# Build chemical components database (this binary was installed by uv sync). +RUN uv run build_data # To work around a known XLA issue causing the compilation time to greatly # increase, the following environment variable setting XLA flags must be enabled @@ -77,4 +94,4 @@ ENV XLA_FLAGS="--xla_gpu_enable_triton_gemm=false" ENV XLA_PYTHON_CLIENT_PREALLOCATE=true ENV XLA_CLIENT_MEM_FRACTION=0.95 -CMD ["python3", "run_alphafold.py"] +CMD ["uv", "run", "python3", "run_alphafold.py"] diff --git a/docs/community_tools.md b/docs/community_tools.md new file mode 100644 index 0000000..c783e7b --- /dev/null +++ b/docs/community_tools.md @@ -0,0 +1,31 @@ +# Community Tools + +## JAAG: a JSON input file Assembler for AlphaFold 3 (with Glycan Integration) + +JAAG is a lightweight, web-based GUI tool that helps generate AlphaFold 3 input +JSON files with integrated glycan support. It automates the creation of correct +glycan syntax (including `bondedAtomPairs` + CCD), reducing manual errors when +preparing glycoprotein or glycan–protein complexes. + +* Web app: https://biofgreat.org/JAAG +* Source code: https://github.com/chinchc/JAAG +* Paper: https://doi.org/10.1093/glycob/cwaf083 + +Note: JAAG is compatible with standalone AlphaFold 3, but not with the AlphaFold +3 server. + +## Modeling glycans with AlphaFold 3: capabilities, caveats, and limitations + +Paper on modeling glycans (and other ligands) with AF3 that modeled and assessed +major glycan classes and provides: + +* Step-by-step tutorial for building ligand inputs (applicable beyond glycans) +* Ready-to-run scripts for each glycan class +* Comprehensive CCD table for all SNFG monosaccharides +* Discussion of caveats and limitations of AF3 +* Full AF3 inputs/outputs archived on ModelArchive for reproducibility + +Useful resource if your AF3 ligand models appear stereochemically off. + +* Paper: https://doi.org/10.1093/glycob/cwaf048 +* ModelArchive: https://doi.org/10.5452/ma-af3glycan diff --git a/docs/installation.md b/docs/installation.md index ac71ba3..1f91b26 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -285,6 +285,14 @@ python dependencies: docker build -t alphafold3 -f docker/Dockerfile . ``` +If you hit `No file descriptors available (os error 24)` on systems like +AlmaLinux/Rocky/RHEL, you need to manually expand the file descriptor limits +during the build by appending `--ulimit nofile=65535:65535`: + +```sh +docker build --ulimit nofile=65535:65535 -t alphafold3 -f docker/Dockerfile . +``` + Create an input JSON file, using either the example in the [README](https://github.com/google-deepmind/alphafold3?tab=readme-ov-file#installation-and-running-your-first-prediction) or a diff --git a/docs/known_issues.md b/docs/known_issues.md index 2f0be4f..5e745fb 100644 --- a/docs/known_issues.md +++ b/docs/known_issues.md @@ -13,3 +13,50 @@ Between commits https://github.com/google-deepmind/alphafold3/commit/f8df1c7 and https://github.com/google-deepmind/alphafold3/commit/4e4023c, AlphaFold 3 handled incorrectly any two-letter atoms (e.g. Cl, Br) in ligands defined using SMILES strings. + +## MSA discrepancy between AlphaFold 3 and AlphaFold Server + +### The root cause of the problem + +The released AlphaFold 3 and AlphaFold Server use the same model weights and +equivalent featurisation and model code. However, the way they run genetic +search is slightly different. The released AlphaFold 3 searches each database in +one go, while AlphaFold Server has a sharded version of each database (split +into multiple smaller FASTA files) and searches all of the shards in parallel. +The results of these parallel searches are then merged together at the end. + +The discrepancy is caused by a different (deeper) MSA on AlphaFold Server in +some cases. We discovered that the issue is caused by running sharded Jackhmmer +in AlphaFold Server without the `--domZ` flag (has to be set together with the +`--Z` flag and set to the same value) which means that effectively the AlphaFold +Server is running with roughly 100× more permissive `--domE` filter. This means +more sequences are sometimes included in the MSA. + +We are keeping behaviour unchanged in both the released AlphaFold 3 and in the +AlphaFold Server, however, we are giving users with local installs an option to +replicate AlphaFold Server behaviour locally. In our large scale tests the +difference did not matter, it is only very specific inputs that get better +accuracy with the deeper MSA. + +See https://github.com/google-deepmind/alphafold3/issues/492 for an example +input where a protein-DNA complex gets significantly higher ipTM and pTM with +AlphaFold Server compared to a local run. + +### Replicating AlphaFold Server behaviour locally + +If you want to replicate AlphaFold Server behaviour (i.e. better folding +accuracy in some cases), you can increase the value of the Jackhmmer/Nhmmer +`--domE` flag by 100× compared to its default value. + +Alternatively, you can run the sharded MSA search while not setting the `--domZ` +value – you would have to modify the code to do it. We added support for +searching against sharded databases in AlphaFold 3 in +https://github.com/google-deepmind/alphafold3/commit/805adc3863841d83d631ccd18136ad58ce3ecb34 +and the way to run AlphaFold 3 with sharded databases is documented in +https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md#sharded-genetic-databases. +It can provide 10–30× speedup (potentially even more, depending on hardware) of +the genetic search. + +In general, we recommend experimenting with MSA if you are seeing a prediction +with low predicted confidence. Typically adding more *relevant* sequences in the +MSA will increase AlphaFold prediction accuracy and model confidence scores. diff --git a/docs/model_parameters.md b/docs/model_parameters.md new file mode 100644 index 0000000..d0aabbf --- /dev/null +++ b/docs/model_parameters.md @@ -0,0 +1,439 @@ +# Model Parameters + +AlphaFold 3 layer names, shapes, and dtypes are documented in the table below. +This can be used for example to generate random parameters for AlphaFold 3 +performance optimisation on new accelerators without having to obtain the +official parameters. It is important to not generate zero-only parameters for +performance optimisations as accelerators often have shortcuts for zero-only +arguments (e.g. `0 * tensor` can be optimised to a no-op). + +Producing random parameters could be done similarly to the following snippet: + +```py +from alphafold3.model import params +import numpy as np +import zstandard + +parameters = ... # Data from the parameters schema. + +with zstandard.open('random_weights.bin.zst', 'wb') as compressed: + for scope_name, shape, dtype in parameters: + if scope_name == '__meta__:__identifier__': + # The identifier can be all zeros. + arr = np.zeros(shape=shape, dtype=dtype) + else: + # Do not use all-zero params, instead sample uniformly between -1 and 1. + arr = np.random.uniform(low=-1, high=1, size=shape).astype(dtype) + scope_name = scope_name.split(':') + compressed.write(params.encode_record(*scope_name, arr)) +``` + +## Parameters Schema + +``` +name=__meta__:__identifier__ dtype=uint8 shape=(64,) +name=diffuser/~/diffusion_head/diffusion_atom_features_layer_norm:scale dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/diffusion_atom_features_to_position_update:weights dtype=float32 shape=(128, 3) +name=diffuser/~/diffusion_head/diffusion_atom_positions_to_features:weights dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderadaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderadaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_adaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_adaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_single_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_single_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_single_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_single_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_transition1:weights dtype=float32 shape=(3, 128, 512) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderffw_transition2:weights dtype=float32 shape=(3, 256, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decodergating_query:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderk_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderksingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderksingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderksingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderksingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderq_projection:bias dtype=float32 shape=(3, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderq_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderqsingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderqsingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderqsingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderqsingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decodertransition2:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/__layer_stack_with_per_layer/diffusion_atom_transformer_decoderv_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/pair_input_layer_norm:scale dtype=float32 shape=(16,) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_decoder/pair_logits_projection:weights dtype=float32 shape=(16, 3, 4) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderadaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderadaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_adaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_adaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_single_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_single_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_single_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_single_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_transition1:weights dtype=float32 shape=(3, 128, 512) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderffw_transition2:weights dtype=float32 shape=(3, 256, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encodergating_query:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderk_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderksingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderksingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderksingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderksingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderq_projection:bias dtype=float32 shape=(3, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderq_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderqsingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderqsingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderqsingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderqsingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encodertransition2:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/__layer_stack_with_per_layer/diffusion_atom_transformer_encoderv_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/pair_input_layer_norm:scale dtype=float32 shape=(16,) +name=diffuser/~/diffusion_head/diffusion_atom_transformer_encoder/pair_logits_projection:weights dtype=float32 shape=(16, 3, 4) +name=diffuser/~/diffusion_head/diffusion_embed_pair_distances_1:weights dtype=float32 shape=(1, 16) +name=diffuser/~/diffusion_head/diffusion_embed_pair_distances:weights dtype=float32 shape=(1, 16) +name=diffuser/~/diffusion_head/diffusion_embed_pair_offsets_1:weights dtype=float32 shape=(3, 16) +name=diffuser/~/diffusion_head/diffusion_embed_pair_offsets_valid:weights dtype=float32 shape=(1, 16) +name=diffuser/~/diffusion_head/diffusion_embed_pair_offsets:weights dtype=float32 shape=(3, 16) +name=diffuser/~/diffusion_head/diffusion_embed_ref_atom_name:weights dtype=float32 shape=(256, 128) +name=diffuser/~/diffusion_head/diffusion_embed_ref_charge:weights dtype=float32 shape=(1, 128) +name=diffuser/~/diffusion_head/diffusion_embed_ref_element:weights dtype=float32 shape=(128, 128) +name=diffuser/~/diffusion_head/diffusion_embed_ref_mask:weights dtype=float32 shape=(1, 128) +name=diffuser/~/diffusion_head/diffusion_embed_ref_pos:weights dtype=float32 shape=(3, 128) +name=diffuser/~/diffusion_head/diffusion_embed_trunk_pair_cond:weights dtype=float32 shape=(128, 16) +name=diffuser/~/diffusion_head/diffusion_embed_trunk_single_cond:weights dtype=float32 shape=(384, 128) +name=diffuser/~/diffusion_head/diffusion_lnorm_trunk_pair_cond:scale dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/diffusion_lnorm_trunk_single_cond:scale dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/diffusion_pair_mlp_1:weights dtype=float32 shape=(16, 16) +name=diffuser/~/diffusion_head/diffusion_pair_mlp_2:weights dtype=float32 shape=(16, 16) +name=diffuser/~/diffusion_head/diffusion_pair_mlp_3:weights dtype=float32 shape=(16, 16) +name=diffuser/~/diffusion_head/diffusion_project_atom_features_for_aggr:weights dtype=float32 shape=(128, 768) +name=diffuser/~/diffusion_head/diffusion_project_token_features_for_broadcast:weights dtype=float32 shape=(768, 128) +name=diffuser/~/diffusion_head/diffusion_single_to_pair_cond_col_1:weights dtype=float32 shape=(128, 16) +name=diffuser/~/diffusion_head/diffusion_single_to_pair_cond_col:weights dtype=float32 shape=(128, 16) +name=diffuser/~/diffusion_head/diffusion_single_to_pair_cond_row_1:weights dtype=float32 shape=(128, 16) +name=diffuser/~/diffusion_head/diffusion_single_to_pair_cond_row:weights dtype=float32 shape=(128, 16) +name=diffuser/~/diffusion_head/noise_embedding_initial_norm:scale dtype=float32 shape=(256,) +name=diffuser/~/diffusion_head/noise_embedding_initial_projection:weights dtype=float32 shape=(256, 384) +name=diffuser/~/diffusion_head/output_norm:scale dtype=float32 shape=(768,) +name=diffuser/~/diffusion_head/pair_cond_initial_norm:scale dtype=float32 shape=(267,) +name=diffuser/~/diffusion_head/pair_cond_initial_projection:weights dtype=float32 shape=(267, 128) +name=diffuser/~/diffusion_head/pair_transition_0ffw_layer_norm:offset dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/pair_transition_0ffw_layer_norm:scale dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/pair_transition_0ffw_transition1:weights dtype=float32 shape=(128, 512) +name=diffuser/~/diffusion_head/pair_transition_0ffw_transition2:weights dtype=float32 shape=(256, 128) +name=diffuser/~/diffusion_head/pair_transition_1ffw_layer_norm:offset dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/pair_transition_1ffw_layer_norm:scale dtype=float32 shape=(128,) +name=diffuser/~/diffusion_head/pair_transition_1ffw_transition1:weights dtype=float32 shape=(128, 512) +name=diffuser/~/diffusion_head/pair_transition_1ffw_transition2:weights dtype=float32 shape=(256, 128) +name=diffuser/~/diffusion_head/single_cond_embedding_norm:scale dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/single_cond_embedding_projection:weights dtype=float32 shape=(384, 768) +name=diffuser/~/diffusion_head/single_cond_initial_norm:scale dtype=float32 shape=(831,) +name=diffuser/~/diffusion_head/single_cond_initial_projection:weights dtype=float32 shape=(831, 384) +name=diffuser/~/diffusion_head/single_transition_0ffw_layer_norm:offset dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/single_transition_0ffw_layer_norm:scale dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/single_transition_0ffw_transition1:weights dtype=float32 shape=(384, 1536) +name=diffuser/~/diffusion_head/single_transition_0ffw_transition2:weights dtype=float32 shape=(768, 384) +name=diffuser/~/diffusion_head/single_transition_1ffw_layer_norm:offset dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/single_transition_1ffw_layer_norm:scale dtype=float32 shape=(384,) +name=diffuser/~/diffusion_head/single_transition_1ffw_transition1:weights dtype=float32 shape=(384, 1536) +name=diffuser/~/diffusion_head/single_transition_1ffw_transition2:weights dtype=float32 shape=(768, 384) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformeradaptive_zero_cond:bias dtype=float32 shape=(6, 4, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformeradaptive_zero_cond:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_adaptive_zero_cond:bias dtype=float32 shape=(6, 4, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_adaptive_zero_cond:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_single_cond_bias:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_single_cond_layer_norm:scale dtype=float32 shape=(6, 4, 384) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_single_cond_scale:bias dtype=float32 shape=(6, 4, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_single_cond_scale:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_transition1:weights dtype=float32 shape=(6, 4, 768, 3072) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerffw_transition2:weights dtype=float32 shape=(6, 4, 1536, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformergating_query:weights dtype=float32 shape=(6, 4, 768, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerk_projection:weights dtype=float32 shape=(6, 4, 768, 16, 48) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerq_projection:bias dtype=float32 shape=(6, 4, 16, 48) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerq_projection:weights dtype=float32 shape=(6, 4, 768, 16, 48) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformersingle_cond_bias:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformersingle_cond_layer_norm:scale dtype=float32 shape=(6, 4, 384) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformersingle_cond_scale:bias dtype=float32 shape=(6, 4, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformersingle_cond_scale:weights dtype=float32 shape=(6, 4, 384, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformertransition2:weights dtype=float32 shape=(6, 4, 768, 768) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/__layer_stack_with_per_layer/transformerv_projection:weights dtype=float32 shape=(6, 4, 768, 16, 48) +name=diffuser/~/diffusion_head/transformer/__layer_stack_with_per_layer/pair_logits_projection:weights dtype=float32 shape=(6, 128, 4, 16) +name=diffuser/~/diffusion_head/transformer/pair_input_layer_norm:scale dtype=float32 shape=(128,) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/act_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/act_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/gating_query:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/k_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/pair_bias_projection:weights dtype=bfloat16 shape=(4, 128, 4) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/q_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention1/v_projection:weights dtype=bfloat16 shape=(4, 128, 4, 32) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/act_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/act_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/gating_query:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/k_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/pair_bias_projection:weights dtype=bfloat16 shape=(4, 128, 4) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/q_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_attention2/v_projection:weights dtype=bfloat16 shape=(4, 128, 4, 32) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_transition/input_layer_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_transition/input_layer_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_transition/transition1:weights dtype=bfloat16 shape=(4, 128, 1024) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/pair_transition/transition2:weights dtype=bfloat16 shape=(4, 512, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_gating_query:weights dtype=bfloat16 shape=(4, 384, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_k_projection:weights dtype=bfloat16 shape=(4, 384, 16, 24) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_layer_norm:offset dtype=float32 shape=(4, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_layer_norm:scale dtype=float32 shape=(4, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_q_projection:bias dtype=bfloat16 shape=(4, 16, 24) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_q_projection:weights dtype=bfloat16 shape=(4, 384, 16, 24) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_transition2:weights dtype=bfloat16 shape=(4, 384, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_attention_v_projection:weights dtype=bfloat16 shape=(4, 384, 16, 24) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_pair_logits_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_pair_logits_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_pair_logits_projection:weights dtype=bfloat16 shape=(4, 128, 16) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_transition/input_layer_norm:offset dtype=float32 shape=(4, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_transition/input_layer_norm:scale dtype=float32 shape=(4, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_transition/transition1:weights dtype=bfloat16 shape=(4, 384, 3072) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/single_transition/transition2:weights dtype=bfloat16 shape=(4, 1536, 384) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/center_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/center_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/gate:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/gating_linear:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/left_norm_input:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/left_norm_input:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_incoming/projection:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/center_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/center_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/gate:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/gating_linear:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/left_norm_input:offset dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/left_norm_input:scale dtype=float32 shape=(4, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/confidence_head/__layer_stack_no_per_layer/confidence_pairformer/triangle_multiplication_outgoing/projection:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/confidence_head/~_embed_features/distogram_feat_project:weights dtype=bfloat16 shape=(39, 128) +name=diffuser/confidence_head/~_embed_features/left_target_feat_project:weights dtype=bfloat16 shape=(447, 128) +name=diffuser/confidence_head/~_embed_features/right_target_feat_project:weights dtype=bfloat16 shape=(447, 128) +name=diffuser/confidence_head/experimentally_resolved_ln:offset dtype=float32 shape=(384,) +name=diffuser/confidence_head/experimentally_resolved_ln:scale dtype=float32 shape=(384,) +name=diffuser/confidence_head/experimentally_resolved_logits:weights dtype=float32 shape=(384, 24, 2) +name=diffuser/confidence_head/left_half_distance_logits:weights dtype=float32 shape=(128, 64) +name=diffuser/confidence_head/logits_ln:offset dtype=float32 shape=(128,) +name=diffuser/confidence_head/logits_ln:scale dtype=float32 shape=(128,) +name=diffuser/confidence_head/pae_logits_ln:offset dtype=float32 shape=(128,) +name=diffuser/confidence_head/pae_logits_ln:scale dtype=float32 shape=(128,) +name=diffuser/confidence_head/pae_logits:weights dtype=float32 shape=(128, 64) +name=diffuser/confidence_head/plddt_logits_ln:offset dtype=float32 shape=(384,) +name=diffuser/confidence_head/plddt_logits_ln:scale dtype=float32 shape=(384,) +name=diffuser/confidence_head/plddt_logits:weights dtype=float32 shape=(384, 24, 50) +name=diffuser/distogram_head/half_logits:weights dtype=float32 shape=(128, 64) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderadaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderadaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_adaptive_zero_cond:bias dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_adaptive_zero_cond:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_single_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_single_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_single_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_single_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_transition1:weights dtype=float32 shape=(3, 128, 512) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderffw_transition2:weights dtype=float32 shape=(3, 256, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encodergating_query:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderk_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderksingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderksingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderksingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderksingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderq_projection:bias dtype=float32 shape=(3, 4, 32) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderq_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderqsingle_cond_bias:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderqsingle_cond_layer_norm:scale dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderqsingle_cond_scale:bias dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderqsingle_cond_scale:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encodertransition2:weights dtype=float32 shape=(3, 128, 128) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/__layer_stack_with_per_layer/evoformer_conditioning_atom_transformer_encoderv_projection:weights dtype=float32 shape=(3, 128, 4, 32) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/pair_input_layer_norm:scale dtype=float32 shape=(16,) +name=diffuser/evoformer_conditioning_atom_transformer_encoder/pair_logits_projection:weights dtype=float32 shape=(16, 3, 4) +name=diffuser/evoformer_conditioning_embed_pair_distances_1:weights dtype=float32 shape=(1, 16) +name=diffuser/evoformer_conditioning_embed_pair_distances:weights dtype=float32 shape=(1, 16) +name=diffuser/evoformer_conditioning_embed_pair_offsets_1:weights dtype=float32 shape=(3, 16) +name=diffuser/evoformer_conditioning_embed_pair_offsets_valid:weights dtype=float32 shape=(1, 16) +name=diffuser/evoformer_conditioning_embed_pair_offsets:weights dtype=float32 shape=(3, 16) +name=diffuser/evoformer_conditioning_embed_ref_atom_name:weights dtype=float32 shape=(256, 128) +name=diffuser/evoformer_conditioning_embed_ref_charge:weights dtype=float32 shape=(1, 128) +name=diffuser/evoformer_conditioning_embed_ref_element:weights dtype=float32 shape=(128, 128) +name=diffuser/evoformer_conditioning_embed_ref_mask:weights dtype=float32 shape=(1, 128) +name=diffuser/evoformer_conditioning_embed_ref_pos:weights dtype=float32 shape=(3, 128) +name=diffuser/evoformer_conditioning_pair_mlp_1:weights dtype=float32 shape=(16, 16) +name=diffuser/evoformer_conditioning_pair_mlp_2:weights dtype=float32 shape=(16, 16) +name=diffuser/evoformer_conditioning_pair_mlp_3:weights dtype=float32 shape=(16, 16) +name=diffuser/evoformer_conditioning_project_atom_features_for_aggr:weights dtype=float32 shape=(128, 384) +name=diffuser/evoformer_conditioning_single_to_pair_cond_col_1:weights dtype=float32 shape=(128, 16) +name=diffuser/evoformer_conditioning_single_to_pair_cond_col:weights dtype=float32 shape=(128, 16) +name=diffuser/evoformer_conditioning_single_to_pair_cond_row_1:weights dtype=float32 shape=(128, 16) +name=diffuser/evoformer_conditioning_single_to_pair_cond_row:weights dtype=float32 shape=(128, 16) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/act_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/act_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/gating_query:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/k_projection:weights dtype=bfloat16 shape=(48, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/output_projection:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/pair_bias_projection:weights dtype=bfloat16 shape=(48, 128, 4) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/q_projection:weights dtype=bfloat16 shape=(48, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention1/v_projection:weights dtype=bfloat16 shape=(48, 128, 4, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/act_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/act_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/gating_query:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/k_projection:weights dtype=bfloat16 shape=(48, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/output_projection:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/pair_bias_projection:weights dtype=bfloat16 shape=(48, 128, 4) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/q_projection:weights dtype=bfloat16 shape=(48, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_attention2/v_projection:weights dtype=bfloat16 shape=(48, 128, 4, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_transition/input_layer_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_transition/input_layer_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_transition/transition1:weights dtype=bfloat16 shape=(48, 128, 1024) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/pair_transition/transition2:weights dtype=bfloat16 shape=(48, 512, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_gating_query:weights dtype=bfloat16 shape=(48, 384, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_k_projection:weights dtype=bfloat16 shape=(48, 384, 16, 24) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_layer_norm:offset dtype=float32 shape=(48, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_layer_norm:scale dtype=float32 shape=(48, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_q_projection:bias dtype=bfloat16 shape=(48, 16, 24) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_q_projection:weights dtype=bfloat16 shape=(48, 384, 16, 24) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_transition2:weights dtype=bfloat16 shape=(48, 384, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_attention_v_projection:weights dtype=bfloat16 shape=(48, 384, 16, 24) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_pair_logits_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_pair_logits_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_pair_logits_projection:weights dtype=bfloat16 shape=(48, 128, 16) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_transition/input_layer_norm:offset dtype=float32 shape=(48, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_transition/input_layer_norm:scale dtype=float32 shape=(48, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_transition/transition1:weights dtype=bfloat16 shape=(48, 384, 3072) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/single_transition/transition2:weights dtype=bfloat16 shape=(48, 1536, 384) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/center_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/center_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/gate:weights dtype=bfloat16 shape=(48, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/gating_linear:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/left_norm_input:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/left_norm_input:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/output_projection:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_incoming/projection:weights dtype=bfloat16 shape=(48, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/center_norm:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/center_norm:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/gate:weights dtype=bfloat16 shape=(48, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/gating_linear:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/left_norm_input:offset dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/left_norm_input:scale dtype=float32 shape=(48, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/output_projection:weights dtype=bfloat16 shape=(48, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer_1/trunk_pairformer/triangle_multiplication_outgoing/projection:weights dtype=bfloat16 shape=(48, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/act_norm:offset dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/act_norm:scale dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/gating_query:weights dtype=bfloat16 shape=(4, 64, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/output_projection:weights dtype=bfloat16 shape=(4, 64, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/pair_logits:weights dtype=bfloat16 shape=(4, 128, 8) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/pair_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/pair_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_attention1/v_projection:weights dtype=bfloat16 shape=(4, 64, 8, 8) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_transition/input_layer_norm:offset dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_transition/input_layer_norm:scale dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_transition/transition1:weights dtype=bfloat16 shape=(4, 64, 512) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/msa_transition/transition2:weights dtype=bfloat16 shape=(4, 256, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean:output_b dtype=bfloat16 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean:output_w dtype=bfloat16 shape=(4, 32, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean/layer_norm_input:offset dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean/layer_norm_input:scale dtype=float32 shape=(4, 64) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean/left_projection:weights dtype=bfloat16 shape=(4, 64, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/outer_product_mean/right_projection:weights dtype=bfloat16 shape=(4, 64, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/act_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/act_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/gating_query:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/k_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/pair_bias_projection:weights dtype=bfloat16 shape=(4, 128, 4) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/q_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention1/v_projection:weights dtype=bfloat16 shape=(4, 128, 4, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/act_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/act_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/gating_query:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/k_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/pair_bias_projection:weights dtype=bfloat16 shape=(4, 128, 4) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/q_projection:weights dtype=bfloat16 shape=(4, 4, 32, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_attention2/v_projection:weights dtype=bfloat16 shape=(4, 128, 4, 32) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_transition/input_layer_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_transition/input_layer_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_transition/transition1:weights dtype=bfloat16 shape=(4, 128, 1024) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/pair_transition/transition2:weights dtype=bfloat16 shape=(4, 512, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/center_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/center_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/gate:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/gating_linear:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/left_norm_input:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/left_norm_input:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_incoming/projection:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/center_norm:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/center_norm:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/gate:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/gating_linear:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/left_norm_input:offset dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/left_norm_input:scale dtype=float32 shape=(4, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/output_projection:weights dtype=bfloat16 shape=(4, 128, 128) +name=diffuser/evoformer/__layer_stack_no_per_layer/msa_stack/triangle_multiplication_outgoing/projection:weights dtype=bfloat16 shape=(4, 128, 256) +name=diffuser/evoformer/~_relative_encoding/position_activations:weights dtype=bfloat16 shape=(139, 128) +name=diffuser/evoformer/bond_embedding:weights dtype=bfloat16 shape=(1, 128) +name=diffuser/evoformer/extra_msa_target_feat:weights dtype=bfloat16 shape=(447, 64) +name=diffuser/evoformer/left_single:weights dtype=bfloat16 shape=(447, 128) +name=diffuser/evoformer/msa_activations:weights dtype=bfloat16 shape=(34, 64) +name=diffuser/evoformer/prev_embedding_layer_norm:offset dtype=float32 shape=(128,) +name=diffuser/evoformer/prev_embedding_layer_norm:scale dtype=float32 shape=(128,) +name=diffuser/evoformer/prev_embedding:weights dtype=bfloat16 shape=(128, 128) +name=diffuser/evoformer/prev_single_embedding_layer_norm:offset dtype=float32 shape=(384,) +name=diffuser/evoformer/prev_single_embedding_layer_norm:scale dtype=float32 shape=(384,) +name=diffuser/evoformer/prev_single_embedding:weights dtype=bfloat16 shape=(384, 384) +name=diffuser/evoformer/right_single:weights dtype=bfloat16 shape=(447, 128) +name=diffuser/evoformer/single_activations:weights dtype=bfloat16 shape=(447, 384) +name=diffuser/evoformer/template_embedding/output_linear:weights dtype=bfloat16 shape=(64, 128) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/act_norm:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/act_norm:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/gating_query:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/k_projection:weights dtype=bfloat16 shape=(2, 4, 16, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/output_projection:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/pair_bias_projection:weights dtype=bfloat16 shape=(2, 64, 4) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/q_projection:weights dtype=bfloat16 shape=(2, 4, 16, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention1/v_projection:weights dtype=bfloat16 shape=(2, 64, 4, 16) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/act_norm:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/act_norm:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/gating_query:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/k_projection:weights dtype=bfloat16 shape=(2, 4, 16, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/output_projection:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/pair_bias_projection:weights dtype=bfloat16 shape=(2, 64, 4) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/q_projection:weights dtype=bfloat16 shape=(2, 4, 16, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_attention2/v_projection:weights dtype=bfloat16 shape=(2, 64, 4, 16) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_transition/input_layer_norm:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_transition/input_layer_norm:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_transition/transition1:weights dtype=bfloat16 shape=(2, 64, 256) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/pair_transition/transition2:weights dtype=bfloat16 shape=(2, 128, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/center_norm:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/center_norm:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/gate:weights dtype=bfloat16 shape=(2, 64, 128) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/gating_linear:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/left_norm_input:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/left_norm_input:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/output_projection:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_incoming/projection:weights dtype=bfloat16 shape=(2, 64, 128) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/center_norm:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/center_norm:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/gate:weights dtype=bfloat16 shape=(2, 64, 128) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/gating_linear:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/left_norm_input:offset dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/left_norm_input:scale dtype=float32 shape=(2, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/output_projection:weights dtype=bfloat16 shape=(2, 64, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/__layer_stack_no_per_layer/template_embedding_iteration/triangle_multiplication_outgoing/projection:weights dtype=bfloat16 shape=(2, 64, 128) +name=diffuser/evoformer/template_embedding/single_template_embedding/output_layer_norm:offset dtype=float32 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/output_layer_norm:scale dtype=float32 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/query_embedding_norm:offset dtype=float32 shape=(128,) +name=diffuser/evoformer/template_embedding/single_template_embedding/query_embedding_norm:scale dtype=float32 shape=(128,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_0:weights dtype=bfloat16 shape=(39, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_1:weights dtype=bfloat16 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_2:weights dtype=bfloat16 shape=(31, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_3:weights dtype=bfloat16 shape=(31, 64) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_4:weights dtype=bfloat16 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_5:weights dtype=bfloat16 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_6:weights dtype=bfloat16 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_7:weights dtype=bfloat16 shape=(64,) +name=diffuser/evoformer/template_embedding/single_template_embedding/template_pair_embedding_8:weights dtype=bfloat16 shape=(128, 64) +``` diff --git a/docs/performance.md b/docs/performance.md index 55caf7d..d8ac9f4 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -18,11 +18,11 @@ utilisation. This can be useful for: ### Data Pipeline Only -Launch `run_alphafold.py` with `--norun_inference` to generate Multiple Sequence -Alignments (MSAs) and templates, without running featurisation and model -inference. This stage can be quite costly in terms of runtime, CPU, and RAM use. -The output will be JSON files augmented with MSAs and templates that can then be -directly used as input for running inference. +Launch `run_alphafold.py` with `--run_inference=false` to generate Multiple +Sequence Alignments (MSAs) and templates, without running featurisation and +model inference. This stage can be quite costly in terms of runtime, CPU, and +RAM use. The output will be JSON files augmented with MSAs and templates that +can then be directly used as input for running inference. ### Pre-computing and reusing MSA and templates @@ -62,10 +62,10 @@ similar issues). ### Featurisation and Model Inference Only -Launch `run_alphafold.py` with `--norun_data_pipeline` to skip the data pipeline -and run only featurisation and model inference. This stage requires the input -JSON file to contain pre-computed MSAs and templates (or they must be explicitly -set to empty if you want to run MSA and template free). +Launch `run_alphafold.py` with `--run_data_pipeline=false` to skip the data +pipeline and run only featurisation and model inference. This stage requires the +input JSON file to contain pre-computed MSAs and templates (or they must be +explicitly set to empty if you want to run MSA and template free). ## Data Pipeline diff --git a/fetch_databases.sh b/fetch_databases.sh index 0373abc..3fa2188 100644 --- a/fetch_databases.sh +++ b/fetch_databases.sh @@ -1,8 +1,17 @@ #!/bin/bash # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -16,6 +25,7 @@ readonly db_dir=${1:-$HOME/public_databases} for cmd in wget tar zstd ; do if ! command -v "${cmd}" > /dev/null 2>&1; then echo "${cmd} is not installed. Please install it." + exit 1 fi done @@ -32,7 +42,8 @@ wget --quiet --output-document=- \ for NAME in mgy_clusters_2022_05.fa \ bfd-first_non_consensus_sequences.fasta \ - uniref90_2022_05.fa uniprot_all_2021_04.fa \ + uniref90_2022_05.fa \ + uniprot_all_2021_04.fa \ pdb_seqres_2022_09_28.fasta \ rnacentral_active_seq_id_90_cov_80_linclust.fasta \ nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fasta \ diff --git a/legal/WEIGHTS_TERMS_OF_USE-Bahasa-Indonesia.md b/legal/WEIGHTS_TERMS_OF_USE-Bahasa-Indonesia.md index c39f4de..0d57e19 100644 --- a/legal/WEIGHTS_TERMS_OF_USE-Bahasa-Indonesia.md +++ b/legal/WEIGHTS_TERMS_OF_USE-Bahasa-Indonesia.md @@ -67,9 +67,8 @@ Sebagaimana digunakan dalam Persyaratan ini: "**AlphaFold 3**" adalah: (a) kode sumber AlphaFold 3 yang disediakan [di sini](https://github.com/google-deepmind/alphafold3/) dan yang dilisensikan -berdasarkan persyaratan lisensi Creative Commons -Attribution-NonCommercial-Sharealike 4.0 International (CC-BY-NC-SA 4.0) dan -kode sumber turunan apa pun, serta (b) Parameter Model. +berdasarkan persyaratan lisensi Apache 2.0 dan kode sumber turunan apa pun, +serta (b) Parameter Model. "**Aset AlphaFold 3**" adalah Parameter dan output Model. diff --git a/legal/WEIGHTS_TERMS_OF_USE-Espanol-Latinoamerica.md b/legal/WEIGHTS_TERMS_OF_USE-Espanol-Latinoamerica.md index e2cbd24..94f9d21 100644 --- a/legal/WEIGHTS_TERMS_OF_USE-Espanol-Latinoamerica.md +++ b/legal/WEIGHTS_TERMS_OF_USE-Espanol-Latinoamerica.md @@ -64,11 +64,10 @@ Google LLC y sus afiliadas. Según su uso en estas Condiciones: -"**AlphaFold 3**" significa: (a) el código fuente de AlphaFold 3 disponible -[aquí](https://github.com/google-deepmind/alphafold3/) y con licencia en virtud -de las condiciones de la Atribución/Reconocimiento-NoComercial-CompartirIgual -4.0 Internacional (CC-BY-NC-SA 4.0) de Creative Commons, y cualquier código -fuente derivado, y (b) los Parámetros del Modelo. +Por "**AlphaFold 3**" se entiende: (a) el código fuente de AlphaFold 3 puesto a +disposición [aquí](https://github.com/google-deepmind/alphafold3/) y licenciado +bajo los términos de la licencia Apache 2.0, así como cualquier código fuente +derivado; y (b) los Parámetros del Modelo. "**Recursos de AlphaFold 3**" hace referencia a los Resultados y los Parámetros del Modelo. diff --git a/legal/WEIGHTS_TERMS_OF_USE-Francais-Canada.md b/legal/WEIGHTS_TERMS_OF_USE-Francais-Canada.md index bacdc18..57d55b2 100644 --- a/legal/WEIGHTS_TERMS_OF_USE-Francais-Canada.md +++ b/legal/WEIGHTS_TERMS_OF_USE-Francais-Canada.md @@ -69,9 +69,8 @@ c'est-à-dire Google LLC et ses filiales. Telle qu’utilisées dans ces Conditions: \ « **AlphaFold 3** » désigne: (a) le code source d'AlphaFold 3 rendu accessible [ici](https://github.com/google-deepmind/alphafold3/) et sous les conditions de -la licence « Creative Commons Attribution-NonCommercial-Sharealike 4.0 -International (CC-BY-NC-SA 4.0) » ainsi que tout code source d'œuvres dérivées -et (b) les Paramètres du modèle. +la licence « Apache 2.0 » ainsi que tout code source d'œuvres dérivées et (b) +les Paramètres du modèle. « **Éléments d'AlphaFold 3** » signifie les Paramètres du modèle et les Résultats. diff --git a/legal/WEIGHTS_TERMS_OF_USE-Portugues-Brazil.md b/legal/WEIGHTS_TERMS_OF_USE-Portugues-Brazil.md index 2d049da..448241d 100644 --- a/legal/WEIGHTS_TERMS_OF_USE-Portugues-Brazil.md +++ b/legal/WEIGHTS_TERMS_OF_USE-Portugues-Brazil.md @@ -66,8 +66,7 @@ Conforme usado nestes Termos: "**AlphaFold 3**" significa: (a) o código-fonte do AlphaFold 3 disponível [neste link](https://github.com/google-deepmind/alphafold3/) e licenciado nos -termos da licença Creative Commons Attribution-NonCommercial-Sharealike 4.0 -International (CC-BY-NC-SA 4.0), bem como qualquer código-fonte derivado, e (b) +termos da licença Apache 2.0, bem como qualquer código-fonte derivado, e (b) Parâmetros do modelo. "**Recursos do AlphaFold 3**" significam as Saídas e os Parâmetros do modelo. diff --git a/pyproject.toml b/pyproject.toml index 45412b6..b1b8e47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,28 +10,33 @@ build-backend = "scikit_build_core.build" [project] name = "alphafold3" -version = "3.0.1" -requires-python = ">=3.11" +dynamic = ["version"] +requires-python = ">=3.12" readme = "README.md" license = {file = "LICENSE"} dependencies = [ - "absl-py", - "dm-haiku==0.0.13", - "dm-tree", - "jax==0.4.34", - "jax[cuda12]==0.4.34", - "jax-triton==0.2.0", - "jaxtyping==0.2.34", + "absl-py>=2.3.1", + "dm-haiku==0.0.16", + "jax==0.9.1", + "jax[cuda12]==0.9.1", "numpy", - "rdkit==2024.3.5", - "triton==3.1.0", + "rdkit==2025.9.4", + "tokamax==0.0.11", "tqdm", - "typeguard==2.13.3", "zstandard", ] -[project.optional-dependencies] -test = ["pytest>=6.0"] +[dependency-groups] +dev = [ + "pytest>=6.0", +] + +[tool.uv] +package = true +environments = [ + "sys_platform == 'linux' and platform_machine == 'x86_64'", + "sys_platform == 'linux' and platform_machine == 'aarch64'", +] [tool.scikit-build] wheel.exclude = [ @@ -46,6 +51,10 @@ sdist.include = [ "WEIGHTS_PROHIBITED_USE_POLICY.md", "WEIGHTS_TERMS_OF_USE.md", ] +metadata.version.provider = "scikit_build_core.metadata.setuptools_scm" + +[tool.setuptools_scm] +fallback_version = "3.0.2" [tool.cibuildwheel] build = "cp3*-manylinux_x86_64" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 9b94e0b..0000000 --- a/requirements.txt +++ /dev/null @@ -1,538 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.12 -# by the following command: -# -# pip-compile --generate-hashes -# -absl-py==2.1.0 \ - --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \ - --hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff - # via - # alphafold3 (pyproject.toml) - # dm-haiku - # jax-triton -dm-haiku==0.0.13 \ - --hash=sha256:029bb91b5b1edb0d3fe23304d3bf12a545ea6e485041f7f5d8c8d85ebcf6e17d \ - --hash=sha256:ee9562c68a059f146ad07f555ca591cb8c11ef751afecc38353863562bd23f43 - # via alphafold3 (pyproject.toml) -dm-tree==0.1.8 \ - --hash=sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6 \ - --hash=sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760 \ - --hash=sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c \ - --hash=sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf \ - --hash=sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430 \ - --hash=sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de \ - --hash=sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317 \ - --hash=sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca \ - --hash=sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913 \ - --hash=sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf \ - --hash=sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef \ - --hash=sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426 \ - --hash=sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1 \ - --hash=sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e \ - --hash=sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60 \ - --hash=sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5 \ - --hash=sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f \ - --hash=sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b \ - --hash=sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410 \ - --hash=sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134 \ - --hash=sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb \ - --hash=sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b \ - --hash=sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7 \ - --hash=sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393 \ - --hash=sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571 \ - --hash=sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368 \ - --hash=sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22 \ - --hash=sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715 \ - --hash=sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80 \ - --hash=sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7 \ - --hash=sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d \ - --hash=sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a \ - --hash=sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d \ - --hash=sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6 \ - --hash=sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5 \ - --hash=sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e \ - --hash=sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68 \ - --hash=sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8 \ - --hash=sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f \ - --hash=sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436 \ - --hash=sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee \ - --hash=sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb \ - --hash=sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144 \ - --hash=sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8 \ - --hash=sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb \ - --hash=sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d - # via alphafold3 (pyproject.toml) -filelock==3.16.1 \ - --hash=sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0 \ - --hash=sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435 - # via triton -jax[cuda12]==0.4.34 \ - --hash=sha256:44196854f40c5f9cea3142824b9f1051f85afc3fcf7593ec5479fc8db01c58db \ - --hash=sha256:b957ca1fc91f7343f91a186af9f19c7f342c946f95a8c11c7f1e5cdfe2e58d9e - # via - # alphafold3 (pyproject.toml) - # jax-triton -jax-cuda12-pjrt==0.4.34 \ - --hash=sha256:0c7cc98f962cc7fc8e0a5ea6331b42a0cee516f202f1c3019f6aa5cd9530cca0 \ - --hash=sha256:6247459827550e7d11480f5707d7d2843cbf57bcda5d833565493456075cd143 - # via jax-cuda12-plugin -jax-cuda12-plugin[with-cuda]==0.4.34 \ - --hash=sha256:15fa8ff678fad34df424e01f1a9151abc966700670d3b22f52b87a2740028602 \ - --hash=sha256:1b037a408b9d9de59367c9513a9445b705da526657a0c4d593e84234312261e7 \ - --hash=sha256:b2099a4407225122ff76f6dcdc8dbdae47e6f29343bdfd21460ad337dc34a209 \ - --hash=sha256:d035ea72bd9b8a65a6ea621bca1affdd33127fa3a52e7bded7692670d360adab \ - --hash=sha256:dadb507949b70ba148567219bf0e7752eb63baec551f7ec860f8fcbb4ee48c48 \ - --hash=sha256:db988b7ba5063483a936ddbf162f04d1b4412e0d64340f11788c7bbc877e8a43 \ - --hash=sha256:e23721d1654b311b47cd6b35768520284bd036f8c7e6b11600143258b4a0409a \ - --hash=sha256:fb72bda44ae9f8a1e5ee9c5f82e2d273abe019874b0110b5c2d8395f1fe966b4 - # via jax -jax-triton==0.2.0 \ - --hash=sha256:6db5f6147327f2462c35becab612d415cf9ee70c96e43020c3c7ea873eb8576e \ - --hash=sha256:ed564a5ffb9e404557dc8d296e7eb30e501da5a4d7b03408acd0837c1c618c21 - # via alphafold3 (pyproject.toml) -jaxlib==0.4.34 \ - --hash=sha256:096f0ca309d41fa692a9d1f2f9baab1c5c8ca0749876ebb3f748e738a27c7ff4 \ - --hash=sha256:133070d4fec5525ffea4dc72956398c1cf647a04dcb37f8a935ee82af78d9965 \ - --hash=sha256:1a30771d85fa77f9ab8f18e63240f455ab3a3f87660ed7b8d5eea6ceecbe5c1e \ - --hash=sha256:3bcfa639ca3cfaf86c8ceebd5fc0d47300fd98a078014a1d0cc03133e1523d5f \ - --hash=sha256:3e60bc826933082e99b19b87c21818a8d26fcdb01f418d47cedff554746fd6cc \ - --hash=sha256:45d719a2ce0ebf21255a277b71d756f3609b7b5be70cddc5d88fd58c35219de0 \ - --hash=sha256:48272e9034ff868d4328cf0055a07882fd2be93f59dfb6283af7de491f9d1290 \ - --hash=sha256:571ef03259835458111596a71a2f4a6fabf4ec34595df4cea555035362ac5bf0 \ - --hash=sha256:6b43a974c5d91a19912d138f2658dd8dbb7d30dcdff5c961d896c673e872b611 \ - --hash=sha256:72e22e99a5dc890a64443c3fc12f13f20091f578c405a76de077ba42b4c62cd7 \ - --hash=sha256:7be673a876ebd1aef440fb7e3ebaf99a91abeb550c9728c644b7d7c7b5d7c108 \ - --hash=sha256:87f25a477cd279840e53718403f97092eba0e8a945fcab47bcf435b6f9119dda \ - --hash=sha256:8ee3f93836e53c86556ccd9449a4ea43516ee05184d031a71dd692e81259f7d9 \ - --hash=sha256:901cb4040ed24eae40071d8114ea8d10dff436277fa74a1a5b9e7206f641151c \ - --hash=sha256:b0001c8f0e2b1c7bc99e4f314b524a340d25653505c1a1484d4041a9d3617f6f \ - --hash=sha256:b7a212a3cb5c6acc201c32ae4f4b5f5a9ac09457fbb77ba8db5ce7e7d4adc214 \ - --hash=sha256:c303f5acaf6c56ce5ff133a923c9b6247bdebedde15bd2c893c24be4d8f71306 \ - --hash=sha256:c7b3e724a30426a856070aba0192b5d199e95b4411070e7ad96ad8b196877b10 \ - --hash=sha256:c9d3adcae43a33aad4332be9c2aedc5ef751d1e755f917a5afb30c7872eacaa8 \ - --hash=sha256:d840e64b85f8865404d6d225b9bb340e158df1457152a361b05680e24792b232 - # via jax -jaxtyping==0.2.34 \ - --hash=sha256:2f81fb6d1586e497a6ea2d28c06dcab37b108a096cbb36ea3fe4fa2e1c1f32e5 \ - --hash=sha256:eed9a3458ec8726c84ea5457ebde53c964f65d2c22c0ec40d0555ae3fed5bbaf - # via alphafold3 (pyproject.toml) -jmp==0.0.4 \ - --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \ - --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d - # via dm-haiku -ml-dtypes==0.5.0 \ - --hash=sha256:099e09edd54e676903b4538f3815b5ab96f5b119690514602d96bfdb67172cbe \ - --hash=sha256:2e7534392682c3098bc7341648c650864207169c654aed83143d7a19c67ae06f \ - --hash=sha256:3e7d3a380fe73a63c884f06136f8baa7a5249cc8e9fdec677997dd78549f8128 \ - --hash=sha256:54415257f00eb44fbcc807454efac3356f75644f1cbfc2d4e5522a72ae1dacab \ - --hash=sha256:5f2b59233a0dbb6a560b3137ed6125433289ccba2f8d9c3695a52423a369ed15 \ - --hash=sha256:60275f2b51b56834e840c4809fca840565f9bf8e9a73f6d8c94f5b5935701215 \ - --hash=sha256:76942f6aeb5c40766d5ea62386daa4148e6a54322aaf5b53eae9e7553240222f \ - --hash=sha256:7ee9c320bb0f9ffdf9f6fa6a696ef2e005d1f66438d6f1c1457338e00a02e8cf \ - --hash=sha256:8c32138975797e681eb175996d64356bcfa124bdbb6a70460b9768c2b35a6fa4 \ - --hash=sha256:968fede07d1f9b926a63df97d25ac656cac1a57ebd33701734eaf704bc55d8d8 \ - --hash=sha256:a03fc861b86cc586728e3d093ba37f0cc05e65330c3ebd7688e7bae8290f8859 \ - --hash=sha256:a38df8df61194aeaae1ab7579075779b4ad32cd1cffd012c28be227fa7f2a70a \ - --hash=sha256:a988bac6572630e1e9c2edd9b1277b4eefd1c86209e52b0d061b775ac33902ff \ - --hash=sha256:ab046f2ff789b1f11b2491909682c5d089934835f9a760fafc180e47dcb676b8 \ - --hash=sha256:afa08343069874a30812871d639f9c02b4158ace065601406a493a8511180c02 \ - --hash=sha256:c7a9152f5876fef565516aa5dd1dccd6fc298a5891b2467973905103eb5c7856 \ - --hash=sha256:cb5cc7b25acabd384f75bbd78892d0c724943f3e2e1986254665a1aa10982e07 \ - --hash=sha256:d3b3db9990c3840986a0e70524e122cfa32b91139c3653df76121ba7776e015f \ - --hash=sha256:d4b1a70a3e5219790d6b55b9507606fc4e02911d1497d16c18dd721eb7efe7d0 \ - --hash=sha256:dc74fd9995513d33eac63d64e436240f5494ec74d522a9f0920194942fc3d2d7 \ - --hash=sha256:e04fde367b2fe901b1d47234426fe8819909bd1dd862a5adb630f27789c20599 - # via - # jax - # jaxlib -numpy==2.1.3 \ - --hash=sha256:016d0f6f5e77b0f0d45d77387ffa4bb89816b57c835580c3ce8e099ef830befe \ - --hash=sha256:02135ade8b8a84011cbb67dc44e07c58f28575cf9ecf8ab304e51c05528c19f0 \ - --hash=sha256:08788d27a5fd867a663f6fc753fd7c3ad7e92747efc73c53bca2f19f8bc06f48 \ - --hash=sha256:0d30c543f02e84e92c4b1f415b7c6b5326cbe45ee7882b6b77db7195fb971e3a \ - --hash=sha256:0fa14563cc46422e99daef53d725d0c326e99e468a9320a240affffe87852564 \ - --hash=sha256:13138eadd4f4da03074851a698ffa7e405f41a0845a6b1ad135b81596e4e9958 \ - --hash=sha256:14e253bd43fc6b37af4921b10f6add6925878a42a0c5fe83daee390bca80bc17 \ - --hash=sha256:15cb89f39fa6d0bdfb600ea24b250e5f1a3df23f901f51c8debaa6a5d122b2f0 \ - --hash=sha256:17ee83a1f4fef3c94d16dc1802b998668b5419362c8a4f4e8a491de1b41cc3ee \ - --hash=sha256:2312b2aa89e1f43ecea6da6ea9a810d06aae08321609d8dc0d0eda6d946a541b \ - --hash=sha256:2564fbdf2b99b3f815f2107c1bbc93e2de8ee655a69c261363a1172a79a257d4 \ - --hash=sha256:3522b0dfe983a575e6a9ab3a4a4dfe156c3e428468ff08ce582b9bb6bd1d71d4 \ - --hash=sha256:4394bc0dbd074b7f9b52024832d16e019decebf86caf909d94f6b3f77a8ee3b6 \ - --hash=sha256:45966d859916ad02b779706bb43b954281db43e185015df6eb3323120188f9e4 \ - --hash=sha256:4d1167c53b93f1f5d8a139a742b3c6f4d429b54e74e6b57d0eff40045187b15d \ - --hash=sha256:4f2015dfe437dfebbfce7c85c7b53d81ba49e71ba7eadbf1df40c915af75979f \ - --hash=sha256:50ca6aba6e163363f132b5c101ba078b8cbd3fa92c7865fd7d4d62d9779ac29f \ - --hash=sha256:50d18c4358a0a8a53f12a8ba9d772ab2d460321e6a93d6064fc22443d189853f \ - --hash=sha256:5641516794ca9e5f8a4d17bb45446998c6554704d888f86df9b200e66bdcce56 \ - --hash=sha256:576a1c1d25e9e02ed7fa5477f30a127fe56debd53b8d2c89d5578f9857d03ca9 \ - --hash=sha256:6a4825252fcc430a182ac4dee5a505053d262c807f8a924603d411f6718b88fd \ - --hash=sha256:72dcc4a35a8515d83e76b58fdf8113a5c969ccd505c8a946759b24e3182d1f23 \ - --hash=sha256:747641635d3d44bcb380d950679462fae44f54b131be347d5ec2bce47d3df9ed \ - --hash=sha256:762479be47a4863e261a840e8e01608d124ee1361e48b96916f38b119cfda04a \ - --hash=sha256:78574ac2d1a4a02421f25da9559850d59457bac82f2b8d7a44fe83a64f770098 \ - --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ - --hash=sha256:8637dcd2caa676e475503d1f8fdb327bc495554e10838019651b76d17b98e512 \ - --hash=sha256:96fe52fcdb9345b7cd82ecd34547fca4321f7656d500eca497eb7ea5a926692f \ - --hash=sha256:973faafebaae4c0aaa1a1ca1ce02434554d67e628b8d805e61f874b84e136b09 \ - --hash=sha256:996bb9399059c5b82f76b53ff8bb686069c05acc94656bb259b1d63d04a9506f \ - --hash=sha256:a38c19106902bb19351b83802531fea19dee18e5b37b36454f27f11ff956f7fc \ - --hash=sha256:a6b46587b14b888e95e4a24d7b13ae91fa22386c199ee7b418f449032b2fa3b8 \ - --hash=sha256:a9f7f672a3388133335589cfca93ed468509cb7b93ba3105fce780d04a6576a0 \ - --hash=sha256:aa08e04e08aaf974d4458def539dece0d28146d866a39da5639596f4921fd761 \ - --hash=sha256:b0df3635b9c8ef48bd3be5f862cf71b0a4716fa0e702155c45067c6b711ddcef \ - --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ - --hash=sha256:baed7e8d7481bfe0874b566850cb0b85243e982388b7b23348c6db2ee2b2ae8e \ - --hash=sha256:bc6f24b3d1ecc1eebfbf5d6051faa49af40b03be1aaa781ebdadcbc090b4539b \ - --hash=sha256:c006b607a865b07cd981ccb218a04fc86b600411d83d6fc261357f1c0966755d \ - --hash=sha256:c181ba05ce8299c7aa3125c27b9c2167bca4a4445b7ce73d5febc411ca692e43 \ - --hash=sha256:c7662f0e3673fe4e832fe07b65c50342ea27d989f92c80355658c7f888fcc83c \ - --hash=sha256:c80e4a09b3d95b4e1cac08643f1152fa71a0a821a2d4277334c88d54b2219a41 \ - --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff \ - --hash=sha256:d7aac50327da5d208db2eec22eb11e491e3fe13d22653dce51b0f4109101b408 \ - --hash=sha256:d89dd2b6da69c4fff5e39c28a382199ddedc3a5be5390115608345dec660b9e2 \ - --hash=sha256:d9beb777a78c331580705326d2367488d5bc473b49a9bc3036c154832520aca9 \ - --hash=sha256:dc258a761a16daa791081d026f0ed4399b582712e6fc887a95af09df10c5ca57 \ - --hash=sha256:e14e26956e6f1696070788252dcdff11b4aca4c3e8bd166e0df1bb8f315a67cb \ - --hash=sha256:e6988e90fcf617da2b5c78902fe8e668361b43b4fe26dbf2d7b0f8034d4cafb9 \ - --hash=sha256:e711e02f49e176a01d0349d82cb5f05ba4db7d5e7e0defd026328e5cfb3226d3 \ - --hash=sha256:ea4dedd6e394a9c180b33c2c872b92f7ce0f8e7ad93e9585312b0c5a04777a4a \ - --hash=sha256:ecc76a9ba2911d8d37ac01de72834d8849e55473457558e12995f4cd53e778e0 \ - --hash=sha256:f55ba01150f52b1027829b50d70ef1dafd9821ea82905b63936668403c3b471e \ - --hash=sha256:f653490b33e9c3a4c1c01d41bc2aef08f9475af51146e4a7710c450cf9761598 \ - --hash=sha256:fa2d1337dc61c8dc417fbccf20f6d1e139896a30721b7f1e832b2bb6ef4eb6c4 - # via - # alphafold3 (pyproject.toml) - # dm-haiku - # jax - # jaxlib - # jmp - # ml-dtypes - # rdkit - # scipy -nvidia-cublas-cu12==12.6.3.3 \ - --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \ - --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \ - --hash=sha256:f33fb68e101d99470c82d17f92a0dd9f74de2a21685c217f4716cdd63b1316eb - # via - # jax-cuda12-plugin - # nvidia-cudnn-cu12 - # nvidia-cusolver-cu12 -nvidia-cuda-cupti-cu12==12.6.80 \ - --hash=sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4 \ - --hash=sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73 \ - --hash=sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a - # via jax-cuda12-plugin -nvidia-cuda-nvcc-cu12==12.6.77 \ - --hash=sha256:26e118b4dae4ebd316ab6f5ee3f3c4cb929f8265f717a8e016ba57e254c0b4f1 \ - --hash=sha256:9a2aad43b8f4e4be9c237ece04f64450df77871eeebffad5a5a872714a97ca77 \ - --hash=sha256:b66bb5dd6b8ae62262586691977d3b4a425e91db61b4b9bc2f6b42bcd4154b96 - # via jax-cuda12-plugin -nvidia-cuda-runtime-cu12==12.6.77 \ - --hash=sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f \ - --hash=sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8 \ - --hash=sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e - # via jax-cuda12-plugin -nvidia-cudnn-cu12==9.5.1.17 \ - --hash=sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2 \ - --hash=sha256:9fd4584468533c61873e5fda8ca41bac3a38bcb2d12350830c69b0a96a7e4def \ - --hash=sha256:d7af0f8a4f3b4b9dbb3122f2ef553b45694ed9c384d5a75bab197b8eefb79ab8 - # via jax-cuda12-plugin -nvidia-cufft-cu12==11.3.0.4 \ - --hash=sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464 \ - --hash=sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca \ - --hash=sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb - # via jax-cuda12-plugin -nvidia-cusolver-cu12==11.7.1.2 \ - --hash=sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0 \ - --hash=sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7 \ - --hash=sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6 - # via jax-cuda12-plugin -nvidia-cusparse-cu12==12.5.4.2 \ - --hash=sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f \ - --hash=sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20 \ - --hash=sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1 - # via - # jax-cuda12-plugin - # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.23.4 \ - --hash=sha256:aa946c8327e22ced28e7cef508a334673abc42064ec85f02d005ba1785ea4cec \ - --hash=sha256:b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1 - # via jax-cuda12-plugin -nvidia-nvjitlink-cu12==12.6.77 \ - --hash=sha256:3bf10d85bb1801e9c894c6e197e44dd137d2a0a9e43f8450e9ad13f2df0dd52d \ - --hash=sha256:410718cd44962bed862a31dd0318620f6f9a8b28a6291967bcfcb446a6516771 \ - --hash=sha256:9ae346d16203ae4ea513be416495167a0101d33d2d14935aa9c1829a3fb45142 - # via - # jax-cuda12-plugin - # nvidia-cufft-cu12 - # nvidia-cusolver-cu12 - # nvidia-cusparse-cu12 -opt-einsum==3.4.0 \ - --hash=sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd \ - --hash=sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac - # via jax -pillow==11.0.0 \ - --hash=sha256:00177a63030d612148e659b55ba99527803288cea7c75fb05766ab7981a8c1b7 \ - --hash=sha256:006bcdd307cc47ba43e924099a038cbf9591062e6c50e570819743f5607404f5 \ - --hash=sha256:084a07ef0821cfe4858fe86652fffac8e187b6ae677e9906e192aafcc1b69903 \ - --hash=sha256:0ae08bd8ffc41aebf578c2af2f9d8749d91f448b3bfd41d7d9ff573d74f2a6b2 \ - --hash=sha256:0e038b0745997c7dcaae350d35859c9715c71e92ffb7e0f4a8e8a16732150f38 \ - --hash=sha256:1187739620f2b365de756ce086fdb3604573337cc28a0d3ac4a01ab6b2d2a6d2 \ - --hash=sha256:16095692a253047fe3ec028e951fa4221a1f3ed3d80c397e83541a3037ff67c9 \ - --hash=sha256:1a61b54f87ab5786b8479f81c4b11f4d61702830354520837f8cc791ebba0f5f \ - --hash=sha256:1c1d72714f429a521d8d2d018badc42414c3077eb187a59579f28e4270b4b0fc \ - --hash=sha256:1e2688958a840c822279fda0086fec1fdab2f95bf2b717b66871c4ad9859d7e8 \ - --hash=sha256:20ec184af98a121fb2da42642dea8a29ec80fc3efbaefb86d8fdd2606619045d \ - --hash=sha256:21a0d3b115009ebb8ac3d2ebec5c2982cc693da935f4ab7bb5c8ebe2f47d36f2 \ - --hash=sha256:224aaa38177597bb179f3ec87eeefcce8e4f85e608025e9cfac60de237ba6316 \ - --hash=sha256:2679d2258b7f1192b378e2893a8a0a0ca472234d4c2c0e6bdd3380e8dfa21b6a \ - --hash=sha256:27a7860107500d813fcd203b4ea19b04babe79448268403172782754870dac25 \ - --hash=sha256:290f2cc809f9da7d6d622550bbf4c1e57518212da51b6a30fe8e0a270a5b78bd \ - --hash=sha256:2e46773dc9f35a1dd28bd6981332fd7f27bec001a918a72a79b4133cf5291dba \ - --hash=sha256:3107c66e43bda25359d5ef446f59c497de2b5ed4c7fdba0894f8d6cf3822dafc \ - --hash=sha256:375b8dd15a1f5d2feafff536d47e22f69625c1aa92f12b339ec0b2ca40263273 \ - --hash=sha256:45c566eb10b8967d71bf1ab8e4a525e5a93519e29ea071459ce517f6b903d7fa \ - --hash=sha256:499c3a1b0d6fc8213519e193796eb1a86a1be4b1877d678b30f83fd979811d1a \ - --hash=sha256:4ad70c4214f67d7466bea6a08061eba35c01b1b89eaa098040a35272a8efb22b \ - --hash=sha256:4b60c9520f7207aaf2e1d94de026682fc227806c6e1f55bba7606d1c94dd623a \ - --hash=sha256:5178952973e588b3f1360868847334e9e3bf49d19e169bbbdfaf8398002419ae \ - --hash=sha256:52a2d8323a465f84faaba5236567d212c3668f2ab53e1c74c15583cf507a0291 \ - --hash=sha256:598b4e238f13276e0008299bd2482003f48158e2b11826862b1eb2ad7c768b97 \ - --hash=sha256:5bd2d3bdb846d757055910f0a59792d33b555800813c3b39ada1829c372ccb06 \ - --hash=sha256:5c39ed17edea3bc69c743a8dd3e9853b7509625c2462532e62baa0732163a904 \ - --hash=sha256:5d203af30149ae339ad1b4f710d9844ed8796e97fda23ffbc4cc472968a47d0b \ - --hash=sha256:5ddbfd761ee00c12ee1be86c9c0683ecf5bb14c9772ddbd782085779a63dd55b \ - --hash=sha256:607bbe123c74e272e381a8d1957083a9463401f7bd01287f50521ecb05a313f8 \ - --hash=sha256:61b887f9ddba63ddf62fd02a3ba7add935d053b6dd7d58998c630e6dbade8527 \ - --hash=sha256:6619654954dc4936fcff82db8eb6401d3159ec6be81e33c6000dfd76ae189947 \ - --hash=sha256:674629ff60030d144b7bca2b8330225a9b11c482ed408813924619c6f302fdbb \ - --hash=sha256:6ec0d5af64f2e3d64a165f490d96368bb5dea8b8f9ad04487f9ab60dc4bb6003 \ - --hash=sha256:6f4dba50cfa56f910241eb7f883c20f1e7b1d8f7d91c750cd0b318bad443f4d5 \ - --hash=sha256:70fbbdacd1d271b77b7721fe3cdd2d537bbbd75d29e6300c672ec6bb38d9672f \ - --hash=sha256:72bacbaf24ac003fea9bff9837d1eedb6088758d41e100c1552930151f677739 \ - --hash=sha256:7326a1787e3c7b0429659e0a944725e1b03eeaa10edd945a86dead1913383944 \ - --hash=sha256:73853108f56df97baf2bb8b522f3578221e56f646ba345a372c78326710d3830 \ - --hash=sha256:73e3a0200cdda995c7e43dd47436c1548f87a30bb27fb871f352a22ab8dcf45f \ - --hash=sha256:75acbbeb05b86bc53cbe7b7e6fe00fbcf82ad7c684b3ad82e3d711da9ba287d3 \ - --hash=sha256:8069c5179902dcdce0be9bfc8235347fdbac249d23bd90514b7a47a72d9fecf4 \ - --hash=sha256:846e193e103b41e984ac921b335df59195356ce3f71dcfd155aa79c603873b84 \ - --hash=sha256:8594f42df584e5b4bb9281799698403f7af489fba84c34d53d1c4bfb71b7c4e7 \ - --hash=sha256:86510e3f5eca0ab87429dd77fafc04693195eec7fd6a137c389c3eeb4cfb77c6 \ - --hash=sha256:8853a3bf12afddfdf15f57c4b02d7ded92c7a75a5d7331d19f4f9572a89c17e6 \ - --hash=sha256:88a58d8ac0cc0e7f3a014509f0455248a76629ca9b604eca7dc5927cc593c5e9 \ - --hash=sha256:8ba470552b48e5835f1d23ecb936bb7f71d206f9dfeee64245f30c3270b994de \ - --hash=sha256:8c676b587da5673d3c75bd67dd2a8cdfeb282ca38a30f37950511766b26858c4 \ - --hash=sha256:8ec4a89295cd6cd4d1058a5e6aec6bf51e0eaaf9714774e1bfac7cfc9051db47 \ - --hash=sha256:94f3e1780abb45062287b4614a5bc0874519c86a777d4a7ad34978e86428b8dd \ - --hash=sha256:9a0f748eaa434a41fccf8e1ee7a3eed68af1b690e75328fd7a60af123c193b50 \ - --hash=sha256:a5629742881bcbc1f42e840af185fd4d83a5edeb96475a575f4da50d6ede337c \ - --hash=sha256:a65149d8ada1055029fcb665452b2814fe7d7082fcb0c5bed6db851cb69b2086 \ - --hash=sha256:b3c5ac4bed7519088103d9450a1107f76308ecf91d6dabc8a33a2fcfb18d0fba \ - --hash=sha256:b4fd7bd29610a83a8c9b564d457cf5bd92b4e11e79a4ee4716a63c959699b306 \ - --hash=sha256:bcd1fb5bb7b07f64c15618c89efcc2cfa3e95f0e3bcdbaf4642509de1942a699 \ - --hash=sha256:c12b5ae868897c7338519c03049a806af85b9b8c237b7d675b8c5e089e4a618e \ - --hash=sha256:c26845094b1af3c91852745ae78e3ea47abf3dbcd1cf962f16b9a5fbe3ee8488 \ - --hash=sha256:c6a660307ca9d4867caa8d9ca2c2658ab685de83792d1876274991adec7b93fa \ - --hash=sha256:c809a70e43c7977c4a42aefd62f0131823ebf7dd73556fa5d5950f5b354087e2 \ - --hash=sha256:c8b2351c85d855293a299038e1f89db92a2f35e8d2f783489c6f0b2b5f3fe8a3 \ - --hash=sha256:cb929ca942d0ec4fac404cbf520ee6cac37bf35be479b970c4ffadf2b6a1cad9 \ - --hash=sha256:d2c0a187a92a1cb5ef2c8ed5412dd8d4334272617f532d4ad4de31e0495bd923 \ - --hash=sha256:d69bfd8ec3219ae71bcde1f942b728903cad25fafe3100ba2258b973bd2bc1b2 \ - --hash=sha256:daffdf51ee5db69a82dd127eabecce20729e21f7a3680cf7cbb23f0829189790 \ - --hash=sha256:e58876c91f97b0952eb766123bfef372792ab3f4e3e1f1a2267834c2ab131734 \ - --hash=sha256:eda2616eb2313cbb3eebbe51f19362eb434b18e3bb599466a1ffa76a033fb916 \ - --hash=sha256:ee217c198f2e41f184f3869f3e485557296d505b5195c513b2bfe0062dc537f1 \ - --hash=sha256:f02541ef64077f22bf4924f225c0fd1248c168f86e4b7abdedd87d6ebaceab0f \ - --hash=sha256:f1b82c27e89fffc6da125d5eb0ca6e68017faf5efc078128cfaa42cf5cb38798 \ - --hash=sha256:fba162b8872d30fea8c52b258a542c5dfd7b235fb5cb352240c8d63b414013eb \ - --hash=sha256:fbbcb7b57dc9c794843e3d1258c0fbf0f48656d46ffe9e09b63bbd6e8cd5d0a2 \ - --hash=sha256:fcb4621042ac4b7865c179bb972ed0da0218a076dc1820ffc48b1d74c1e37fe9 - # via rdkit -rdkit==2024.3.5 \ - --hash=sha256:01fd323f5902a720925293c0ce08a33e630422374cf263d04ef173c106c63c36 \ - --hash=sha256:04079aef5f91965bd2a48bf432b32c36ea1bf088e26efb2c9a43b625b6a75d15 \ - --hash=sha256:1cb7cdb29956649d4c875790b94142868c08c0735dee4d90b2d70dfd378e9d2e \ - --hash=sha256:292141f0ba2359be07d8566d9a39ce697c1d6e7eb406031a2bf0387026503a47 \ - --hash=sha256:2affa1b4cf601d27b157bbc44255c640baac9bc04cfdf5fa88f43c34c3f360d5 \ - --hash=sha256:316bb5273bff5be76d4e5b0def998455e9b1a20898df44548b0e299f90786cb2 \ - --hash=sha256:455d510beff8806e62e48b977d7acbfbc351474fa132124738a07223440c1b9a \ - --hash=sha256:53638d0d94f7c511273ff5765abc6a1fe60fe6c8633db13cad03bc527550a7a9 \ - --hash=sha256:6af38c62d4af9d669f8996551d3ae809de9ab245570b861be3f7a2b1a536f79f \ - --hash=sha256:73fa2fe3a9a2334edf8448a5a3093a6f949ee715de77d3baa685751e2cf9317b \ - --hash=sha256:7c91c5ccbf8d51cd3dd894bf770ca99da5fe37714f6e2a66016e8af4f57d2b56 \ - --hash=sha256:7d5576bad6672959d0c1939e8d7e2fcd8656f8daf9788ce58e5c313db719b5ec \ - --hash=sha256:89f7a72466ffbd45320d845ee1bcde1bddb1a322352836835d26ebb712228a41 \ - --hash=sha256:8d4ebafb7e845bdc55eb6e2e188123ec84fb2bf724d3a7eb2da9bcfade7acc65 \ - --hash=sha256:90376467f4988ee44f3f47d1e558ffc27646254a99fe65d19a7d9a123e832ef2 \ - --hash=sha256:913494a03fcf6121c103e5db43fa33bca0fcc65ebd5a6c46661baa6851f79d91 \ - --hash=sha256:98dac94a5e414ed1880b215eacec63e5fd77b0ea351e0a18fc5229acb8e886f5 \ - --hash=sha256:a51cc403792ba37d1ee8ad84ad829177a4453547a4e085e1115ae2b237084ddb \ - --hash=sha256:a98099cef42c1eb52512b343beb68a8b7cc4da79924a67181f58692a6376d91b \ - --hash=sha256:c971eba2a67f8ae301b9b8c66b3431cd62472f648cdec9694e5413d8ebc59720 \ - --hash=sha256:cf2747e7c82d6c90a8a631fe063d24d77b31f535e0f8a17ac9dbfbe4ab456886 \ - --hash=sha256:d8714c03558086d022649617197c0c45c75c301f6a9fb67205695493a7e5b826 \ - --hash=sha256:e5b0dfb28aa0949152ff816fba0a2f6267154f5c25159a9ec84c27cf18f72322 \ - --hash=sha256:e69cec14cbb42c5ca912346b594d3f6a7aac572463949c55c7749da257e68123 \ - --hash=sha256:f662ed10e14c7874d0e011c1f8260a6cf4a5547100806f9f4540f8550a3f1d36 - # via alphafold3 (pyproject.toml) -scipy==1.14.1 \ - --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \ - --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \ - --hash=sha256:278266012eb69f4a720827bdd2dc54b2271c97d84255b2faaa8f161a158c3b37 \ - --hash=sha256:2843f2d527d9eebec9a43e6b406fb7266f3af25a751aa91d62ff416f54170bc5 \ - --hash=sha256:2da0469a4ef0ecd3693761acbdc20f2fdeafb69e6819cc081308cc978153c675 \ - --hash=sha256:2ff0a7e01e422c15739ecd64432743cf7aae2b03f3084288f399affcefe5222d \ - --hash=sha256:2ff38e22128e6c03ff73b6bb0f85f897d2362f8c052e3b8ad00532198fbdae3f \ - --hash=sha256:30ac8812c1d2aab7131a79ba62933a2a76f582d5dbbc695192453dae67ad6310 \ - --hash=sha256:3a1b111fac6baec1c1d92f27e76511c9e7218f1695d61b59e05e0fe04dc59617 \ - --hash=sha256:4079b90df244709e675cdc8b93bfd8a395d59af40b72e339c2287c91860deb8e \ - --hash=sha256:5149e3fd2d686e42144a093b206aef01932a0059c2a33ddfa67f5f035bdfe13e \ - --hash=sha256:5a275584e726026a5699459aa72f828a610821006228e841b94275c4a7c08417 \ - --hash=sha256:631f07b3734d34aced009aaf6fedfd0eb3498a97e581c3b1e5f14a04164a456d \ - --hash=sha256:716e389b694c4bb564b4fc0c51bc84d381735e0d39d3f26ec1af2556ec6aad94 \ - --hash=sha256:8426251ad1e4ad903a4514712d2fa8fdd5382c978010d1c6f5f37ef286a713ad \ - --hash=sha256:8475230e55549ab3f207bff11ebfc91c805dc3463ef62eda3ccf593254524ce8 \ - --hash=sha256:8bddf15838ba768bb5f5083c1ea012d64c9a444e16192762bd858f1e126196d0 \ - --hash=sha256:8e32dced201274bf96899e6491d9ba3e9a5f6b336708656466ad0522d8528f69 \ - --hash=sha256:8f9ea80f2e65bdaa0b7627fb00cbeb2daf163caa015e59b7516395fe3bd1e066 \ - --hash=sha256:97c5dddd5932bd2a1a31c927ba5e1463a53b87ca96b5c9bdf5dfd6096e27efc3 \ - --hash=sha256:a49f6ed96f83966f576b33a44257d869756df6cf1ef4934f59dd58b25e0327e5 \ - --hash=sha256:af29a935803cc707ab2ed7791c44288a682f9c8107bc00f0eccc4f92c08d6e07 \ - --hash=sha256:b05d43735bb2f07d689f56f7b474788a13ed8adc484a85aa65c0fd931cf9ccd2 \ - --hash=sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389 \ - --hash=sha256:b99722ea48b7ea25e8e015e8341ae74624f72e5f21fc2abd45f3a93266de4c5d \ - --hash=sha256:baff393942b550823bfce952bb62270ee17504d02a1801d7fd0719534dfb9c84 \ - --hash=sha256:c0ee987efa6737242745f347835da2cc5bb9f1b42996a4d97d5c7ff7928cb6f2 \ - --hash=sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3 \ - --hash=sha256:e0cf28db0f24a38b2a0ca33a85a54852586e43cf6fd876365c86e0657cfe7d73 \ - --hash=sha256:e4f5a7c49323533f9103d4dacf4e4f07078f360743dec7f7596949149efeec06 \ - --hash=sha256:eb58ca0abd96911932f688528977858681a59d61a7ce908ffd355957f7025cfc \ - --hash=sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1 \ - --hash=sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2 - # via - # jax - # jaxlib -tabulate==0.9.0 \ - --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \ - --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f - # via dm-haiku -tqdm==4.67.0 \ - --hash=sha256:0cd8af9d56911acab92182e88d763100d4788bdf421d251616040cc4d44863be \ - --hash=sha256:fe5a6f95e6fe0b9755e9469b77b9c3cf850048224ecaa8293d7d2d31f97d869a - # via alphafold3 (pyproject.toml) -triton==3.1.0 \ - --hash=sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c \ - --hash=sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8 \ - --hash=sha256:6dadaca7fc24de34e180271b5cf864c16755702e9f63a16f62df714a8099126a \ - --hash=sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11 \ - --hash=sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc - # via - # alphafold3 (pyproject.toml) - # jax-triton -typeguard==2.13.3 \ - --hash=sha256:00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4 \ - --hash=sha256:5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1 - # via - # alphafold3 (pyproject.toml) - # jaxtyping -zstandard==0.23.0 \ - --hash=sha256:034b88913ecc1b097f528e42b539453fa82c3557e414b3de9d5632c80439a473 \ - --hash=sha256:0a7f0804bb3799414af278e9ad51be25edf67f78f916e08afdb983e74161b916 \ - --hash=sha256:11e3bf3c924853a2d5835b24f03eeba7fc9b07d8ca499e247e06ff5676461a15 \ - --hash=sha256:12a289832e520c6bd4dcaad68e944b86da3bad0d339ef7989fb7e88f92e96072 \ - --hash=sha256:1516c8c37d3a053b01c1c15b182f3b5f5eef19ced9b930b684a73bad121addf4 \ - --hash=sha256:157e89ceb4054029a289fb504c98c6a9fe8010f1680de0201b3eb5dc20aa6d9e \ - --hash=sha256:1bfe8de1da6d104f15a60d4a8a768288f66aa953bbe00d027398b93fb9680b26 \ - --hash=sha256:1e172f57cd78c20f13a3415cc8dfe24bf388614324d25539146594c16d78fcc8 \ - --hash=sha256:1fd7e0f1cfb70eb2f95a19b472ee7ad6d9a0a992ec0ae53286870c104ca939e5 \ - --hash=sha256:203d236f4c94cd8379d1ea61db2fce20730b4c38d7f1c34506a31b34edc87bdd \ - --hash=sha256:27d3ef2252d2e62476389ca8f9b0cf2bbafb082a3b6bfe9d90cbcbb5529ecf7c \ - --hash=sha256:29a2bc7c1b09b0af938b7a8343174b987ae021705acabcbae560166567f5a8db \ - --hash=sha256:2ef230a8fd217a2015bc91b74f6b3b7d6522ba48be29ad4ea0ca3a3775bf7dd5 \ - --hash=sha256:2ef3775758346d9ac6214123887d25c7061c92afe1f2b354f9388e9e4d48acfc \ - --hash=sha256:2f146f50723defec2975fb7e388ae3a024eb7151542d1599527ec2aa9cacb152 \ - --hash=sha256:2fb4535137de7e244c230e24f9d1ec194f61721c86ebea04e1581d9d06ea1269 \ - --hash=sha256:32ba3b5ccde2d581b1e6aa952c836a6291e8435d788f656fe5976445865ae045 \ - --hash=sha256:34895a41273ad33347b2fc70e1bff4240556de3c46c6ea430a7ed91f9042aa4e \ - --hash=sha256:379b378ae694ba78cef921581ebd420c938936a153ded602c4fea612b7eaa90d \ - --hash=sha256:38302b78a850ff82656beaddeb0bb989a0322a8bbb1bf1ab10c17506681d772a \ - --hash=sha256:3aa014d55c3af933c1315eb4bb06dd0459661cc0b15cd61077afa6489bec63bb \ - --hash=sha256:4051e406288b8cdbb993798b9a45c59a4896b6ecee2f875424ec10276a895740 \ - --hash=sha256:40b33d93c6eddf02d2c19f5773196068d875c41ca25730e8288e9b672897c105 \ - --hash=sha256:43da0f0092281bf501f9c5f6f3b4c975a8a0ea82de49ba3f7100e64d422a1274 \ - --hash=sha256:445e4cb5048b04e90ce96a79b4b63140e3f4ab5f662321975679b5f6360b90e2 \ - --hash=sha256:48ef6a43b1846f6025dde6ed9fee0c24e1149c1c25f7fb0a0585572b2f3adc58 \ - --hash=sha256:50a80baba0285386f97ea36239855f6020ce452456605f262b2d33ac35c7770b \ - --hash=sha256:519fbf169dfac1222a76ba8861ef4ac7f0530c35dd79ba5727014613f91613d4 \ - --hash=sha256:53dd9d5e3d29f95acd5de6802e909ada8d8d8cfa37a3ac64836f3bc4bc5512db \ - --hash=sha256:53ea7cdc96c6eb56e76bb06894bcfb5dfa93b7adcf59d61c6b92674e24e2dd5e \ - --hash=sha256:576856e8594e6649aee06ddbfc738fec6a834f7c85bf7cadd1c53d4a58186ef9 \ - --hash=sha256:59556bf80a7094d0cfb9f5e50bb2db27fefb75d5138bb16fb052b61b0e0eeeb0 \ - --hash=sha256:5d41d5e025f1e0bccae4928981e71b2334c60f580bdc8345f824e7c0a4c2a813 \ - --hash=sha256:61062387ad820c654b6a6b5f0b94484fa19515e0c5116faf29f41a6bc91ded6e \ - --hash=sha256:61f89436cbfede4bc4e91b4397eaa3e2108ebe96d05e93d6ccc95ab5714be512 \ - --hash=sha256:62136da96a973bd2557f06ddd4e8e807f9e13cbb0bfb9cc06cfe6d98ea90dfe0 \ - --hash=sha256:64585e1dba664dc67c7cdabd56c1e5685233fbb1fc1966cfba2a340ec0dfff7b \ - --hash=sha256:65308f4b4890aa12d9b6ad9f2844b7ee42c7f7a4fd3390425b242ffc57498f48 \ - --hash=sha256:66b689c107857eceabf2cf3d3fc699c3c0fe8ccd18df2219d978c0283e4c508a \ - --hash=sha256:6a41c120c3dbc0d81a8e8adc73312d668cd34acd7725f036992b1b72d22c1772 \ - --hash=sha256:6f77fa49079891a4aab203d0b1744acc85577ed16d767b52fc089d83faf8d8ed \ - --hash=sha256:72c68dda124a1a138340fb62fa21b9bf4848437d9ca60bd35db36f2d3345f373 \ - --hash=sha256:752bf8a74412b9892f4e5b58f2f890a039f57037f52c89a740757ebd807f33ea \ - --hash=sha256:76e79bc28a65f467e0409098fa2c4376931fd3207fbeb6b956c7c476d53746dd \ - --hash=sha256:774d45b1fac1461f48698a9d4b5fa19a69d47ece02fa469825b442263f04021f \ - --hash=sha256:77da4c6bfa20dd5ea25cbf12c76f181a8e8cd7ea231c673828d0386b1740b8dc \ - --hash=sha256:77ea385f7dd5b5676d7fd943292ffa18fbf5c72ba98f7d09fc1fb9e819b34c23 \ - --hash=sha256:80080816b4f52a9d886e67f1f96912891074903238fe54f2de8b786f86baded2 \ - --hash=sha256:80a539906390591dd39ebb8d773771dc4db82ace6372c4d41e2d293f8e32b8db \ - --hash=sha256:82d17e94d735c99621bf8ebf9995f870a6b3e6d14543b99e201ae046dfe7de70 \ - --hash=sha256:837bb6764be6919963ef41235fd56a6486b132ea64afe5fafb4cb279ac44f259 \ - --hash=sha256:84433dddea68571a6d6bd4fbf8ff398236031149116a7fff6f777ff95cad3df9 \ - --hash=sha256:8c24f21fa2af4bb9f2c492a86fe0c34e6d2c63812a839590edaf177b7398f700 \ - --hash=sha256:8ed7d27cb56b3e058d3cf684d7200703bcae623e1dcc06ed1e18ecda39fee003 \ - --hash=sha256:9206649ec587e6b02bd124fb7799b86cddec350f6f6c14bc82a2b70183e708ba \ - --hash=sha256:983b6efd649723474f29ed42e1467f90a35a74793437d0bc64a5bf482bedfa0a \ - --hash=sha256:98da17ce9cbf3bfe4617e836d561e433f871129e3a7ac16d6ef4c680f13a839c \ - --hash=sha256:9c236e635582742fee16603042553d276cca506e824fa2e6489db04039521e90 \ - --hash=sha256:9da6bc32faac9a293ddfdcb9108d4b20416219461e4ec64dfea8383cac186690 \ - --hash=sha256:a05e6d6218461eb1b4771d973728f0133b2a4613a6779995df557f70794fd60f \ - --hash=sha256:a0817825b900fcd43ac5d05b8b3079937073d2b1ff9cf89427590718b70dd840 \ - --hash=sha256:a4ae99c57668ca1e78597d8b06d5af837f377f340f4cce993b551b2d7731778d \ - --hash=sha256:a8c86881813a78a6f4508ef9daf9d4995b8ac2d147dcb1a450448941398091c9 \ - --hash=sha256:a8fffdbd9d1408006baaf02f1068d7dd1f016c6bcb7538682622c556e7b68e35 \ - --hash=sha256:a9b07268d0c3ca5c170a385a0ab9fb7fdd9f5fd866be004c4ea39e44edce47dd \ - --hash=sha256:ab19a2d91963ed9e42b4e8d77cd847ae8381576585bad79dbd0a8837a9f6620a \ - --hash=sha256:ac184f87ff521f4840e6ea0b10c0ec90c6b1dcd0bad2f1e4a9a1b4fa177982ea \ - --hash=sha256:b0e166f698c5a3e914947388c162be2583e0c638a4703fc6a543e23a88dea3c1 \ - --hash=sha256:b2170c7e0367dde86a2647ed5b6f57394ea7f53545746104c6b09fc1f4223573 \ - --hash=sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09 \ - --hash=sha256:b4567955a6bc1b20e9c31612e615af6b53733491aeaa19a6b3b37f3b65477094 \ - --hash=sha256:b69bb4f51daf461b15e7b3db033160937d3ff88303a7bc808c67bbc1eaf98c78 \ - --hash=sha256:b8c0bd73aeac689beacd4e7667d48c299f61b959475cdbb91e7d3d88d27c56b9 \ - --hash=sha256:be9b5b8659dff1f913039c2feee1aca499cfbc19e98fa12bc85e037c17ec6ca5 \ - --hash=sha256:bf0a05b6059c0528477fba9054d09179beb63744355cab9f38059548fedd46a9 \ - --hash=sha256:c16842b846a8d2a145223f520b7e18b57c8f476924bda92aeee3a88d11cfc391 \ - --hash=sha256:c363b53e257246a954ebc7c488304b5592b9c53fbe74d03bc1c64dda153fb847 \ - --hash=sha256:c7c517d74bea1a6afd39aa612fa025e6b8011982a0897768a2f7c8ab4ebb78a2 \ - --hash=sha256:d20fd853fbb5807c8e84c136c278827b6167ded66c72ec6f9a14b863d809211c \ - --hash=sha256:d2240ddc86b74966c34554c49d00eaafa8200a18d3a5b6ffbf7da63b11d74ee2 \ - --hash=sha256:d477ed829077cd945b01fc3115edd132c47e6540ddcd96ca169facff28173057 \ - --hash=sha256:d50d31bfedd53a928fed6707b15a8dbeef011bb6366297cc435accc888b27c20 \ - --hash=sha256:dc1d33abb8a0d754ea4763bad944fd965d3d95b5baef6b121c0c9013eaf1907d \ - --hash=sha256:dc5d1a49d3f8262be192589a4b72f0d03b72dcf46c51ad5852a4fdc67be7b9e4 \ - --hash=sha256:e2d1a054f8f0a191004675755448d12be47fa9bebbcffa3cdf01db19f2d30a54 \ - --hash=sha256:e7792606d606c8df5277c32ccb58f29b9b8603bf83b48639b7aedf6df4fe8171 \ - --hash=sha256:ed1708dbf4d2e3a1c5c69110ba2b4eb6678262028afd6c6fbcc5a8dac9cda68e \ - --hash=sha256:f2d4380bf5f62daabd7b751ea2339c1a21d1c9463f1feb7fc2bdcea2c29c3160 \ - --hash=sha256:f3513916e8c645d0610815c257cbfd3242adfd5c4cfa78be514e5a3ebb42a41b \ - --hash=sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58 \ - --hash=sha256:f83fa6cae3fff8e98691248c9320356971b59678a17f20656a9e59cd32cee6d8 \ - --hash=sha256:fa6ce8b52c5987b3e34d5674b0ab529a4602b632ebab0a93b07bfb4dfc8f8a33 \ - --hash=sha256:fb2b1ecfef1e67897d336de3a0e3f52478182d6a47eda86cbd42504c5cbd009a \ - --hash=sha256:fc9ca1c9718cb3b06634c7c8dec57d24e9438b2aa9a0f02b8bb36bf478538880 \ - --hash=sha256:fd30d9c67d13d891f2360b2a120186729c111238ac63b43dbd37a5a40670b8ca \ - --hash=sha256:fd7699e8fd9969f455ef2926221e0233f81a2542921471382e77a9e2f2b57f4b \ - --hash=sha256:fe3b385d996ee0822fd46528d9f0443b880d4d05528fd26a9119a54ec3f91c69 - # via alphafold3 (pyproject.toml) diff --git a/run_alphafold.py b/run_alphafold.py index e8d2a1b..0bdd8a5 100644 --- a/run_alphafold.py +++ b/run_alphafold.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -10,8 +19,8 @@ """AlphaFold 3 structure prediction script. -AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of -this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/ +AlphaFold 3 source code is licensed under Apache License, Version 2.0. To view a +copy of this license, visit http://www.apache.org/licenses/LICENSE-2.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 @@ -42,7 +51,6 @@ from alphafold3.data import featurisation from alphafold3.data import pipeline from alphafold3.data.tools import shards -from alphafold3.jax.attention import attention from alphafold3.model import features from alphafold3.model import model from alphafold3.model import params @@ -52,9 +60,9 @@ import jax from jax import numpy as jnp import numpy as np +import tokamax - -_HOME_DIR = pathlib.Path(os.environ.get('HOME')) +_HOME_DIR = pathlib.Path.home() _DEFAULT_MODEL_DIR = _HOME_DIR / 'models' _DEFAULT_DB_DIR = _HOME_DIR / 'public_databases' @@ -127,7 +135,6 @@ 'Path to the directory containing the databases. Can be specified multiple' ' times to search multiple directories in order.', ) - _SMALL_BFD_DATABASE_PATH = flags.DEFINE_string( 'small_bfd_database_path', '${DB_DIR}/bfd-first_non_consensus_sequences.fasta', @@ -285,6 +292,14 @@ 'conformer search.', lower_bound=0, ) +_FIX_STANDALONE_GLYCANS = flags.DEFINE_bool( + 'fix_standalone_glycans', + False, + 'AlphaFold 3 model training and evaluation filtered out leaving atoms from' + ' glycan ligands even if they were not bonded to anything ("standalone"' + ' glycans). Setting this flag to True fixes this undesirable behavior, but' + ' moves away from the regime where AlphaFold 3 was trained and evaluated.', +) # JAX inference performance tuning. _JAX_COMPILATION_CACHE_DIR = flags.DEFINE_string( @@ -369,11 +384,18 @@ ' and is non-empty. Useful to set this to True to run the data pipeline and' ' the inference separately, but use the same output directory.', ) +_COMPRESS_LARGE_OUTPUT_FILES = flags.DEFINE_bool( + 'compress_large_output_files', + False, + 'If True, compresses the output mmCIF and confidences JSON files (the two' + ' largest files) using zstandard. Note that embeddings and distogram, if' + ' saved, are already stored in a compressed format.', +) def make_model_config( *, - flash_attention_implementation: attention.Implementation = 'triton', + flash_attention_implementation: tokamax.DotProductAttentionImplementation = 'triton', num_diffusion_samples: int = 5, num_recycles: int = 10, return_embeddings: bool = False, @@ -506,10 +528,12 @@ class ResultsForSeed: def predict_structure( fold_input: folding_input.Input, model_runner: ModelRunner, + *, buckets: Sequence[int] | None = None, ref_max_modified_date: datetime.date | None = None, conformer_max_iterations: int | None = None, resolve_msa_overlaps: bool = True, + fix_standalone_glycans: bool = False, ) -> Sequence[ResultsForSeed]: """Runs the full inference pipeline to predict structures for each seed.""" @@ -524,6 +548,7 @@ def predict_structure( ref_max_modified_date=ref_max_modified_date, conformer_max_iterations=conformer_max_iterations, resolve_msa_overlaps=resolve_msa_overlaps, + fix_standalone_glycans=fix_standalone_glycans, ) print( f'Featurising data with {len(fold_input.rng_seeds)} seed(s) took' @@ -594,6 +619,7 @@ def write_outputs( all_inference_results: Sequence[ResultsForSeed], output_dir: os.PathLike[str] | str, job_name: str, + compress_large_output_files: bool = False, ) -> None: """Writes outputs to the specified output directory.""" ranking_scores = [] @@ -614,6 +640,7 @@ def write_outputs( inference_result=result, output_dir=sample_dir, name=f'{job_name}_seed-{seed}_sample-{sample_idx}', + compress=compress_large_output_files, ) ranking_score = float(result.metadata['ranking_score']) ranking_scores.append((seed, sample_idx, ranking_score)) @@ -646,6 +673,7 @@ def write_outputs( # The output terms of use are the same for all seeds/samples. terms_of_use=output_terms, name=job_name, + compress=compress_large_output_files, ) # Save csv of ranking scores with seeds and sample indices, to allow easier # comparison of ranking scores across different runs. @@ -681,13 +709,16 @@ def replace_db_dir(path_with_db_dir: str, db_dirs: Sequence[str]) -> str: def process_fold_input( fold_input: folding_input.Input, data_pipeline_config: pipeline.DataPipelineConfig | None, + *, model_runner: None, output_dir: os.PathLike[str] | str, buckets: Sequence[int] | None = None, ref_max_modified_date: datetime.date | None = None, conformer_max_iterations: int | None = None, resolve_msa_overlaps: bool = True, + fix_standalone_glycans: bool = False, force_output_dir: bool = False, + compress_large_output_files: bool = False, ) -> folding_input.Input: ... @@ -696,13 +727,16 @@ def process_fold_input( def process_fold_input( fold_input: folding_input.Input, data_pipeline_config: pipeline.DataPipelineConfig | None, + *, model_runner: ModelRunner, output_dir: os.PathLike[str] | str, buckets: Sequence[int] | None = None, ref_max_modified_date: datetime.date | None = None, conformer_max_iterations: int | None = None, resolve_msa_overlaps: bool = True, + fix_standalone_glycans: bool = False, force_output_dir: bool = False, + compress_large_output_files: bool = False, ) -> Sequence[ResultsForSeed]: ... @@ -710,13 +744,16 @@ def process_fold_input( def process_fold_input( fold_input: folding_input.Input, data_pipeline_config: pipeline.DataPipelineConfig | None, + *, model_runner: ModelRunner | None, output_dir: os.PathLike[str] | str, buckets: Sequence[int] | None = None, ref_max_modified_date: datetime.date | None = None, conformer_max_iterations: int | None = None, resolve_msa_overlaps: bool = True, + fix_standalone_glycans: bool = False, force_output_dir: bool = False, + compress_large_output_files: bool = False, ) -> folding_input.Input | Sequence[ResultsForSeed]: """Runs data pipeline and/or inference on a single fold input. @@ -743,10 +780,17 @@ def process_fold_input( paper. Set this to false if providing custom paired MSA using the unpaired MSA field to keep it exactly as is as deduplication against the paired MSA could break the manually crafted pairing between MSA sequences. + fix_standalone_glycans: If True, standalone glycans are preserved when + filter_leaving_atoms is True. This is False by default to match the + AlphaFold 3 paper. Note that the model has been trained with the default + setting, so setting this to True may cause non-standard behaviour of the + model. force_output_dir: If True, do not create a new output directory even if the existing one is non-empty. Instead use the existing output directory and potentially overwrite existing files. If False, create a new timestamped output directory instead if the existing one is non-empty. + compress_large_output_files: If True, compress large output files (mmCIF and + confidences JSON) using zstandard. Returns: The processed fold input, or the inference results for each seed. @@ -797,12 +841,14 @@ def process_fold_input( ref_max_modified_date=ref_max_modified_date, conformer_max_iterations=conformer_max_iterations, resolve_msa_overlaps=resolve_msa_overlaps, + fix_standalone_glycans=fix_standalone_glycans, ) print(f'Writing outputs with {len(fold_input.rng_seeds)} seed(s)...') write_outputs( all_inference_results=all_inference_results, output_dir=output_dir, job_name=fold_input.sanitised_name(), + compress_large_output_files=compress_large_output_files, ) output = all_inference_results @@ -840,6 +886,9 @@ def main(_): 'Exactly one of --json_path or --input_dir must be specified.' ) + if _OUTPUT_DIR.value is None: + raise ValueError('Output directory must be specified with --output_dir.') + # Make sure we can create the output directory before running anything. try: os.makedirs(_OUTPUT_DIR.value, exist_ok=True) @@ -935,7 +984,8 @@ def main(_): model_runner = ModelRunner( config=make_model_config( flash_attention_implementation=typing.cast( - attention.Implementation, _FLASH_ATTENTION_IMPLEMENTATION.value + tokamax.DotProductAttentionImplementation, + _FLASH_ATTENTION_IMPLEMENTATION.value, ), num_diffusion_samples=_NUM_DIFFUSION_SAMPLES.value, num_recycles=_NUM_RECYCLES.value, @@ -965,7 +1015,9 @@ def main(_): ref_max_modified_date=max_template_date, conformer_max_iterations=_CONFORMER_MAX_ITERATIONS.value, resolve_msa_overlaps=_RESOLVE_MSA_OVERLAPS.value, + fix_standalone_glycans=_FIX_STANDALONE_GLYCANS.value, force_output_dir=_FORCE_OUTPUT_DIR.value, + compress_large_output_files=_COMPRESS_LARGE_OUTPUT_FILES.value, ) num_fold_inputs += 1 diff --git a/run_alphafold_data_test.py b/run_alphafold_data_test.py index 5465e88..bde41a0 100644 --- a/run_alphafold_data_test.py +++ b/run_alphafold_data_test.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/run_alphafold_test.py b/run_alphafold_test.py index d97a1db..420d7f6 100644 --- a/run_alphafold_test.py +++ b/run_alphafold_test.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -208,7 +217,7 @@ def test_inference(self, bucket, seed): actual = run_alphafold.process_fold_input( fold_input, self._data_pipeline_config, - run_alphafold.ModelRunner( + model_runner=run_alphafold.ModelRunner( config=self._model_config, device=jax.local_devices(backend='gpu')[0], model_dir=pathlib.Path(run_alphafold.MODEL_DIR.value), diff --git a/src/alphafold3/__init__.py b/src/alphafold3/__init__.py index c36629d..2ce1e6f 100644 --- a/src/alphafold3/__init__.py +++ b/src/alphafold3/__init__.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/build_data.py b/src/alphafold3/build_data.py index 0c33338..030134d 100644 --- a/src/alphafold3/build_data.py +++ b/src/alphafold3/build_data.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -11,6 +20,7 @@ """Script for building intermediate data.""" from importlib import resources +import os import pathlib import site @@ -21,13 +31,21 @@ def build_data(): """Builds intermediate data.""" - for site_path in site.getsitepackages(): - path = pathlib.Path(site_path) / 'share/libcifpp/components.cif' - if path.exists(): - cif_path = path - break + libcifpp_data_dir = os.environ.get('LIBCIFPP_DATA_DIR') + if libcifpp_data_dir: + cif_path = pathlib.Path(libcifpp_data_dir) / 'components.cif' else: - raise ValueError('Could not find components.cif') + for site_path in site.getsitepackages(): + path = pathlib.Path(site_path) / 'share/libcifpp/components.cif' + if path.exists(): + cif_path = path + break + else: + raise ValueError( + 'Could not find components.cif. If libcifpp is installed in a' + ' non-standard location, please set the LIBCIFPP_DATA_DIR environment' + ' variable to the directory where libcifpp is installed.' + ) out_root = resources.files(alphafold3.constants.converters) ccd_pickle_path = out_root.joinpath('ccd.pickle') diff --git a/src/alphafold3/common/base_config.py b/src/alphafold3/common/base_config.py index 29cc1d8..58f4933 100644 --- a/src/alphafold3/common/base_config.py +++ b/src/alphafold3/common/base_config.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/common/folding_input.py b/src/alphafold3/common/folding_input.py index 3e3925d..5ab04c3 100644 --- a/src/alphafold3/common/folding_input.py +++ b/src/alphafold3/common/folding_input.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -129,7 +138,9 @@ def mmcif(self) -> str: def __hash__(self) -> int: return hash((self._mmcif, tuple(sorted(self._query_to_template)))) - def __eq__(self, other: Self) -> bool: + def __eq__(self, other: object) -> bool: + if not isinstance(other, Template): + return NotImplemented mmcifs_equal = self._mmcif == other._mmcif maps_equal = sorted(self._query_to_template) == sorted( other._query_to_template @@ -252,7 +263,9 @@ def templates(self) -> Sequence[Template] | None: def __len__(self) -> int: return len(self._sequence) - def __eq__(self, other: Self) -> bool: + def __eq__(self, other: object) -> bool: + if not isinstance(other, ProteinChain): + return NotImplemented return ( self._id == other._id and self._sequence == other._sequence @@ -572,7 +585,9 @@ def unpaired_msa(self) -> str | None: def __len__(self) -> int: return len(self._sequence) - def __eq__(self, other: Self) -> bool: + def __eq__(self, other: object) -> bool: + if not isinstance(other, RnaChain): + return NotImplemented return ( self._id == other._id and self._sequence == other._sequence @@ -703,7 +718,7 @@ def fill_missing_fields(self) -> Self: sequence=self.sequence, modifications=self.modifications, residue_ids=self._residue_ids, - description=self._description, + description=self.description, unpaired_msa=self._unpaired_msa or '', ) @@ -784,7 +799,9 @@ def residue_ids(self) -> Sequence[int] | None: def __len__(self) -> int: return len(self._sequence) - def __eq__(self, other: Self) -> bool: + def __eq__(self, other: object) -> bool: + if not isinstance(other, DnaChain): + return NotImplemented return ( self._id == other._id and self._sequence == other._sequence diff --git a/src/alphafold3/common/resources.py b/src/alphafold3/common/resources.py index bd2ed98..edbc307 100644 --- a/src/alphafold3/common/resources.py +++ b/src/alphafold3/common/resources.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/common/safe_pickle.py b/src/alphafold3/common/safe_pickle.py index 14a19ca..cb2eb10 100644 --- a/src/alphafold3/common/safe_pickle.py +++ b/src/alphafold3/common/safe_pickle.py @@ -1,7 +1,16 @@ # Copyright 2025 DeepMind Technologies Limited # -# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of -# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/common/testing/data.py b/src/alphafold3/common/testing/data.py index 5695d45..a67a869 100644 --- a/src/alphafold3/common/testing/data.py +++ b/src/alphafold3/common/testing/data.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/atom_types.py b/src/alphafold3/constants/atom_types.py index fae5109..22f9a48 100644 --- a/src/alphafold3/constants/atom_types.py +++ b/src/alphafold3/constants/atom_types.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/chemical_component_sets.py b/src/alphafold3/constants/chemical_component_sets.py index 077ed8c..fcb97b3 100644 --- a/src/alphafold3/constants/chemical_component_sets.py +++ b/src/alphafold3/constants/chemical_component_sets.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/chemical_components.py b/src/alphafold3/constants/chemical_components.py index b93cb32..d219ff5 100644 --- a/src/alphafold3/constants/chemical_components.py +++ b/src/alphafold3/constants/chemical_components.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -67,15 +76,17 @@ def __init__( if not user_ccd: raise ValueError('User CCD cannot be an empty string.') user_ccd_cifs = { - key: {k: tuple(v) for k, v in value.items()} + key: value.to_dict() for key, value in cif_dict.parse_multi_data_cif(user_ccd).items() } self._dict.update(user_ccd_cifs) - def __getitem__(self, key: str) -> Mapping[str, Sequence[str]]: + def __getitem__(self, key: object) -> Mapping[str, Sequence[str]]: + if not isinstance(key, str): + raise TypeError(f'The CCD key must be a string, got {type(key)}') return self._dict[key] - def __contains__(self, key: str) -> bool: + def __contains__(self, key: object) -> bool: return key in self._dict def __iter__(self) -> Iterator[str]: @@ -87,7 +98,7 @@ def __len__(self) -> int: def __hash__(self) -> int: return id(self) # Ok since this is immutable. - def get( + def get( # pyrefly: ignore[bad-override] self, key: str, default: None | Mapping[str, Sequence[str]] = None ) -> Mapping[str, Sequence[str]] | None: return self._dict.get(key, default) diff --git a/src/alphafold3/constants/converters/ccd_pickle_gen.py b/src/alphafold3/constants/converters/ccd_pickle_gen.py index 6d1dce5..b5b3d17 100644 --- a/src/alphafold3/constants/converters/ccd_pickle_gen.py +++ b/src/alphafold3/constants/converters/ccd_pickle_gen.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -34,7 +43,7 @@ def main(argv: Sequence[str]) -> None: with opener(input_file, 'rb') as f: whole_file = f.read() result = { - key: {k: tuple(v) for k, v in value.items()} + key: value.to_dict() for key, value in tqdm.tqdm( cif_dict.parse_multi_data_cif(whole_file).items(), disable=None ) diff --git a/src/alphafold3/constants/converters/chemical_component_sets_gen.py b/src/alphafold3/constants/converters/chemical_component_sets_gen.py index 046a4e8..d9ca895 100644 --- a/src/alphafold3/constants/converters/chemical_component_sets_gen.py +++ b/src/alphafold3/constants/converters/chemical_component_sets_gen.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/mmcif_names.py b/src/alphafold3/constants/mmcif_names.py index fd191ed..9684dbb 100644 --- a/src/alphafold3/constants/mmcif_names.py +++ b/src/alphafold3/constants/mmcif_names.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/periodic_table.py b/src/alphafold3/constants/periodic_table.py index d4e8771..6e9ca0b 100644 --- a/src/alphafold3/constants/periodic_table.py +++ b/src/alphafold3/constants/periodic_table.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/residue_names.py b/src/alphafold3/constants/residue_names.py index 3a5643f..677c1a1 100644 --- a/src/alphafold3/constants/residue_names.py +++ b/src/alphafold3/constants/residue_names.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/constants/side_chains.py b/src/alphafold3/constants/side_chains.py index 9511830..7dfb216 100644 --- a/src/alphafold3/constants/side_chains.py +++ b/src/alphafold3/constants/side_chains.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -62,18 +71,18 @@ 'VAL': [('N', 'CA', 'CB', 'CG1')], } -CHI_GROUPS_FOR_ATOM = {} +_CHI_GROUPS_FOR_ATOM: dict[tuple[str, str], list[tuple[int, int]]] = {} for res_name, chi_angle_atoms_for_res in CHI_ANGLES_ATOMS.items(): for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res): for atom_i, atom in enumerate(chi_group): - CHI_GROUPS_FOR_ATOM.setdefault((res_name, atom), []).append( + _CHI_GROUPS_FOR_ATOM.setdefault((res_name, atom), []).append( (chi_group_i, atom_i) ) # Mapping from (residue_name, atom_name) pairs to the atom's chi group index # and atom index within that group. CHI_GROUPS_FOR_ATOM: Mapping[tuple[str, str], Sequence[tuple[int, int]]] = ( - CHI_GROUPS_FOR_ATOM + _CHI_GROUPS_FOR_ATOM ) MAX_NUM_CHI_ANGLES: int = 4 diff --git a/src/alphafold3/cpp.cc b/src/alphafold3/cpp.cc index bc3a927..5bb4c0f 100644 --- a/src/alphafold3/cpp.cc +++ b/src/alphafold3/cpp.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 @@ -9,6 +18,7 @@ // https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md #include "alphafold3/data/cpp/msa_profile_pybind.h" +#include "alphafold3/model/json_serialize_pybind.h" #include "alphafold3/model/mkdssp_pybind.h" #include "alphafold3/parsers/cpp/cif_dict_pybind.h" #include "alphafold3/parsers/cpp/fasta_iterator_pybind.h" @@ -37,6 +47,7 @@ PYBIND11_MODULE(cpp, m) { RegisterModuleAggregation(m.def_submodule("aggregation")); RegisterModuleStringArray(m.def_submodule("string_array")); RegisterModuleMmcifAtomSite(m.def_submodule("mmcif_atom_site")); + RegisterModuleJsonSerialize(m.def_submodule("json_serialize")); RegisterModuleMkdssp(m.def_submodule("mkdssp")); RegisterModuleMsaProfile(m.def_submodule("msa_profile")); } diff --git a/src/alphafold3/data/cpp/msa_profile_pybind.cc b/src/alphafold3/data/cpp/msa_profile_pybind.cc index 10ba588..aa7fbbb 100644 --- a/src/alphafold3/data/cpp/msa_profile_pybind.cc +++ b/src/alphafold3/data/cpp/msa_profile_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/data/cpp/msa_profile_pybind.h b/src/alphafold3/data/cpp/msa_profile_pybind.h index 19ac96d..c88e7b1 100644 --- a/src/alphafold3/data/cpp/msa_profile_pybind.h +++ b/src/alphafold3/data/cpp/msa_profile_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_DATA_PYTHON_MSA_PROFILE_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_DATA_PYTHON_MSA_PROFILE_PYBIND_H_ diff --git a/src/alphafold3/data/featurisation.py b/src/alphafold3/data/featurisation.py index 98e8591..501e615 100644 --- a/src/alphafold3/data/featurisation.py +++ b/src/alphafold3/data/featurisation.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -42,6 +51,7 @@ def featurise_input( ref_max_modified_date: datetime.date | None = None, conformer_max_iterations: int | None = None, resolve_msa_overlaps: bool = True, + fix_standalone_glycans: bool = False, verbose: bool = False, ) -> Sequence[features.BatchDict]: """Featurise the folding input. @@ -66,6 +76,12 @@ def featurise_input( paper. Set this to false if providing custom paired MSA using the unpaired MSA field to keep it exactly as is as deduplication against the paired MSA could break the manually crafted pairing between MSA sequences. + fix_standalone_glycans: AlphaFold 3 model training and evaluation filtered + out leaving atoms from glycan ligands even if they were not bonded to + anything ("standalone" glycans). Setting this flag to True fixes this + undesirable behavior, but moves away from the regime where AlphaFold 3 was + trained and evaluated. This has only an effect if filter_leaving_atoms is + True in the WholePdbPipeline.Config. verbose: Whether to print progress messages. Returns: @@ -80,6 +96,7 @@ def featurise_input( ref_max_modified_date=ref_max_modified_date, conformer_max_iterations=conformer_max_iterations, resolve_msa_overlaps=resolve_msa_overlaps, + fix_standalone_glycans=fix_standalone_glycans, ), ) diff --git a/src/alphafold3/data/msa.py b/src/alphafold3/data/msa.py index b07a210..9b59fb1 100644 --- a/src/alphafold3/data/msa.py +++ b/src/alphafold3/data/msa.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -26,7 +35,7 @@ class Error(Exception): - """Error indicatating a problem with MSA Search.""" + """Error indicating a problem with MSA Search.""" def _featurize(seq: str, chain_poly_type: str) -> str | list[int]: @@ -298,7 +307,9 @@ def get_msa_tool( n_iter=msa_tool_config.n_iter, e_value=msa_tool_config.e_value, z_value=msa_tool_config.z_value, + dom_z_value=msa_tool_config.dom_z_value, max_sequences=msa_tool_config.max_sequences, + max_threads=msa_tool_config.max_parallel_shards, ) case msa_config.NhmmerConfig(): return nhmmer.Nhmmer( @@ -308,8 +319,10 @@ def get_msa_tool( database_path=msa_tool_config.database_config.path, n_cpu=msa_tool_config.n_cpu, e_value=msa_tool_config.e_value, + z_value=msa_tool_config.z_value, max_sequences=msa_tool_config.max_sequences, alphabet=msa_tool_config.alphabet, + max_threads=msa_tool_config.max_parallel_shards, ) case _: raise ValueError(f'Unknown MSA tool: {msa_tool_config}.') diff --git a/src/alphafold3/data/msa_config.py b/src/alphafold3/data/msa_config.py index c2b0333..5dd655c 100644 --- a/src/alphafold3/data/msa_config.py +++ b/src/alphafold3/data/msa_config.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/msa_features.py b/src/alphafold3/data/msa_features.py index 2c385a7..379333f 100644 --- a/src/alphafold3/data/msa_features.py +++ b/src/alphafold3/data/msa_features.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/msa_identifiers.py b/src/alphafold3/data/msa_identifiers.py index bd939ef..76debed 100644 --- a/src/alphafold3/data/msa_identifiers.py +++ b/src/alphafold3/data/msa_identifiers.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/msa_store.py b/src/alphafold3/data/msa_store.py deleted file mode 100644 index bb79e92..0000000 --- a/src/alphafold3/data/msa_store.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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 - -"""Interface and implementations for fetching MSA data.""" - -from collections.abc import Sequence -from typing import Protocol, TypeAlias - -from alphafold3.data import msa -from alphafold3.data import msa_config - - -MsaErrors: TypeAlias = Sequence[tuple[msa_config.RunConfig, str]] - - -class MsaProvider(Protocol): - """Interface for providing Multiple Sequence Alignments.""" - - def __call__( - self, - query_sequence: str, - chain_polymer_type: str, - ) -> tuple[msa.Msa, MsaErrors]: - """Retrieve MSA for the given polymer query_sequence. - - Args: - query_sequence: The residue sequence of the polymer to search for. - chain_polymer_type: The polymer type of the query_sequence. This must - match the chain_polymer_type of the provider. - - Returns: - A tuple containing the MSA and MsaErrors. MsaErrors is a Sequence - containing a tuple for each msa_query that failed. Each tuple contains - the failing query and the associated error message. - """ - - -class EmptyMsaProvider: - """MSA provider that returns just the query sequence, useful for testing.""" - - def __init__(self, chain_polymer_type: str): - self._chain_polymer_type = chain_polymer_type - - def __call__( - self, query_sequence: str, chain_polymer_type: str - ) -> tuple[msa.Msa, MsaErrors]: - """Returns an MSA containing just the query sequence, never errors.""" - if chain_polymer_type != self._chain_polymer_type: - raise ValueError( - f'EmptyMsaProvider of type {self._chain_polymer_type} called with ' - f'sequence of {chain_polymer_type=}, {query_sequence=}.' - ) - return ( - msa.Msa.from_empty( - query_sequence=query_sequence, - chain_poly_type=self._chain_polymer_type, - ), - (), - ) diff --git a/src/alphafold3/data/parsers.py b/src/alphafold3/data/parsers.py index a48d0e1..1155a9c 100644 --- a/src/alphafold3/data/parsers.py +++ b/src/alphafold3/data/parsers.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/pipeline.py b/src/alphafold3/data/pipeline.py index db48d03..d518317 100644 --- a/src/alphafold3/data/pipeline.py +++ b/src/alphafold3/data/pipeline.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -155,9 +164,9 @@ def _get_protein_msa_and_templates( @functools.cache def _get_rna_msa( sequence: str, - nt_rna_msa_config: msa_config.NhmmerConfig, - rfam_msa_config: msa_config.NhmmerConfig, - rnacentral_msa_config: msa_config.NhmmerConfig, + nt_rna_msa_config: msa_config.RunConfig, + rfam_msa_config: msa_config.RunConfig, + rnacentral_msa_config: msa_config.RunConfig, ) -> msa.Msa: """Processes a single RNA chain.""" logging.info('Getting RNA MSAs for sequence %s', sequence) @@ -271,11 +280,11 @@ class DataPipelineConfig: uniref90_z_value: int | None = None # Nhmmer databases. ntrna_database_path: str - ntrna_z_value: int | None = None + ntrna_z_value: float | None = None rfam_database_path: str - rfam_z_value: int | None = None + rfam_z_value: float | None = None rna_central_database_path: str - rna_central_z_value: int | None = None + rna_central_z_value: float | None = None # Template search databases. seqres_database_path: str pdb_database_path: str @@ -538,6 +547,7 @@ def process_protein_chain( id=chain.id, sequence=chain.sequence, ptms=chain.ptms, + description=chain.description, unpaired_msa=unpaired_msa, paired_msa=paired_msa, templates=templates, @@ -570,6 +580,7 @@ def process_rna_chain( id=chain.id, sequence=chain.sequence, modifications=chain.modifications, + description=chain.description, unpaired_msa=unpaired_msa, ) diff --git a/src/alphafold3/data/structure_stores.py b/src/alphafold3/data/structure_stores.py index aed02eb..58904a0 100644 --- a/src/alphafold3/data/structure_stores.py +++ b/src/alphafold3/data/structure_stores.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/template_realign.py b/src/alphafold3/data/template_realign.py index 51963f0..d25d40d 100644 --- a/src/alphafold3/data/template_realign.py +++ b/src/alphafold3/data/template_realign.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/template_store.py b/src/alphafold3/data/template_store.py deleted file mode 100644 index b5cf84b..0000000 --- a/src/alphafold3/data/template_store.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -"""Interface and implementations for fetching templates data.""" - -from collections.abc import Mapping -import datetime -from typing import Any, Protocol, TypeAlias - - -TemplateFeatures: TypeAlias = Mapping[str, Any] - - -class TemplateFeatureProvider(Protocol): - """Interface for providing Template Features.""" - - def __call__( - self, - sequence: str, - release_date: datetime.date | None, - include_ligand_features: bool = True, - ) -> TemplateFeatures: - """Retrieve template features for the given sequence and release_date. - - Args: - sequence: The residue sequence of the query. - release_date: The release_date of the template query, this is used to - filter templates for training, ensuring that they do not leak structure - information from the future. - include_ligand_features: Whether to include ligand features. - - Returns: - Template features: A mapping of template feature labels to features, which - may be numpy arrays, bytes objects, or for the special case of label - `ligand_features`, a nested feature map of labels to numpy arrays. - - Raises: - TemplateRetrievalError if the template features were not found. - """ diff --git a/src/alphafold3/data/templates.py b/src/alphafold3/data/templates.py index 9ab9d7e..3f6477d 100644 --- a/src/alphafold3/data/templates.py +++ b/src/alphafold3/data/templates.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/hmmalign.py b/src/alphafold3/data/tools/hmmalign.py index 81cd421..fa8aa67 100644 --- a/src/alphafold3/data/tools/hmmalign.py +++ b/src/alphafold3/data/tools/hmmalign.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/hmmbuild.py b/src/alphafold3/data/tools/hmmbuild.py index 3aa0cb1..09c9311 100644 --- a/src/alphafold3/data/tools/hmmbuild.py +++ b/src/alphafold3/data/tools/hmmbuild.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/hmmsearch.py b/src/alphafold3/data/tools/hmmsearch.py index b62c775..48c7387 100644 --- a/src/alphafold3/data/tools/hmmsearch.py +++ b/src/alphafold3/data/tools/hmmsearch.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/jackhmmer.py b/src/alphafold3/data/tools/jackhmmer.py index bb2fec1..5acc8b5 100644 --- a/src/alphafold3/data/tools/jackhmmer.py +++ b/src/alphafold3/data/tools/jackhmmer.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -299,12 +308,14 @@ def _merged_a3m_tbl_iter(a3m: str) -> Iterable[tuple[str, str, str, str]]: # Skip sequences for which we don't have tbl information. yield sequence, description, tbl_info, name - def sort_key(seq_data: tuple[str, str, str, str]) -> tuple[float, str]: + def sort_key(seq_data: tuple[str, str, str, str]) -> tuple[float, float, str]: unused_seq, unused_description, tbl_info, name = seq_data # Tblout lines have 19 whitespace delimited columns. "-" used if no value - # present. We want e-value in column with index 4, so do only 5 splits. - # Use the name in case of a e-value tie. - return float(tbl_info.split(maxsplit=5)[4]), name + # present. We want e-value (column 5) and bit score (column 6), so do only 6 + # splits. E-value and bit score are equivalent, but bit score might have + # higher resolution. Use the name in case of a tie. + e_value, bit_score = tbl_info.split(maxsplit=6)[4:6] + return float(e_value), -float(bit_score), name # A3M/TBL is sorted by e-value and name, hence we can merge them efficiently. merged_a3m_and_tblout = heapq.merge( diff --git a/src/alphafold3/data/tools/msa_tool.py b/src/alphafold3/data/tools/msa_tool.py index c6c076c..29b79c8 100644 --- a/src/alphafold3/data/tools/msa_tool.py +++ b/src/alphafold3/data/tools/msa_tool.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/nhmmer.py b/src/alphafold3/data/tools/nhmmer.py index f833ac7..9b9ac87 100644 --- a/src/alphafold3/data/tools/nhmmer.py +++ b/src/alphafold3/data/tools/nhmmer.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -54,7 +63,7 @@ def __init__( NOTE: The MSA obtained by running against sharded dbs won't be always exactly the same as the MSA obtained by running against an unsharded db. - This is because of Jackhmmer deduplication logic, which won't spot duplicate + This is because of Nhmmer deduplication logic, which won't spot duplicate hits across multiple shards. Usually this means that the sharded search finds more hits (likely bounded by the number of shards), but this should not pose an issue given how the results are used downstream. The problem is diff --git a/src/alphafold3/data/tools/rdkit_utils.py b/src/alphafold3/data/tools/rdkit_utils.py index a9b61a1..0458c1c 100644 --- a/src/alphafold3/data/tools/rdkit_utils.py +++ b/src/alphafold3/data/tools/rdkit_utils.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/shards.py b/src/alphafold3/data/tools/shards.py index 09e2387..1af49cf 100644 --- a/src/alphafold3/data/tools/shards.py +++ b/src/alphafold3/data/tools/shards.py @@ -1,7 +1,16 @@ # Copyright 2025 DeepMind Technologies Limited # -# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of -# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/data/tools/subprocess_utils.py b/src/alphafold3/data/tools/subprocess_utils.py index 0ab15b9..168c7d7 100644 --- a/src/alphafold3/data/tools/subprocess_utils.py +++ b/src/alphafold3/data/tools/subprocess_utils.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/jax/attention/attention.py b/src/alphafold3/jax/attention/attention.py deleted file mode 100644 index 2632003..0000000 --- a/src/alphafold3/jax/attention/attention.py +++ /dev/null @@ -1,139 +0,0 @@ -# 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 - -"""Scaled dot-product attention.""" - -import typing -from typing import Literal, TypeAlias - -from alphafold3.jax.attention import attention_base as base -from alphafold3.jax.attention import flash_attention as attention_triton -from alphafold3.jax.attention import xla_attention -from 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"] - - -@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", - ) - - 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) diff --git a/src/alphafold3/jax/attention/attention_base.py b/src/alphafold3/jax/attention/attention_base.py deleted file mode 100644 index 1b12869..0000000 --- a/src/alphafold3/jax/attention/attention_base.py +++ /dev/null @@ -1,363 +0,0 @@ -# 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 - -"""Common types and utilities for attention kernels.""" - -import abc -import dataclasses -import enum -import functools -import math -from typing import Any, Self - -from alphafold3.jax.common import array_view -from alphafold3.jax.common import precision as precision_lib -import jax -import jax.numpy as jnp -from jax.typing import DTypeLike # pylint: disable=g-importing-member -import jaxtyping -from jaxtyping import Array, Bool, Float, Int # pylint: disable=g-multiple-import,g-importing-member -import typeguard - - -class AUTO: # Used as a sentinel value. - pass - - -DotPrecisionLike = jax.lax.Precision | precision_lib.DotPrecision - - -@jax.tree_util.register_pytree_node_class -@dataclasses.dataclass(frozen=True) -class Mask: - """An attention mask. - - `k_start` (inclusive) and `k_end` (exclusive) define range of enabled - k-sequence values for each row of logits. - - For example, a local attention mask could be defined as follows: - ``` - seq_len_q = seq_len_k = 4 - window_size = 2 - k_start = jnp.maximum(0, jnp.arange(seq_len_q) + 1 - window_size) - mask = Mask(k_start=k_start, is_causal=True) - assert mask.as_array(seq_len_q, seq_len_k) == jnp.array( - [[1, 0, 0, 0], - [1, 1, 0, 0], - [0, 1, 1, 0], - [0, 0, 1, 1]], dtype=bool) - ``` - Or equivalently (but less efficiently): - ``` - k_end = jnp.arange(seq_len_q) + 1 - k_start = jnp.maximum(0, k_end - window_size) - mask = Mask(k_start=k_start, k_end=k_end) - assert mask.as_array(seq_len_q, seq_len_k) == jnp.array( - [[1, 0, 0, 0], - [1, 1, 0, 0], - [0, 1, 1, 0], - [0, 0, 1, 1]], dtype=bool) - ``` - - A mask for two independent causal sequences could be defined as follows: - ``` - k_start = jnp.array([0, 0, 2, 2]) - mask = Mask(k_start=k_start, is_causal=True) - assert mask.as_array(seq_len_q, seq_len_k) == jnp.array( - [[1, 0, 0, 0], - [1, 1, 0, 0], - [0, 0, 1, 0], - [0, 0, 1, 1]], dtype=bool) - ``` - """ - - bool_mask: Bool[Array, "*#B #T #t"] | None = None - _: dataclasses.KW_ONLY - q_start: Int[Array, "*#B #t"] | None = None - q_end: Int[Array, "*#B #t"] | None = None - k_start: Int[Array, "*#B #T"] | None = None - k_end: Int[Array, "*#B #T"] | None = None - is_causal: bool = False - - def tree_flatten(self): - return ( - self.bool_mask, - self.q_start, - self.q_end, - self.k_start, - self.k_end, - ), (self.is_causal,) - - @classmethod - def tree_unflatten(cls, aux, children) -> Self: - (is_causal,) = aux - bool_mask, q_start, q_end, k_start, k_end = children - return cls( - bool_mask, - q_start=q_start, - q_end=q_end, - k_start=k_start, - k_end=k_end, - is_causal=is_causal, - ) - - def as_array( - self, - q_len_or_indices: int | Int[Array, "*#B T"], - k_len_or_indices: int | Int[Array, "*#B t"], - ) -> Bool[Array, "*#B #T #t"] | None: - """Returns the mask as a boolean array.""" - if isinstance(q_len_or_indices, int): - q_indices = jnp.arange(q_len_or_indices) - else: - q_indices = q_len_or_indices - - if isinstance(k_len_or_indices, int): - k_indices = jnp.arange(k_len_or_indices) - else: - k_indices = k_len_or_indices - - q_indices = q_indices[..., None] - k_indices = k_indices[..., None, :] - - mask = [] - if self.bool_mask is not None: - mask.append(self.bool_mask) - # Check `bool_mask` shape is compatible with `{q,kv}_indices`. - _ = jnp.broadcast_shapes( - q_indices.shape, k_indices.shape, self.bool_mask.shape - ) - - if self.q_start is not None: - mask.append(q_indices >= self.q_start[..., None, :]) - - if self.q_end is not None: - mask.append(q_indices < self.q_end[..., None, :]) - - if self.k_start is not None: - mask.append(k_indices >= self.k_start[..., None]) - - if self.k_end is not None: - mask.append(k_indices < self.k_end[..., None]) - - if self.is_causal: - mask.append(q_indices >= k_indices) - - logical_and = functools.partial(functools.reduce, jnp.logical_and) - return jax.lax.broadcast_to_rank(logical_and(mask), 3) if mask else None - - def take(self, *attrs: str) -> tuple[Any, ...]: - """Returns a mask with attrs removed and the removed attrs.""" - default_mask = type(self)() - replacements = {attr: getattr(default_mask, attr) for attr in attrs} - values = (getattr(self, attr) for attr in attrs) - return dataclasses.replace(self, **replacements), *values - - def __and__(self, other: "Bool[Array, '*#B #T #t'] | Mask") -> "Mask": # pylint: disable=g-inconsistent-quotes - """Returns the intersection of two masks.""" - if not isinstance(other, Mask): - other = Mask(other) - - def combine(op): - return lambda a, b: b if a is None else a if b is None else op(a, b) - - return Mask( - bool_mask=combine(jnp.logical_and)(self.bool_mask, other.bool_mask), - q_end=combine(jnp.minimum)(self.q_end, other.q_end), - k_start=combine(jnp.maximum)(self.k_start, other.k_start), - k_end=combine(jnp.minimum)(self.k_end, other.k_end), - is_causal=self.is_causal or other.is_causal, - ) - - -CAUSAL_MASK = Mask(is_causal=True) - - -SoftmaxResidual = ( - tuple[Float[Array, "*B H T"], Float[Array, "*B H T"]] - | Float[Array, "*B H T"] -) - - -@enum.unique -class SoftmaxResidualMode(enum.Enum): - """The mode of storing softmax residuals for the backwards pass. - - The stable softmax calculation performs two reductions calculating: - - the maximum input value (`x_max`), - - the sum of exponentiated values (`denom`). - - We can store these values as residuals to avoid the need to recompute them - in the backwards pass. - - It is also possible to combine the two residuals into a single residual, - `res = x_max + log(denom)`, as `exp(x - res) === exp(x - x_max - log(denom)) - === exp(x - x_max) / denom`. Combining the residuals reduces the memory usage - of the residuals, but will reduce the accuracy of the backwards pass if - `abs(x_max) >> log(denom)`. - """ - - SEPARATE = "separate" - COMBINED = "combined" - - def conform(self, aux: SoftmaxResidual) -> SoftmaxResidual | None: - match self, aux: - case None, _: - return None - case SoftmaxResidualMode.SEPARATE, (_, _): - return aux - case SoftmaxResidualMode.SEPARATE, _: # pytype: disable=redundant-match # b/300135240 - raise ValueError("`aux` has been combined.") - case SoftmaxResidualMode.COMBINED, (x_max, denom): - return x_max + jnp.log(denom) - case SoftmaxResidualMode.COMBINED, _: # pytype: disable=redundant-match # b/300135240 - return aux - - -class DotProductAttention(abc.ABC): - """Dot product attention function.""" - - @jaxtyping.jaxtyped(typechecker=typeguard.typechecked) - def __call__( - self, - query: Float[Array | array_view.ArrayView, "*B T H D"], - key: Float[Array | array_view.ArrayView, "*B t h D"], - value: Float[Array | array_view.ArrayView, "*B t h D"], - *, - precision: ( - DotPrecisionLike | tuple[DotPrecisionLike, DotPrecisionLike] - ) = jax.lax.Precision.DEFAULT, - logits_dtype: DTypeLike | type[AUTO] = AUTO, - bias: Float[Array, "*#B #H #T #t"] | None = None, - mask: Bool[Array, "*#B #H #T #t"] | Mask | None = None, - q_indices: Int[Array, "*#B #H T"] | None = None, - k_indices: Int[Array, "*#B #H t"] | None = None, - ) -> Float[Array, "*B T H D"]: - """Performs scaled dot-product attention. - - Scaled dot-product attention from "Attention is all you need" - https://arxiv.org/abs/1706.03762. - - Computes self- or cross-attention. The following is computed: - softmax(qk_scale * query @ key^T + bias) @ value. - - Supports both multi-head and multi-query attention - (https://arxiv.org/abs/1911.02150). - - Arguments: - query: Query array of shape `[batch, seq_len_q, num_heads_q, head_dim]`. - It must be a multiple of num_heads_kv. - Here's an example of how q/kv heads are interleaved: - For 8 key/value heads and 4 query heads: - - key/value heads [0, 1] see query head 0 - - key/value heads [2, 3] see query head 1 - - key/value heads [4, 5] see query head 2 - key: Key array of shape `[batch, seq_len_kv, num_heads_kv, head_dim]`. It - must be divisible by num_heads_q. - value: Value array of shape `[batch, seq_len_kv, num_heads_kv, head_dim]`. - precision: The precision for the dot products. Either a tuple `( - query_key_dot_precision, weights_value_dot_precision)` or a single - precision applied to both dot products. - logits_dtype: Data type for attention logits (`query @ key^T`). If `AUTO` - is passed (the default), the accumulator type from the `query @ key^T` - dot product will be used. - bias: Optional bias array, broadcastable to shape `[batch, num_heads, - seq_len_q, seq_len_kv]`. - mask: Optional boolean mask, broadcastable to `[batch, num_heads, - seq_len_q, seq_len_kv]`. Attention weights are masked out if the - corresponding mask value is `False`. - q_indices: Optional indices for each token in query sequence. - k_indices: Optional indices for each token in key/value sequence. - - Returns: - An array with the same shape as `query`. - """ # fmt: skip - return self.fwd( - query, - key, - value, - precision=precision, - logits_dtype=logits_dtype, - bias=bias, - mask=mask, - q_indices=q_indices, - k_indices=k_indices, - ) - - @jaxtyping.jaxtyped(typechecker=typeguard.typechecked) - def fwd( - self, - query: Float[Array | array_view.ArrayView, "*B T H D"], - key: Float[Array | array_view.ArrayView, "*B t h D"], - value: Float[Array | array_view.ArrayView, "*B t h D"], - *, - precision: ( - DotPrecisionLike | tuple[DotPrecisionLike, DotPrecisionLike] - ) = jax.lax.Precision.DEFAULT, - logits_dtype: DTypeLike | type[AUTO] = AUTO, - bias: Float[Array, "*#B #H #T #t"] | None = None, - mask: Bool[Array, "*#B #H #T #t"] | Mask | None = None, - q_indices: Int[Array, "*#B #H T"] | None = None, - k_indices: Int[Array, "*#B #H t"] | None = None, - ) -> Float[Array, "*B T H D"]: - """Performs attention.""" - if not isinstance(precision, tuple): - precision = (precision, precision) - - q_k_dot_precision, weights_v_dot_precision = precision - - if not isinstance(q_k_dot_precision, precision_lib.DotPrecision): - q_k_dot_precision = precision_lib.get_equivalent_dot_precision( - query.dtype, key.dtype, q_k_dot_precision - ) - - if not isinstance(weights_v_dot_precision, precision_lib.DotPrecision): - weights_v_dot_precision = precision_lib.get_equivalent_dot_precision( - value.dtype, value.dtype, weights_v_dot_precision - ) - - if logits_dtype is AUTO: - logits_dtype = q_k_dot_precision.accumulator_dtype - - if not isinstance(mask, Mask): - mask = Mask(mask) - - return self._fwd( - array_view.as_array_view(query), - array_view.as_array_view(key), - array_view.as_array_view(value), - q_k_dot_precision=q_k_dot_precision, - logits_dtype=jnp.dtype(logits_dtype), - logits_scale=1 / math.sqrt(query.shape[-1]), - bias=bias, - mask=mask, - weights_v_dot_precision=weights_v_dot_precision, - q_indices=q_indices, - k_indices=k_indices, - ) - - @abc.abstractmethod - def _fwd( - self, - q: Float[array_view.ArrayView, "*B T H D"], - k: Float[array_view.ArrayView, "*B t h D"], - v: Float[array_view.ArrayView, "*B t h D"], - *, - q_k_dot_precision: precision_lib.DotPrecision, - logits_dtype: jnp.dtype, - logits_scale: float, - bias: Float[Array, "*#B #H #T #t"] | None, - mask: Mask | None, - weights_v_dot_precision: precision_lib.DotPrecision, - q_indices: Int[Array, "*#B #H T"] | None = None, - k_indices: Int[Array, "*#B #H t"] | None = None, - ) -> Float[Array, "*B T H D"]: - """Performs attention.""" - ... diff --git a/src/alphafold3/jax/attention/attention_call_arg_specs.py b/src/alphafold3/jax/attention/attention_call_arg_specs.py deleted file mode 100644 index 9f0cfc3..0000000 --- a/src/alphafold3/jax/attention/attention_call_arg_specs.py +++ /dev/null @@ -1,62 +0,0 @@ -# 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 - -"""Attention call argument specifications. - -Attention argument specifications used by users of the library. -They are the most important test cases, and also cases for optimize -performance of via autotuning. -""" - -from typing import Any - -import jax - -ShapedArray = jax.ShapeDtypeStruct - - -def _make_argspec( - *, - q_shape, - dtype, - k_shape=None, - v_shape=None, - bias_shape=None, - mask_shape=None, - **kwargs, -) -> dict[str, Any]: - """Make argspec from shapes and kwargs.""" - if k_shape is None: - k_shape = q_shape - if v_shape is None: - v_shape = k_shape - - return dict( - query=ShapedArray(q_shape, dtype), - key=ShapedArray(k_shape, dtype), - value=ShapedArray(v_shape, dtype), - bias=ShapedArray(bias_shape, dtype) if bias_shape is not None else None, - mask=ShapedArray(mask_shape, 'bool_') if mask_shape is not None else None, - **kwargs, - ) - - -# A subset of the full set of argument specifications. Useful for tap-tests and -# microbenchmarks. -CALL_ARG_SPECS = dict( - vanilla_f32=_make_argspec(q_shape=(8, 1024, 4, 128), dtype='float32'), - vanilla_bf16=_make_argspec(q_shape=(8, 1024, 4, 128), dtype='bfloat16'), - alphafold=_make_argspec( - q_shape=(384, 384, 4, 32), - bias_shape=(1, 4, 384, 384), - mask_shape=(384, 1, 1, 384), - dtype='bfloat16', - ), -) diff --git a/src/alphafold3/jax/attention/flash_attention.py b/src/alphafold3/jax/attention/flash_attention.py deleted file mode 100644 index e5befb1..0000000 --- a/src/alphafold3/jax/attention/flash_attention.py +++ /dev/null @@ -1,703 +0,0 @@ -# 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 - -"""Triton FlashAttention implementation.""" - -import dataclasses -import functools - -from alphafold3.jax.attention import attention_base as base -from alphafold3.jax.common import array_view -from alphafold3.jax.common import precision as precision_lib -from alphafold3.jax.common import triton_utils -import jax -import jax.numpy as jnp -import jax_triton as jt -import jaxtyping -from jaxtyping import Array, Bool, Float, Int # pylint: disable=g-multiple-import,g-importing-member -import triton -import triton.language as tl -import typeguard - - -@triton.jit -def _fwd_kernel_inner( - start_loop, - end_loop, - q, - span_q, - k_block_ptr, - v_block_ptr, - bias_block_ptr, - mask_block_ptr, - k_start, - k_end, - seq_len_k, - acc, - m_i, - l_i, - bias_advance: tl.constexpr, - mask_advance: tl.constexpr, - is_causal: tl.constexpr, - use_attention_mask: tl.constexpr, - use_k_start: tl.constexpr, - use_k_end: tl.constexpr, - use_bias: tl.constexpr, - block_k: tl.constexpr, - use_mask_k: tl.constexpr, - k_boundary_check: tl.constexpr, - v_boundary_check: tl.constexpr, - dot_fn_qk: tl.constexpr, - dot_fn_kv: tl.constexpr, -): - """Triton MHA forward kernel's inner loop.""" - - for start_k in range(start_loop, end_loop, block_k): - start_k = tl.multiple_of(start_k, block_k) - span_k = start_k + tl.arange(0, block_k) - - k = tl.load( - k_block_ptr, - boundary_check=k_boundary_check, - padding_option="zero" if len(k_boundary_check.value) else "", - ) - v = tl.load( - v_block_ptr, - boundary_check=v_boundary_check, - padding_option="zero" if len(v_boundary_check.value) else "", - ) - - if use_bias: - bias = tl.load(bias_block_ptr) - - qk = dot_fn_qk(q.to(k.dtype), k) # [block_q, block_k] - - if use_bias: - # Prevent dot accumulating into the bias tensor. It appears that Triton - # doesn't pipeline the bias load as it does the `k` load, so the bias load - # blocks the matmul if the add is merged. - qk = qk.to(tl.uint32, bitcast=True) & 0xFFFFFFFF - qk = qk.to(tl.float32, bitcast=True) - qk += bias - - if use_attention_mask | use_k_start | use_k_end: - mask_value = float(jnp.finfo(jnp.float32).min) - - if use_attention_mask: - mask = tl.load(mask_block_ptr) - qk = tl.where(mask, qk, mask_value) - - if use_k_start: - # This check is there to work around a triton compiler bug, but it - # shouldn't be strictly needed. - if tl.sum(k_start) != 0: - qk = tl.where(k_start[:, None] <= span_k[None, :], qk, mask_value) - if is_causal: - qk = tl.where(span_q[:, None] >= span_k[None, :], qk, float("-inf")) - elif use_k_end: - # When called with k_end and is_causal=True, the causal mask gets folded - # into k_end and is_causal is set to False. - qk = tl.where(k_end[:, None] > span_k[None, :], qk, mask_value) - - if use_mask_k: - qk = tl.where((span_k < seq_len_k)[None, :], qk, float("-inf")) - - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) # Shape [block_q]. - p = tl.exp(qk - m_ij[:, None]) # Shape [block_q, block_k]. - alpha = tl.exp(m_i - m_ij) - m_i = m_ij - acc *= alpha[:, None] - l_i *= alpha - l_i += tl.sum(p, axis=1) - - # Add the new block of attention weights. - acc += dot_fn_kv(p.to(v.dtype), v) - - k_block_ptr = tl.advance(k_block_ptr, (0, block_k)) - v_block_ptr = tl.advance(v_block_ptr, (block_k, 0)) - bias_block_ptr = tl.advance(bias_block_ptr, bias_advance.value) - mask_block_ptr = tl.advance(mask_block_ptr, mask_advance.value) - - return ( - k_block_ptr, - v_block_ptr, - bias_block_ptr, - mask_block_ptr, - acc, - m_i, - l_i, - ) - - -# Based on Algorithm 1 of https://arxiv.org/abs/2205.14135. -# Inspired by the official Triton tutorial implementation -# https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html -@triton.jit -def _fwd_kernel( - # Input arrays. - q_ptr, - k_ptr, - v_ptr, - bias_ptr, - mask_ptr, - k_start_ptr, - k_end_ptr, - # Scalar inputs. - q_offset, - k_offset, - v_offset, - q_stride_b, - q_stride_s, - q_stride_h, - q_stride_d, - k_stride_b, - k_stride_s, - k_stride_h, - k_stride_d, - v_stride_b, - v_stride_s, - v_stride_h, - v_stride_d, - bias_stride_b, - bias_stride_h, - bias_stride_sq, - bias_stride_sk, - mask_stride_b, - mask_stride_h, - mask_stride_sq, - mask_stride_sk, - k_start_stride_b, - k_start_stride_h, - k_start_stride_sq, - k_end_stride_b, - k_end_stride_h, - k_end_stride_sq, - o_stride_b, - o_stride_s, - o_stride_h, - o_stride_d, - num_heads_q, - num_heads_k, - seq_len_q, - seq_len_k, - # Output arrays. - o_ptr, - # Compile-time constants. - is_causal: tl.constexpr, - use_attention_mask: tl.constexpr, - use_k_start: tl.constexpr, - use_k_end: tl.constexpr, - use_bias: tl.constexpr, - sm_scale: tl.constexpr, - block_q: tl.constexpr, - block_k: tl.constexpr, - head_dim: tl.constexpr, - use_mask_q: tl.constexpr, - use_mask_k: tl.constexpr, - bias_bcast_sq: tl.constexpr, - mask_bcast_sq: tl.constexpr, - dot_fn_qk: tl.constexpr, - dot_fn_kv: tl.constexpr, -): - """Triton MHA forward kernel.""" - # pytype: disable=annotation-type-mismatch,unsupported-operands - block_d: tl.constexpr = jt.utils.next_power_of_2(head_dim.value) - - # Each thread block processes one batch element (b) and one head (h). - start_q = tl.program_id(1) * block_q - off_h = tl.program_id(0) # int in [0, num_heads_o). - off_b = tl.program_id(2) # int in [0, batch_size) - - off_h_k = off_h // (num_heads_q // num_heads_k) - - q_ptr += off_h * q_stride_h + off_b * q_stride_b + q_offset - k_ptr += off_h_k * k_stride_h + off_b * k_stride_b + k_offset - v_ptr += off_h_k * v_stride_h + off_b * v_stride_b + v_offset - o_ptr += off_h * o_stride_h + off_b * o_stride_b - - if use_bias: - bias_ptr += off_b * bias_stride_b + off_h * bias_stride_h - if use_attention_mask: - mask_ptr += off_b * mask_stride_b + off_h * mask_stride_h - if use_k_start: - k_start_ptr += off_b * k_start_stride_b + off_h * k_start_stride_h - if use_k_end: - k_end_ptr += off_b * k_end_stride_b + off_h * k_end_stride_h - - q_block_ptr = tl.make_block_ptr( - q_ptr, - shape=(seq_len_q, head_dim), - strides=(q_stride_s, q_stride_d), - offsets=(start_q, 0), - block_shape=(block_q, block_d), - order=(1, 0), - ) - k_block_ptr = tl.make_block_ptr( - k_ptr, - shape=(head_dim, seq_len_k), - strides=(k_stride_d, k_stride_s), - offsets=(0, 0), - block_shape=(block_d, block_k), - order=(0, 1), - ) - v_block_ptr = tl.make_block_ptr( - v_ptr, - shape=(seq_len_k, head_dim), - strides=(v_stride_s, v_stride_d), - offsets=(0, 0), - block_shape=(block_k, block_d), - order=(1, 0), - ) - - q_boundary_check0: tl.constexpr = (0,) if use_mask_q else () - q_boundary_check1: tl.constexpr = (1,) if head_dim != block_d else () - q_boundary_check: tl.constexpr = q_boundary_check0 + q_boundary_check1 - q_padding_option: tl.constexpr = "zero" if len(q_boundary_check.value) else "" - k_boundary_check: tl.constexpr = (0,) if head_dim != block_d else () - v_boundary_check: tl.constexpr = (0,) if use_mask_k else () - - # If broadcasting in a given dim, use a 1D block (observed to be faster). - bias_start_dim: tl.constexpr = 1 if bias_bcast_sq else 0 - bias_block_ptr = tl.make_block_ptr( - bias_ptr, - shape=(seq_len_q, seq_len_k)[bias_start_dim:], - strides=(bias_stride_sq, bias_stride_sk)[bias_start_dim:], - offsets=(start_q, 0)[bias_start_dim:], - block_shape=(block_q, block_k)[bias_start_dim:], - order=(1, 0)[bias_start_dim:], - ) - bias_advance: tl.constexpr = (0, block_k)[bias_start_dim:] - - mask_start_dim: tl.constexpr = 1 if mask_bcast_sq else 0 - mask_block_ptr = tl.make_block_ptr( - mask_ptr, - shape=(seq_len_q, seq_len_k)[mask_start_dim:], - strides=(mask_stride_sq, mask_stride_sk)[mask_start_dim:], - offsets=(start_q, 0)[mask_start_dim:], - block_shape=(block_q, block_k)[mask_start_dim:], - order=(1, 0)[mask_start_dim:], - ) - mask_advance: tl.constexpr = (0, block_k)[mask_start_dim:] - - k_start_block_ptr = tl.make_block_ptr( - k_start_ptr, - shape=(seq_len_q,), - strides=(k_start_stride_sq,), - offsets=(start_q,), - block_shape=(block_q,), - order=(0,), - ) - k_end_block_ptr = tl.make_block_ptr( - k_end_ptr, - shape=(seq_len_q,), - strides=(k_end_stride_sq,), - offsets=(start_q,), - block_shape=(block_q,), - order=(0,), - ) - # pytype: enable=annotation-type-mismatch,unsupported-operands - - # Each thread block processes a block of block_q queries. - span_q = start_q + tl.arange(0, block_q) - - # m_i and l_i (see FlashAttention paper) are updated during the k,v loop. - m_i = tl.full([block_q], float("-inf"), dtype=tl.float32) - l_i = tl.zeros([block_q], dtype=tl.float32) - # acc is the buffer where we accumulate the output on sram. - acc = tl.zeros([block_q, block_d], dtype=tl.float32) - - # Load q: it will stay in smem throughout. Indices form a matrix because we - # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index. - q = tl.load( - q_block_ptr, - boundary_check=q_boundary_check, - padding_option=q_padding_option, - ) - q *= sm_scale - - # In FlashAttention algorithm 1 there are 2 loops: slow over tiles of kv (size - # (Bc == block_k here), and fast over blocks of q (size Br == block_q here). - # Here we only loop over blocks of kv to process entire seq_len, the loop over - # blocks of q is carried out by the grid. - k_start = None - if use_k_start: - k_start = tl.load(k_start_block_ptr) - start_loop = tl.maximum(tl.min(k_start), 0) - blocks_to_skip = start_loop // block_k - start_loop = block_k * blocks_to_skip # Floor to multiple of block_k. - for _ in range(blocks_to_skip): - # Advance all block pointers to the first valid block. - k_block_ptr = tl.advance(k_block_ptr, (0, block_k)) - v_block_ptr = tl.advance(v_block_ptr, (block_k, 0)) - bias_block_ptr = tl.advance(bias_block_ptr, bias_advance.value) - mask_block_ptr = tl.advance(mask_block_ptr, mask_advance.value) - else: - start_loop = 0 - - k_end = None - if is_causal: - end_loop = tl.minimum((start_q // block_k) * block_k, seq_len_k) - elif use_k_end: - k_end = tl.load(k_end_block_ptr) - end_loop = tl.minimum(tl.max(k_end), seq_len_k) - else: - end_loop = seq_len_k - - ( - k_block_ptr, - v_block_ptr, - bias_block_ptr, - mask_block_ptr, - acc, - m_i, - l_i, - ) = _fwd_kernel_inner( - start_loop, - end_loop, - q, - span_q, - k_block_ptr, - v_block_ptr, - bias_block_ptr, - mask_block_ptr, - k_start, - k_end, - seq_len_k, - acc, - m_i, - l_i, - bias_advance, - mask_advance, - False, # is_causal - use_attention_mask, - use_k_start, - use_k_end, - use_bias, - block_k, - use_mask_k, - k_boundary_check, - v_boundary_check, - dot_fn_qk, - dot_fn_kv, - ) - - if is_causal: - tl.debug_barrier() # Help compiler schedule loops independently. - start_loop, end_loop = end_loop, tl.minimum(end_loop + block_k, seq_len_k) - - _, _, _, _, acc, _, l_i = _fwd_kernel_inner( - start_loop, - end_loop, - q, - span_q, - k_block_ptr, - v_block_ptr, - bias_block_ptr, - mask_block_ptr, - k_start, - k_end, - seq_len_k, - acc, - m_i, - l_i, - bias_advance, - mask_advance, - True, # is_causal - use_attention_mask, - use_k_start, - use_k_end, - use_bias, - block_k, - use_mask_k, - k_boundary_check, - v_boundary_check, - dot_fn_qk, - dot_fn_kv, - ) - - # It is possible that every value in a row was masked to f32 min or that the - # main loop has been completely optimised out, and that `l_i` is `0` for that - # row. Add epsilon value to avoid NaNs from `0 / 0`. - l_i += float(jnp.finfo(jnp.float32).tiny) - - acc /= l_i[:, None] - - # Write output to dram. - o_block_ptr = tl.make_block_ptr( - o_ptr, - shape=(seq_len_q, head_dim), - strides=(o_stride_s, o_stride_d), - offsets=(start_q, 0), - block_shape=(block_q, block_d), - order=(1, 0), - ) - acc = acc.to(o_ptr.dtype.element_ty) - tl.store(o_block_ptr, acc, boundary_check=q_boundary_check) - - -@jaxtyping.jaxtyped(typechecker=typeguard.typechecked) -def _fwd( - q: Float[array_view.ArrayView, "*B T H D"], - k: Float[array_view.ArrayView, "*B t h D"], - v: Float[array_view.ArrayView, "*B t h D"], - bias: Float[Array, "*#B #H #T #t"] | None, - mask: Bool[Array, "*#B #H #T #t"] | None, - k_start: Int[Array, "*#B #H #T"] | None, - k_end: Int[Array, "*#B #H #T"] | None, - *, - logits_scale: float, - is_causal: bool, - q_k_dot_precision: precision_lib.DotPrecision, - weights_v_dot_precision: precision_lib.DotPrecision, -) -> Float[Array, "*B T H D"]: - """Forward pass of Triton FlashAttention.""" - - orig_q_shape = q.shape - q = q.collapse(0, -3, allow_copy=True) - batch_size, seq_len_q, num_heads_q, head_dim = q.shape - *_, seq_len_k, num_heads_kv, _ = k.shape - # Maybe broadcast `k`/`v` heads dimension. - kv_shape = (batch_size, seq_len_k, num_heads_kv, head_dim) - k = k.collapse(0, -3, allow_copy=True).broadcast_to(kv_shape) - v = v.collapse(0, -3, allow_copy=True).broadcast_to(kv_shape) - - def get_bias_mask_view(x, dtype): - if x is None: - x = jnp.array([], dtype=dtype) - return array_view.ArrayView(x, shape=(0, 0, 0, 0), strides=(0, 0, 0, 0)) - - shape = orig_q_shape[:-3] + (num_heads_q, seq_len_q, seq_len_k) - return ( - array_view.ArrayView(x) - .broadcast_to(shape) - .collapse(0, -3, allow_copy=True) - ) - - bias = get_bias_mask_view(bias, dtype=q.dtype) - mask = get_bias_mask_view(mask, dtype=jnp.bool_) - - def get_range_view(x, seq_len): - if x is None: - x = jnp.array([], dtype=jnp.int32) - return array_view.ArrayView(x, shape=(0, 0, 0), strides=(0, 0, 0)) - - shape = orig_q_shape[:-3] + (num_heads_q, seq_len) - return ( - array_view.ArrayView(x) - .broadcast_to(shape) - .collapse(0, -2, allow_copy=True) - ) - - k_start = get_range_view(k_start, seq_len_q) - k_end = get_range_view(k_end, seq_len_q) - - block_q = 64 - block_k = 64 - - return jt.triton_call( - q.base, - k.base, - v.base, - bias.base, - mask.base, - k_start.base, - k_end.base, - q.offset, - k.offset, - v.offset, - *q.strides, - *k.strides, - *v.strides, - *bias.strides, - *mask.strides, - k_start.strides, - k_end.strides, - *jt.utils.strides_from_shape(q.shape), # out strides. - num_heads_q, - num_heads_kv, - seq_len_q, - seq_len_k, - kernel=_fwd_kernel, - name="triton_flash_attention", - out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype), - grid=(num_heads_q, triton.cdiv(seq_len_q, block_q), batch_size), - num_stages=2, - num_warps=4, - is_causal=is_causal, - use_attention_mask=(mask.size != 0), - use_k_start=(k_start.size != 0), - use_k_end=(k_end.size != 0), - use_bias=(bias.size != 0), - sm_scale=logits_scale, - block_q=block_q, - block_k=block_k, - head_dim=head_dim, - use_mask_q=(seq_len_q % block_q != 0), - use_mask_k=(seq_len_k % block_q != 0), - bias_bcast_sq=(bias.strides[-2] == 0), - mask_bcast_sq=(mask.strides[-2] == 0), - dot_fn_qk=triton_utils.get_tl_dot_fn(q_k_dot_precision), - dot_fn_kv=triton_utils.get_tl_dot_fn(weights_v_dot_precision), - ).reshape(orig_q_shape) - - -def _as_batched_array_view(x, axis_size): - batched_shape = (axis_size,) + x.shape - batched_strides = (x.base.size // axis_size,) + x.strides - return dataclasses.replace(x, shape=batched_shape, strides=batched_strides) - - -def _fwd_vmap_rule( - axis_size, in_batched, *args, fn: jax.custom_batching.custom_vmap -): - """`vmap` rule for Triton FlashAttention forward op.""" - q, k, v, bias, mask, k_start, k_end = args - ( - q_batched, - k_batched, - v_batched, - bias_batched, - mask_batched, - k_start_batched, - k_end_batched, - ) = in_batched - - if q_batched.base: - q = _as_batched_array_view(q, axis_size) - if k_batched.base: - k = _as_batched_array_view(k, axis_size) - if v_batched.base: - v = _as_batched_array_view(v, axis_size) - - # Triton op requires `q`, `k`, `v` batch dims to be identical. - if q_batched.base and k_batched.base and v_batched.base: - if bias is not None and not bias_batched: - bias = jax.lax.broadcast_to_rank(bias, bias.ndim + 1) - if mask is not None and not mask_batched: - mask = jax.lax.broadcast_to_rank(mask, mask.ndim + 1) - if k_start is not None and not k_start_batched: - k_start = jax.lax.broadcast_to_rank(k_start, k_start.ndim + 1) - if k_end is not None and not k_end_batched: - k_end = jax.lax.broadcast_to_rank(k_end, k_end.ndim + 1) - out = fn(q, k, v, bias, mask, k_start, k_end) - out_batched = True - return out, out_batched - - # Fallback to sequential loop. - q, k, v = map(jnp.asarray, (q, k, v)) - in_batched = [ - q_batched.base, - k_batched.base, - v_batched.base, - bias_batched, - mask_batched, - k_start_batched, - k_end_batched, - ] - - def f(q, k, v, *args, **kwargs): - q, k, v = map(array_view.ArrayView, (q, k, v)) - return fn.fun(q, k, v, *args, **kwargs) - - sequential_vmap = jax.custom_batching.sequential_vmap(f) - return sequential_vmap.vmap_rule(axis_size, in_batched, q, k, v, *args[3:]) - - -def _decompose_mask(mask, q, k, q_indices, k_indices): - """Decomposes `mask` into a mask array, `is_causal`, `k_start` and `k_end`.""" - if mask is None: - return None, False, None, None - - is_causal = False - k_start = None - k_end = None - if q_indices is None and k_indices is None: - mask, is_causal, k_start, k_end = mask.take("is_causal", "k_start", "k_end") - if k_start is not None: - k_start = jax.lax.broadcast_to_rank(k_start, 2) - if k_end is not None: - k_end = jax.lax.broadcast_to_rank(k_end, 2) - if is_causal: # Fold is_causal into k_end - k_end = jnp.minimum(k_end, jnp.arange(1, q.shape[-3] + 1)) - is_causal = False - - q_len_or_indices = q.shape[-3] if q_indices is None else q_indices - k_len_or_indices = k.shape[-3] if k_indices is None else k_indices - return ( - mask.as_array(q_len_or_indices, k_len_or_indices), - is_causal, - k_start, - k_end, - ) - - -@dataclasses.dataclass(frozen=True) -class TritonFlashAttention(base.DotProductAttention): - """Triton FlashAttention implementation.""" - - @jaxtyping.jaxtyped(typechecker=typeguard.typechecked) - def _fwd( - self, - q: Float[array_view.ArrayView, "*B T H D"], - k: Float[array_view.ArrayView, "*B t h D"], - v: Float[array_view.ArrayView, "*B t h D"], - bias: Float[Array, "*#B #H #T #t"] | None, - *, - q_k_dot_precision: precision_lib.DotPrecision, - logits_dtype: jnp.dtype, - logits_scale: float, - mask: base.Mask | None, - weights_v_dot_precision: precision_lib.DotPrecision, - q_indices: Int[Array, "*#B #H T"] | None = None, - k_indices: Int[Array, "*#B #H t"] | None = None, - ) -> Float[Array, "*B T H D"]: - if logits_dtype != jnp.float32: - raise ValueError("`logits_dtype` must be float32.") - - kwargs = dict( - logits_scale=logits_scale, - q_k_dot_precision=q_k_dot_precision, - weights_v_dot_precision=weights_v_dot_precision, - ) - - def attend_fwd( - q, - k, - v, - bias, - mask_, - q_indices, - k_indices, - ): - - mask, is_causal, k_start, k_end = _decompose_mask( - mask_, q, k, q_indices, k_indices - ) - - fwd_closed_kwargs = dict( - is_causal=is_causal, - **kwargs, - ) - - fwd_closed = functools.partial(_fwd, **fwd_closed_kwargs) - fwd_closed = jax.custom_batching.custom_vmap(fwd_closed) - fwd_closed.def_vmap(functools.partial(_fwd_vmap_rule, fn=fwd_closed)) - - return fwd_closed(q, k, v, bias, mask, k_start, k_end) - - return attend_fwd( - q, - k, - v, - bias, - mask, - q_indices, - k_indices, - ) diff --git a/src/alphafold3/jax/attention/xla_attention.py b/src/alphafold3/jax/attention/xla_attention.py deleted file mode 100644 index e095f45..0000000 --- a/src/alphafold3/jax/attention/xla_attention.py +++ /dev/null @@ -1,140 +0,0 @@ -# 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 - -"""XLA implementation of scaled dot-product attention.""" - -import dataclasses - -from alphafold3.jax.attention import attention_base as base -from alphafold3.jax.common import array_view -from alphafold3.jax.common import precision as precision_lib -import jax -import jax.numpy as jnp -import jaxtyping -from jaxtyping import Array, Float, Int # pylint: disable=g-multiple-import,g-importing-member -import typeguard - - -def _get_precision( - backend: str, precision: precision_lib.DotPrecision -) -> jax.lax.Precision: - if backend == "gpu" and precision == precision_lib.DotPrecision.F32_F32: - return jax.lax.Precision.HIGHEST - return jax.lax.Precision.DEFAULT - - -def einsum_with_dot_precision( - subscript: str, - a: jax.Array, - b: jax.Array, - *, - precision: precision_lib.DotPrecision, -) -> jax.Array: - """Evaluate `fn` with the given precision.""" - result = jnp.einsum( - subscript, - a.astype(precision.operand_dtype), - b.astype(precision.operand_dtype), - precision=_get_precision(jax.default_backend().lower(), precision), - preferred_element_type=precision.accumulator_dtype, - ) - assert result.dtype == precision.accumulator_dtype - return result - - -def _softmax(x: jax.Array) -> jax.Array: - """Computes softmax.""" - # Always perform reductions in at least f32 precision. - dtype = jnp.promote_types(x.dtype, jnp.float32) - x_max = jnp.max(x.astype(dtype), axis=-1, keepdims=True) - unnormalized = jnp.exp(x - x_max) - denom = jnp.sum(unnormalized, axis=-1, keepdims=True) - return unnormalized / denom - - -@jaxtyping.jaxtyped(typechecker=typeguard.typechecked) -def _attend( - q: Float[array_view.ArrayView, "*B T H D"], - k: Float[array_view.ArrayView, "*B t #H D"], - v: Float[array_view.ArrayView, "*B t #H D"], - *, - q_k_dot_precision: precision_lib.DotPrecision, - logits_dtype: jnp.dtype, - logits_scale: float, - bias: Float[Array, "*#B #H #T #t"] | None, - mask: base.Mask | None, - weights_v_dot_precision: precision_lib.DotPrecision, - q_indices: Int[Array, "*#B #H T"] | None, - k_indices: Int[Array, "*#B #H t"] | None, -) -> Float[Array, "*B T H D"]: - """Computes attention.""" - logits = einsum_with_dot_precision( - "...qhd,...khd->...hqk", q, k, precision=q_k_dot_precision - ).astype(logits_dtype) - - logits *= logits_scale - - if bias is not None: - logits += bias - - if mask is not None: - q_len_or_indices = q.shape[-3] if q_indices is None else q_indices - k_len_or_indices = k.shape[-3] if k_indices is None else k_indices - mask = mask.as_array(q_len_or_indices, k_len_or_indices) - - if mask is not None: - mask_value = float(jnp.finfo(logits.dtype).min) - - logits = jnp.where(jnp.asarray(mask), logits, mask_value) - - weights = _softmax(logits) - - weights = weights.astype(v.dtype) - out = einsum_with_dot_precision( - "...hqk,...khd->...qhd", weights, v, precision=weights_v_dot_precision - ).astype(q.dtype) - return out - - -@dataclasses.dataclass(frozen=True) -class XlaDotProductAttention(base.DotProductAttention): - """XLA dot product attention function.""" - - _: dataclasses.KW_ONLY - - def _fwd( - self, - q: Float[array_view.ArrayView, "*B T H D"], - k: Float[array_view.ArrayView, "*B t #H D"], - v: Float[array_view.ArrayView, "*B t #H D"], - *, - q_k_dot_precision: precision_lib.DotPrecision, - logits_dtype: jnp.dtype, - logits_scale: float, - bias: Float[Array, "*#B #H #T #t"] | None, - mask: base.Mask | None, - weights_v_dot_precision: precision_lib.DotPrecision, - q_indices: Int[Array, "*#B #H T"] | None = None, - k_indices: Int[Array, "*#B #H t"] | None = None, - ) -> Float[Array, "*B T H D"]: - - return _attend( - q, - k, - v, - bias=bias, - mask=mask, - q_indices=q_indices, - k_indices=k_indices, - q_k_dot_precision=q_k_dot_precision, - logits_dtype=logits_dtype, - logits_scale=logits_scale, - weights_v_dot_precision=weights_v_dot_precision, - ) diff --git a/src/alphafold3/jax/common/array_view.py b/src/alphafold3/jax/common/array_view.py deleted file mode 100644 index c37c299..0000000 --- a/src/alphafold3/jax/common/array_view.py +++ /dev/null @@ -1,405 +0,0 @@ -# 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 - -"""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 diff --git a/src/alphafold3/jax/common/precision.py b/src/alphafold3/jax/common/precision.py deleted file mode 100644 index e39aa84..0000000 --- a/src/alphafold3/jax/common/precision.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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 - -"""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)] diff --git a/src/alphafold3/jax/common/triton_utils.py b/src/alphafold3/jax/common/triton_utils.py deleted file mode 100644 index ebf67fc..0000000 --- a/src/alphafold3/jax/common/triton_utils.py +++ /dev/null @@ -1,125 +0,0 @@ -# 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 - -"""Triton utils.""" - -from collections.abc import Callable, Mapping - -from alphafold3.jax.common import precision as precision_lib -import jax -import jax.numpy as jnp -import triton -import triton.language as tl - - -_JNP_TO_TL_DTYPES: Mapping[jnp.dtype, tl.dtype] = { - jnp.bool_: tl.int1, - jnp.int8: tl.int8, - jnp.int16: tl.int16, - jnp.int32: tl.int32, - jnp.int64: tl.int64, - jnp.uint8: tl.uint8, - jnp.uint16: tl.uint16, - jnp.uint32: tl.uint32, - jnp.uint64: tl.uint64, - jnp.float16: tl.float16, - jnp.bfloat16: tl.bfloat16, - jnp.float32: tl.float32, - jnp.float64: tl.float64, -} - - -def jnp_to_tl_dtype(jnp_dtype: jnp.dtype) -> tl.dtype: - return _JNP_TO_TL_DTYPES[jnp_dtype] - - -def get_tl_dot_fn( - precision: precision_lib.DotPrecision, -) -> Callable[..., tl.tensor]: - """Returns a tl `dot` implementation with the specified precision. - - Args: - precision: The `dot` precision. - """ - if not is_precision_supported(precision): - raise ValueError(f'Unsupported dot precision: {precision}') - - if precision == precision_lib.DotPrecision.TF32_F32_3X: - return _dot_tf32_f32_3x - - in_dtype = jnp_to_tl_dtype(precision.operand_dtype) - out_dtype = jnp_to_tl_dtype(precision.accumulator_dtype) - allow_tf32 = precision == precision_lib.DotPrecision.TF32_F32 - - @tl.core.extern - def _dot_fn( - a: tl.core.tensor, - b: tl.core.tensor, - *, - trans_a: bool = False, - trans_b: bool = False, - _builder, - ): - if in_dtype == tl.float32: - tl.static_assert(a.dtype == tl.float32, _builder=_builder) - tl.static_assert(b.dtype == tl.float32, _builder=_builder) - else: - tl.static_assert(a.dtype.is_standard_floating(), _builder=_builder) - tl.static_assert(b.dtype.is_standard_floating(), _builder=_builder) - a = a.to(in_dtype, _builder=_builder) - b = b.to(in_dtype, _builder=_builder) - a = tl.trans(a, _builder=_builder) if trans_a else a - b = tl.trans(b, _builder=_builder) if trans_b else b - return tl.dot( - a, b, allow_tf32=allow_tf32, out_dtype=out_dtype, _builder=_builder - ) - - return _dot_fn - - -def is_precision_supported(precision: precision_lib.DotPrecision) -> bool: - return precision in { - precision_lib.DotPrecision.F32_F32, - precision_lib.DotPrecision.TF32_F32, - precision_lib.DotPrecision.F16_F32, - precision_lib.DotPrecision.BF16_F32, - precision_lib.DotPrecision.TF32_F32_3X, - } - - -@triton.jit -def _dot_tf32_f32_3x(a, b, trans_a=False, trans_b=False): - """Perform the 3-pass tf32 dot function.""" - tl.static_assert(a.dtype == tl.float32) - tl.static_assert(b.dtype == tl.float32) - a_ = (a.to(tl.uint32, bitcast=True) & 0xFFFFE000).to(tl.float32, bitcast=True) - b_ = (b.to(tl.uint32, bitcast=True) & 0xFFFFE000).to(tl.float32, bitcast=True) - a_err = a - a_ - b_err = b - b_ - if trans_a: - a_ = tl.trans(a_) - a_err = tl.trans(a_err) - if trans_b: - b_ = tl.trans(b_) - b_err = tl.trans(b_err) - # Add smallest terms first for better accuracy. - return tl.dot(a_, b_, out_dtype=tl.float32) + ( - tl.dot(a_, b_err, out_dtype=tl.float32) - + tl.dot(a_err, b_, out_dtype=tl.float32) - ) - - -def has_triton_support() -> bool: - """Returns True if Triton is supported by the default JAX device.""" - if jax.default_backend() != 'gpu': - return False - - # Only currently supported for Ampere and above. - return float(jax.devices()[0].compute_capability) >= 8.0 diff --git a/src/alphafold3/jax/gated_linear_unit/block.py b/src/alphafold3/jax/gated_linear_unit/block.py deleted file mode 100644 index 43c7e79..0000000 --- a/src/alphafold3/jax/gated_linear_unit/block.py +++ /dev/null @@ -1,126 +0,0 @@ -# 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 - -"""Pallas block load / store utilities.""" - -from collections.abc import Sequence -from typing import Any, TypeAlias - -from alphafold3.jax.common import array_view -import jax -import jax.experimental -from jax.experimental import pallas as pl -import jax.numpy as jnp -import jaxtyping -from jaxtyping import Int # pylint: disable=g-importing-member -import numpy as np -import typeguard - -ArrayT: TypeAlias = Any -ScalarInt: TypeAlias = ( - Int[ArrayT, ""] | Int[np.generic, ""] | Int[jnp.generic, ""] -) - - -@jaxtyping.jaxtyped(typechecker=typeguard.typechecked) -def load_block( - ref, - idx: Sequence[int | ScalarInt], - *, - block_shape: Sequence[int | None], - other=None, - **kwargs, -) -> jax.Array: - """Loads a block from the given `ref`, masking where necessary.""" - idx, mask = _get_block_indexer_and_mask(ref, idx, block_shape=block_shape) - if isinstance(ref, array_view.ArrayView): - idx = ref[idx].offsets - ref = ref.base - other = None if mask is None else other - with jax.experimental.enable_x64(): - return pl.load(ref, idx, mask=mask, other=other, **kwargs) - - -@jaxtyping.jaxtyped(typechecker=typeguard.typechecked) -def store_block( - ref, - val: jax.Array, - idx: Sequence[int | ScalarInt], - *, - block_shape: Sequence[int | None] | None = None, - **kwargs, -): - """Stores a block from the given `ref`, masking where necessary.""" - if block_shape is None: - block_shape = val.shape - idx, mask = _get_block_indexer_and_mask(ref, idx, block_shape=block_shape) - if isinstance(ref, array_view.ArrayView): - idx = ref[idx].offsets - ref = ref.base - with jax.experimental.enable_x64(): - pl.store(ref, idx, val.astype(ref.dtype), mask=mask, **kwargs) - - -def in_bounds_mask( - idx: Sequence[int | slice | pl.Slice | jax.Array], - shape: Sequence[int], - *, - check: Sequence[bool] | None = None, -) -> jax.Array | None: - """Returns a boolean mask denoting which indices are within bounds. - - Args: - idx: Indices for each dimension. - shape: Shape designating the valid bounds. - check: Whether or not to check bounds in each dimension. Useful for ignoring - indices known to be in bounds. Defaults to all True. - """ - if check is None: - check = [True] * len(shape) - - # Remove `int` indexed dims (mask shape must match slice result shape). - shape = [dim for i, dim in enumerate(shape) if not isinstance(idx[i], int)] - check = [chk for i, chk in enumerate(check) if not isinstance(idx[i], int)] - idx = [idx for idx in idx if not isinstance(idx, int)] - - mask = None - for i, (dim_idx, dim, chk) in enumerate(zip(idx, shape, check, strict=True)): - if not chk: - continue - - if isinstance(dim_idx, slice): - dim_idx = pl.Slice.from_slice(dim_idx, dim) - if isinstance(dim_idx, pl.Slice): - dim_idx = dim_idx.start + dim_idx.stride * jnp.arange(dim_idx.size) - if dim_idx.ndim != 1: - raise NotImplementedError("Only one-dimensional indices are supported.") - - bcast_axes = [a for a in range(len(shape)) if a != i] - dim_mask = jnp.expand_dims(dim_idx < dim, bcast_axes) - mask = dim_mask if mask is None else (mask & dim_mask) - return mask - - -def _get_block_indexer_and_mask( - ref, idx: Sequence[int | ScalarInt], *, block_shape: Sequence[int | None] -) -> tuple[tuple[int | slice | pl.Slice, ...], jax.Array | None]: - """Return indices and mask for loading / storing a block.""" - shape = ref.shape - idxs = [] - check = [] - for dim, block_idx, block_dim in zip(shape, idx, block_shape, strict=True): - if block_dim is None: - idxs.append(block_idx) - check.append(False) - else: - idxs.append(pl.dslice(block_dim * block_idx, block_dim)) - check.append(dim % block_dim != 0) - - return tuple(idxs), in_bounds_mask(idxs, shape, check=check) diff --git a/src/alphafold3/jax/gated_linear_unit/gated_linear_unit.py b/src/alphafold3/jax/gated_linear_unit/gated_linear_unit.py deleted file mode 100644 index 304d813..0000000 --- a/src/alphafold3/jax/gated_linear_unit/gated_linear_unit.py +++ /dev/null @@ -1,124 +0,0 @@ -# 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 - -"""Public API for gated linear unit functions.""" - -from collections.abc import Callable -import typing -from typing import Literal, TypeAlias - -from alphafold3.jax.common import array_view -from alphafold3.jax.common import triton_utils -from alphafold3.jax.gated_linear_unit import gated_linear_unit_base -from alphafold3.jax.gated_linear_unit import matmul_ext -import jax -import jaxtyping -from jaxtyping import Array, Float # pylint: disable=g-importing-member,g-multiple-import -import typeguard - -Implementation: TypeAlias = Literal['xla', 'triton'] - - -class PallasGatedLinearUnit(gated_linear_unit_base.GatedLinearUnit): - """Pallas gated linear unit.""" - - def _fwd(self, x, weight, *, activation, precision): - weight_view = array_view.ArrayView(weight) - return self.apply_vmap_rule_forward( - matmul_ext.gated_linear_unit, - activation=activation, - precision=precision, - )( - x, - weight_view[:, 1], - weight_view[:, 0], - ) - - -@jaxtyping.jaxtyped(typechecker=typeguard.typechecked) -def gated_linear_unit( - x: Float[Array, '*B M K'], - weight: Float[Array, 'K 2 N'], - *, - activation: Callable[[jax.Array], jax.Array] | None = None, - precision: jax.lax.Precision | None = None, - implementation: Implementation | None = None, -) -> Float[Array, '*B M N']: - """Applies a gated linear unit (https://arxiv.org/abs/1612.08083). - - Computes `activation(x @ weight[:, 0]) * x @ weight[:, 1]`. - - This is SwiGLU when `activation=jax.nn.swish`, GEGLU when - `activation=jax.nn.gelu`, REGLU when `activation=jax.nn.relu`, and GLU when - `activation=jax.nn.sigmoid` (https://arxiv.org/abs/2002.05202). - - Args: - x: the input array. - weight: the combined weight array. - activation: optional activation function. - precision: specifies the matrix multiplication precision. Either `None` - (default), which means the default precision for the backend, or a - `jax.lax.Precision` enum. - implementation: if `None` (default), an implementation is automatically - chosen. 'xla' will use standard XLA and work on any platform, and 'triton' - will use a fused Triton GPU kernel. Only a subset of data types, shapes - and GPUs are supported by 'triton', with an exception thrown in this case. - - Raises: - NotImplementedError: if `implementation='triton'` does not support a given - input or device. - ValueError: if the arguments are invalid. - - Returns: - The output array. - """ - - match implementation: - case 'triton': - if not triton_utils.has_triton_support(): - raise NotImplementedError('Triton not supported on this platform.') - case _: - ... - - if x.dtype.name != weight.dtype.name: - raise ValueError( - f'Input and weight must have the same dtype. {x.dtype} !=' - f' {weight.dtype}' - ) - - if implementation is not None: - named_args = typing.get_args(Implementation) - if implementation not in named_args: - raise ValueError( - f'Unsupported named implementation. Must be one of {named_args}.' - ) - - if implementation is None or implementation == 'triton': - try: - return PallasGatedLinearUnit()( - x=x, - weight=weight, - activation=activation, - precision=precision, - ) - # When `implementation=None`, we must catch any exception, and use XLA - # as a fallback. As we rely on a third-party library (Triton), it might - # not be possible to enumerate all possible exceptions that could be - # thrown, hence catching the broadest possible one. - except Exception as e: # pylint: disable=broad-exception-caught - if implementation == 'triton': - raise e - - return gated_linear_unit_base.gated_linear_unit_xla( - x=x, - weight=weight, - activation=activation, - precision=precision, - ) diff --git a/src/alphafold3/jax/gated_linear_unit/gated_linear_unit_base.py b/src/alphafold3/jax/gated_linear_unit/gated_linear_unit_base.py deleted file mode 100644 index 164bcf7..0000000 --- a/src/alphafold3/jax/gated_linear_unit/gated_linear_unit_base.py +++ /dev/null @@ -1,130 +0,0 @@ -# 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 - -"""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 diff --git a/src/alphafold3/jax/gated_linear_unit/matmul_config.py b/src/alphafold3/jax/gated_linear_unit/matmul_config.py deleted file mode 100644 index 062855f..0000000 --- a/src/alphafold3/jax/gated_linear_unit/matmul_config.py +++ /dev/null @@ -1,78 +0,0 @@ -# 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 - -"""Auto-tuned configs for matmul.""" - -import dataclasses -import functools -import math - -import jax -from jax.experimental import pallas as pl - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class Config: - block_m: int - block_n: int - block_k: int - num_warps: int - num_stages: int - - -@functools.cache -def _get_best_block_size( - m: int, n: int, k: int, core_count: int -) -> tuple[int, int, int]: - """Returns the best block size for the given shape.""" - min_block_dim = 32 - block_m = min(max(min_block_dim, pl.next_power_of_2(m)), 128) - block_n = min(max(min_block_dim, pl.next_power_of_2(n)), 256) - block_n = min(block_n, (128 * 128) // block_m) - block_k = 32 - split_k = 1 - num_blocks = pl.cdiv(m, block_m) * pl.cdiv(n, block_n) - while num_blocks < core_count: - if block_m > min_block_dim: - block_m //= 2 - num_blocks = pl.cdiv(m, block_m) * pl.cdiv(n, block_n) - elif split_k * block_k < pl.next_power_of_2(k): - split_k *= 2 - num_blocks *= 2 - else: - break - return block_m, block_n, block_k - - -def _abstractify(x): - return jax.api_util.shaped_abstractify(x) if isinstance(x, jax.Array) else x - - -def get_config( - x: jax.Array, w: jax.Array, core_count: int | None = None -) -> Config: - """Returns a config for the given args.""" - if core_count is None: - core_count = jax.devices()[0].core_count - x = _abstractify(x) - w = _abstractify(w) - m, k = math.prod(x.shape[:-1]), x.shape[-1] - n = w.shape[1] - if n >= m: # Prefer `block_n` > `block_m`. - block_m, block_n, block_k = _get_best_block_size(m, n, k, core_count) - else: - block_n, block_m, block_k = _get_best_block_size(n, m, k, core_count) - return Config( - block_m=block_m, - block_n=block_n // 2, # Halve `block_n` as we read two `w` blocks. - block_k=block_k, - num_warps=4, - num_stages=4, - ) diff --git a/src/alphafold3/jax/gated_linear_unit/matmul_ext.py b/src/alphafold3/jax/gated_linear_unit/matmul_ext.py deleted file mode 100644 index 9970a7c..0000000 --- a/src/alphafold3/jax/gated_linear_unit/matmul_ext.py +++ /dev/null @@ -1,273 +0,0 @@ -# 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 - -"""Extended matmul ops.""" - -from collections.abc import Callable -import functools -from typing import Any, TypeAlias - -from alphafold3.jax.common import array_view -from alphafold3.jax.common import triton_utils -from alphafold3.jax.gated_linear_unit import block -from 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]) diff --git a/src/alphafold3/jax/geometry/__init__.py b/src/alphafold3/jax/geometry/__init__.py index 334d07b..df9cd74 100644 --- a/src/alphafold3/jax/geometry/__init__.py +++ b/src/alphafold3/jax/geometry/__init__.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/jax/geometry/rigid_matrix_vector.py b/src/alphafold3/jax/geometry/rigid_matrix_vector.py index e6abb93..4761e55 100644 --- a/src/alphafold3/jax/geometry/rigid_matrix_vector.py +++ b/src/alphafold3/jax/geometry/rigid_matrix_vector.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/jax/geometry/rotation_matrix.py b/src/alphafold3/jax/geometry/rotation_matrix.py index 4d9c837..070bf63 100644 --- a/src/alphafold3/jax/geometry/rotation_matrix.py +++ b/src/alphafold3/jax/geometry/rotation_matrix.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/jax/geometry/struct_of_array.py b/src/alphafold3/jax/geometry/struct_of_array.py index 916111e..a4bed29 100644 --- a/src/alphafold3/jax/geometry/struct_of_array.py +++ b/src/alphafold3/jax/geometry/struct_of_array.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/jax/geometry/utils.py b/src/alphafold3/jax/geometry/utils.py index 7f2fc66..24aed8f 100644 --- a/src/alphafold3/jax/geometry/utils.py +++ b/src/alphafold3/jax/geometry/utils.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -13,72 +22,9 @@ 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) @@ -93,18 +39,6 @@ def angdiff(alpha: jnp.ndarray, beta: jnp.ndarray) -> jnp.ndarray: 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, diff --git a/src/alphafold3/jax/geometry/vector.py b/src/alphafold3/jax/geometry/vector.py index 6d3bf84..091ab2d 100644 --- a/src/alphafold3/jax/geometry/vector.py +++ b/src/alphafold3/jax/geometry/vector.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/atom_layout/atom_layout.py b/src/alphafold3/model/atom_layout/atom_layout.py index 541fe7f..52093a3 100644 --- a/src/alphafold3/model/atom_layout/atom_layout.py +++ b/src/alphafold3/model/atom_layout/atom_layout.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -28,7 +37,6 @@ import numpy as np from rdkit import Chem - xnp_ndarray: TypeAlias = np.ndarray | jnp.ndarray # pylint: disable=invalid-name NumpyIndex: TypeAlias = Any @@ -641,6 +649,7 @@ def get_link_drop_atoms( is_end_terminus: bool, bonded_atoms: set[str], drop_ligand_leaving_atoms: bool = False, + fix_standalone_glycans: bool = False, ) -> set[str]: """Returns set of atoms that are dropped when this res_name gets linked. @@ -651,6 +660,12 @@ def get_link_drop_atoms( 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. + fix_standalone_glycans: AlphaFold 3 model training and evaluation filtered + out leaving atoms from glycan ligands even if they were not bonded to + anything ("standalone" glycans). Setting this flag to True fixes this + undesirable behavior, but moves away from the regime where AlphaFold 3 was + trained and evaluated. This has only an effect if + drop_ligand_leaving_atoms is True. Returns: Set of atoms that are dropped when this amino acid gets linked. @@ -677,8 +692,12 @@ def get_link_drop_atoms( *chemical_component_sets.GLYCAN_OTHER_LIGANDS, *chemical_component_sets.GLYCAN_LINKING_LIGANDS, }: - if 'O1' not in bonded_atoms: - drop_atoms.update({'O1'}) + if fix_standalone_glycans: + if bonded_atoms and 'O1' not in bonded_atoms: + drop_atoms.update({'O1'}) + else: + if 'O1' not in bonded_atoms: + drop_atoms.update({'O1'}) return drop_atoms @@ -743,6 +762,7 @@ def make_flat_atom_layout( with_hydrogens: bool = False, skip_unk_residues: bool = True, drop_ligand_leaving_atoms: bool = False, + fix_standalone_glycans: bool = False, ) -> AtomLayout: """Make a flat atom layout for given residues. @@ -761,6 +781,12 @@ def make_flat_atom_layout( 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. + fix_standalone_glycans: AlphaFold 3 model training and evaluation filtered + out leaving atoms from glycan ligands even if they were not bonded to + anything ("standalone" glycans). Setting this flag to True fixes this + undesirable behavior, but moves away from the regime where AlphaFold 3 was + trained and evaluated. This has only an effect if + drop_ligand_leaving_atoms is True. Returns: an `AtomLayout` object @@ -834,6 +860,7 @@ def make_flat_atom_layout( is_end_terminus=residues.is_end_terminus[idx], bonded_atoms=bonded_atoms, drop_ligand_leaving_atoms=drop_ligand_leaving_atoms, + fix_standalone_glycans=fix_standalone_glycans, ) # If deprotonation info is available, remove the specific atoms. diff --git a/src/alphafold3/model/components/haiku_modules.py b/src/alphafold3/model/components/haiku_modules.py index 0c20b69..91a6d94 100644 --- a/src/alphafold3/model/components/haiku_modules.py +++ b/src/alphafold3/model/components/haiku_modules.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/components/mapping.py b/src/alphafold3/model/components/mapping.py index 44ab109..c02f31b 100644 --- a/src/alphafold3/model/components/mapping.py +++ b/src/alphafold3/model/components/mapping.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/components/utils.py b/src/alphafold3/model/components/utils.py index b5ff4db..ee4a8d5 100644 --- a/src/alphafold3/model/components/utils.py +++ b/src/alphafold3/model/components/utils.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -11,6 +20,7 @@ """Utility functions for training AlphaFold and similar models.""" from collections import abc +from collections.abc import Sequence import contextlib import numbers @@ -19,7 +29,6 @@ import jax.numpy as jnp import numpy as np - VALID_DTYPES = [np.float32, np.float64, np.int8, np.int32, np.int64, bool] @@ -48,23 +57,26 @@ def bfloat16_context(): yield -def mask_mean(mask, value, axis=None, keepdims=False, eps=1e-10): +def mask_mean( + mask: jnp.ndarray, + value: jnp.ndarray, + axis: int | Sequence[int] | None = None, + keepdims: bool = False, + eps: float = 1e-10, +) -> jnp.ndarray: """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 len(mask_shape) != len(value_shape): + raise ValueError(f'Incompatible shapes: {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"' + assert isinstance(axis, abc.Sequence) broadcast_factor = 1.0 for axis_ in axis: @@ -73,8 +85,8 @@ def mask_mean(mask, value, axis=None, keepdims=False, eps=1e-10): 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 + if mask_size != value_size: + raise ValueError(f'Incompatible shapes: {mask_shape=}, {value_shape=}') return jnp.sum(mask * value, keepdims=keepdims, axis=axis) / ( jnp.maximum( diff --git a/src/alphafold3/model/confidence_types.py b/src/alphafold3/model/confidence_types.py index 5ad948b..dfdead6 100644 --- a/src/alphafold3/model/confidence_types.py +++ b/src/alphafold3/model/confidence_types.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -16,47 +25,12 @@ from typing import Any, Self from absl import logging +from alphafold3.cpp import json_serialize from 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( @@ -298,4 +272,11 @@ def from_json(cls, json_string: str) -> Self: def to_json(self) -> str: """Converts StructureConfidenceFull to json string.""" - return json.dumps(self, cls=StructureConfidenceFullEncoder) + return json_serialize.structure_confidence_full_to_json( + pae=self.pae, + token_chain_ids=self.token_chain_ids, + token_res_ids=self.token_res_ids, + atom_plddts=self.atom_plddts, + atom_chain_ids=self.atom_chain_ids, + contact_probs=self.contact_probs, + ) diff --git a/src/alphafold3/model/confidences.py b/src/alphafold3/model/confidences.py index 35a3585..0119821 100644 --- a/src/alphafold3/model/confidences.py +++ b/src/alphafold3/model/confidences.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/data3.py b/src/alphafold3/model/data3.py index 22503f1..f2944d3 100644 --- a/src/alphafold3/model/data3.py +++ b/src/alphafold3/model/data3.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/data_constants.py b/src/alphafold3/model/data_constants.py index eabdcfd..387b8eb 100644 --- a/src/alphafold3/model/data_constants.py +++ b/src/alphafold3/model/data_constants.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/feat_batch.py b/src/alphafold3/model/feat_batch.py index e2f8de4..1ce8efe 100644 --- a/src/alphafold3/model/feat_batch.py +++ b/src/alphafold3/model/feat_batch.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/features.py b/src/alphafold3/model/features.py index dfc13e3..79b2ffc 100644 --- a/src/alphafold3/model/features.py +++ b/src/alphafold3/model/features.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -10,6 +19,7 @@ """Data-side of the input features processing.""" +from collections.abc import Mapping, Sequence import dataclasses import datetime import itertools @@ -419,7 +429,8 @@ def compute_features( all_tokens: atom_layout.AtomLayout, standard_token_idxs: np.ndarray, padding_shapes: PaddingShapes, - fold_input: folding_input.Input, + unpaired_msa_by_chain_id: Mapping[str, str], + paired_msa_by_chain_id: Mapping[str, str], logging_name: str, max_paired_sequence_per_species: int, resolve_msa_overlaps: bool = True, @@ -437,12 +448,10 @@ def compute_features( 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] @@ -480,14 +489,19 @@ def compute_features( if chain_type in mmcif_names.STANDARD_POLYMER_CHAIN_TYPES: unpaired_a3m = '' + if not skip_chain and chain_type in ( + mmcif_names.RNA_CHAIN, + mmcif_names.PROTEIN_CHAIN, + ): + unpaired_a3m = unpaired_msa_by_chain_id[b_chain_id] + 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 ( + not skip_chain + and need_msa_pairing + and chain_type == mmcif_names.PROTEIN_CHAIN + ): + paired_a3m = paired_msa_by_chain_id[b_chain_id] # 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( @@ -718,7 +732,7 @@ def compute_features( all_tokens: atom_layout.AtomLayout, standard_token_idxs: np.ndarray, padding_shapes: PaddingShapes, - fold_input: folding_input.Input, + templates_by_chain_id: Mapping[str, Sequence[folding_input.Template]], max_templates: int, logging_name: str, ) -> Self: @@ -734,13 +748,12 @@ def compute_features( ) np_chains_list = [] - input_chains_by_id = {chain.id: chain for chain in fold_input.chains} + sequences = substruct.chain_single_letter_sequence() 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] @@ -764,10 +777,9 @@ def compute_features( 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: + for template in templates_by_chain_id[chain_id]: struc = structure.from_mmcif( template.mmcif, fix_mse_residues=True, @@ -779,7 +791,7 @@ def compute_features( hit_features = templates.get_polymer_features( chain=struc, chain_poly_type=mmcif_names.PROTEIN_CHAIN, - query_sequence_length=len(chain.sequence), + query_sequence_length=len(sequences[chain_id]), query_to_hit_mapping=dict(template.query_to_template_map), ) sorted_features.append(hit_features) @@ -790,7 +802,8 @@ def compute_features( ) template_features = data3.fix_template_features( - template_features=template_features, num_res=len(chain.sequence) + template_features=template_features, + num_res=len(sequences[chain_id]), ) template_features = _reduce_template_features( @@ -1821,9 +1834,19 @@ def as_data_dict(self) -> BatchDict: @dataclasses.dataclass(frozen=True) class ConvertModelOutput: - """Contains atom layout info.""" + """Contains information needed to convert model flat output back to structure. + + Attributes: + token_atoms_layout: Layout mapping tokens to their corresponding atoms. + flat_output_layout: Flat layout of all atoms present in the structure. + empty_output_struc: Container structure without coordinates used to store + the converted model predictions. + polymer_ligand_bonds: Atom layout containing information on the bonds + between polymers and ligands. + ligand_ligand_bonds: Atom layout containing information on the intra- and + inter-ligand covalent bonds. + """ - cleaned_struc: structure.Structure token_atoms_layout: atom_layout.AtomLayout flat_output_layout: atom_layout.AtomLayout empty_output_struc: structure.Structure @@ -1835,7 +1858,6 @@ 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, @@ -1848,7 +1870,6 @@ def compute_features( ) return cls( - cleaned_struc=cleaned_struc, token_atoms_layout=token_atoms_layout, flat_output_layout=flat_output_layout, empty_output_struc=empty_output_struc, @@ -1861,7 +1882,6 @@ 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)), @@ -1871,7 +1891,6 @@ def from_data_dict(cls, batch: BatchDict) -> Self: 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), diff --git a/src/alphafold3/model/json_serialize_pybind.cc b/src/alphafold3/model/json_serialize_pybind.cc new file mode 100644 index 0000000..6e6f5a1 --- /dev/null +++ b/src/alphafold3/model/json_serialize_pybind.cc @@ -0,0 +1,273 @@ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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 +#include +#include +#include +#include +#include + +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "pybind11/cast.h" +#include "pybind11/numpy.h" +#include "pybind11/pybind11.h" +#include "pybind11_abseil/absl_casters.h" + +namespace alphafold3 { +namespace { + +namespace py = pybind11; + +constexpr std::array, 10'000> MakeStringFrom2DigitNum() { + std::array, 10'000> decimals = {}; + for (int i = 0; i < decimals.size(); ++i) { + decimals[i][0] = (i / 1000 > 0) ? i / 1000 + '0' : ' '; + decimals[i][1] = (i / 100 % 10) + '0'; + decimals[i][2] = '.'; + decimals[i][3] = (i / 10 % 10) + '0'; + decimals[i][4] = (i % 10) + '0'; + } + return decimals; +} + +static constexpr auto kStringFrom2DigitNum = MakeStringFrom2DigitNum(); + +// Converts floats in the range [0, 1] to strings with 2 decimal places. +// Returns "null" for NaN otherwise clamps to [0.00, 1.00]. +constexpr absl::string_view StringFromFraction(float value) { + if (std::isnan(value)) { + return "null"; + } + float value_clamped = std::clamp(value, 0.f, 1.f); + const int index = static_cast(std::round(value_clamped * 100.0f)); + const auto& result_array = kStringFrom2DigitNum[index]; + absl::string_view result(result_array.data() + 1, result_array.size() - 1); + if (result.back() == '0') { + result.remove_suffix(1); + } + return result; +} + +// Converts floats in the range [0, 99.9] to to strings with 1 decimal places. +// Returns "null" for NaN otherwise clamps to [0, 99.9]. +constexpr absl::string_view StringFrom2DigitNum1Decimal(float value) { + if (std::isnan(value)) { + return "null"; + } + float value_clamped = std::clamp(value, 0.f, 99.9f); + const int index = static_cast(std::round(value_clamped * 10.0f)) * 10; + const auto& result_array = kStringFrom2DigitNum[index]; + absl::string_view result(result_array.data(), result_array.size() - 1); + if (result[0] == ' ') { + result.remove_prefix(1); + } + return result; +} + +// Converts floats in the range [0, 99.9] to to strings with 2 decimal places. +// Returns "null" for NaN otherwise clamps to [0, 99.99]. +constexpr absl::string_view StringFrom2DigitNum2Decimal(float value) { + if (std::isnan(value)) { + return "null"; + } + float value_clamped = std::clamp(value, 0.f, 99.99f); + const int index = static_cast(std::round(value_clamped * 100.0f)); + const auto& result_array = kStringFrom2DigitNum[index]; + absl::string_view result(result_array.data(), result_array.size()); + if (result[0] == ' ') { + result.remove_prefix(1); + } + if (result.back() == '0') { + result.remove_suffix(1); + } + return result; +} + +constexpr absl::string_view kAtomChainIdsJson = R"( "atom_chain_ids": [")"; +constexpr absl::string_view kAtomPlddtsJson = R"( "atom_plddts": [)"; +constexpr absl::string_view kContactProbsJson = R"( "contact_probs": [)"; +constexpr absl::string_view kPaeJson = R"( "pae": [)"; +constexpr absl::string_view kTokenChainIdsJson = R"( "token_chain_ids": [")"; +constexpr absl::string_view kTokenResIdsJson = R"( "token_res_ids": [)"; + +using NpFloatArray = + py::array_t; +using NpIntArray = + py::array_t; + +std::string StructureConfidenceFullToJson( + NpFloatArray pae_array, // + const std::vector& token_chain_ids, // + NpIntArray token_res_ids_array, // + NpFloatArray atom_plddts_array, // + const std::vector& atom_chain_ids, // + NpFloatArray contact_probs_array) { + if (pae_array.ndim() != 2) { + throw py::value_error("pae_array must be 2D"); + } + + if (contact_probs_array.ndim() != 2) { + throw py::value_error("contact_probs_array must be 2D"); + } + + const size_t pae_num_rows = pae_array.shape()[0]; + const size_t pae_num_cols = pae_array.shape()[1]; + const size_t contact_probs_num_rows = contact_probs_array.shape()[0]; + const size_t contact_probs_num_cols = contact_probs_array.shape()[1]; + + absl::Span pae(pae_array); + absl::Span token_res_ids(token_res_ids_array); + absl::Span atom_plddts(atom_plddts_array); + absl::Span contact_probs(contact_probs_array); + + py::gil_scoped_release gil_release; + + // \b deletes the previous character comma. + size_t capacity = 2; // {\n + capacity += kAtomChainIdsJson.size() + 1; // [ + for (const auto& s : atom_chain_ids) { + capacity += s.size() + 3; // '","' + } + // \b\b], + + capacity += kAtomPlddtsJson.size() + 1; // [ + capacity += atom_plddts.size() * 6; // 10.00, + capacity += 1; // \b], + + capacity += kContactProbsJson.size() + 1; // [ + capacity += contact_probs.size() * 5; // 0.05, + capacity += 2 * contact_probs_num_rows; // \b],[ + capacity += 1; // \b\b]], + + capacity += kPaeJson.size() + 2; // [[ + capacity += pae.size() * 5; // 10.0, + capacity += 2 * pae_num_rows; // \b],[ + // \b\b]] + + capacity += kTokenChainIdsJson.size() + 1; // [ + for (const auto& s : token_chain_ids) { + capacity += s.size() + 3; // '","' + } + // \b] + + capacity += kTokenResIdsJson.size() + 1; // [ + capacity += token_res_ids.size() * 2; // 0, + for (int res_id : token_res_ids) { + while ((res_id /= 10) > 0) { + capacity += 1; // Bump for each extra digit. + } + } + // \b] + + capacity += 1; // "}" + + std::string json; + json.reserve(capacity); + json.append("{\n"); + + // atom_chain_ids + json.append(kAtomChainIdsJson); + for (size_t i = 0; i + 1 < atom_chain_ids.size(); ++i) { + json.append(atom_chain_ids[i]); + json.append("\",\""); + } + json.append(atom_chain_ids.back()); + json.append("\"],\n"); + + // atom_plddts + json.append(kAtomPlddtsJson); + for (size_t i = 0; i + 1 < atom_plddts.size(); ++i) { + json.append(StringFrom2DigitNum2Decimal(atom_plddts[i])); + json.append(","); + } + json.append(StringFrom2DigitNum2Decimal(atom_plddts.back())); + json.append("],\n"); + + // contact_probs + json.append(kContactProbsJson); + size_t pos = 0; + for (size_t row = 0; row < contact_probs_num_rows; ++row) { + json.append("["); + for (size_t col = 0; col < contact_probs_num_cols - 1; ++col) { + json.append(StringFromFraction(contact_probs[pos++])); + json.append(","); + } + json.append(StringFromFraction(contact_probs[pos++])); + json.append(row != contact_probs_num_rows - 1 ? "]," : "]"); + } + json.append("],\n"); + + // pae + json.append(kPaeJson); + pos = 0; + for (size_t row = 0; row < pae_num_rows; ++row) { + json.append("["); + for (size_t col = 0; col + 1 < pae_num_cols; ++col) { + json.append(StringFrom2DigitNum1Decimal(pae[pos++])); + json.append(","); + } + json.append(StringFrom2DigitNum1Decimal(pae[pos++])); + json.append(row != pae_num_rows - 1 ? "]," : "]"); + } + json.append("],\n"); + + // token_chain_ids + json.append(kTokenChainIdsJson); + for (size_t col = 0; col + 1 < token_chain_ids.size(); ++col) { + json.append(token_chain_ids[col]); + json.append("\",\""); + } + json.append(token_chain_ids.back()); + json.append("\"],\n"); + + // token_res_ids + json.append(kTokenResIdsJson); + for (size_t col = 0; col + 1 < token_res_ids.size(); ++col) { + json.append(absl::StrCat(token_res_ids[col])); + json.append(","); + } + json.append(absl::StrCat(token_res_ids.back())); + json.append("]\n"); + + json.append("}"); + return json; +} + +constexpr char kStructureConfidenceFullToJson[] = R"( + Serializes StructureConfidenceFull to JSON faster than Python (yes, low bar). + )"; +} // namespace + +void RegisterModuleJsonSerialize(pybind11::module m) { + m.def("structure_confidence_full_to_json", // + &StructureConfidenceFullToJson, // + py::arg("pae"), // + py::arg("token_chain_ids"), // + py::arg("token_res_ids"), // + py::arg("atom_plddts"), // + py::arg("atom_chain_ids"), // + py::arg("contact_probs"), // + py::doc(kStructureConfidenceFullToJson + 1) // + ); +} + +} // namespace alphafold3 diff --git a/src/alphafold3/model/json_serialize_pybind.h b/src/alphafold3/model/json_serialize_pybind.h new file mode 100644 index 0000000..77a3c01 --- /dev/null +++ b/src/alphafold3/model/json_serialize_pybind.h @@ -0,0 +1,31 @@ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_JSON_SERIALIZE_PYBIND_H_ +#define ALPHAFOLD3_SRC_ALPHAFOLD3_MODEL_JSON_SERIALIZE_PYBIND_H_ + +#include "pybind11/pybind11.h" + +namespace alphafold3 { + +void RegisterModuleJsonSerialize(pybind11::module m); + +} + +#endif // ALPHAFOLD3_SRC_ALPHAFOLD3_MODEL_JSON_SERIALIZE_PYBIND_H_ diff --git a/src/alphafold3/model/merging_features.py b/src/alphafold3/model/merging_features.py index 81bc48e..619eb92 100644 --- a/src/alphafold3/model/merging_features.py +++ b/src/alphafold3/model/merging_features.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/mkdssp_pybind.cc b/src/alphafold3/model/mkdssp_pybind.cc index 27ca7b5..4434863 100644 --- a/src/alphafold3/model/mkdssp_pybind.cc +++ b/src/alphafold3/model/mkdssp_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 @@ -25,22 +34,25 @@ namespace alphafold3 { namespace py = pybind11; void RegisterModuleMkdssp(pybind11::module m) { - py::module site = py::module::import("site"); - py::list paths = py::cast(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(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 (!getenv("LIBCIFPP_DATA_DIR")) { + py::module site = py::module::import("site"); + py::list paths = py::cast(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(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."); } - } - if (!found) { - throw py::type_error("Could not find the libcifpp components.cif file."); } m.def( "get_dssp", @@ -59,5 +71,4 @@ void RegisterModuleMkdssp(pybind11::module m) { py::arg("calculate_surface_accessibility") = false, py::doc("Gets secondary structure from an mmCIF file.")); } - } // namespace alphafold3 diff --git a/src/alphafold3/model/mkdssp_pybind.h b/src/alphafold3/model/mkdssp_pybind.h index a1e4832..3b38243 100644 --- a/src/alphafold3/model/mkdssp_pybind.h +++ b/src/alphafold3/model/mkdssp_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_ diff --git a/src/alphafold3/model/mmcif_metadata.py b/src/alphafold3/model/mmcif_metadata.py index 5239046..0004d7f 100644 --- a/src/alphafold3/model/mmcif_metadata.py +++ b/src/alphafold3/model/mmcif_metadata.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/model.py b/src/alphafold3/model/model.py index b099be4..52412ff 100644 --- a/src/alphafold3/model/model.py +++ b/src/alphafold3/model/model.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -39,10 +48,9 @@ ModelResult: TypeAlias = Mapping[str, Any] -_ScalarNumberOrArray: TypeAlias = Mapping[str, float | int | np.ndarray] -@dataclasses.dataclass(frozen=True) +@dataclasses.dataclass(frozen=True, kw_only=True) class InferenceResult: """Postprocessed model result. @@ -58,8 +66,12 @@ class InferenceResult: """ predicted_structure: structure.Structure = dataclasses.field() - numerical_data: _ScalarNumberOrArray = dataclasses.field(default_factory=dict) - metadata: _ScalarNumberOrArray = dataclasses.field(default_factory=dict) + numerical_data: Mapping[str, float | int | np.ndarray] = dataclasses.field( + default_factory=dict + ) + metadata: Mapping[str, float | int | np.ndarray] = dataclasses.field( + default_factory=dict + ) debug_outputs: Mapping[str, Any] = dataclasses.field(default_factory=dict) model_id: bytes = b'' @@ -464,9 +476,8 @@ def get_inference_result( # 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 + max_workers=min(len(pred_structures), 32) ) as executor: has_clash = list(executor.map(confidences.has_clash, pred_structures)) fraction_disordered = list( diff --git a/src/alphafold3/model/model_config.py b/src/alphafold3/model/model_config.py index 2040d8f..3bf720c 100644 --- a/src/alphafold3/model/model_config.py +++ b/src/alphafold3/model/model_config.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -14,13 +23,14 @@ from typing import Literal, TypeAlias from alphafold3.common import base_config -from alphafold3.jax.attention import attention - +import tokamax _Shape2DType: TypeAlias = tuple[int | None, int | None] class GlobalConfig(base_config.BaseConfig): + """Global configuration for the AlphaFold3 model.""" + bfloat16: Literal['all', 'none', 'intermediate'] = 'all' final_init: Literal['zeros', 'linear'] = 'zeros' pair_attention_chunk_size: Sequence[_Shape2DType] = ((1536, 128), (None, 32)) @@ -29,4 +39,6 @@ class GlobalConfig(base_config.BaseConfig): (None, 1024), ) # Note: flash_attention_implementation = 'xla' means no flash attention. - flash_attention_implementation: attention.Implementation = 'triton' + flash_attention_implementation: tokamax.DotProductAttentionImplementation = ( + 'triton' + ) diff --git a/src/alphafold3/model/msa_pairing.py b/src/alphafold3/model/msa_pairing.py index eaed479..ae81c6d 100644 --- a/src/alphafold3/model/msa_pairing.py +++ b/src/alphafold3/model/msa_pairing.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -23,7 +32,7 @@ - Is deduplicated such that it doesn't contain any sequences in the paired MSA. """ -from typing import Mapping, MutableMapping, Sequence +from collections.abc import Mapping, MutableMapping, Sequence from alphafold3.model import data_constants import numpy as np diff --git a/src/alphafold3/model/network/atom_cross_attention.py b/src/alphafold3/model/network/atom_cross_attention.py index 5833313..6667a10 100644 --- a/src/alphafold3/model/network/atom_cross_attention.py +++ b/src/alphafold3/model/network/atom_cross_attention.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/confidence_head.py b/src/alphafold3/model/network/confidence_head.py index 35fc86c..151de70 100644 --- a/src/alphafold3/model/network/confidence_head.py +++ b/src/alphafold3/model/network/confidence_head.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/diffusion_head.py b/src/alphafold3/model/network/diffusion_head.py index d6fe609..14fe909 100644 --- a/src/alphafold3/model/network/diffusion_head.py +++ b/src/alphafold3/model/network/diffusion_head.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/diffusion_transformer.py b/src/alphafold3/model/network/diffusion_transformer.py index 58823e8..153137c 100644 --- a/src/alphafold3/model/network/diffusion_transformer.py +++ b/src/alphafold3/model/network/diffusion_transformer.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -11,13 +20,13 @@ """Diffusion transformer model.""" from alphafold3.common import base_config -from alphafold3.jax.gated_linear_unit import gated_linear_unit from alphafold3.model import model_config from alphafold3.model.atom_layout import atom_layout from alphafold3.model.components import haiku_modules as hm import haiku as hk import jax from jax import numpy as jnp +import tokamax def adaptive_layernorm(x, single_cond, name): @@ -97,9 +106,7 @@ def transition_block( 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 - ) + c = tokamax.gated_linear_unit(x=x, weights=weights, activation=jax.nn.swish) else: x = hm.Linear( num_intermediates * 2, initializer='relu', name=f'{name}ffw_transition1' diff --git a/src/alphafold3/model/network/distogram_head.py b/src/alphafold3/model/network/distogram_head.py index 57c374e..15492f6 100644 --- a/src/alphafold3/model/network/distogram_head.py +++ b/src/alphafold3/model/network/distogram_head.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/evoformer.py b/src/alphafold3/model/network/evoformer.py index 2e01bb2..0fd6a4a 100644 --- a/src/alphafold3/model/network/evoformer.py +++ b/src/alphafold3/model/network/evoformer.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/featurization.py b/src/alphafold3/model/network/featurization.py index eeb1606..fc1a349 100644 --- a/src/alphafold3/model/network/featurization.py +++ b/src/alphafold3/model/network/featurization.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/modules.py b/src/alphafold3/model/network/modules.py index 2c1d8a6..6ef3ec9 100644 --- a/src/alphafold3/model/network/modules.py +++ b/src/alphafold3/model/network/modules.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -14,8 +23,6 @@ from typing import Literal from alphafold3.common import base_config -from alphafold3.jax.attention import attention -from alphafold3.jax.gated_linear_unit import gated_linear_unit from alphafold3.model import model_config from alphafold3.model.components import haiku_modules as hm from alphafold3.model.components import mapping @@ -23,6 +30,7 @@ import haiku as hk import jax import jax.numpy as jnp +import tokamax def get_shard_size( @@ -68,8 +76,8 @@ def __call__(self, act, broadcast_dim=0): name='transition1', ) weights = jnp.reshape(weights, (len(weights), 2, num_intermediate)) - c = gated_linear_unit.gated_linear_unit( - x=act, weight=weights, implementation=None, activation=jax.nn.swish + c = tokamax.gated_linear_unit( + x=act, weights=weights, activation=jax.nn.swish ) else: act = hm.Linear( @@ -172,7 +180,7 @@ def _attention( # Dot product attention requires the bias term to have a batch dimension. bias = jnp.expand_dims(bias, 0) - weighted_avg = attention.dot_product_attention( + weighted_avg = tokamax.dot_product_attention( q, k, v, @@ -289,11 +297,8 @@ def __call__(self, act, mask): ) weights_glu = jnp.stack([weights_gate, weights_projection], axis=1) - projection = gated_linear_unit.gated_linear_unit( - x=act, - weight=weights_glu, - activation=jax.nn.sigmoid, - implementation=None, + projection = tokamax.gated_linear_unit( + act, weights_glu, activation=jax.nn.sigmoid ) projection = jnp.transpose(projection, (2, 0, 1)) projection *= mask diff --git a/src/alphafold3/model/network/noise_level_embeddings.py b/src/alphafold3/model/network/noise_level_embeddings.py index dc0381f..f6ee585 100644 --- a/src/alphafold3/model/network/noise_level_embeddings.py +++ b/src/alphafold3/model/network/noise_level_embeddings.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/network/template_modules.py b/src/alphafold3/model/network/template_modules.py index 64a0d76..1081b87 100644 --- a/src/alphafold3/model/network/template_modules.py +++ b/src/alphafold3/model/network/template_modules.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/params.py b/src/alphafold3/model/params.py index 2fa3407..97fb84c 100644 --- a/src/alphafold3/model/params.py +++ b/src/alphafold3/model/params.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -110,6 +119,7 @@ def _abs_to_rel(self, pos: int) -> tuple[int, int]: def close(self): self._stack.close() + @property def closed(self) -> bool: return all(handle.closed for handle in self._handles) @@ -154,7 +164,8 @@ def readinto(self, b: bytearray | memoryview) -> int: 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) + buffered = io.BufferedReader(f) + yield zstandard.ZstdDecompressor().stream_reader(buffered) else: yield f diff --git a/src/alphafold3/model/pipeline/inter_chain_bonds.py b/src/alphafold3/model/pipeline/inter_chain_bonds.py index 4038f90..6a8c2d4 100644 --- a/src/alphafold3/model/pipeline/inter_chain_bonds.py +++ b/src/alphafold3/model/pipeline/inter_chain_bonds.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/pipeline/pipeline.py b/src/alphafold3/model/pipeline/pipeline.py index a7e395f..2b39931 100644 --- a/src/alphafold3/model/pipeline/pipeline.py +++ b/src/alphafold3/model/pipeline/pipeline.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -11,11 +20,12 @@ """The main featurizer.""" import bisect -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import datetime import itertools from absl import logging +from alphafold3 import structure from alphafold3.common import base_config from alphafold3.common import folding_input from alphafold3.constants import chemical_components @@ -102,8 +112,6 @@ class Config(base_config.BaseConfig): date the model coordinates can be used as a fallback. max_templates: The maximum number of templates to send through the network set to 0 to switch off templates. - filter_clashes: If true then will remove clashing chains. - filter_crystal_aids: If true ligands in the cryal aid list are removed. max_paired_sequence_per_species: The maximum number of sequences per species that will be used for MSA pairing. drop_ligand_leaving_atoms: Flag for handling leaving atoms for ligands. @@ -114,8 +122,6 @@ class Config(base_config.BaseConfig): atom_cross_att_keys_subset_size: keys subset size in atom cross attention flatten_non_standard_residues: Whether to expand non-standard polymer residues into flat-atom format. - remove_nonsymmetric_bonds: Whether to remove nonsymmetric bonds from - symmetric polymer chains. deterministic_frames: Whether to use fixed-seed reference positions to construct deterministic frames. resolve_msa_overlaps: Whether to deduplicate unpaired MSA against paired @@ -124,6 +130,12 @@ class Config(base_config.BaseConfig): unpaired MSA field to keep it exactly as is as deduplication against the paired MSA could break the manually crafted pairing between MSA sequences. + fix_standalone_glycans: AlphaFold 3 model training and evaluation filtered + out leaving atoms from glycan ligands even if they were not bonded to + anything ("standalone" glycans). Setting this flag to True fixes this + undesirable behavior, but moves away from the regime where AlphaFold 3 + was trained and evaluated. This has only an effect if + drop_ligand_leaving_atoms is True. """ max_atoms_per_token: int = 24 @@ -135,18 +147,16 @@ class Config(base_config.BaseConfig): max_template_date: datetime.date | None = None ref_max_modified_date: datetime.date | None = None max_templates: int = 4 - filter_clashes: bool = False - filter_crystal_aids: bool = False max_paired_sequence_per_species: int = 600 drop_ligand_leaving_atoms: bool = True average_num_atoms_per_token: int = 24 atom_cross_att_queries_subset_size: int = 32 atom_cross_att_keys_subset_size: int = 128 flatten_non_standard_residues: bool = True - remove_nonsymmetric_bonds: bool = False deterministic_frames: bool = True conformer_max_iterations: int | None = None resolve_msa_overlaps: bool = True + fix_standalone_glycans: bool = False def __init__(self, *, config: Config): """Initializes WholePdb data pipeline. @@ -156,31 +166,30 @@ def __init__(self, *, config: Config): """ self._config = config - def process_item( + def process_structure( self, - fold_input: folding_input.Input, + struct: structure.Structure, random_state: np.random.RandomState, ccd: chemical_components.Ccd, + unpaired_msa_by_chain_id: Mapping[str, str], + paired_msa_by_chain_id: Mapping[str, str], + templates_by_chain_id: Mapping[str, Sequence[folding_input.Template]], random_seed: int | None = None, - ) -> features.BatchDict: - """Takes requests from in_queue, adds (key, serialized ex) to out_queue.""" + ) -> feat_batch.Batch: + """Computes features for a structure and associated MSAs/templates.""" if random_seed is None: random_seed = random_state.randint(2**31) random_state = np.random.RandomState(seed=random_seed) - - logging_name = f'{fold_input.name}, random_seed={random_seed}' - logging.info('processing %s', logging_name) - struct = fold_input.to_structure(ccd=ccd) + logging_name = f'{struct.name}, random_seed={random_seed}' + logging.info('Processing %s', logging_name) # Clean structure. - cleaned_struc, cleaning_metadata = structure_cleaning.clean_structure( + cleaned_struc = structure_cleaning.clean_structure( struct, ccd=ccd, drop_non_standard_atoms=True, drop_missing_sequence=True, - filter_clashes=self._config.filter_clashes, - filter_crystal_aids=self._config.filter_crystal_aids, filter_waters=True, filter_hydrogens=True, filter_leaving_atoms=self._config.drop_ligand_leaving_atoms, @@ -188,21 +197,10 @@ def process_item( covalent_bonds_only=True, remove_polymer_polymer_bonds=True, remove_bad_bonds=True, - remove_nonsymmetric_bonds=self._config.remove_nonsymmetric_bonds, + fix_standalone_glycans=self._config.fix_standalone_glycans, ) - num_clashing_chains_removed = cleaning_metadata[ - 'num_clashing_chains_removed' - ] - - if num_clashing_chains_removed: - logging.info( - 'Removed %d clashing chains from %s', - num_clashing_chains_removed, - logging_name, - ) - - # No chains after fixes + # No chains after cleaning. if cleaned_struc.num_chains == 0: raise MmcifNumChainsError(f'{logging_name}: No chains in structure!') @@ -228,6 +226,7 @@ def process_item( polymer_ligand_bonds=polymer_ligand_bonds, ligand_ligand_bonds=ligand_ligand_bonds, drop_ligand_leaving_atoms=self._config.drop_ligand_leaving_atoms, + fix_standalone_glycans=self._config.fix_standalone_glycans, ) ) @@ -321,7 +320,6 @@ def process_item( batch_convert_model_output = features.ConvertModelOutput.compute_features( all_token_atoms_layout=all_token_atoms_layout, padding_shapes=padding_shapes, - cleaned_struc=cleaned_struc, flat_output_layout=flat_output_layout, empty_output_struc=empty_output_struc, polymer_ligand_bonds=polymer_ligand_bonds, @@ -342,7 +340,8 @@ def process_item( all_tokens=all_tokens, standard_token_idxs=standard_token_idxs, padding_shapes=padding_shapes, - fold_input=fold_input, + unpaired_msa_by_chain_id=unpaired_msa_by_chain_id, + paired_msa_by_chain_id=paired_msa_by_chain_id, logging_name=logging_name, max_paired_sequence_per_species=self._config.max_paired_sequence_per_species, resolve_msa_overlaps=self._config.resolve_msa_overlaps, @@ -353,7 +352,7 @@ def process_item( all_tokens=all_tokens, standard_token_idxs=standard_token_idxs, padding_shapes=padding_shapes, - fold_input=fold_input, + templates_by_chain_id=templates_by_chain_id, max_templates=self._config.max_templates, logging_name=logging_name, ) @@ -436,9 +435,39 @@ def process_item( frames=batch_frames, ) + return batch + + def process_item( + self, + fold_input: folding_input.Input, + random_state: np.random.RandomState, + ccd: chemical_components.Ccd, + random_seed: int | None = None, + ) -> features.BatchDict: + """Takes requests from in_queue, adds (key, serialized ex) to out_queue.""" + struct = fold_input.to_structure(ccd=ccd) + unpaired_msa_by_chain_id = {} + paired_msa_by_chain_id = {} + templates_by_chain_id = {} + for chain in fold_input.chains: + if isinstance( + chain, (folding_input.ProteinChain, folding_input.RnaChain) + ): + unpaired_msa_by_chain_id[chain.id] = chain.unpaired_msa + if isinstance(chain, folding_input.ProteinChain): + paired_msa_by_chain_id[chain.id] = chain.paired_msa + templates_by_chain_id[chain.id] = list(chain.templates) + + batch = self.process_structure( + struct=struct, + random_state=random_state, + ccd=ccd, + unpaired_msa_by_chain_id=unpaired_msa_by_chain_id, + paired_msa_by_chain_id=paired_msa_by_chain_id, + templates_by_chain_id=templates_by_chain_id, + random_seed=random_seed, + ) np_example = batch.as_data_dict() - if 'num_iter_recycling' in np_example: - del np_example['num_iter_recycling'] # that does not belong here for name, value in np_example.items(): if ( @@ -447,9 +476,7 @@ def process_item( and np.isnan(np.sum(value)) ): raise NanDataError( - 'The output of the data pipeline contained nans. ' - f'nan feature: {name}, fold input name: {fold_input.name}, ' - f'random_seed {random_seed}' + f'Data pipeline output for struct.name={struct.name},' + f' random_seed={random_seed} contains NaNs. NaN feature: {name}' ) - return np_example diff --git a/src/alphafold3/model/pipeline/structure_cleaning.py b/src/alphafold3/model/pipeline/structure_cleaning.py index 4043a3f..23acb7f 100644 --- a/src/alphafold3/model/pipeline/structure_cleaning.py +++ b/src/alphafold3/model/pipeline/structure_cleaning.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -10,8 +19,6 @@ """Prepare PDB structure for training or inference.""" -from typing import Any - from absl import logging from alphafold3 import structure from alphafold3.constants import chemical_component_sets @@ -19,8 +26,6 @@ from alphafold3.constants import mmcif_names from alphafold3.model.atom_layout import atom_layout from alphafold3.model.pipeline import inter_chain_bonds -from alphafold3.model.scoring import covalent_bond_cleaning -from alphafold3.structure import sterics import numpy as np @@ -32,6 +37,7 @@ def _get_leaving_atom_mask( chain_type: str, res_id: int, res_name: str, + fix_standalone_glycans: bool, ) -> np.ndarray: """Updates a drop_leaving_atoms mask with new leaving atom locations.""" bonded_atoms = atom_layout.get_bonded_atoms( @@ -48,6 +54,7 @@ def _get_leaving_atom_mask( is_end_terminus=False, bonded_atoms=bonded_atoms, drop_ligand_leaving_atoms=True, + fix_standalone_glycans=fix_standalone_glycans, ) # Default mask where everything is false, which equates to being kept. drop_atom_filter_atoms = struc.chain_id != struc.chain_id @@ -70,9 +77,7 @@ def clean_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, @@ -80,17 +85,15 @@ def clean_structure( 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. + fix_standalone_glycans: bool, +) -> structure.Structure: + """Returns a cleaned version of the input 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. @@ -99,43 +102,19 @@ def clean_structure( 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. + fix_standalone_glycans: AlphaFold 3 model training and evaluation filtered + out leaving atoms from glycan ligands even if they were not bonded to + anything ("standalone" glycans). Setting this flag to True fixes this + undesirable behavior, but moves away from the regime where AlphaFold 3 was + trained and evaluated. This has only an effect if filter_leaving_atoms is + True. """ - 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: @@ -167,8 +146,10 @@ def clean_structure( *chemical_component_sets.GLYCAN_LINKING_LIGANDS, } # If only glycan ligands and no O1 atoms, we can do parallel drop. + # If we need to keep O1 for standalone glycans, iterate over each residue. if ( - only_glycan_ligands_for_leaving_atoms + not fix_standalone_glycans + and only_glycan_ligands_for_leaving_atoms and (not (ligand_ligand_bonds.atom_name == 'O1').any()) and (not (polymer_ligand_bonds.atom_name == 'O1').any()) ): @@ -192,23 +173,22 @@ def clean_structure( chain_type=res['chain_type'], res_id=res['res_id'], res_name=res_name, + fix_standalone_glycans=fix_standalone_glycans, ) 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, + if np.any(drop_leaving_atoms_all): + logging.info( + 'Dropped %d atoms: chain_id %s, res_id %s, res_name %s, atom_name %s', + struc.num_atoms - np.sum(~drop_leaving_atoms_all), + struc.chain_id[drop_leaving_atoms_all], + struc.res_id[drop_leaving_atoms_all], + struc.res_name[drop_leaving_atoms_all], + struc.atom_name[drop_leaving_atoms_all], ) + struc = struc.filter_out(drop_leaving_atoms_all) # Can filter by bond type without having to iterate over bonds. if struc.bonds and covalent_bonds_only: @@ -268,27 +248,7 @@ def clean_structure( 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 + return struc def create_empty_output_struc_and_layout( @@ -300,6 +260,7 @@ def create_empty_output_struc_and_layout( polymer_ligand_bonds: atom_layout.AtomLayout | None = None, ligand_ligand_bonds: atom_layout.AtomLayout | None = None, drop_ligand_leaving_atoms: bool = False, + fix_standalone_glycans: bool = False, ) -> tuple[structure.Structure, atom_layout.AtomLayout]: """Make zero-coordinate structure from all physical residues. @@ -311,6 +272,8 @@ def create_empty_output_struc_and_layout( 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. + fix_standalone_glycans: If True, standalone glycans are preserved when + drop_ligand_leaving_atoms is True. Returns: Tuple of structure with all bonds, physical residues and coordinates set to @@ -351,6 +314,7 @@ def create_empty_output_struc_and_layout( polymer_ligand_bonds=polymer_ligand_bonds, ligand_ligand_bonds=ligand_ligand_bonds, drop_ligand_leaving_atoms=drop_ligand_leaving_atoms, + fix_standalone_glycans=fix_standalone_glycans, ) empty_output_struc = atom_layout.make_structure( diff --git a/src/alphafold3/model/post_processing.py b/src/alphafold3/model/post_processing.py index 0f91342..22eef2f 100644 --- a/src/alphafold3/model/post_processing.py +++ b/src/alphafold3/model/post_processing.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -19,6 +28,7 @@ from alphafold3.model import mmcif_metadata from alphafold3.model import model import numpy as np +import zstandard @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) @@ -78,7 +88,7 @@ def post_process_inference_result( ) return ProcessedInferenceResult( cif=cif, - mean_confidence_1d=mean_confidence_1d, + mean_confidence_1d=float(mean_confidence_1d), ranking_score=float(inference_result.metadata['ranking_score']), structure_confidence_summary_json=structure_confidence_summary_json, structure_full_data_json=structure_full_data_json, @@ -91,23 +101,34 @@ def write_output( output_dir: os.PathLike[str] | str, terms_of_use: str | None = None, name: str | None = None, + compress: bool = False, ) -> None: """Writes processed inference result to a directory.""" processed_result = post_process_inference_result(inference_result) prefix = f'{name}_' if name is not None else '' - with open(os.path.join(output_dir, f'{prefix}model.cif'), 'wb') as f: - f.write(processed_result.cif) + if compress: + opener = zstandard.open + path_transform = lambda path: f'{path}.zst' + else: + opener = open + path_transform = lambda path: path - with open( - os.path.join(output_dir, f'{prefix}summary_confidences.json'), 'wb' - ) as f: - f.write(processed_result.structure_confidence_summary_json) + mmcif_path = os.path.join(output_dir, f'{prefix}model.cif') + with opener(path_transform(mmcif_path), 'wb') as f: + f.write(processed_result.cif) - with open(os.path.join(output_dir, f'{prefix}confidences.json'), 'wb') as f: + full_confidences_path = os.path.join(output_dir, f'{prefix}confidences.json') + with opener(path_transform(full_confidences_path), 'wb') as f: f.write(processed_result.structure_full_data_json) + summary_confidences_path = os.path.join( + output_dir, f'{prefix}summary_confidences.json' + ) + with open(summary_confidences_path, 'wb') as f: + f.write(processed_result.structure_confidence_summary_json) + if terms_of_use is not None: with open(os.path.join(output_dir, 'TERMS_OF_USE.md'), 'wt') as f: f.write(terms_of_use) diff --git a/src/alphafold3/model/protein_data_processing.py b/src/alphafold3/model/protein_data_processing.py index d2f800a..455034c 100644 --- a/src/alphafold3/model/protein_data_processing.py +++ b/src/alphafold3/model/protein_data_processing.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/scoring/alignment.py b/src/alphafold3/model/scoring/alignment.py index f417007..a611b92 100644 --- a/src/alphafold3/model/scoring/alignment.py +++ b/src/alphafold3/model/scoring/alignment.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/scoring/chirality.py b/src/alphafold3/model/scoring/chirality.py index 0b5c85e..97620f5 100644 --- a/src/alphafold3/model/scoring/chirality.py +++ b/src/alphafold3/model/scoring/chirality.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/model/scoring/covalent_bond_cleaning.py b/src/alphafold3/model/scoring/covalent_bond_cleaning.py deleted file mode 100644 index 3763401..0000000 --- a/src/alphafold3/model/scoring/covalent_bond_cleaning.py +++ /dev/null @@ -1,265 +0,0 @@ -# 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 - -"""Some methods to compute metrics for PTMs.""" - -import collections -from collections.abc import Mapping -import dataclasses - -from alphafold3 import structure -from alphafold3.constants import mmcif_names -from alphafold3.model.atom_layout import atom_layout -import numpy as np - - -@dataclasses.dataclass(frozen=True) -class ResIdMapping: - old_res_ids: np.ndarray - new_res_ids: np.ndarray - - -def _count_symmetric_chains(struc: structure.Structure) -> Mapping[str, int]: - """Returns a dict with each chain ID and count.""" - chain_res_name_sequence_from_chain_id = struc.chain_res_name_sequence( - include_missing_residues=True, fix_non_standard_polymer_res=False - ) - counts_for_chain_res_name_sequence = collections.Counter( - chain_res_name_sequence_from_chain_id.values() - ) - chain_symmetric_count = {} - for chain_id, chain_res_name in chain_res_name_sequence_from_chain_id.items(): - chain_symmetric_count[chain_id] = counts_for_chain_res_name_sequence[ - chain_res_name - ] - return chain_symmetric_count - - -def has_nonsymmetric_bonds_on_symmetric_polymer_chains( - struc: structure.Structure, polymer_ligand_bonds: atom_layout.AtomLayout -) -> bool: - """Returns true if nonsymmetric bonds found on polymer chains.""" - try: - _get_polymer_dim(polymer_ligand_bonds) - except ValueError: - return True - if _has_non_polymer_ligand_ptm_bonds(polymer_ligand_bonds): - return True - if _has_multiple_polymers_bonded_to_one_ligand(polymer_ligand_bonds): - return True - combined_struc, _ = _combine_polymer_ligand_ptm_chains( - struc, polymer_ligand_bonds - ) - struc = struc.filter(chain_type=mmcif_names.POLYMER_CHAIN_TYPES) - combined_struc = combined_struc.filter( - chain_type=mmcif_names.POLYMER_CHAIN_TYPES - ) - return _count_symmetric_chains(struc) != _count_symmetric_chains( - combined_struc - ) - - -def _has_non_polymer_ligand_ptm_bonds( - polymer_ligand_bonds: atom_layout.AtomLayout, -): - """Checks if all bonds are between a polymer chain and a ligand chain type.""" - for start_chain_type, end_chain_type in polymer_ligand_bonds.chain_type: - if ( - start_chain_type in mmcif_names.POLYMER_CHAIN_TYPES - and end_chain_type in mmcif_names.LIGAND_CHAIN_TYPES - ): - continue - elif ( - start_chain_type in mmcif_names.LIGAND_CHAIN_TYPES - and end_chain_type in mmcif_names.POLYMER_CHAIN_TYPES - ): - continue - else: - return True - return False - - -def _combine_polymer_ligand_ptm_chains( - struc: structure.Structure, - polymer_ligand_bonds: atom_layout.AtomLayout, -) -> tuple[structure.Structure, dict[tuple[str, str], ResIdMapping]]: - """Combines the ptm polymer-ligand chains together. - - This will prevent them from being permuted away from each other when chains - are matched to the ground truth. This function also returns the res_id mapping - from the separate ligand res_ids to their res_ids in the combined - polymer-ligand chain; this information is needed to later separate the - combined polymer-ligand chain. - - Args: - struc: Structure to be modified. - polymer_ligand_bonds: AtomLayout with polymer-ligand bond info. - - Returns: - A tuple of a Structure with each ptm polymer-ligand chain relabelled as one - chain and a dict from bond chain pair to the res_id mapping. - """ - if not _has_only_single_bond_from_each_chain(polymer_ligand_bonds): - if _has_multiple_ligands_bonded_to_one_polymer(polymer_ligand_bonds): - # For structures where a polymer chain is connected to multiple ligands, - # we need to sort the multiple bonds from the same chain by res_id to - # ensure that the combined polymer-ligand chain will always be the same - # when you have repeated symmetric polymer-ligand chains. - polymer_ligand_bonds = ( - _sort_polymer_ligand_bonds_by_polymer_chain_and_res_id( - polymer_ligand_bonds - ) - ) - else: - raise ValueError( - 'Code cannot handle multiple bonds from one chain unless' - ' its several ligands bonded to a polymer.' - ) - res_id_mappings_for_bond_chain_pair = dict() - for (start_chain_id, end_chain_id), (start_chain_type, end_chain_type) in zip( - polymer_ligand_bonds.chain_id, polymer_ligand_bonds.chain_type - ): - poly_info, ligand_info = _get_polymer_and_ligand_chain_ids_and_types( - start_chain_id, end_chain_id, start_chain_type, end_chain_type - ) - polymer_chain_id, polymer_chain_type = poly_info - ligand_chain_id, _ = ligand_info - - # Join the ligand chain to the polymer chain. - ligand_res_ids = struc.filter(chain_id=ligand_chain_id).res_id - new_res_ids = ligand_res_ids + len(struc.all_residues[polymer_chain_id]) - res_id_mappings_for_bond_chain_pair[(polymer_chain_id, ligand_chain_id)] = ( - ResIdMapping(old_res_ids=ligand_res_ids, new_res_ids=new_res_ids) - ) - chain_groups = [] - chain_group_ids = [] - chain_group_types = [] - for chain_id, chain_type in zip( - struc.chains_table.id, struc.chains_table.type - ): - if chain_id == ligand_chain_id: - continue - elif chain_id == polymer_chain_id: - chain_groups.append([polymer_chain_id, ligand_chain_id]) - chain_group_ids.append(polymer_chain_id) - chain_group_types.append(polymer_chain_type) - else: - chain_groups.append([chain_id]) - chain_group_ids.append(chain_id) - chain_group_types.append(chain_type) - - struc = struc.merge_chains( - chain_groups=chain_groups, - chain_group_ids=chain_group_ids, - chain_group_types=chain_group_types, - ) - - return struc, res_id_mappings_for_bond_chain_pair - - -def _has_only_single_bond_from_each_chain( - polymer_ligand_bonds: atom_layout.AtomLayout, -) -> bool: - """Checks that there is at most one bond from each chain.""" - chain_ids = [] - for chains in polymer_ligand_bonds.chain_id: - chain_ids.extend(chains) - if len(chain_ids) != len(set(chain_ids)): - return False - return True - - -def _get_polymer_and_ligand_chain_ids_and_types( - start_chain_id: str, - end_chain_id: str, - start_chain_type: str, - end_chain_type: str, -) -> tuple[tuple[str, str], tuple[str, str]]: - """Finds polymer and ligand chain ids from chain types.""" - if ( - start_chain_type in mmcif_names.POLYMER_CHAIN_TYPES - and end_chain_type in mmcif_names.LIGAND_CHAIN_TYPES - ): - return (start_chain_id, start_chain_type), (end_chain_id, end_chain_type) - elif ( - start_chain_type in mmcif_names.LIGAND_CHAIN_TYPES - and end_chain_type in mmcif_names.POLYMER_CHAIN_TYPES - ): - return (end_chain_id, end_chain_type), (start_chain_id, start_chain_type) - else: - raise ValueError( - 'This code only handles PTM-bonds from polymer chain to ligands.' - ) - - -def _get_polymer_dim(polymer_ligand_bonds: atom_layout.AtomLayout) -> int: - """Gets polymer dimension from the polymer-ligand bond layout.""" - start_chain_types = [] - end_chain_types = [] - for start_chain_type, end_chain_type in polymer_ligand_bonds.chain_type: - start_chain_types.append(start_chain_type) - end_chain_types.append(end_chain_type) - if set(start_chain_types).issubset( - set(mmcif_names.POLYMER_CHAIN_TYPES) - ) and set(end_chain_types).issubset(set(mmcif_names.LIGAND_CHAIN_TYPES)): - return 0 - elif set(start_chain_types).issubset(mmcif_names.LIGAND_CHAIN_TYPES) and set( - end_chain_types - ).issubset(set(mmcif_names.POLYMER_CHAIN_TYPES)): - return 1 - else: - raise ValueError( - 'Polymer and ligand dimensions are not consistent within the structure.' - ) - - -def _has_multiple_ligands_bonded_to_one_polymer(polymer_ligand_bonds): - """Checks if there are multiple ligands bonded to one polymer.""" - polymer_dim = _get_polymer_dim(polymer_ligand_bonds) - polymer_chain_ids = [ - chains[polymer_dim] for chains in polymer_ligand_bonds.chain_id - ] - if len(polymer_chain_ids) != len(set(polymer_chain_ids)): - return True - return False - - -def _has_multiple_polymers_bonded_to_one_ligand(polymer_ligand_bonds): - """Checks if there are multiple polymer chains bonded to one ligand.""" - polymer_dim = _get_polymer_dim(polymer_ligand_bonds) - ligand_dim = 1 - polymer_dim - ligand_chain_ids = [ - chains[ligand_dim] for chains in polymer_ligand_bonds.chain_id - ] - if len(ligand_chain_ids) != len(set(ligand_chain_ids)): - return True - return False - - -def _sort_polymer_ligand_bonds_by_polymer_chain_and_res_id( - polymer_ligand_bonds, -): - """Sorts bonds by res_id (for when a polymer chain has multiple bonded ligands).""" - - polymer_dim = _get_polymer_dim(polymer_ligand_bonds) - - polymer_chain_ids = [ - chains[polymer_dim] for chains in polymer_ligand_bonds.chain_id - ] - polymer_res_ids = [res[polymer_dim] for res in polymer_ligand_bonds.res_id] - - polymer_chain_and_res_id = zip(polymer_chain_ids, polymer_res_ids) - sorted_indices = [ - idx - for idx, _ in sorted( - enumerate(polymer_chain_and_res_id), key=lambda x: x[1] - ) - ] - return polymer_ligand_bonds[sorted_indices] diff --git a/src/alphafold3/model/scoring/scoring.py b/src/alphafold3/model/scoring/scoring.py index b19108f..3ee354a 100644 --- a/src/alphafold3/model/scoring/scoring.py +++ b/src/alphafold3/model/scoring/scoring.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/parsers/cpp/cif_dict.pyi b/src/alphafold3/parsers/cpp/cif_dict.pyi index 09d915c..b342f78 100644 --- a/src/alphafold3/parsers/cpp/cif_dict.pyi +++ b/src/alphafold3/parsers/cpp/cif_dict.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -109,6 +118,7 @@ class CifDict: def items(self) -> CifDict.ItemView: ... def keys(self) -> CifDict.KeyView: ... def to_string(self) -> str: ... + def to_dict(self) -> dict[str, list[str]]: ... def value_length(self, key: str) -> int: ... def values(self) -> CifDict.ValueView: ... def __bool__(self) -> bool: ... diff --git a/src/alphafold3/parsers/cpp/cif_dict_lib.cc b/src/alphafold3/parsers/cpp/cif_dict_lib.cc index 4eb51e7..01598d5 100644 --- a/src/alphafold3/parsers/cpp/cif_dict_lib.cc +++ b/src/alphafold3/parsers/cpp/cif_dict_lib.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 @@ -384,13 +393,17 @@ absl::StatusOr CifDict::FromString(absl::string_view cif_string) { cif["data_"].emplace_back(first_token); // Counters for CIF loop_ regions. - int loop_token_index = 0; int num_loop_keys = 0; // Loops have usually O(10) columns but could have up to O(10^6) rows. It is // therefore wasteful to look up the cif vector where to add a loop value // since that means doing `columns * rows` map lookups. If we save pointers to // these loop column fields instead, we need only 1 cif lookup per column. std::vector*> loop_column_values; + // Use a resetting column index instead of computing modulo on every token. + // This is the hot path for _atom_site tables with millions of values. + int token_column_index = 0; + // Total number of loop values seen (for column size validation). + int loop_token_count = 0; // Skip the first element since we already processed it above. for (auto token_itr = tokens->begin() + 1; token_itr != tokens->end(); @@ -399,23 +412,22 @@ absl::StatusOr CifDict::FromString(absl::string_view cif_string) { if (absl::EqualsIgnoreCase(token, "loop_")) { // A new loop started, check the previous loop and get rid of its data. absl::Status loop_status = - CheckLoopColumnSizes(num_loop_keys, loop_token_index); + CheckLoopColumnSizes(num_loop_keys, loop_token_count); if (!loop_status.ok()) { return loop_status; } loop_flag = true; loop_column_values.clear(); - loop_token_index = 0; + token_column_index = 0; + loop_token_count = 0; num_loop_keys = 0; continue; } else if (loop_flag) { // The second condition checks we are in the first column. Some mmCIF // files (e.g. 4q9r) have values in later columns starting with an // underscore and we don't want to read these as keys. - int token_column_index = - num_loop_keys == 0 ? 0 : loop_token_index % num_loop_keys; if (token_column_index == 0 && !token.empty() && token[0] == '_') { - if (loop_token_index > 0) { + if (loop_token_count > 0) { // We are out of the loop. loop_flag = false; } else { @@ -449,7 +461,11 @@ absl::StatusOr CifDict::FromString(absl::string_view cif_string) { " expected at most: ", loop_column_values.size())); } loop_column_values[token_column_index]->emplace_back(token); - loop_token_index++; + loop_token_count++; + if (++token_column_index == num_loop_keys) { + // We have completed a row, reset the column index for the next row. + token_column_index = 0; + } continue; } } @@ -470,7 +486,7 @@ absl::StatusOr CifDict::FromString(absl::string_view cif_string) { } } absl::Status loop_status = - CheckLoopColumnSizes(num_loop_keys, loop_token_index); + CheckLoopColumnSizes(num_loop_keys, loop_token_count); if (!loop_status.ok()) { return loop_status; } diff --git a/src/alphafold3/parsers/cpp/cif_dict_lib.h b/src/alphafold3/parsers/cpp/cif_dict_lib.h index 5c16eaa..58bc3ef 100644 --- a/src/alphafold3/parsers/cpp/cif_dict_lib.h +++ b/src/alphafold3/parsers/cpp/cif_dict_lib.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// To request access to the AlphaFold 3 model parameters, follow the process set +// out at https://github.com/google-deepmind/alphafold3. You may only use these +// if received directly from Google. Use is subject to terms of use available at +// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md // A C++ implementation of a CIF parser. For the format specification see // https://www.iucr.org/resources/cif/spec/version1.1/cifsyntax diff --git a/src/alphafold3/parsers/cpp/cif_dict_pybind.cc b/src/alphafold3/parsers/cpp/cif_dict_pybind.cc index 130a821..9589ead 100644 --- a/src/alphafold3/parsers/cpp/cif_dict_pybind.cc +++ b/src/alphafold3/parsers/cpp/cif_dict_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 @@ -460,6 +469,16 @@ void RegisterModuleCifDict(pybind11::module m) { return *result; }, "Serialize to a string", py::call_guard()) + .def( + "to_dict", + [](const CifDict& self) { + py::dict result; + for (const auto& [key, value] : *self.dict()) { + result[py::cast(key)] = py::cast(value); + } + return result; + }, + "Returns the CIF data as a Python dict[str, list[str]].") .def("value_length", &CifDict::ValueLength, py::arg("key"), "Num elements in value") .def("__len__", @@ -606,10 +625,12 @@ void RegisterModuleCifDict(pybind11::module m) { py::keep_alive<0, 1>()); py::class_(cif_dict, "KeyView") - .def("__contains__", - [](const KeyView& v, absl::string_view k) { - return v.map.dict()->find(k) != v.map.dict()->end(); - }) + .def( + "__contains__", + [](const KeyView& v, absl::string_view k) { + return v.map.dict()->find(k) != v.map.dict()->end(); + }, + py::call_guard()) .def("__contains__", [](const KeyView&, py::handle) { return false; }) .def("__len__", [](const KeyView& v) { return v.map.dict()->size(); }) .def( @@ -633,19 +654,19 @@ void RegisterModuleCifDict(pybind11::module m) { cif_dict .def( "__iter__", - [](CifDict& self) { + [](const CifDict& self) { return py::make_key_iterator(self.dict()->begin(), self.dict()->end()); }, py::keep_alive<0, 1>()) .def( - "keys", [](CifDict& self) { return KeyView{self}; }, + "keys", [](const CifDict& self) { return KeyView{self}; }, "Returns an iterable view of the map's keys.") .def( - "values", [](CifDict& self) { return ValueView{self}; }, + "values", [](const CifDict& self) { return ValueView{self}; }, "Returns an iterable view of the map's values.") .def( - "items", [](CifDict& self) { return ItemView{self}; }, + "items", [](const CifDict& self) { return ItemView{self}; }, "Returns an iterable view of the map's items."); } diff --git a/src/alphafold3/parsers/cpp/cif_dict_pybind.h b/src/alphafold3/parsers/cpp/cif_dict_pybind.h index ca4f947..4dacaa5 100644 --- a/src/alphafold3/parsers/cpp/cif_dict_pybind.h +++ b/src/alphafold3/parsers/cpp/cif_dict_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_PARSERS_PYTHON_CIF_DICT_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_PARSERS_PYTHON_CIF_DICT_PYBIND_H_ diff --git a/src/alphafold3/parsers/cpp/fasta_iterator.pyi b/src/alphafold3/parsers/cpp/fasta_iterator.pyi index d5da60e..b1b4c49 100755 --- a/src/alphafold3/parsers/cpp/fasta_iterator.pyi +++ b/src/alphafold3/parsers/cpp/fasta_iterator.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/parsers/cpp/fasta_iterator_lib.cc b/src/alphafold3/parsers/cpp/fasta_iterator_lib.cc index 82cac93..6f6cb68 100644 --- a/src/alphafold3/parsers/cpp/fasta_iterator_lib.cc +++ b/src/alphafold3/parsers/cpp/fasta_iterator_lib.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/parsers/cpp/fasta_iterator_lib.h b/src/alphafold3/parsers/cpp/fasta_iterator_lib.h index 486d05f..eceb411 100644 --- a/src/alphafold3/parsers/cpp/fasta_iterator_lib.h +++ b/src/alphafold3/parsers/cpp/fasta_iterator_lib.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// To request access to the AlphaFold 3 model parameters, follow the process set +// out at https://github.com/google-deepmind/alphafold3. You may only use these +// if received directly from Google. Use is subject to terms of use available at +// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md // A C++ implementation of a FASTA parser. #ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_PARSERS_PYTHON_FASTA_ITERATOR_LIB_H_ diff --git a/src/alphafold3/parsers/cpp/fasta_iterator_pybind.cc b/src/alphafold3/parsers/cpp/fasta_iterator_pybind.cc index 0b47933..e3486d4 100644 --- a/src/alphafold3/parsers/cpp/fasta_iterator_pybind.cc +++ b/src/alphafold3/parsers/cpp/fasta_iterator_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/parsers/cpp/fasta_iterator_pybind.h b/src/alphafold3/parsers/cpp/fasta_iterator_pybind.h index 091ea3f..e2e3c6b 100644 --- a/src/alphafold3/parsers/cpp/fasta_iterator_pybind.h +++ b/src/alphafold3/parsers/cpp/fasta_iterator_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_PARSERS_PYTHON_FASTA_ITERATOR_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_PARSERS_PYTHON_FASTA_ITERATOR_PYBIND_H_ diff --git a/src/alphafold3/parsers/cpp/msa_conversion.pyi b/src/alphafold3/parsers/cpp/msa_conversion.pyi index 3602032..87e862d 100755 --- a/src/alphafold3/parsers/cpp/msa_conversion.pyi +++ b/src/alphafold3/parsers/cpp/msa_conversion.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/parsers/cpp/msa_conversion_pybind.cc b/src/alphafold3/parsers/cpp/msa_conversion_pybind.cc index c192052..023f1d4 100644 --- a/src/alphafold3/parsers/cpp/msa_conversion_pybind.cc +++ b/src/alphafold3/parsers/cpp/msa_conversion_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/parsers/cpp/msa_conversion_pybind.h b/src/alphafold3/parsers/cpp/msa_conversion_pybind.h index 65f5fe9..11e0d03 100644 --- a/src/alphafold3/parsers/cpp/msa_conversion_pybind.h +++ b/src/alphafold3/parsers/cpp/msa_conversion_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_PARSERS_PYTHON_MSA_CONVERSION_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_PARSERS_PYTHON_MSA_CONVERSION_PYBIND_H_ diff --git a/src/alphafold3/scripts/copy_to_ssd.sh b/src/alphafold3/scripts/copy_to_ssd.sh index 9d53716..e1b6431 100755 --- a/src/alphafold3/scripts/copy_to_ssd.sh +++ b/src/alphafold3/scripts/copy_to_ssd.sh @@ -1,8 +1,17 @@ #!/bin/bash # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/scripts/gcp_mount_ssd.sh b/src/alphafold3/scripts/gcp_mount_ssd.sh index 18f9788..c4b9246 100755 --- a/src/alphafold3/scripts/gcp_mount_ssd.sh +++ b/src/alphafold3/scripts/gcp_mount_ssd.sh @@ -1,8 +1,17 @@ #!/bin/bash # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/__init__.py b/src/alphafold3/structure/__init__.py index 94da943..20ff1c9 100644 --- a/src/alphafold3/structure/__init__.py +++ b/src/alphafold3/structure/__init__.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/bioassemblies.py b/src/alphafold3/structure/bioassemblies.py index 9eeba29..b5ae2f9 100644 --- a/src/alphafold3/structure/bioassemblies.py +++ b/src/alphafold3/structure/bioassemblies.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/bonds.py b/src/alphafold3/structure/bonds.py index 176d79c..af0f1d0 100644 --- a/src/alphafold3/structure/bonds.py +++ b/src/alphafold3/structure/bonds.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -108,59 +117,59 @@ def to_mmcif_dict_from_atom_arrays( auth_seq_id: A (num_atom,) array of auth_seq_id strings. insertion_code: A (num_atom,) array of insertion code strings. """ - mmcif_dict = collections.defaultdict(list) - ptnr1_indices, ptnr2_indices = self.get_atom_indices(atom_key) - - mmcif_dict['_struct_conn.ptnr1_label_asym_id'] = chain_id[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_label_asym_id'] = chain_id[ptnr2_indices] - mmcif_dict['_struct_conn.ptnr1_label_comp_id'] = res_name[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_label_comp_id'] = res_name[ptnr2_indices] - mmcif_dict['_struct_conn.ptnr1_label_seq_id'] = res_id[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_label_seq_id'] = res_id[ptnr2_indices] - mmcif_dict['_struct_conn.ptnr1_label_atom_id'] = atom_name[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_label_atom_id'] = atom_name[ptnr2_indices] - - mmcif_dict['_struct_conn.ptnr1_auth_asym_id'] = auth_asym_id[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_auth_asym_id'] = auth_asym_id[ptnr2_indices] - mmcif_dict['_struct_conn.ptnr1_auth_seq_id'] = auth_seq_id[ptnr1_indices] - mmcif_dict['_struct_conn.ptnr2_auth_seq_id'] = auth_seq_id[ptnr2_indices] - mmcif_dict['_struct_conn.pdbx_ptnr1_PDB_ins_code'] = insertion_code[ - ptnr1_indices - ] - mmcif_dict['_struct_conn.pdbx_ptnr2_PDB_ins_code'] = insertion_code[ - ptnr2_indices - ] + cif = collections.defaultdict(list) + p1_indices, p2_indices = self.get_atom_indices(atom_key) + + cif['_struct_conn.ptnr1_label_asym_id'] = chain_id[p1_indices].tolist() + cif['_struct_conn.ptnr2_label_asym_id'] = chain_id[p2_indices].tolist() + cif['_struct_conn.ptnr1_label_comp_id'] = res_name[p1_indices].tolist() + cif['_struct_conn.ptnr2_label_comp_id'] = res_name[p2_indices].tolist() + cif['_struct_conn.ptnr1_label_seq_id'] = res_id[p1_indices].tolist() + cif['_struct_conn.ptnr2_label_seq_id'] = res_id[p2_indices].tolist() + cif['_struct_conn.ptnr1_label_atom_id'] = atom_name[p1_indices].tolist() + cif['_struct_conn.ptnr2_label_atom_id'] = atom_name[p2_indices].tolist() + + cif['_struct_conn.ptnr1_auth_asym_id'] = auth_asym_id[p1_indices].tolist() + cif['_struct_conn.ptnr2_auth_asym_id'] = auth_asym_id[p2_indices].tolist() + cif['_struct_conn.ptnr1_auth_seq_id'] = auth_seq_id[p1_indices].tolist() + cif['_struct_conn.ptnr2_auth_seq_id'] = auth_seq_id[p2_indices].tolist() + cif['_struct_conn.pdbx_ptnr1_PDB_ins_code'] = insertion_code[ + p1_indices + ].tolist() + cif['_struct_conn.pdbx_ptnr2_PDB_ins_code'] = insertion_code[ + p2_indices + ].tolist() label_alt_id = ['?'] * self.size - mmcif_dict['_struct_conn.pdbx_ptnr1_label_alt_id'] = label_alt_id - mmcif_dict['_struct_conn.pdbx_ptnr2_label_alt_id'] = label_alt_id + cif['_struct_conn.pdbx_ptnr1_label_alt_id'] = label_alt_id + cif['_struct_conn.pdbx_ptnr2_label_alt_id'] = label_alt_id # We need to set this to make visualisation work in NGL/PyMOL. - mmcif_dict['_struct_conn.pdbx_value_order'] = ['?'] * self.size + cif['_struct_conn.pdbx_value_order'] = ['?'] * self.size # We use a symmetry of 1_555 which is the no-op transformation. Other # values are used when bonds involve atoms that only exist after expanding # the bioassembly, but we don't support this kind of bond at the moment. symmetry = ['1_555'] * self.size - mmcif_dict['_struct_conn.ptnr1_symmetry'] = symmetry - mmcif_dict['_struct_conn.ptnr2_symmetry'] = symmetry + cif['_struct_conn.ptnr1_symmetry'] = symmetry + cif['_struct_conn.ptnr2_symmetry'] = symmetry bond_type_counter = collections.Counter() for bond_row in self.iterrows(): bond_type = bond_row['type'] bond_type_counter[bond_type] += 1 - mmcif_dict['_struct_conn.id'].append( + cif['_struct_conn.id'].append( f'{bond_type}{bond_type_counter[bond_type]}' ) - mmcif_dict['_struct_conn.pdbx_role'].append(bond_row['role']) - mmcif_dict['_struct_conn.conn_type_id'].append(bond_type) + cif['_struct_conn.pdbx_role'].append(bond_row['role']) + cif['_struct_conn.conn_type_id'].append(bond_type) - bond_types = np.unique(self.type) - mmcif_dict['_struct_conn_type.id'] = bond_types + bond_types = np.unique(self.type).tolist() + cif['_struct_conn_type.id'] = bond_types unknown = ['?'] * len(bond_types) - mmcif_dict['_struct_conn_type.criteria'] = unknown - mmcif_dict['_struct_conn_type.reference'] = unknown + cif['_struct_conn_type.criteria'] = unknown + cif['_struct_conn_type.reference'] = unknown - return dict(mmcif_dict) + return dict(cif) def concat_with_atom_keys( diff --git a/src/alphafold3/structure/chemical_components.py b/src/alphafold3/structure/chemical_components.py index e036560..2d73955 100644 --- a/src/alphafold3/structure/chemical_components.py +++ b/src/alphafold3/structure/chemical_components.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -10,17 +19,76 @@ """Utilities for manipulating chemical components data.""" -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Collection, Iterable, Mapping, Sequence import dataclasses import functools -from typing import Self +from typing import Any, Self from alphafold3.constants import chemical_components from alphafold3.constants import residue_names from alphafold3.structure import mmcif +import numpy as np import rdkit.Chem as rd_chem +def _value_is_missing(value: Collection[Any] | str | None) -> bool: + return value in ('.', '?', '', None) + + +def _to_optional_int(values: Sequence[str | None]) -> Sequence[int | None]: + return [None if _value_is_missing(x) else int(x) for x in values] + + +def _to_optional_float(values: Sequence[str | None]) -> Sequence[float | None]: + return [None if _value_is_missing(x) else float(x) for x in values] + + +@dataclasses.dataclass(frozen=True, kw_only=True, slots=True) +class ChemCompAtom: + """Items of _chem_comp_atom category. + + See mmcif.wwpdb.org/dictionaries/mmcif_ma.dic/Categories/chem_comp_atom.html + + Attributes: + type_symbol: _chem_comp_atom.type_symbol + ordinal: _chem_comp_atom.pdbx_ordinal, can be optional + charge: _chem_comp_atom.charge + leaving_atom_flag: _chem_comp_atom.pdbx_leaving_atom_flag + model_ideal_x: _chem_comp_atom.pdbx_model_Cartn_x_ideal + model_ideal_y: _chem_comp_atom.pdbx_model_Cartn_y_ideal + model_ideal_z: _chem_comp_atom.pdbx_model_Cartn_z_ideal + """ + + type_symbol: str + ordinal: int | None = None + charge: int | None = None + leaving_atom_flag: str + model_ideal_x: float | None = None + model_ideal_y: float | None = None + model_ideal_z: float | None = None + + +@dataclasses.dataclass(frozen=True, kw_only=True, slots=True) +class ChemCompBond: + """Items of _chem_comp_bond category. + + See mmcif.wwpdb.org/dictionaries/mmcif_ma.dic/Categories/_chem_comp_bond.html + + Attributes: + atom_id_1: _chem_comp_bond.atom_id_1 + atom_id_2: _chem_comp_bond.atom_id_2 + value_order: _chem_comp_bond.value_order + aromatic_flag: _chem_comp_bond.aromatic_flag + stereo_config: _chem_comp_bond.stereo_config + """ + + atom_id_1: str + atom_id_2: str + value_order: str + aromatic_flag: str + stereo_config: str + + @dataclasses.dataclass(frozen=True) class ChemCompEntry: """Items of _chem_comp category. @@ -36,6 +104,9 @@ class ChemCompEntry: formula_weight: str = '?' mon_nstd_flag: str = '?' pdbx_smiles: str | None = None + # Mapping from _chem_comp_atom.atom_id to ChemCompAtom. + chem_comp_atoms: Mapping[str, ChemCompAtom] | None = None + chem_comp_bonds: list[ChemCompBond] | None = None def __post_init__(self): for field, value in vars(self).items(): @@ -97,12 +168,20 @@ def from_mmcif( smiles = cif.get('_chem_comp.pdbx_smiles', ['?'] * len(id_)) smiles = [None if s == '?' else s for s in smiles] - chem_comp = { - component_name: ChemCompEntry(*entry) - for component_name, *entry in zip( - id_, type_, name, synonyms, formula, weight, mon_nstd_flag, smiles - ) - } + chem_comp_atoms_mapping = parse_atom_data(cif) + chem_comp_bonds_mapping = parse_bond_data(cif) + + chem_comp = {} + for component_name, *entry in zip( + id_, type_, name, synonyms, formula, weight, mon_nstd_flag, smiles + ): + chem_comp_atoms = chem_comp_atoms_mapping.get(component_name) + chem_comp_bonds = chem_comp_bonds_mapping.get(component_name) + chem_comp[component_name] = ChemCompEntry( + *entry, + chem_comp_atoms=chem_comp_atoms, + chem_comp_bonds=chem_comp_bonds, + ) if fix_mse and 'MSE' in chem_comp: if 'MET' not in chem_comp: @@ -135,26 +214,214 @@ def to_mmcif_dict(self) -> Mapping[str, Sequence[str]]: """Returns chemical components data as a dict suitable for `mmcif.Mmcif`.""" mmcif_dict = {} - mmcif_fields = set() - for entry in self.chem_comp.values(): - for field, value in vars(entry).items(): - if value: - mmcif_fields.add(field) chem_comp_ids = [] + chem_comp_types = [] + chem_comp_names = [] + chem_comp_pdbx_synonyms = [] + chem_comp_formulas = [] + chem_comp_formula_weights = [] + chem_comp_mon_nstd_flags = [] + chem_comp_pdbx_smiles = [] + + # _chem_comp_atom category + chem_comp_atom_comp_ids = [] + chem_comp_atom_atom_ids = [] + chem_comp_atom_type_symbols = [] + chem_comp_atom_charges = [] + chem_comp_atom_leaving_atom_flags = [] + chem_comp_atom_model_ideal_xs = [] + chem_comp_atom_model_ideal_ys = [] + chem_comp_atom_model_ideal_zs = [] + chem_comp_atom_ordinals = [] + + # _chem_comp_bond category + chem_comp_bond_comp_ids = [] + chem_comp_bond_atom_id_1s = [] + chem_comp_bond_atom_id_2s = [] + chem_comp_bond_value_orders = [] + chem_comp_bond_aromatic_flags = [] + chem_comp_bond_stereo_configs = [] + for component_id in sorted(self.chem_comp): entry = self.chem_comp[component_id] chem_comp_ids.append(component_id) - for field in mmcif_fields: - mmcif_dict.setdefault(f'_chem_comp.{field}', []).append( - getattr(entry, field) or '?' - ) + chem_comp_types.append(entry.type) + chem_comp_names.append(entry.name) + chem_comp_pdbx_synonyms.append(entry.pdbx_synonyms) + chem_comp_formulas.append(entry.formula) + chem_comp_formula_weights.append(entry.formula_weight) + chem_comp_mon_nstd_flags.append(entry.mon_nstd_flag) + chem_comp_pdbx_smiles.append(entry.pdbx_smiles) + + if entry.chem_comp_atoms: + xs = [] + ys = [] + zs = [] + for atom_id, chem_comp_atom in entry.chem_comp_atoms.items(): + chem_comp_atom_comp_ids.append(component_id) + chem_comp_atom_atom_ids.append(atom_id) + chem_comp_atom_type_symbols.append(chem_comp_atom.type_symbol) + chem_comp_atom_charges.append( + str(chem_comp_atom.charge) + if chem_comp_atom.charge is not None + else '?' + ) + chem_comp_atom_leaving_atom_flags.append( + chem_comp_atom.leaving_atom_flag + ) + xs.append(chem_comp_atom.model_ideal_x) + ys.append(chem_comp_atom.model_ideal_y) + zs.append(chem_comp_atom.model_ideal_z) + chem_comp_atom_ordinals.append( + str(chem_comp_atom.ordinal) + if chem_comp_atom.ordinal is not None + else '?' + ) + + def format_coords(coords: Sequence[float | None]) -> Sequence[str]: + if any(c is None for c in coords): + return [f'{c:.3f}' if c is not None else '?' for c in coords] + else: + return mmcif.format_float_array( + np.array(coords, dtype=np.float64), num_decimal_places=3 + ) + + chem_comp_atom_model_ideal_xs.extend(format_coords(xs)) + chem_comp_atom_model_ideal_ys.extend(format_coords(ys)) + chem_comp_atom_model_ideal_zs.extend(format_coords(zs)) + + if entry.chem_comp_bonds: + for bond in entry.chem_comp_bonds: + chem_comp_bond_atom_id_1s.append(bond.atom_id_1) + chem_comp_bond_atom_id_2s.append(bond.atom_id_2) + chem_comp_bond_comp_ids.append(component_id) + chem_comp_bond_value_orders.append(bond.value_order or '?') + chem_comp_bond_aromatic_flags.append(bond.aromatic_flag or '?') + chem_comp_bond_stereo_configs.append(bond.stereo_config or '?') + if chem_comp_ids: mmcif_dict['_chem_comp.id'] = chem_comp_ids + mmcif_dict['_chem_comp.type'] = chem_comp_types + mmcif_dict['_chem_comp.name'] = chem_comp_names + mmcif_dict['_chem_comp.pdbx_synonyms'] = chem_comp_pdbx_synonyms + mmcif_dict['_chem_comp.formula'] = chem_comp_formulas + mmcif_dict['_chem_comp.formula_weight'] = chem_comp_formula_weights + mmcif_dict['_chem_comp.mon_nstd_flag'] = chem_comp_mon_nstd_flags + + if not all((v is None for v in chem_comp_pdbx_smiles)): + mmcif_dict['_chem_comp.pdbx_smiles'] = [ + v or '?' for v in chem_comp_pdbx_smiles + ] + + if chem_comp_atom_comp_ids: + mmcif_dict['_chem_comp_atom.comp_id'] = chem_comp_atom_comp_ids + mmcif_dict['_chem_comp_atom.atom_id'] = chem_comp_atom_atom_ids + mmcif_dict['_chem_comp_atom.type_symbol'] = chem_comp_atom_type_symbols + mmcif_dict['_chem_comp_atom.charge'] = chem_comp_atom_charges + + mmcif_dict['_chem_comp_atom.pdbx_leaving_atom_flag'] = ( + chem_comp_atom_leaving_atom_flags + ) + + mmcif_dict['_chem_comp_atom.pdbx_model_Cartn_x_ideal'] = ( + chem_comp_atom_model_ideal_xs + ) + mmcif_dict['_chem_comp_atom.pdbx_model_Cartn_y_ideal'] = ( + chem_comp_atom_model_ideal_ys + ) + mmcif_dict['_chem_comp_atom.pdbx_model_Cartn_z_ideal'] = ( + chem_comp_atom_model_ideal_zs + ) + mmcif_dict['_chem_comp_atom.pdbx_ordinal'] = chem_comp_atom_ordinals + + if chem_comp_bond_comp_ids: + mmcif_dict['_chem_comp_bond.comp_id'] = chem_comp_bond_comp_ids + mmcif_dict['_chem_comp_bond.atom_id_1'] = chem_comp_bond_atom_id_1s + mmcif_dict['_chem_comp_bond.atom_id_2'] = chem_comp_bond_atom_id_2s + mmcif_dict['_chem_comp_bond.value_order'] = chem_comp_bond_value_orders + mmcif_dict['_chem_comp_bond.pdbx_aromatic_flag'] = ( + chem_comp_bond_aromatic_flags + ) + mmcif_dict['_chem_comp_bond.pdbx_stereo_config'] = ( + chem_comp_bond_stereo_configs + ) + return mmcif_dict -def _value_is_missing(value: str) -> bool: - return not value or value in ('.', '?') +def parse_atom_data( + cif: mmcif.Mmcif, +) -> Mapping[str, Mapping[str, ChemCompAtom]]: + """Parses _chem_comp_atom data from an Mmcif object.""" + + atom_comps = cif.get('_chem_comp_atom.comp_id', []) + num_atoms = len(atom_comps) + if not num_atoms: + return {} + + atom_ids = cif['_chem_comp_atom.atom_id'] + atom_types = cif['_chem_comp_atom.type_symbol'] + atom_leaving_atom_flags = cif.get( + '_chem_comp_atom.pdbx_leaving_atom_flag', ['?'] * num_atoms + ) + + nones = [None] * num_atoms + atom_charges = _to_optional_int(cif.get('_chem_comp_atom.charge', nones)) + atom_model_ideal_xs = _to_optional_float( + cif.get('_chem_comp_atom.pdbx_model_Cartn_x_ideal', nones) + ) + atom_model_ideal_ys = _to_optional_float( + cif.get('_chem_comp_atom.pdbx_model_Cartn_y_ideal', nones) + ) + atom_model_ideal_zs = _to_optional_float( + cif.get('_chem_comp_atom.pdbx_model_Cartn_z_ideal', nones) + ) + ordinal = _to_optional_int(cif.get('_chem_comp_atom.pdbx_ordinal', nones)) + + chem_comp_atoms = {} + for i, atom_comp in enumerate(atom_comps): + comp_atom_vals = chem_comp_atoms.setdefault(atom_comp, {}) + comp_atom_vals[atom_ids[i]] = ChemCompAtom( + type_symbol=atom_types[i], + ordinal=ordinal[i], + charge=atom_charges[i], + leaving_atom_flag=atom_leaving_atom_flags[i], + model_ideal_x=atom_model_ideal_xs[i], + model_ideal_y=atom_model_ideal_ys[i], + model_ideal_z=atom_model_ideal_zs[i], + ) + return chem_comp_atoms + + +def parse_bond_data( + cif: mmcif.Mmcif, +) -> Mapping[str, list[ChemCompBond]]: + """Parses _chem_comp_bond data from an Mmcif object.""" + bond_comps = cif.get('_chem_comp_bond.comp_id', []) + len_bonds = len(bond_comps) + if not len_bonds: + return {} + + bond_atom_ids_1 = cif['_chem_comp_bond.atom_id_1'] + bond_atom_ids_2 = cif['_chem_comp_bond.atom_id_2'] + + unknowns = ['?'] * len_bonds + bond_value_orders = cif.get('_chem_comp_bond.value_order', unknowns) + bond_aromatic_flags = cif.get('_chem_comp_bond.pdbx_aromatic_flag', unknowns) + bond_stereo_configs = cif.get('_chem_comp_bond.pdbx_stereo_config', unknowns) + + chem_comp_bonds = {} + for i, bond_comp in enumerate(bond_comps): + chem_comp_bonds.setdefault(bond_comp, []).append( + ChemCompBond( + atom_id_1=bond_atom_ids_1[i], + atom_id_2=bond_atom_ids_2[i], + value_order=bond_value_orders[i], + aromatic_flag=bond_aromatic_flags[i], + stereo_config=bond_stereo_configs[i], + ) + ) + return chem_comp_bonds def get_data_for_ccd_components( @@ -170,6 +437,7 @@ def get_data_for_ccd_components( ) if not chem_data: continue + chem_comp[chemical_component_id] = ChemCompEntry( type=chem_data.type, name=chem_data.name, diff --git a/src/alphafold3/structure/cpp/aggregation.pyi b/src/alphafold3/structure/cpp/aggregation.pyi index 8f4a8b3..706df09 100755 --- a/src/alphafold3/structure/cpp/aggregation.pyi +++ b/src/alphafold3/structure/cpp/aggregation.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/aggregation_pybind.cc b/src/alphafold3/structure/cpp/aggregation_pybind.cc index 5ac46d6..161c73c 100644 --- a/src/alphafold3/structure/cpp/aggregation_pybind.cc +++ b/src/alphafold3/structure/cpp/aggregation_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/aggregation_pybind.h b/src/alphafold3/structure/cpp/aggregation_pybind.h index 9547b94..b0f124a 100644 --- a/src/alphafold3/structure/cpp/aggregation_pybind.h +++ b/src/alphafold3/structure/cpp/aggregation_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_AGGREGATION_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_AGGREGATION_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/membership.pyi b/src/alphafold3/structure/cpp/membership.pyi index 305f366..a77ef46 100755 --- a/src/alphafold3/structure/cpp/membership.pyi +++ b/src/alphafold3/structure/cpp/membership.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/membership_pybind.cc b/src/alphafold3/structure/cpp/membership_pybind.cc index 2b3faf8..31a6012 100644 --- a/src/alphafold3/structure/cpp/membership_pybind.cc +++ b/src/alphafold3/structure/cpp/membership_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/membership_pybind.h b/src/alphafold3/structure/cpp/membership_pybind.h index d224fb1..6e3cc1c 100644 --- a/src/alphafold3/structure/cpp/membership_pybind.h +++ b/src/alphafold3/structure/cpp/membership_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MEMBERSHIP_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MEMBERSHIP_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_altlocs.cc b/src/alphafold3/structure/cpp/mmcif_altlocs.cc index cea9a1b..656c7b4 100644 --- a/src/alphafold3/structure/cpp/mmcif_altlocs.cc +++ b/src/alphafold3/structure/cpp/mmcif_altlocs.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_altlocs.h b/src/alphafold3/structure/cpp/mmcif_altlocs.h index fab5781..c1e61be 100644 --- a/src/alphafold3/structure/cpp/mmcif_altlocs.h +++ b/src/alphafold3/structure/cpp/mmcif_altlocs.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_ALTLOCS_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ALTLOCS_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_atom_site.pyi b/src/alphafold3/structure/cpp/mmcif_atom_site.pyi index 5f0ba34..3861cf3 100755 --- a/src/alphafold3/structure/cpp/mmcif_atom_site.pyi +++ b/src/alphafold3/structure/cpp/mmcif_atom_site.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.cc b/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.cc index 6037fe0..67a1e13 100644 --- a/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.cc +++ b/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.h b/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.h index 1f2104e..d09bffe 100644 --- a/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.h +++ b/src/alphafold3/structure/cpp/mmcif_atom_site_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_ATOM_SITE_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ATOM_SITE_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_layout.h b/src/alphafold3/structure/cpp/mmcif_layout.h index 51c67c5..70d1129 100644 --- a/src/alphafold3/structure/cpp/mmcif_layout.h +++ b/src/alphafold3/structure/cpp/mmcif_layout.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_LAYOUT_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_layout.pyi b/src/alphafold3/structure/cpp/mmcif_layout.pyi index add1b05..b7d1c30 100755 --- a/src/alphafold3/structure/cpp/mmcif_layout.pyi +++ b/src/alphafold3/structure/cpp/mmcif_layout.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/mmcif_layout_lib.cc b/src/alphafold3/structure/cpp/mmcif_layout_lib.cc index 91ad70c..493aff7 100644 --- a/src/alphafold3/structure/cpp/mmcif_layout_lib.cc +++ b/src/alphafold3/structure/cpp/mmcif_layout_lib.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_layout_pybind.cc b/src/alphafold3/structure/cpp/mmcif_layout_pybind.cc index 8eb69be..7d1d298 100644 --- a/src/alphafold3/structure/cpp/mmcif_layout_pybind.cc +++ b/src/alphafold3/structure/cpp/mmcif_layout_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_layout_pybind.h b/src/alphafold3/structure/cpp/mmcif_layout_pybind.h index c79b2dd..9b22280 100644 --- a/src/alphafold3/structure/cpp/mmcif_layout_pybind.h +++ b/src/alphafold3/structure/cpp/mmcif_layout_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_LAYOUT_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_struct_conn.h b/src/alphafold3/structure/cpp/mmcif_struct_conn.h index 821be65..1aac3e2 100644 --- a/src/alphafold3/structure/cpp/mmcif_struct_conn.h +++ b/src/alphafold3/structure/cpp/mmcif_struct_conn.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_struct_conn.pyi b/src/alphafold3/structure/cpp/mmcif_struct_conn.pyi index d293e66..c0c8ac7 100755 --- a/src/alphafold3/structure/cpp/mmcif_struct_conn.pyi +++ b/src/alphafold3/structure/cpp/mmcif_struct_conn.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc b/src/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc index afb930f..adf196f 100644 --- a/src/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc +++ b/src/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.cc b/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.cc index 111715a..df370d0 100644 --- a/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.cc +++ b/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.h b/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.h index acdbf7b..b443aa1 100644 --- a/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.h +++ b/src/alphafold3/structure/cpp/mmcif_struct_conn_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/mmcif_utils.pyi b/src/alphafold3/structure/cpp/mmcif_utils.pyi index 0a8967b..662eb09 100755 --- a/src/alphafold3/structure/cpp/mmcif_utils.pyi +++ b/src/alphafold3/structure/cpp/mmcif_utils.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/mmcif_utils_pybind.cc b/src/alphafold3/structure/cpp/mmcif_utils_pybind.cc index 17f90e1..967e0f7 100644 --- a/src/alphafold3/structure/cpp/mmcif_utils_pybind.cc +++ b/src/alphafold3/structure/cpp/mmcif_utils_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 diff --git a/src/alphafold3/structure/cpp/mmcif_utils_pybind.h b/src/alphafold3/structure/cpp/mmcif_utils_pybind.h index 7ba1942..5721277 100644 --- a/src/alphafold3/structure/cpp/mmcif_utils_pybind.h +++ b/src/alphafold3/structure/cpp/mmcif_utils_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_MMCIF_UTILS_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_UTILS_PYBIND_H_ diff --git a/src/alphafold3/structure/cpp/string_array.pyi b/src/alphafold3/structure/cpp/string_array.pyi index b4b76c2..4524aa2 100755 --- a/src/alphafold3/structure/cpp/string_array.pyi +++ b/src/alphafold3/structure/cpp/string_array.pyi @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/cpp/string_array_pybind.cc b/src/alphafold3/structure/cpp/string_array_pybind.cc index 29fac72..5790504 100644 --- a/src/alphafold3/structure/cpp/string_array_pybind.cc +++ b/src/alphafold3/structure/cpp/string_array_pybind.cc @@ -1,7 +1,16 @@ // 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/ +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. // // 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 @@ -174,6 +183,9 @@ py::array_t IsIn( py::array RemapMultipleArrays( const std::vector>& arrays, const py::dict& mapping) { + if (arrays.empty()) { + return py::array_t(0); + } size_t array_size = arrays[0].size(); for (const auto& array : arrays) { if (array.size() != array_size) { @@ -184,10 +196,6 @@ py::array RemapMultipleArrays( // Create a result buffer. auto result = py::array_t(array_size); absl::Span result_buffer(result.mutable_data(), array_size); - PyObject* entry = PyTuple_New(arrays.size()); - if (entry == nullptr) { - throw py::error_already_set(); - } std::vector> array_spans; array_spans.reserve(arrays.size()); for (const auto& array : arrays) { @@ -197,34 +205,38 @@ py::array RemapMultipleArrays( // Iterate over arrays and look up elements in the `py_dict`. bool fail = false; for (size_t i = 0; i < array_size; ++i) { + PyObject* entry = PyTuple_New(arrays.size()); + if (entry == nullptr) { + fail = true; + break; + } for (size_t j = 0; j < array_spans.size(); ++j) { - PyTuple_SET_ITEM(entry, j, array_spans[j][i]); + PyTuple_SET_ITEM(entry, j, Py_NewRef(array_spans[j][i])); } PyObject* result = PyDict_GetItem(mapping.ptr(), entry); if (result != nullptr) { int64_t result_value = PyLong_AsLongLong(result); if (result_value == -1 && PyErr_Occurred()) { fail = true; + Py_DECREF(entry); break; } if (result_value > std::numeric_limits::max() || result_value < std::numeric_limits::lowest()) { PyErr_SetString(PyExc_OverflowError, "Result value too large."); fail = true; + Py_DECREF(entry); break; } result_buffer[i] = result_value; } else { PyErr_Format(PyExc_KeyError, "%R", entry); fail = true; + Py_DECREF(entry); break; } + Py_DECREF(entry); } - - for (size_t j = 0; j < array_spans.size(); ++j) { - PyTuple_SET_ITEM(entry, j, nullptr); - } - Py_XDECREF(entry); if (fail) { throw py::error_already_set(); } diff --git a/src/alphafold3/structure/cpp/string_array_pybind.h b/src/alphafold3/structure/cpp/string_array_pybind.h index 85790dd..7ffb60d 100644 --- a/src/alphafold3/structure/cpp/string_array_pybind.h +++ b/src/alphafold3/structure/cpp/string_array_pybind.h @@ -1,14 +1,21 @@ -/* - * 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 - */ +// Copyright 2024 DeepMind Technologies Limited +// +// AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with the +// License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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_STRUCTURE_PYTHON_STRING_ARRAY_PYBIND_H_ #define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_STRING_ARRAY_PYBIND_H_ diff --git a/src/alphafold3/structure/mmcif.py b/src/alphafold3/structure/mmcif.py index 78fd267..d205dae 100644 --- a/src/alphafold3/structure/mmcif.py +++ b/src/alphafold3/structure/mmcif.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/parsing.py b/src/alphafold3/structure/parsing.py index fc67617..18c35f2 100644 --- a/src/alphafold3/structure/parsing.py +++ b/src/alphafold3/structure/parsing.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/structure/sterics.py b/src/alphafold3/structure/sterics.py deleted file mode 100644 index 34f99a7..0000000 --- a/src/alphafold3/structure/sterics.py +++ /dev/null @@ -1,140 +0,0 @@ -# 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 - -"""Functions relating to spatial locations of atoms within a structure.""" - -from collections.abc import Collection, Sequence - -from alphafold3 import structure -from alphafold3.structure import mmcif -import numpy as np -import scipy - - -def _make_atom_has_clash_mask( - kd_query_result: np.ndarray, - struc: structure.Structure, - ignore_chains: Collection[str], -) -> np.ndarray: - """Returns a boolean NumPy array representing whether each atom has a clash. - - Args: - kd_query_result: NumPy array containing N-atoms arrays, each array - containing indices to atoms that clash with the N'th atom. - struc: Structure over which clashes were detected. - ignore_chains: Collection of chains that should not be considered clashing. - A boolean NumPy array of length N atoms. - """ - atom_is_clashing = np.zeros((struc.num_atoms,), dtype=bool) - for atom_index, clashes in enumerate(kd_query_result): - chain_i = struc.chain_id[atom_index] - if chain_i in ignore_chains: - continue - islig_i = struc.is_ligand_mask[atom_index] - for clashing_atom_index in clashes: - chain_c = struc.chain_id[clashing_atom_index] - if chain_c in ignore_chains: - continue - islig_c = struc.is_ligand_mask[clashing_atom_index] - if ( - clashing_atom_index == atom_index - or chain_i == chain_c - or islig_i != islig_c - ): - # Ignore clashes within chain or between ligand and polymer. - continue - atom_is_clashing[atom_index] = True - return atom_is_clashing - - -def find_clashing_chains( - struc: structure.Structure, - clash_thresh_angstrom: float = 1.7, - clash_thresh_fraction: float = 0.3, -) -> Sequence[str]: - """Finds chains that clash with others. - - Clashes are defined by polymer backbone atoms and all ligand atoms. - Ligand-polymer clashes are not dropped. - - Will not find clashes if all coordinates are 0. Coordinates are all 0s if - the structure is generated from sequences only, as done for inference in - dendro for example. - - Args: - struc: The structure defining the chains and atom positions. - clash_thresh_angstrom: Below this distance, atoms are considered clashing. - clash_thresh_fraction: Chains with more than this fraction of their atoms - considered clashing will be dropped. This value should be in the range (0, - 1]. - - Returns: - A sequence of chain ids for chains that clash. - - Raises: - ValueError: If `clash_thresh_fraction` is not in range (0,1]. - """ - if not 0 < clash_thresh_fraction <= 1: - raise ValueError('clash_thresh_fraction must be in range (0,1]') - - struc_backbone = struc.filter_polymers_to_single_atom_per_res() - if struc_backbone.num_chains == 0: - return [] - - # If the coordinates are all 0, do not search for clashes. - if not np.any(struc_backbone.coords): - return [] - - coord_kdtree = scipy.spatial.cKDTree(struc_backbone.coords) - - # For each atom coordinate, find all atoms within the clash thresh radius. - clashing_per_atom = coord_kdtree.query_ball_point( - struc_backbone.coords, r=clash_thresh_angstrom - ) - chain_ids = struc_backbone.chains - if struc_backbone.atom_occupancy is not None: - chain_occupancy = np.array([ - np.mean(struc_backbone.atom_occupancy[start:end]) - for start, end in struc_backbone.iter_chain_ranges() - ]) - else: - chain_occupancy = None - - # Remove chains until no more significant clashing. - chains_to_remove = set() - for _ in range(len(chain_ids)): - # Calculate maximally clashing. - atom_has_clash = _make_atom_has_clash_mask( - clashing_per_atom, struc_backbone, chains_to_remove - ) - clashes_per_chain = np.array([ - atom_has_clash[start:end].mean() - for start, end in struc_backbone.iter_chain_ranges() - ]) - max_clash = np.max(clashes_per_chain) - if max_clash <= clash_thresh_fraction: - # None of the remaining chains exceed the clash fraction threshold, so - # we can exit. - break - - # Greedily remove worst with the lowest occupancy. - most_clashes = np.nonzero(clashes_per_chain == max_clash)[0] - if chain_occupancy is not None: - occupancy_clashing = chain_occupancy[most_clashes] - last_lowest_occupancy = ( - len(occupancy_clashing) - np.argmin(occupancy_clashing[::-1]) - 1 - ) - worst_and_last = most_clashes[last_lowest_occupancy] - else: - worst_and_last = most_clashes[-1] - - chains_to_remove.add(chain_ids[worst_and_last]) - - return sorted(chains_to_remove, key=mmcif.str_id_to_int_id) diff --git a/src/alphafold3/structure/structure.py b/src/alphafold3/structure/structure.py index d8abc5a..d54ecf9 100644 --- a/src/alphafold3/structure/structure.py +++ b/src/alphafold3/structure/structure.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -50,6 +59,7 @@ class _UnsetSentinel(enum.Enum): UNSET = object() +_UnsetType = Literal[_UnsetSentinel.UNSET] _UNSET = _UnsetSentinel.UNSET @@ -793,7 +803,7 @@ def iter_residues( yield row | current_chain def _iter_atom_ranges( - self, boundaries: Sequence[int] + self, boundaries: Sequence[int] | np.ndarray ) -> Iterator[tuple[int, int]]: """Iterator for (start, end) pairs from an array of start indices.""" yield from itertools.pairwise(boundaries) @@ -803,7 +813,7 @@ def _iter_atom_ranges( def _iter_residue_ranges( self, - boundaries: Sequence[int], + boundaries: Sequence[int] | np.ndarray, *, count_unresolved: bool, ) -> Iterator[tuple[int, int]]: @@ -1131,20 +1141,20 @@ def _iter_residues(struc: Self) -> Iterable[tuple[str, str]]: def copy_and_update( self, *, - name: str | Literal[_UNSET] = _UNSET, - release_date: datetime.date | None | Literal[_UNSET] = _UNSET, - resolution: float | None | Literal[_UNSET] = _UNSET, - structure_method: str | None | Literal[_UNSET] = _UNSET, + name: str | _UnsetType = _UNSET, + release_date: datetime.date | None | _UnsetType = _UNSET, + resolution: float | None | _UnsetType = _UNSET, + structure_method: str | None | _UnsetType = _UNSET, bioassembly_data: ( - bioassemblies.BioassemblyData | None | Literal[_UNSET] + bioassemblies.BioassemblyData | None | _UnsetType ) = _UNSET, chemical_components_data: ( - struc_chem_comps.ChemicalComponentsData | None | Literal[_UNSET] + struc_chem_comps.ChemicalComponentsData | None | _UnsetType ) = _UNSET, - chains: structure_tables.Chains | None | Literal[_UNSET] = _UNSET, - residues: structure_tables.Residues | None | Literal[_UNSET] = _UNSET, - atoms: structure_tables.Atoms | None | Literal[_UNSET] = _UNSET, - bonds: structure_tables.Bonds | None | Literal[_UNSET] = _UNSET, + chains: structure_tables.Chains | None | _UnsetType = _UNSET, + residues: structure_tables.Residues | None | _UnsetType = _UNSET, + atoms: structure_tables.Atoms | None | _UnsetType = _UNSET, + bonds: structure_tables.Bonds | None | _UnsetType = _UNSET, skip_validation: bool = False, ) -> Self: """Performs a shallow copy but with specified fields updated.""" @@ -1322,15 +1332,15 @@ def copy_and_update_from_res_arrays( def copy_and_update_globals( self, *, - name: str | Literal[_UNSET] = _UNSET, - release_date: datetime.date | Literal[_UNSET] | None = _UNSET, - resolution: float | Literal[_UNSET] | None = _UNSET, - structure_method: str | Literal[_UNSET] | None = _UNSET, + name: str | _UnsetType = _UNSET, + release_date: datetime.date | _UnsetType | None = _UNSET, + resolution: float | _UnsetType | None = _UNSET, + structure_method: str | _UnsetType | None = _UNSET, bioassembly_data: ( - bioassemblies.BioassemblyData | Literal[_UNSET] | None + bioassemblies.BioassemblyData | _UnsetType | None ) = _UNSET, chemical_components_data: ( - struc_chem_comps.ChemicalComponentsData | Literal[_UNSET] | None + struc_chem_comps.ChemicalComponentsData | _UnsetType | None ) = _UNSET, ) -> Self: """Returns a shallow copy with the global columns updated.""" diff --git a/src/alphafold3/structure/structure_tables.py b/src/alphafold3/structure/structure_tables.py index 79f7fd5..35fbe8d 100644 --- a/src/alphafold3/structure/structure_tables.py +++ b/src/alphafold3/structure/structure_tables.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -67,7 +76,9 @@ class AuthorNamingScheme: def _default( - candidate_value: np.ndarray | None, default_value: Sequence[Any], dtype: Any + candidate_value: np.ndarray | None, + default_value: Sequence[Any] | np.ndarray, + dtype: Any, ) -> np.ndarray: if candidate_value is None: return np.array(default_value, dtype=dtype) @@ -417,7 +428,7 @@ def to_mmcif_atom_site_and_bonds_table( coords_decimal_places: int, ) -> Mapping[str, Sequence[str]]: """Returns raw _atom_site and _struct_conn mmCIF tables.""" - raw_mmcif = collections.defaultdict(list) + raw_mmcif: dict[str, Sequence[str]] = {} # Use [value] * num wherever possible since it is about 10x faster than list # comprehension in such cases. Also use f-strings instead of str() - faster. total_atoms = atoms.size * atoms.num_models @@ -521,7 +532,7 @@ def tile_for_models(arr: np.ndarray) -> list[str]: def _flatten_author_naming_scheme_table( - res_table: Mapping[str, Mapping[int, str]], + res_table: Mapping[str, Mapping[int, str | None]], chain_ids: np.ndarray, res_chain_ids: np.ndarray, res_ids: np.ndarray, diff --git a/src/alphafold3/structure/table.py b/src/alphafold3/structure/table.py index 287850b..05755dc 100644 --- a/src/alphafold3/structure/table.py +++ b/src/alphafold3/structure/table.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -478,7 +487,7 @@ def foreign_keys(self) -> Mapping[str, Collection[tuple[str, str]]]: @abc.abstractmethod def copy_and_update( self: _DatabaseT, - **new_field_by_field_name: ..., + **new_field_by_field_name, ) -> _DatabaseT: """Returns a copy of this database with the specified changes applied.""" diff --git a/src/alphafold3/structure/test_utils.py b/src/alphafold3/structure/test_utils.py index fc25ae4..65d87dc 100644 --- a/src/alphafold3/structure/test_utils.py +++ b/src/alphafold3/structure/test_utils.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 diff --git a/src/alphafold3/test_data/featurised_example.json b/src/alphafold3/test_data/featurised_example.json index 257fa69..d2d13e0 100644 --- a/src/alphafold3/test_data/featurised_example.json +++ b/src/alphafold3/test_data/featurised_example.json @@ -2,7 +2,6 @@ { "aatype": "46551530df4e79371352989100ae6af82ccd939009517b788173bdbb323ae40a", "asym_id": "42c714885d17f6c0fcfebebd171b5ed7d7e77a56cf9c4983d9ff9ff90995bd1c", - "cleaned_struc": "<>", "deletion_matrix": "54f980b5b3be8ce80cb6490c527e38d681deade50f239f2cb7d23cf9d0108c37", "deletion_mean": "655a3ef0465a9f30fddf25f4dde0c19a05c6f9069b83961800c1944165955273", "empty_output_struc": "<>", diff --git a/src/alphafold3/version.py b/src/alphafold3/version.py index ff4c5f0..dedb298 100644 --- a/src/alphafold3/version.py +++ b/src/alphafold3/version.py @@ -1,7 +1,16 @@ # 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/ +# AlphaFold 3 source code is licensed under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with the +# License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # # 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 @@ -10,4 +19,4 @@ """Single source of truth for the AlphaFold version.""" -__version__ = '3.0.1' +__version__ = '3.0.2' diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..d527c5b --- /dev/null +++ b/uv.lock @@ -0,0 +1,1049 @@ +version = 1 +revision = 3 +requires-python = ">=3.12" +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'linux'", + "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", + "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'", + "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.13.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", +] +supported-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] + +[[package]] +name = "absl-py" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/2a/c93173ffa1b39c1d0395b7e842bbdc62e556ca9d8d3b5572926f3e4ca752/absl_py-2.3.1.tar.gz", hash = "sha256:a97820526f7fbfd2ec1bce83f3f25e3a14840dac0d8e02a0b71cd75db3f77fc9", size = 116588, upload-time = "2025-07-03T09:31:44.05Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/aa/ba0014cc4659328dc818a28827be78e6d97312ab0cb98105a770924dc11e/absl_py-2.3.1-py3-none-any.whl", hash = "sha256:eeecf07f0c2a93ace0772c92e596ace6d3d3996c042b2128459aaae2a76de11d", size = 135811, upload-time = "2025-07-03T09:31:42.253Z" }, +] + +[[package]] +name = "aiofiles" +version = "25.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/c3/534eac40372d8ee36ef40df62ec129bee4fdb5ad9706e58a29be53b2c970/aiofiles-25.1.0.tar.gz", hash = "sha256:a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2", size = 46354, upload-time = "2025-10-09T20:51:04.358Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, +] + +[[package]] +name = "alphafold3" +source = { editable = "." } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "dm-haiku", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", extra = ["cuda12"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rdkit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tokamax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "zstandard", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "absl-py", specifier = ">=2.3.1" }, + { name = "dm-haiku", specifier = "==0.0.16" }, + { name = "jax", specifier = "==0.9.1" }, + { name = "jax", extras = ["cuda12"], specifier = "==0.9.1" }, + { name = "numpy" }, + { name = "rdkit", specifier = "==2025.9.4" }, + { name = "tokamax", specifier = "==0.0.11" }, + { name = "tqdm" }, + { name = "zstandard" }, +] + +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=6.0" }] + +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + +[[package]] +name = "chex" +version = "0.1.91" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "toolz", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/7d/812f01e7b2ddf28a0caa8dde56bd951a2c8f691c9bbfce38d469458d1502/chex-0.1.91.tar.gz", hash = "sha256:65367a521415ada905b8c0222b0a41a68337fcadf79a1fb6fc992dbd95dd9f76", size = 90302, upload-time = "2025-09-01T21:49:32.834Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/0c/96102c01dd02ae740d4afc3644d5c7d7fc51d3feefd67300a2aa1ddbf7cb/chex-0.1.91-py3-none-any.whl", hash = "sha256:6fc4cbfc22301c08d4a7ef706045668410100962eba8ba6af03fa07f4e5dcf9b", size = 100965, upload-time = "2025-09-01T21:49:31.141Z" }, +] + +[[package]] +name = "dm-haiku" +version = "0.0.16" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jmp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tabulate", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/fc/daf4689198f4c0af8b71611f39fcd5d68ce0ae59fa919b9e58192a7d70f5/dm_haiku-0.0.16.tar.gz", hash = "sha256:1830b0ce63c5cef2fb3a63a13033c9d8f612ee7f896f2b0b25a6ba484f5fad28", size = 263092, upload-time = "2025-12-17T15:55:35.145Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/91/0f53835d0292a74e6b37e68125b669827e2a75a26e01c34741d6c13cca6c/dm_haiku-0.0.16-py3-none-any.whl", hash = "sha256:cc355d4d5aaa85af20e5a23ccd278bc751232ac8e5971261bed39318c07d744f", size = 374267, upload-time = "2025-12-17T15:55:33.9Z" }, +] + +[[package]] +name = "einshape" +version = "1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cb/c6/95ad0a036656aec1cb32177a5d5abfcfbf53a01c1416484cacb8c7332a84/einshape-1.0.tar.gz", hash = "sha256:53538d75dd099f4ead4a4f786fafdcb0b729bb587e0b3afeca25ceef18c9ac14", size = 14571, upload-time = "2022-12-19T17:09:34.618Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/bb/34cc02f13b438d550e4709216ee1df9da8e55e15b0cc87a2cb5dee19a729/einshape-1.0-py3-none-any.whl", hash = "sha256:42da4c2dea3a27f87ee45a7cee5072a636b97cb184bb07bf5d6412ba0ff7b965", size = 21392, upload-time = "2022-12-19T17:09:32.904Z" }, +] + +[[package]] +name = "etils" +version = "1.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9b/a0/522bbff0f3cdd37968f90dd7f26c7aa801ed87f5ba335f156de7f2b88a48/etils-1.13.0.tar.gz", hash = "sha256:a5b60c71f95bcd2d43d4e9fb3dc3879120c1f60472bb5ce19f7a860b1d44f607", size = 106368, upload-time = "2025-07-15T10:29:10.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/98/87b5946356095738cb90a6df7b35ff69ac5750f6e783d5fbcc5cb3b6cbd7/etils-1.13.0-py3-none-any.whl", hash = "sha256:d9cd4f40fbe77ad6613b7348a18132cc511237b6c076dbb89105c0b520a4c6bb", size = 170603, upload-time = "2025-07-15T10:29:09.076Z" }, +] + +[package.optional-dependencies] +epath = [ + { name = "fsspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "importlib-resources", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "zipp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +epy = [ + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "flax" +version = "0.12.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "msgpack", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "optax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "orbax-checkpoint", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tensorstore", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "treescope", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6b/7e/c4c66ab9b41149cf7a1961907d9a844832af1e76b121b35235a618c92825/flax-0.12.2.tar.gz", hash = "sha256:e9723b0881e571abe61885bb8770f53fdb3c383b6b3f5a923dcf6f1e9a687905", size = 5008370, upload-time = "2025-12-18T22:36:19.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/6b/7b75508251f4220df8f68e7718b476ee3d614a2a51f9eace97393ee91b46/flax-0.12.2-py3-none-any.whl", hash = "sha256:912fdd8a7c623ec8b2694b28d2827608e7fc82a3a6f8fff17ec5038f2bca66f4", size = 488031, upload-time = "2025-12-18T22:36:18.01Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d5/7d/5df2650c57d47c57232af5ef4b4fdbff182070421e405e0d62c6cdbfaa87/fsspec-2026.1.0.tar.gz", hash = "sha256:e987cb0496a0d81bba3a9d1cee62922fb395e7d4c3b575e57f547953334fe07b", size = 310496, upload-time = "2026-01-09T15:21:35.562Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/c9/97cc5aae1648dcb851958a3ddf73ccd7dbe5650d95203ecb4d7720b4cdbf/fsspec-2026.1.0-py3-none-any.whl", hash = "sha256:cb76aa913c2285a3b49bdd5fc55b1d7c708d7208126b60f2eb8194fe1b4cbdcc", size = 201838, upload-time = "2026-01-09T15:21:34.041Z" }, +] + +[[package]] +name = "humanize" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/66/a3921783d54be8a6870ac4ccffcd15c4dc0dd7fcce51c6d63b8c63935276/humanize-4.15.0.tar.gz", hash = "sha256:1dd098483eb1c7ee8e32eb2e99ad1910baefa4b75c3aff3a82f4d78688993b10", size = 83599, upload-time = "2025-12-20T20:16:13.19Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, +] + +[[package]] +name = "immutabledict" +version = "4.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/12/1da8e1a9050d0603ba65fb1796ed8860a705b906701c96e77f85cc7490be/immutabledict-4.2.2.tar.gz", hash = "sha256:cb6ed3090df593148f94cb407d218ca526fd2639694afdb553dc4f50ce6feeca", size = 6099, upload-time = "2025-10-12T13:32:59.755Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/7b/04ab6afa1ff7eb9ccb09049918c0407b205f5009092c0416147d163e4e2b/immutabledict-4.2.2-py3-none-any.whl", hash = "sha256:97c31d098a2c850e93a958badeef765e4736ed7942ec73e439facd764a3a7217", size = 4736, upload-time = "2025-10-12T13:32:58.326Z" }, +] + +[[package]] +name = "importlib-resources" +version = "6.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cf/8c/f834fbf984f691b4f7ff60f50b514cc3de5cc08abfc3295564dd89c5e2e7/importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c", size = 44693, upload-time = "2025-01-03T18:51:56.698Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jax" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ml-dtypes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opt-einsum", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/4d/f45853fdc2b811e78b866d5f80b8a21a848278361f66c066706132f415cf/jax-0.9.1.tar.gz", hash = "sha256:ce1b82477ee192f0b1d9801b095aa0cf3839bc1fe0cbc071c961a24b3ff30361", size = 2625994, upload-time = "2026-03-02T11:24:18.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/e4/88778c6a23b65224e5088e68fd0924e5bde2196a26e76edb3ea3543fed6a/jax-0.9.1-py3-none-any.whl", hash = "sha256:d11cb53d362912253013e8c4d6926cb9f3a4b59ab5b25a7dc08123567067d088", size = 3062162, upload-time = "2026-03-02T11:22:05.089Z" }, +] + +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", extra = ["with-cuda"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1b/b8/b725d6ae617b5b6d988eb3b8672e997fff52f1749b3b138c371ae6cf7c2d/jax_cuda12_pjrt-0.9.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:31284cd1c279612cebcf2e9383db9cc2e0703ad8ef56c887476b40331fed7c93", size = 149993391, upload-time = "2026-03-02T11:22:08.125Z" }, + { url = "https://files.pythonhosted.org/packages/a7/2c/8ddb471091b46de99bba7eaa7f4e3983f9c8e74e310e585ff08915ce8b7a/jax_cuda12_pjrt-0.9.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:6447fe3d661b1d824ae7329914c4ac05fe54aeab68c7a7ca23c4f5306dec2e1f", size = 156039796, upload-time = "2026-03-02T11:22:13.316Z" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax-cuda12-pjrt", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/3e/df484dc2f9c42ad2344ff677ab5cb28c05420990a2697708ea0c32ef23bd/jax_cuda12_plugin-0.9.1-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:2843ca4278617a199f9f0a3d992d8684fabe8beaf988ff793ec19d836bdcadc4", size = 5652156, upload-time = "2026-03-02T11:22:19.536Z" }, + { url = "https://files.pythonhosted.org/packages/4f/4e/cd2aa23847bb3c631d9c55ff35fda51c0c5d0410159df897b22c6c02ff5c/jax_cuda12_plugin-0.9.1-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:f4c21de4507b1c442f4e10beb123e4b916055bff4c0102729aa2a49f2ed0fc88", size = 5654547, upload-time = "2026-03-02T11:22:21.365Z" }, + { url = "https://files.pythonhosted.org/packages/95/80/0d15d44900bbe2c6276537dc62c8ec49cb08d77b11f8a948b398a3d17eec/jax_cuda12_plugin-0.9.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:d1bf3374819c8a87e90df04b7144272e5ac0727a5e1fc4872531c52e387c0e38", size = 5652459, upload-time = "2026-03-02T11:22:22.728Z" }, + { url = "https://files.pythonhosted.org/packages/ca/25/7257a5e4b94634495c138b57f15e77573c440f5128fd44882ae798c70181/jax_cuda12_plugin-0.9.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:389c896216a6d4af78709a8f4ca75916ab4ff7bbf4b1af5c09dce9823c003b73", size = 5654488, upload-time = "2026-03-02T11:22:23.978Z" }, + { url = "https://files.pythonhosted.org/packages/45/5b/eadf6199ec61d0640c00c41bb6152c67c2b4b116f7c118ab1eec7034d13e/jax_cuda12_plugin-0.9.1-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:6c7a92b7890de3b3cf9f16a017d4ce7f364ef7d3003da86a59d941a2da4a0bda", size = 5665206, upload-time = "2026-03-02T11:22:26.471Z" }, + { url = "https://files.pythonhosted.org/packages/8d/30/302d9331fe44db041efe042d70fa963c88eaa32e92b4c241a2229a6fea76/jax_cuda12_plugin-0.9.1-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:0a5d2c6604c8fec04df575190f13908a4aa588fb715c338839099d7cb8f90ad5", size = 5663128, upload-time = "2026-03-02T11:22:27.809Z" }, + { url = "https://files.pythonhosted.org/packages/cc/8a/d626849b8667c971a95784cf468d272d92823133dbb2311ceadcb07427aa/jax_cuda12_plugin-0.9.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:d04ac5cff48d18f8497d99c25a26f15bb2df3ba287a4e5eb86c61c89f058c2f9", size = 5653063, upload-time = "2026-03-02T11:22:28.998Z" }, + { url = "https://files.pythonhosted.org/packages/17/e9/3632d7eddb8f282b50bf7c095bba9de91aae0de9baea56d1699d982fe5f8/jax_cuda12_plugin-0.9.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:55888255c325f98af961cf5f665ca8223d0e36363eb544fc0286277be2324e8f", size = 5655068, upload-time = "2026-03-02T11:22:30.251Z" }, + { url = "https://files.pythonhosted.org/packages/a9/90/cfa3222b419194524d28b63b3c6d09ce22ddf823701ae73426df6688f26a/jax_cuda12_plugin-0.9.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:6b4533951bca6e03b07bb092241f2b471d7675e6c0fd5a4c7cc371d030e0c95a", size = 5665690, upload-time = "2026-03-02T11:22:31.438Z" }, + { url = "https://files.pythonhosted.org/packages/11/65/c123fd6343008df4366840a68f9d5a5c2ee1c7a88b94f0485d1e5bbfd1f8/jax_cuda12_plugin-0.9.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:88d3714a638e38f80fa91ba4b9ffc4fd767f4c158ef156067801c74b38d6370f", size = 5663278, upload-time = "2026-03-02T11:22:32.687Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-cupti-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-runtime-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cudnn-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufft-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusolver-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nccl-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvshmem-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[[package]] +name = "jaxlib" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/b9/e0419783cbff9fa3bbc053dbe130f9051f60de4f424f650d70aae7f3bdf1/jaxlib-0.9.1-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:f80e8aead3461683657027e14e814e5bdd00be8ce8e05c0a5db86403db297c2e", size = 76828062, upload-time = "2026-03-02T11:23:19.202Z" }, + { url = "https://files.pythonhosted.org/packages/53/6b/b381bda5850f5611822d791cd25dfe36efda2688a68c4dda0f8a92c36dec/jaxlib-0.9.1-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:e2ab8c97be30354a34e64d17066df0fce7d1d0f40f7a48eded19e9e837896f5d", size = 82472923, upload-time = "2026-03-02T11:23:23.352Z" }, + { url = "https://files.pythonhosted.org/packages/fd/5f/d4a79d6802f3cef02773852453d9528569dd0896964117d4401658828aba/jaxlib-0.9.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:9e88c35248b37d5219423ff8ddca60c6a561e665ded5c4fcbc61f0763e03f1e3", size = 76828438, upload-time = "2026-03-02T11:23:34.793Z" }, + { url = "https://files.pythonhosted.org/packages/3e/2e/d84cafbd07e8cdc7701d9f840f4eea0cfcf3487a99ada14507702172da14/jaxlib-0.9.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:da60d967b4ac2084a3e3535ad982392894dd6bdf79c9a56978aba08404a58c82", size = 82473711, upload-time = "2026-03-02T11:23:38.356Z" }, + { url = "https://files.pythonhosted.org/packages/e9/63/0f5acacd3bd6906f2e1f730ceeafac4afc5cc612f43be4820785608cb951/jaxlib-0.9.1-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:469f08a30f6b541557e29c5de61ea6df16ac0ef9225879373bb2b332f1b27d14", size = 76949185, upload-time = "2026-03-02T11:23:49.378Z" }, + { url = "https://files.pythonhosted.org/packages/91/c5/a4dee13627d913c7bd0cf29b7f5c1d6a2605760d08a7cff952f9098ebb61/jaxlib-0.9.1-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:2e2225b80689610cbb472822dadf7cc200aa4bdac813112a3f6e074d96b1458c", size = 82584273, upload-time = "2026-03-02T11:23:52.762Z" }, + { url = "https://files.pythonhosted.org/packages/0c/e7/237ec5f4cd07420ef50d79a048b769664dbe306e31bdb10f9dcb9accabe9/jaxlib-0.9.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:531dff9fae7aea14449ee544cc1415880cc8a346a9287d347dbd1b2b51d8aabd", size = 76846925, upload-time = "2026-03-02T11:23:59.18Z" }, + { url = "https://files.pythonhosted.org/packages/76/fe/67d2c414b0860d42f4a20b1fadbe7aeffb1b3d885efebd7aedf22a4bc2a2/jaxlib-0.9.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:2287a1c891b152c52eb9b73925f57cde01be35d2bab4dad9673d3c83c5982ca8", size = 82484342, upload-time = "2026-03-02T11:24:02.541Z" }, + { url = "https://files.pythonhosted.org/packages/50/d2/35ecc2e92065ac035a954fcb4b752baa72747dcc3a3466525c42c4404958/jaxlib-0.9.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:30fe58e8e4e105dffe364a6f0dccca16d93433576d4a015babc83339ca7f1f38", size = 76948543, upload-time = "2026-03-02T11:24:12.026Z" }, + { url = "https://files.pythonhosted.org/packages/ba/cb/a8de776aee88f42937d07472953cf7980e45f5fb30aa9d5ee652b4acc771/jaxlib-0.9.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:6b6654a20d54e7cc77d1d54c33f1db851ef9d70bb112b627776178221036e720", size = 82585090, upload-time = "2026-03-02T11:24:15.783Z" }, +] + +[[package]] +name = "jaxtyping" +version = "0.3.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/30/0e29eab664d33127a7a30aebd7a845a95fbfa1887ede692676769820acea/jaxtyping-0.3.5.tar.gz", hash = "sha256:8150ad5b72b62fa63f573d492a79e9e455f070abe3b260f7dc15270b3eb9bba6", size = 45482, upload-time = "2026-01-05T10:51:34.486Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/19/a29de61c7ada0edf4eacfc19355ab8560091b65ea5d2450bcb984e203259/jaxtyping-0.3.5-py3-none-any.whl", hash = "sha256:862c39fa2e526274e82dc96ee8dbe9369dadb651ab1e05d95bd685acb4e2ef02", size = 55824, upload-time = "2026-01-05T10:51:33.523Z" }, +] + +[[package]] +name = "jmp" +version = "0.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ab/b0/e90fbbffef4b345329c878a69f0336d3edc5a1f9fcba193931aca2132d62/jmp-0.0.4.tar.gz", hash = "sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730", size = 18582, upload-time = "2023-01-30T12:47:13.634Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/e5/cce82de2831e5aff9332d8d624bb57188f1b2af6ccf6979caf898a8a4348/jmp-0.0.4-py3-none-any.whl", hash = "sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d", size = 18274, upload-time = "2023-01-30T12:47:11.931Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/04/f9/067b84365c7e83bda15bba2b06c6ca250ce27b20630b1128c435fb7a09aa/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465", size = 5036145, upload-time = "2025-11-17T22:32:12.783Z" }, + { url = "https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f", size = 5010230, upload-time = "2025-11-17T22:32:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/5a/e7/85cb99fe80a7a5513253ec7faa88a65306be071163485e9a626fce1b6e84/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7", size = 5355358, upload-time = "2025-11-17T22:32:19.7Z" }, + { url = "https://files.pythonhosted.org/packages/79/2b/a826ba18d2179a56e144aef69e57fb2ab7c464ef0b2111940ee8a3a223a2/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf", size = 5366332, upload-time = "2025-11-17T22:32:21.193Z" }, +] + +[[package]] +name = "msgpack" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, + { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, + { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, + { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, + { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, + { url = "https://files.pythonhosted.org/packages/8e/a9/3536e385167b88c2cc8f4424c49e28d49a6fc35206d4a8060f136e71f94c/msgpack-1.1.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99e2cb7b9031568a2a5c73aa077180f93dd2e95b4f8d3b8e14a73ae94a9e667e", size = 411885, upload-time = "2025-10-08T09:15:27.22Z" }, + { url = "https://files.pythonhosted.org/packages/2f/40/dc34d1a8d5f1e51fc64640b62b191684da52ca469da9cd74e84936ffa4a6/msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:180759d89a057eab503cf62eeec0aa61c4ea1200dee709f3a8e9397dbb3b6931", size = 419658, upload-time = "2025-10-08T09:15:28.4Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ef/2b92e286366500a09a67e03496ee8b8ba00562797a52f3c117aa2b29514b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:04fb995247a6e83830b62f0b07bf36540c213f6eac8e851166d8d86d83cbd014", size = 403290, upload-time = "2025-10-08T09:15:29.764Z" }, + { url = "https://files.pythonhosted.org/packages/78/90/e0ea7990abea5764e4655b8177aa7c63cdfa89945b6e7641055800f6c16b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8e22ab046fa7ede9e36eeb4cfad44d46450f37bb05d5ec482b02868f451c95e2", size = 415234, upload-time = "2025-10-08T09:15:31.022Z" }, + { url = "https://files.pythonhosted.org/packages/fc/6b/62e85ff7193663fbea5c0254ef32f0c77134b4059f8da89b958beb7696f3/msgpack-1.1.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5559d03930d3aa0f3aacb4c42c776af1a2ace2611871c84a75afe436695e6245", size = 435242, upload-time = "2025-10-08T09:15:37.647Z" }, + { url = "https://files.pythonhosted.org/packages/c1/47/5c74ecb4cc277cf09f64e913947871682ffa82b3b93c8dad68083112f412/msgpack-1.1.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70c5a7a9fea7f036b716191c29047374c10721c389c21e9ffafad04df8c52c90", size = 432509, upload-time = "2025-10-08T09:15:38.794Z" }, + { url = "https://files.pythonhosted.org/packages/24/a4/e98ccdb56dc4e98c929a3f150de1799831c0a800583cde9fa022fa90602d/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f2cb069d8b981abc72b41aea1c580ce92d57c673ec61af4c500153a626cb9e20", size = 415957, upload-time = "2025-10-08T09:15:40.238Z" }, + { url = "https://files.pythonhosted.org/packages/da/28/6951f7fb67bc0a4e184a6b38ab71a92d9ba58080b27a77d3e2fb0be5998f/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d62ce1f483f355f61adb5433ebfd8868c5f078d1a52d042b0a998682b4fa8c27", size = 422910, upload-time = "2025-10-08T09:15:41.505Z" }, +] + +[[package]] +name = "nest-asyncio" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/83/f8/51569ac65d696c8ecbee95938f89d4abf00f47d58d48f6fbabfe8f0baefe/nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe", size = 7418, upload-time = "2024-01-21T14:25:19.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/62/ae72ff66c0f1fd959925b4c11f8c2dea61f47f6acaea75a08512cdfe3fed/numpy-2.4.1.tar.gz", hash = "sha256:a1ceafc5042451a858231588a104093474c6a5c57dcc724841f5c888d237d690", size = 20721320, upload-time = "2026-01-10T06:44:59.619Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/e3/5f115fae982565771be994867c89bcd8d7208dbfe9469185497d70de5ddf/numpy-2.4.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0093e85df2960d7e4049664b26afc58b03236e967fb942354deef3208857a04c", size = 14404859, upload-time = "2026-01-10T06:42:49.947Z" }, + { url = "https://files.pythonhosted.org/packages/d9/7d/9c8a781c88933725445a859cac5d01b5871588a15969ee6aeb618ba99eee/numpy-2.4.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ad270f438cbdd402c364980317fb6b117d9ec5e226fff5b4148dd9aa9fc6e02", size = 16371419, upload-time = "2026-01-10T06:42:52.409Z" }, + { url = "https://files.pythonhosted.org/packages/a6/d2/8aa084818554543f17cf4162c42f162acbd3bb42688aefdba6628a859f77/numpy-2.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:297c72b1b98100c2e8f873d5d35fb551fce7040ade83d67dd51d38c8d42a2162", size = 16182131, upload-time = "2026-01-10T06:42:54.694Z" }, + { url = "https://files.pythonhosted.org/packages/60/db/0425216684297c58a8df35f3284ef56ec4a043e6d283f8a59c53562caf1b/numpy-2.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:cf6470d91d34bf669f61d515499859fa7a4c2f7c36434afb70e82df7217933f9", size = 18295342, upload-time = "2026-01-10T06:42:56.991Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a4/77333f4d1e4dac4395385482557aeecf4826e6ff517e32ca48e1dafbe42a/numpy-2.4.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0faba4a331195bfa96f93dd9dfaa10b2c7aa8cda3a02b7fd635e588fe821bf5", size = 14402871, upload-time = "2026-01-10T06:43:17.324Z" }, + { url = "https://files.pythonhosted.org/packages/ba/87/d341e519956273b39d8d47969dd1eaa1af740615394fe67d06f1efa68773/numpy-2.4.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3e3087f53e2b4428766b54932644d148613c5a595150533ae7f00dab2f319a8", size = 16359305, upload-time = "2026-01-10T06:43:19.376Z" }, + { url = "https://files.pythonhosted.org/packages/32/91/789132c6666288eaa20ae8066bb99eba1939362e8f1a534949a215246e97/numpy-2.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:49e792ec351315e16da54b543db06ca8a86985ab682602d90c60ef4ff4db2a9c", size = 16181909, upload-time = "2026-01-10T06:43:21.808Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b8/090b8bd27b82a844bb22ff8fdf7935cb1980b48d6e439ae116f53cdc2143/numpy-2.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:79e9e06c4c2379db47f3f6fc7a8652e7498251789bf8ff5bd43bf478ef314ca2", size = 18284380, upload-time = "2026-01-10T06:43:23.957Z" }, + { url = "https://files.pythonhosted.org/packages/50/96/9fa8635ed9d7c847d87e30c834f7109fac5e88549d79ef3324ab5c20919f/numpy-2.4.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:941c2a93313d030f219f3a71fd3d91a728b82979a5e8034eb2e60d394a2b83f9", size = 14462352, upload-time = "2026-01-10T06:43:39.479Z" }, + { url = "https://files.pythonhosted.org/packages/03/d1/8cf62d8bb2062da4fb82dd5d49e47c923f9c0738032f054e0a75342faba7/numpy-2.4.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:529050522e983e00a6c1c6b67411083630de8b57f65e853d7b03d9281b8694d2", size = 16407279, upload-time = "2026-01-10T06:43:41.93Z" }, + { url = "https://files.pythonhosted.org/packages/86/1c/95c86e17c6b0b31ce6ef219da00f71113b220bcb14938c8d9a05cee0ff53/numpy-2.4.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:2302dc0224c1cbc49bb94f7064f3f923a971bfae45c33870dcbff63a2a550505", size = 16248316, upload-time = "2026-01-10T06:43:44.121Z" }, + { url = "https://files.pythonhosted.org/packages/30/b4/e7f5ff8697274c9d0fa82398b6a372a27e5cef069b37df6355ccb1f1db1a/numpy-2.4.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:9171a42fcad32dcf3fa86f0a4faa5e9f8facefdb276f54b8b390d90447cff4e2", size = 18329884, upload-time = "2026-01-10T06:43:46.613Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c4/2e7908915c0e32ca636b92e4e4a3bdec4cb1e7eb0f8aedf1ed3c68a0d8cd/numpy-2.4.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5d558123217a83b2d1ba316b986e9248a1ed1971ad495963d555ccd75dcb1556", size = 14418963, upload-time = "2026-01-10T06:44:04.047Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c0/3ed5083d94e7ffd7c404e54619c088e11f2e1939a9544f5397f4adb1b8ba/numpy-2.4.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2f44de05659b67d20499cbc96d49f2650769afcb398b79b324bb6e297bfe3844", size = 16363811, upload-time = "2026-01-10T06:44:06.207Z" }, + { url = "https://files.pythonhosted.org/packages/0e/68/42b66f1852bf525050a67315a4fb94586ab7e9eaa541b1bef530fab0c5dd/numpy-2.4.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:69e7419c9012c4aaf695109564e3387f1259f001b4326dfa55907b098af082d3", size = 16197643, upload-time = "2026-01-10T06:44:08.33Z" }, + { url = "https://files.pythonhosted.org/packages/d2/40/e8714fc933d85f82c6bfc7b998a0649ad9769a32f3494ba86598aaf18a48/numpy-2.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2ffd257026eb1b34352e749d7cc1678b5eeec3e329ad8c9965a797e08ccba205", size = 18289601, upload-time = "2026-01-10T06:44:10.841Z" }, + { url = "https://files.pythonhosted.org/packages/c8/e5/0989b44ade47430be6323d05c23207636d67d7362a1796ccbccac6773dd2/numpy-2.4.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:423797bdab2eeefbe608d7c1ec7b2b4fd3c58d51460f1ee26c7500a1d9c9ee93", size = 14464653, upload-time = "2026-01-10T06:44:26.706Z" }, + { url = "https://files.pythonhosted.org/packages/10/a7/cfbe475c35371cae1358e61f20c5f075badc18c4797ab4354140e1d283cf/numpy-2.4.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:52b5f61bdb323b566b528899cc7db2ba5d1015bda7ea811a8bcf3c89c331fa42", size = 16405144, upload-time = "2026-01-10T06:44:29.378Z" }, + { url = "https://files.pythonhosted.org/packages/f8/a3/0c63fe66b534888fa5177cc7cef061541064dbe2b4b60dcc60ffaf0d2157/numpy-2.4.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:42d7dd5fa36d16d52a84f821eb96031836fd405ee6955dd732f2023724d0aa01", size = 16247425, upload-time = "2026-01-10T06:44:31.721Z" }, + { url = "https://files.pythonhosted.org/packages/6b/2b/55d980cfa2c93bd40ff4c290bf824d792bd41d2fe3487b07707559071760/numpy-2.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e7b6b5e28bbd47b7532698e5db2fe1db693d84b58c254e4389d99a27bb9b8f6b", size = 18330053, upload-time = "2026-01-10T06:44:34.617Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.9.1.4" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/6c/90d3f532f608a03a13c1d6c16c266ffa3828e8011b1549d3b61db2ad59f5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7a950dae01add3b415a5a5cdc4ec818fb5858263e9cca59004bb99fdbbd3a5d6", size = 575006342, upload-time = "2025-06-05T20:04:16.902Z" }, + { url = "https://files.pythonhosted.org/packages/77/3c/aa88abe01f3be3d1f8f787d1d33dc83e76fec05945f9a28fbb41cfb99cd5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:453611eb21a7c1f2c2156ed9f3a45b691deda0440ec550860290dc901af5b4c2", size = 581242350, upload-time = "2025-06-05T20:04:51.979Z" }, +] + +[[package]] +name = "nvidia-cuda-cccl-cu12" +version = "12.9.27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/7e/82e49956b046bdc506c789235c587d9b3ef58b8bc1782258c1e247229647/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7898b38aa68beaa234d48f0868273702342a196d6e2e9d0ef058dca2390ebea", size = 3152245, upload-time = "2025-05-01T19:32:04.802Z" }, + { url = "https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37869e17ce2e1ecec6eddf1927cca0f8c34e64fd848d40453df559091e2d7117", size = 3152243, upload-time = "2025-05-01T19:32:13.955Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/78/351b5c8cdbd9a6b4fb0d6ee73fb176dcdc1b6b6ad47c2ffff5ae8ca4a1f7/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:791853b030602c6a11d08b5578edfb957cadea06e9d3b26adbf8d036135a4afe", size = 10077166, upload-time = "2025-06-05T20:01:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, + { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, + { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.17.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/55/9987152bd99746b6d5d899a3e920f72f565f5adb4835dc44899382482e2c/nvidia_cudnn_cu12-9.17.1.4-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:d18d61bdd596fca0dbe459c129f6eb7a24ed2e6de1d7988b0a37ac63184ee05d", size = 646648394, upload-time = "2025-12-22T16:42:08.8Z" }, + { url = "https://files.pythonhosted.org/packages/91/8b/b15aef578b64ff8b72f68f682b6a8254ee6aa6d3d236b4282a4d32dbdf1a/nvidia_cudnn_cu12-9.17.1.4-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b5083ced291edb2baf8eab09951c6bc58b79b2023c4ec885657a63acdf51d0bd", size = 647835891, upload-time = "2025-12-22T16:43:28.317Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.4.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.5.82" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, + { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.10.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, + { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.29.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/b2/e4dc7b33020645746710040cb2a6ac0de8332687d3ce902156dd3d7c351a/nvidia_nccl_cu12-2.29.2-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:0712e55c067965c6093cc793a9bbcc5f37b5b47248e9ebf8ae3af06867757587", size = 289707761, upload-time = "2026-01-07T00:21:30.514Z" }, + { url = "https://files.pythonhosted.org/packages/23/2d/609d0392d992259c6dc39881688a7fc13b1397a668bc360fbd68d1396f85/nvidia_nccl_cu12-2.29.2-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:3a9a0bf4142126e0d0ed99ec202579bef8d007601f9fab75af60b10324666b12", size = 289762233, upload-time = "2026-01-07T00:21:56.124Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.5.19" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-cccl-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/1b/06a698b584b67670e7f108381ecdf680d3a497c0567f5183644034f4212e/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:333d91f68038539f20a591ef4d5cee25a85d440537b31809c4e42ae1700fccce", size = 152515031, upload-time = "2026-01-02T04:24:25.119Z" }, + { url = "https://files.pythonhosted.org/packages/64/b9/6ab941001c23cfb43499b5b0b7417b0bb4dfba3a29ffa2b06985422dad50/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f29a23c5ee75461ec4342e17757f9e010369ada8eb0d441070049db787d9e51b", size = 152681200, upload-time = "2026-01-02T04:24:46.524Z" }, +] + +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004, upload-time = "2024-09-26T14:33:24.483Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, +] + +[[package]] +name = "optax" +version = "0.2.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "chex", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/3b/90c11f740a3538200b61cd2b7d9346959cb9e31e0bdea3d2f886b7262203/optax-0.2.6.tar.gz", hash = "sha256:ba8d1e12678eba2657484d6feeca4fb281b8066bdfd5efbfc0f41b87663109c0", size = 269660, upload-time = "2025-09-15T22:41:24.76Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/ec/19c6cc6064c7fc8f0cd6d5b37c4747849e66040c6ca98f86565efc2c227c/optax-0.2.6-py3-none-any.whl", hash = "sha256:f875251a5ab20f179d4be57478354e8e21963373b10f9c3b762b94dcb8c36d91", size = 367782, upload-time = "2025-09-15T22:41:22.825Z" }, +] + +[[package]] +name = "orbax-checkpoint" +version = "0.11.31" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "aiofiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "etils", extra = ["epath", "epy"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "humanize", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "msgpack", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nest-asyncio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "psutil", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "simplejson", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tensorstore", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4b/1e/c65e35ab5ef9d380f4a4ce7f983a0dd360d229eec22204aacb80d2b91aca/orbax_checkpoint-0.11.31.tar.gz", hash = "sha256:f021193a619782655798bc4a285f40612f6fe647ddeb303d1f49cdbc5645e319", size = 406137, upload-time = "2025-12-11T18:09:17.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/3a/abbc3c5cac2e082e88cfec2161bf837f18fef786caaa3f007594c839fc8c/orbax_checkpoint-0.11.31-py3-none-any.whl", hash = "sha256:b00e39cd61cbd6c7c78b091ccac0ed1bbf3cf7788e761618e7070761195bfcc0", size = 602358, upload-time = "2025-12-11T18:09:15.667Z" }, +] + +[[package]] +name = "packaging" +version = "25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a1/d4/1fc4078c65507b51b96ca8f8c3ba19e6a61c8253c72794544580a7b6c24d/packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f", size = 165727, upload-time = "2025-04-19T11:48:59.673Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/12/38679034af332785aac8774540895e234f4d07f7545804097de4b666afd8/packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484", size = 66469, upload-time = "2025-04-19T11:48:57.875Z" }, +] + +[[package]] +name = "pillow" +version = "12.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/02/d52c733a2452ef1ffcc123b68e6606d07276b0e358db70eabad7e40042b7/pillow-12.1.0.tar.gz", hash = "sha256:5c5ae0a06e9ea030ab786b0251b32c7e4ce10e58d983c0d5c56029455180b5b9", size = 46977283, upload-time = "2026-01-02T09:13:29.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ad/26/7b82c0ab7ef40ebede7a97c72d473bda5950f609f8e0c77b04af574a0ddb/pillow-12.1.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:efdc140e7b63b8f739d09a99033aa430accce485ff78e6d311973a67b6bf3208", size = 6234878, upload-time = "2026-01-02T09:11:14.096Z" }, + { url = "https://files.pythonhosted.org/packages/76/25/27abc9792615b5e886ca9411ba6637b675f1b77af3104710ac7353fe5605/pillow-12.1.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bef9768cab184e7ae6e559c032e95ba8d07b3023c289f79a2bd36e8bf85605a5", size = 8044868, upload-time = "2026-01-02T09:11:15.903Z" }, + { url = "https://files.pythonhosted.org/packages/0a/ea/f200a4c36d836100e7bc738fc48cd963d3ba6372ebc8298a889e0cfc3359/pillow-12.1.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:742aea052cf5ab5034a53c3846165bc3ce88d7c38e954120db0ab867ca242661", size = 6349468, upload-time = "2026-01-02T09:11:17.631Z" }, + { url = "https://files.pythonhosted.org/packages/11/8f/48d0b77ab2200374c66d344459b8958c86693be99526450e7aee714e03e4/pillow-12.1.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a6dfc2af5b082b635af6e08e0d1f9f1c4e04d17d4e2ca0ef96131e85eda6eb17", size = 7041518, upload-time = "2026-01-02T09:11:19.389Z" }, + { url = "https://files.pythonhosted.org/packages/1d/23/c281182eb986b5d31f0a76d2a2c8cd41722d6fb8ed07521e802f9bba52de/pillow-12.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:609e89d9f90b581c8d16358c9087df76024cf058fa693dd3e1e1620823f39670", size = 6462829, upload-time = "2026-01-02T09:11:21.28Z" }, + { url = "https://files.pythonhosted.org/packages/25/ef/7018273e0faac099d7b00982abdcc39142ae6f3bd9ceb06de09779c4a9d6/pillow-12.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:43b4899cfd091a9693a1278c4982f3e50f7fb7cff5153b05174b4afc9593b616", size = 7166756, upload-time = "2026-01-02T09:11:23.559Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c7/2530a4aa28248623e9d7f27316b42e27c32ec410f695929696f2e0e4a778/pillow-12.1.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:7b5dd7cbae20285cdb597b10eb5a2c13aa9de6cde9bb64a3c1317427b1db1ae1", size = 4062543, upload-time = "2026-01-02T09:11:31.566Z" }, + { url = "https://files.pythonhosted.org/packages/8f/1f/40b8eae823dc1519b87d53c30ed9ef085506b05281d313031755c1705f73/pillow-12.1.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:29a4cef9cb672363926f0470afc516dbf7305a14d8c54f7abbb5c199cd8f8179", size = 4138373, upload-time = "2026-01-02T09:11:33.367Z" }, + { url = "https://files.pythonhosted.org/packages/d4/77/6fa60634cf06e52139fd0e89e5bbf055e8166c691c42fb162818b7fda31d/pillow-12.1.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:681088909d7e8fa9e31b9799aaa59ba5234c58e5e4f1951b4c4d1082a2e980e0", size = 3601241, upload-time = "2026-01-02T09:11:35.011Z" }, + { url = "https://files.pythonhosted.org/packages/1d/fd/f5a0896839762885b3376ff04878f86ab2b097c2f9a9cdccf4eda8ba8dc0/pillow-12.1.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6b7a9d1db5dad90e2991645874f708e87d9a3c370c243c2d7684d28f7e133e6b", size = 6232605, upload-time = "2026-01-02T09:11:40.602Z" }, + { url = "https://files.pythonhosted.org/packages/98/aa/938a09d127ac1e70e6ed467bd03834350b33ef646b31edb7452d5de43792/pillow-12.1.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6258f3260986990ba2fa8a874f8b6e808cf5abb51a94015ca3dc3c68aa4f30ea", size = 8041617, upload-time = "2026-01-02T09:11:42.721Z" }, + { url = "https://files.pythonhosted.org/packages/17/e8/538b24cb426ac0186e03f80f78bc8dc7246c667f58b540bdd57c71c9f79d/pillow-12.1.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e115c15e3bc727b1ca3e641a909f77f8ca72a64fff150f666fcc85e57701c26c", size = 6346509, upload-time = "2026-01-02T09:11:44.955Z" }, + { url = "https://files.pythonhosted.org/packages/01/9a/632e58ec89a32738cabfd9ec418f0e9898a2b4719afc581f07c04a05e3c9/pillow-12.1.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6741e6f3074a35e47c77b23a4e4f2d90db3ed905cb1c5e6e0d49bff2045632bc", size = 7038117, upload-time = "2026-01-02T09:11:46.736Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a2/d40308cf86eada842ca1f3ffa45d0ca0df7e4ab33c83f81e73f5eaed136d/pillow-12.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:935b9d1aed48fcfb3f838caac506f38e29621b44ccc4f8a64d575cb1b2a88644", size = 6460151, upload-time = "2026-01-02T09:11:48.625Z" }, + { url = "https://files.pythonhosted.org/packages/f1/88/f5b058ad6453a085c5266660a1417bdad590199da1b32fb4efcff9d33b05/pillow-12.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5fee4c04aad8932da9f8f710af2c1a15a83582cfb884152a9caa79d4efcdbf9c", size = 7164534, upload-time = "2026-01-02T09:11:50.445Z" }, + { url = "https://files.pythonhosted.org/packages/2b/c0/5273a99478956a099d533c4f46cbaa19fd69d606624f4334b85e50987a08/pillow-12.1.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2479c7f02f9d505682dc47df8c0ea1fc5e264c4d1629a5d63fe3e2334b89554", size = 6268974, upload-time = "2026-01-02T09:12:02.572Z" }, + { url = "https://files.pythonhosted.org/packages/b4/26/0bf714bc2e73d5267887d47931d53c4ceeceea6978148ed2ab2a4e6463c4/pillow-12.1.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f188d580bd870cda1e15183790d1cc2fa78f666e76077d103edf048eed9c356e", size = 8073070, upload-time = "2026-01-02T09:12:04.75Z" }, + { url = "https://files.pythonhosted.org/packages/43/cf/1ea826200de111a9d65724c54f927f3111dc5ae297f294b370a670c17786/pillow-12.1.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0fde7ec5538ab5095cc02df38ee99b0443ff0e1c847a045554cf5f9af1f4aa82", size = 6380176, upload-time = "2026-01-02T09:12:06.626Z" }, + { url = "https://files.pythonhosted.org/packages/03/e0/7938dd2b2013373fd85d96e0f38d62b7a5a262af21ac274250c7ca7847c9/pillow-12.1.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0ed07dca4a8464bada6139ab38f5382f83e5f111698caf3191cb8dbf27d908b4", size = 7067061, upload-time = "2026-01-02T09:12:08.624Z" }, + { url = "https://files.pythonhosted.org/packages/86/ad/a2aa97d37272a929a98437a8c0ac37b3cf012f4f8721e1bd5154699b2518/pillow-12.1.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:f45bd71d1fa5e5749587613037b172e0b3b23159d1c00ef2fc920da6f470e6f0", size = 6491824, upload-time = "2026-01-02T09:12:10.488Z" }, + { url = "https://files.pythonhosted.org/packages/a4/44/80e46611b288d51b115826f136fb3465653c28f491068a72d3da49b54cd4/pillow-12.1.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:277518bf4fe74aa91489e1b20577473b19ee70fb97c374aa50830b279f25841b", size = 7190911, upload-time = "2026-01-02T09:12:12.772Z" }, + { url = "https://files.pythonhosted.org/packages/8c/87/bdf971d8bbcf80a348cc3bacfcb239f5882100fe80534b0ce67a784181d8/pillow-12.1.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:5cb7bc1966d031aec37ddb9dcf15c2da5b2e9f7cc3ca7c54473a20a927e1eb91", size = 4062533, upload-time = "2026-01-02T09:12:20.791Z" }, + { url = "https://files.pythonhosted.org/packages/ff/4f/5eb37a681c68d605eb7034c004875c81f86ec9ef51f5be4a63eadd58859a/pillow-12.1.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:97e9993d5ed946aba26baf9c1e8cf18adbab584b99f452ee72f7ee8acb882796", size = 4138546, upload-time = "2026-01-02T09:12:23.664Z" }, + { url = "https://files.pythonhosted.org/packages/11/6d/19a95acb2edbace40dcd582d077b991646b7083c41b98da4ed7555b59733/pillow-12.1.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:414b9a78e14ffeb98128863314e62c3f24b8a86081066625700b7985b3f529bd", size = 3601163, upload-time = "2026-01-02T09:12:26.338Z" }, + { url = "https://files.pythonhosted.org/packages/6c/6b/c5742cea0f1ade0cd61485dc3d81f05261fc2276f537fbdc00802de56779/pillow-12.1.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e5dcbe95016e88437ecf33544ba5db21ef1b8dd6e1b434a2cb2a3d605299e643", size = 6232114, upload-time = "2026-01-02T09:12:32.936Z" }, + { url = "https://files.pythonhosted.org/packages/bf/8f/9f521268ce22d63991601aafd3d48d5ff7280a246a1ef62d626d67b44064/pillow-12.1.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d0a7735df32ccbcc98b98a1ac785cc4b19b580be1bdf0aeb5c03223220ea09d5", size = 8042708, upload-time = "2026-01-02T09:12:34.78Z" }, + { url = "https://files.pythonhosted.org/packages/1a/eb/257f38542893f021502a1bbe0c2e883c90b5cff26cc33b1584a841a06d30/pillow-12.1.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c27407a2d1b96774cbc4a7594129cc027339fd800cd081e44497722ea1179de", size = 6347762, upload-time = "2026-01-02T09:12:36.748Z" }, + { url = "https://files.pythonhosted.org/packages/c4/5a/8ba375025701c09b309e8d5163c5a4ce0102fa86bbf8800eb0d7ac87bc51/pillow-12.1.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:15c794d74303828eaa957ff8070846d0efe8c630901a1c753fdc63850e19ecd9", size = 7039265, upload-time = "2026-01-02T09:12:39.082Z" }, + { url = "https://files.pythonhosted.org/packages/cf/dc/cf5e4cdb3db533f539e88a7bbf9f190c64ab8a08a9bc7a4ccf55067872e4/pillow-12.1.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c990547452ee2800d8506c4150280757f88532f3de2a58e3022e9b179107862a", size = 6462341, upload-time = "2026-01-02T09:12:40.946Z" }, + { url = "https://files.pythonhosted.org/packages/d0/47/0291a25ac9550677e22eda48510cfc4fa4b2ef0396448b7fbdc0a6946309/pillow-12.1.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b63e13dd27da389ed9475b3d28510f0f954bca0041e8e551b2a4eb1eab56a39a", size = 7165395, upload-time = "2026-01-02T09:12:42.706Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e1/53ee5163f794aef1bf84243f755ee6897a92c708505350dd1923f4afec48/pillow-12.1.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8e178e3e99d3c0ea8fc64b88447f7cac8ccf058af422a6cedc690d0eadd98c51", size = 6269908, upload-time = "2026-01-02T09:12:54.884Z" }, + { url = "https://files.pythonhosted.org/packages/bc/0b/b4b4106ff0ee1afa1dc599fde6ab230417f800279745124f6c50bcffed8e/pillow-12.1.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:079af2fb0c599c2ec144ba2c02766d1b55498e373b3ac64687e43849fbbef5bc", size = 8074733, upload-time = "2026-01-02T09:12:56.802Z" }, + { url = "https://files.pythonhosted.org/packages/19/9f/80b411cbac4a732439e629a26ad3ef11907a8c7fc5377b7602f04f6fe4e7/pillow-12.1.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bdec5e43377761c5dbca620efb69a77f6855c5a379e32ac5b158f54c84212b14", size = 6381431, upload-time = "2026-01-02T09:12:58.823Z" }, + { url = "https://files.pythonhosted.org/packages/8f/b7/d65c45db463b66ecb6abc17c6ba6917a911202a07662247e1355ce1789e7/pillow-12.1.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:565c986f4b45c020f5421a4cea13ef294dde9509a8577f29b2fc5edc7587fff8", size = 7068529, upload-time = "2026-01-02T09:13:00.885Z" }, + { url = "https://files.pythonhosted.org/packages/50/96/dfd4cd726b4a45ae6e3c669fc9e49deb2241312605d33aba50499e9d9bd1/pillow-12.1.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:43aca0a55ce1eefc0aefa6253661cb54571857b1a7b2964bd8a1e3ef4b729924", size = 6492981, upload-time = "2026-01-02T09:13:03.314Z" }, + { url = "https://files.pythonhosted.org/packages/4d/1c/b5dc52cf713ae46033359c5ca920444f18a6359ce1020dd3e9c553ea5bc6/pillow-12.1.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0deedf2ea233722476b3a81e8cdfbad786f7adbed5d848469fa59fe52396e4ef", size = 7191878, upload-time = "2026-01-02T09:13:05.276Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "protobuf" +version = "6.33.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/5c/f912bdebdd4af4160da6a2c2b1b3aaa1b8c578d0243ba8f694f93c7095f0/protobuf-6.33.3.tar.gz", hash = "sha256:c8794debeb402963fddff41a595e1f649bcd76616ba56c835645cab4539e810e", size = 444318, upload-time = "2026-01-09T23:05:02.79Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/46/551c69b6ff1957bd703654342bfb776bb97db400bc80afc56fbb64e7c11d/protobuf-6.33.3-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:08a6ca12f60ba99097dd3625ef4275280f99c9037990e47ce9368826b159b890", size = 324469, upload-time = "2026-01-09T23:04:54.332Z" }, + { url = "https://files.pythonhosted.org/packages/38/8c/6522b8e543ece46f645911c3cebe361d8460134c0fee02ddcf70ebf32999/protobuf-6.33.3-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:6fa9b5f4baa12257542273e5e6f3c3d3867b30bc2770c14ad9ac8315264bf986", size = 323298, upload-time = "2026-01-09T23:04:56.866Z" }, + { url = "https://files.pythonhosted.org/packages/a6/b9/067b8a843569d5605ba6f7c039b9319720a974f82216cd623e13186d3078/protobuf-6.33.3-py3-none-any.whl", hash = "sha256:c2bf221076b0d463551efa2e1319f08d4cffcc5f0d864614ccd3d0e77a637794", size = 170518, upload-time = "2026-01-09T23:05:01.227Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/73/cb/09e5184fb5fc0358d110fc3ca7f6b1d033800734d34cac10f4136cfac10e/psutil-7.2.1.tar.gz", hash = "sha256:f7583aec590485b43ca601dd9cea0dcd65bd7bb21d30ef4ddbf4ea6b5ed1bdd3", size = 490253, upload-time = "2025-12-29T08:26:00.169Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/6d/ed44901e830739af5f72a85fa7ec5ff1edea7f81bfbf4875e409007149bd/psutil-7.2.1-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ea46c0d060491051d39f0d2cff4f98d5c72b288289f57a21556cc7d504db37fc", size = 180612, upload-time = "2025-12-29T08:26:08.276Z" }, + { url = "https://files.pythonhosted.org/packages/c7/65/b628f8459bca4efbfae50d4bf3feaab803de9a160b9d5f3bd9295a33f0c2/psutil-7.2.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:35630d5af80d5d0d49cfc4d64c1c13838baf6717a13effb35869a5919b854cdf", size = 183201, upload-time = "2025-12-29T08:26:10.622Z" }, + { url = "https://files.pythonhosted.org/packages/77/69/19f1eb0e01d24c2b3eacbc2f78d3b5add8a89bf0bb69465bc8d563cc33de/psutil-7.2.1-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:93f3f7b0bb07711b49626e7940d6fe52aa9940ad86e8f7e74842e73189712129", size = 181518, upload-time = "2025-12-29T08:26:20.241Z" }, + { url = "https://files.pythonhosted.org/packages/e1/6d/7e18b1b4fa13ad370787626c95887b027656ad4829c156bb6569d02f3262/psutil-7.2.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d34d2ca888208eea2b5c68186841336a7f5e0b990edec929be909353a202768a", size = 184348, upload-time = "2025-12-29T08:26:22.215Z" }, + { url = "https://files.pythonhosted.org/packages/29/f8/40e01c350ad9a2b3cb4e6adbcc8a83b17ee50dd5792102b6142385937db5/psutil-7.2.1-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e38404ca2bb30ed7267a46c02f06ff842e92da3bb8c5bfdadbd35a5722314d8", size = 154694, upload-time = "2025-12-29T08:26:32.147Z" }, + { url = "https://files.pythonhosted.org/packages/06/e4/b751cdf839c011a9714a783f120e6a86b7494eb70044d7d81a25a5cd295f/psutil-7.2.1-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ab2b98c9fc19f13f59628d94df5cc4cc4844bc572467d113a8b517d634e362c6", size = 156136, upload-time = "2025-12-29T08:26:34.079Z" }, + { url = "https://files.pythonhosted.org/packages/44/ad/bbf6595a8134ee1e94a4487af3f132cef7fce43aef4a93b49912a48c3af7/psutil-7.2.1-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f78baafb38436d5a128f837fab2d92c276dfb48af01a240b861ae02b2413ada8", size = 148108, upload-time = "2025-12-29T08:26:36.225Z" }, + { url = "https://files.pythonhosted.org/packages/1c/15/dd6fd869753ce82ff64dcbc18356093471a5a5adf4f77ed1f805d473d859/psutil-7.2.1-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:99a4cd17a5fdd1f3d014396502daa70b5ec21bf4ffe38393e152f8e449757d67", size = 147402, upload-time = "2025-12-29T08:26:39.21Z" }, +] + +[[package]] +name = "pydantic" +version = "2.12.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-core", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-inspection", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.41.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, + { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, + { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, + { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, + { url = "https://files.pythonhosted.org/packages/15/df/a4c740c0943e93e6500f9eb23f4ca7ec9bf71b19e608ae5b579678c8d02f/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0cbaad15cb0c90aa221d43c00e77bb33c93e8d36e0bf74760cd00e732d10a6a0", size = 1919307, upload-time = "2025-11-04T13:40:29.806Z" }, + { url = "https://files.pythonhosted.org/packages/cf/4e/35a80cae583a37cf15604b44240e45c05e04e86f9cfd766623149297e971/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:406bf18d345822d6c21366031003612b9c77b3e29ffdb0f612367352aab7d586", size = 2073164, upload-time = "2025-11-04T13:40:40.289Z" }, + { url = "https://files.pythonhosted.org/packages/75/c7/20bd7fc05f0c6ea2056a4565c6f36f8968c0924f19b7d97bbfea55780e73/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:01a3d0ab748ee531f4ea6c3e48ad9dac84ddba4b0d82291f87248f2f9de8d740", size = 2137788, upload-time = "2025-11-04T13:40:44.752Z" }, + { url = "https://files.pythonhosted.org/packages/9c/59/013626bf8c78a5a5d9350d12e7697d3d4de951a75565496abd40ccd46bee/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:915c3d10f81bec3a74fbd4faebe8391013ba61e5a1a8d48c4455b923bdda7858", size = 2324852, upload-time = "2025-11-04T13:40:48.575Z" }, + { url = "https://files.pythonhosted.org/packages/23/04/e89c29e267b8060b40dca97bfc64a19b2a3cf99018167ea1677d96368273/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:25e1c2af0fce638d5f1988b686f3b3ea8cd7de5f244ca147c777769e798a9cd1", size = 1915040, upload-time = "2025-11-04T13:41:00.853Z" }, + { url = "https://files.pythonhosted.org/packages/4c/d2/ef2074dc020dd6e109611a8be4449b98cd25e1b9b8a303c2f0fca2f2bcf7/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22f0fb8c1c583a3b6f24df2470833b40207e907b90c928cc8d3594b76f874375", size = 2064877, upload-time = "2025-11-04T13:41:09.827Z" }, + { url = "https://files.pythonhosted.org/packages/d3/9e/3ce66cebb929f3ced22be85d4c2399b8e85b622db77dad36b73c5387f8f8/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0177272f88ab8312479336e1d777f6b124537d47f2123f89cb37e0accea97f90", size = 2138960, upload-time = "2025-11-04T13:41:14.627Z" }, + { url = "https://files.pythonhosted.org/packages/3c/0d/f05e79471e889d74d3d88f5bd20d0ed189ad94c2423d81ff8d0000aab4ff/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:e56ba91f47764cc14f1daacd723e3e82d1a89d783f0f5afe9c364b8bb491ccdb", size = 2326039, upload-time = "2025-11-04T13:41:18.934Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f0/e5e6b99d4191da102f2b0eb9687aaa7f5bea5d9964071a84effc3e40f997/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3006c3dd9ba34b0c094c544c6006cc79e87d8612999f1a5d43b769b89181f23c", size = 1878092, upload-time = "2025-11-04T13:41:33.21Z" }, + { url = "https://files.pythonhosted.org/packages/aa/4e/2ae1aa85d6af35a39b236b1b1641de73f5a6ac4d5a7509f77b814885760c/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2ff4321e56e879ee8d2a879501c8e469414d948f4aba74a2d4593184eb326660", size = 2041078, upload-time = "2025-11-04T13:41:42.323Z" }, + { url = "https://files.pythonhosted.org/packages/02/7a/f999a6dcbcd0e5660bc348a3991c8915ce6599f4f2c6ac22f01d7a10816c/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:a39455728aabd58ceabb03c90e12f71fd30fa69615760a075b9fec596456ccc3", size = 2129560, upload-time = "2025-11-04T13:41:47.474Z" }, + { url = "https://files.pythonhosted.org/packages/d9/02/3c562f3a51afd4d88fff8dffb1771b30cfdfd79befd9883ee094f5b6c0d8/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:2a5e06546e19f24c6a96a129142a75cee553cc018ffee48a460059b1185f4470", size = 2331955, upload-time = "2025-11-04T13:41:54.079Z" }, + { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, + { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, +] + +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "iniconfig", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pluggy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, +] + +[[package]] +name = "qwix" +version = "0.1.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9a/6d/119d353c5a2597f8122c992466aaa69e102e6d8f59587a55e65517f34edb/qwix-0.1.5.tar.gz", hash = "sha256:935fefd41f2b26d0fe545e433bff658b1ee476c83b7c6e467e31f769d67a74e2", size = 74227, upload-time = "2025-12-12T01:12:03.809Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/70/a64d02681d30cfdce59278968b23945169a37f8eb5fc3c1ba590f809edc6/qwix-0.1.5-py3-none-any.whl", hash = "sha256:21e71c52e22b95b3926b48b90453fcd7b9bd80f5251d52429bf36adbaffaa043", size = 96125, upload-time = "2025-12-12T01:12:02.799Z" }, +] + +[[package]] +name = "rdkit" +version = "2025.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/fe/48a1b8a77c144682be59f2ec0213609b5b55900a4f572c8c25b54d5aed61/rdkit-2025.9.4-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:86f3871783de1b8996832ec02ebcb350260ece4508f77cba32f0a11acdcdce0f", size = 35100276, upload-time = "2026-01-31T12:07:49.263Z" }, + { url = "https://files.pythonhosted.org/packages/be/d1/c687917cf1cdf5c5412437a5f4a91921497c2d0ffeb827b23f5489a10842/rdkit-2025.9.4-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:79a081174c53498efb7d6d48737bc3803aaf04153d0c4d6b0ea132dfeec68b51", size = 36641090, upload-time = "2026-01-31T12:07:54.627Z" }, + { url = "https://files.pythonhosted.org/packages/8e/25/74d28e1bb5f36260f57955119ceef9c7259a01f35e737310b853d3abfa72/rdkit-2025.9.4-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1a65e1a78357638b5533584c1418a1ddafbc1bd7dffde8a8e08135d54495b58e", size = 35098273, upload-time = "2026-01-31T12:08:06.188Z" }, + { url = "https://files.pythonhosted.org/packages/7b/8b/a2bd32153a9c6526eb7db7aa848328927b2e37674ac5591f470ce3578ac0/rdkit-2025.9.4-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:19e3d3bb763c56831a8299974e209634d5f065e9cfb4ca67a79730d11bec233a", size = 36639736, upload-time = "2026-01-31T12:08:10.106Z" }, + { url = "https://files.pythonhosted.org/packages/64/a7/fb688b2ccae058a5226604a52be36d9becf2c85af2041231a3c3594eb082/rdkit-2025.9.4-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:90b5f9ce64e10c550bb313877a7cf27526aad981bb46aa94c2df123fb5fb9535", size = 35131961, upload-time = "2026-01-31T12:08:21.949Z" }, + { url = "https://files.pythonhosted.org/packages/46/33/3e07254b69080b3c03fe931b33af47d5fe823961f294fbc544ce8710a5a9/rdkit-2025.9.4-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:5a1c827cd97488a2bf6d38f15811a03ebeb691f1f0427d063e739b03d8180f8c", size = 36651962, upload-time = "2026-01-31T12:08:26.426Z" }, +] + +[[package]] +name = "rich" +version = "14.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fb/d2/8920e102050a0de7bfabeb4c4614a49248cf8d5d7a8d01885fbb24dc767a/rich-14.2.0.tar.gz", hash = "sha256:73ff50c7c0c1c77c8243079283f4edb376f0f6442433aecb8ce7e6d0b92d1fe4", size = 219990, upload-time = "2025-10-09T14:16:53.064Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/7a/b0178788f8dc6cafce37a212c99565fa1fe7872c70c6c9c1e1a372d9d88f/rich-14.2.0-py3-none-any.whl", hash = "sha256:76bc51fe2e57d2b1be1f96c524b890b816e334ab4c1e45888799bfaab0021edd", size = 243393, upload-time = "2025-10-09T14:16:51.245Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/3e/9cca699f3486ce6bc12ff46dc2031f1ec8eb9ccc9a320fdaf925f1417426/scipy-1.17.0.tar.gz", hash = "sha256:2591060c8e648d8b96439e111ac41fd8342fdeff1876be2e19dea3fe8930454e", size = 30396830, upload-time = "2026-01-10T21:34:23.009Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/69/7c347e857224fcaf32a34a05183b9d8a7aca25f8f2d10b8a698b8388561a/scipy-1.17.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5194c445d0a1c7a6c1a4a4681b6b7c71baad98ff66d96b949097e7513c9d6742", size = 32724197, upload-time = "2026-01-10T21:25:44.084Z" }, + { url = "https://files.pythonhosted.org/packages/d1/fe/66d73b76d378ba8cc2fe605920c0c75092e3a65ae746e1e767d9d020a75a/scipy-1.17.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9eeb9b5f5997f75507814ed9d298ab23f62cf79f5a3ef90031b1ee2506abdb5b", size = 35009148, upload-time = "2026-01-10T21:25:50.591Z" }, + { url = "https://files.pythonhosted.org/packages/af/07/07dec27d9dc41c18d8c43c69e9e413431d20c53a0339c388bcf72f353c4b/scipy-1.17.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:40052543f7bbe921df4408f46003d6f01c6af109b9e2c8a66dd1cf6cf57f7d5d", size = 34798766, upload-time = "2026-01-10T21:25:59.41Z" }, + { url = "https://files.pythonhosted.org/packages/81/61/0470810c8a093cdacd4ba7504b8a218fd49ca070d79eca23a615f5d9a0b0/scipy-1.17.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0cf46c8013fec9d3694dc572f0b54100c28405d55d3e2cb15e2895b25057996e", size = 37405953, upload-time = "2026-01-10T21:26:07.75Z" }, + { url = "https://files.pythonhosted.org/packages/c9/10/be13397a0e434f98e0c79552b2b584ae5bb1c8b2be95db421533bbca5369/scipy-1.17.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe508b5690e9eaaa9467fc047f833af58f1152ae51a0d0aed67aa5801f4dd7d6", size = 32696338, upload-time = "2026-01-10T21:26:55.521Z" }, + { url = "https://files.pythonhosted.org/packages/63/1e/12fbf2a3bb240161651c94bb5cdd0eae5d4e8cc6eaeceb74ab07b12a753d/scipy-1.17.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6680f2dfd4f6182e7d6db161344537da644d1cf85cf293f015c60a17ecf08752", size = 34977201, upload-time = "2026-01-10T21:27:03.501Z" }, + { url = "https://files.pythonhosted.org/packages/19/5b/1a63923e23ccd20bd32156d7dd708af5bbde410daa993aa2500c847ab2d2/scipy-1.17.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:eec3842ec9ac9de5917899b277428886042a93db0b227ebbe3a333b64ec7643d", size = 34777384, upload-time = "2026-01-10T21:27:11.423Z" }, + { url = "https://files.pythonhosted.org/packages/39/22/b5da95d74edcf81e540e467202a988c50fef41bd2011f46e05f72ba07df6/scipy-1.17.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d7425fcafbc09a03731e1bc05581f5fad988e48c6a861f441b7ab729a49a55ea", size = 37379586, upload-time = "2026-01-10T21:27:20.171Z" }, + { url = "https://files.pythonhosted.org/packages/9b/9d/025cccdd738a72140efc582b1641d0dd4caf2e86c3fb127568dc80444e6e/scipy-1.17.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:130d12926ae34399d157de777472bf82e9061c60cc081372b3118edacafe1d00", size = 32815098, upload-time = "2026-01-10T21:27:54.389Z" }, + { url = "https://files.pythonhosted.org/packages/48/5f/09b879619f8bca15ce392bfc1894bd9c54377e01d1b3f2f3b595a1b4d945/scipy-1.17.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e886000eb4919eae3a44f035e63f0fd8b651234117e8f6f29bad1cd26e7bc45", size = 35031342, upload-time = "2026-01-10T21:28:03.012Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9a/f0f0a9f0aa079d2f106555b984ff0fbb11a837df280f04f71f056ea9c6e4/scipy-1.17.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:13c4096ac6bc31d706018f06a49abe0485f96499deb82066b94d19b02f664209", size = 34893199, upload-time = "2026-01-10T21:28:10.832Z" }, + { url = "https://files.pythonhosted.org/packages/90/b8/4f0f5cf0c5ea4d7548424e6533e6b17d164f34a6e2fb2e43ffebb6697b06/scipy-1.17.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:cacbaddd91fcffde703934897c5cd2c7cb0371fac195d383f4e1f1c5d3f3bd04", size = 37438061, upload-time = "2026-01-10T21:28:19.684Z" }, + { url = "https://files.pythonhosted.org/packages/80/5c/ea5d239cda2dd3d31399424967a24d556cf409fbea7b5b21412b0fd0a44f/scipy-1.17.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a38c3337e00be6fd8a95b4ed66b5d988bac4ec888fd922c2ea9fe5fb1603dd67", size = 32757834, upload-time = "2026-01-10T21:29:23.406Z" }, + { url = "https://files.pythonhosted.org/packages/b8/7e/8c917cc573310e5dc91cbeead76f1b600d3fb17cf0969db02c9cf92e3cfa/scipy-1.17.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00fb5f8ec8398ad90215008d8b6009c9db9fa924fd4c7d6be307c6f945f9cd73", size = 34995775, upload-time = "2026-01-10T21:29:31.915Z" }, + { url = "https://files.pythonhosted.org/packages/c5/43/176c0c3c07b3f7df324e7cdd933d3e2c4898ca202b090bd5ba122f9fe270/scipy-1.17.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f2a4942b0f5f7c23c7cd641a0ca1955e2ae83dedcff537e3a0259096635e186b", size = 34841240, upload-time = "2026-01-10T21:29:39.995Z" }, + { url = "https://files.pythonhosted.org/packages/44/8c/d1f5f4b491160592e7f084d997de53a8e896a3ac01cd07e59f43ca222744/scipy-1.17.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:dbf133ced83889583156566d2bdf7a07ff89228fe0c0cb727f777de92092ec6b", size = 37394463, upload-time = "2026-01-10T21:29:48.723Z" }, + { url = "https://files.pythonhosted.org/packages/7c/74/3498563a2c619e8a3ebb4d75457486c249b19b5b04a30600dfd9af06bea5/scipy-1.17.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5fb10d17e649e1446410895639f3385fd2bf4c3c7dfc9bea937bddcbc3d7b9ba", size = 32829770, upload-time = "2026-01-10T21:30:16.359Z" }, + { url = "https://files.pythonhosted.org/packages/48/d1/7b50cedd8c6c9d6f706b4b36fa8544d829c712a75e370f763b318e9638c1/scipy-1.17.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8547e7c57f932e7354a2319fab613981cde910631979f74c9b542bb167a8b9db", size = 35051093, upload-time = "2026-01-10T21:30:22.987Z" }, + { url = "https://files.pythonhosted.org/packages/e2/82/a2d684dfddb87ba1b3ea325df7c3293496ee9accb3a19abe9429bce94755/scipy-1.17.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:33af70d040e8af9d5e7a38b5ed3b772adddd281e3062ff23fec49e49681c38cf", size = 34909905, upload-time = "2026-01-10T21:30:28.704Z" }, + { url = "https://files.pythonhosted.org/packages/ef/5e/e565bd73991d42023eb82bb99e51c5b3d9e2c588ca9d4b3e2cc1d3ca62a6/scipy-1.17.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9eb55bb97d00f8b7ab95cb64f873eb0bf54d9446264d9f3609130381233483f", size = 37457743, upload-time = "2026-01-10T21:30:34.819Z" }, +] + +[[package]] +name = "simplejson" +version = "3.20.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/f4/a1ac5ed32f7ed9a088d62a59d410d4c204b3b3815722e2ccfb491fa8251b/simplejson-3.20.2.tar.gz", hash = "sha256:5fe7a6ce14d1c300d80d08695b7f7e633de6cd72c80644021874d985b3393649", size = 85784, upload-time = "2025-09-26T16:29:36.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/ec/5c6db08e42f380f005d03944be1af1a6bd501cc641175429a1cbe7fb23b9/simplejson-3.20.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:12a6b2816b6cab6c3fd273d43b1948bc9acf708272074c8858f579c394f4cbc9", size = 150335, upload-time = "2025-09-26T16:28:05.027Z" }, + { url = "https://files.pythonhosted.org/packages/20/05/ed9b2571bbf38f1a2425391f18e3ac11cb1e91482c22d644a1640dea9da7/simplejson-3.20.2-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:979ce23ea663895ae39106946ef3d78527822d918a136dbc77b9e2b7f006237e", size = 152367, upload-time = "2025-09-26T16:28:08.921Z" }, + { url = "https://files.pythonhosted.org/packages/81/2c/bad68b05dd43e93f77994b920505634d31ed239418eb6a88997d06599983/simplejson-3.20.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a2ba921b047bb029805726800819675249ef25d2f65fd0edb90639c5b1c3033c", size = 150205, upload-time = "2025-09-26T16:28:10.086Z" }, + { url = "https://files.pythonhosted.org/packages/71/ad/d7f3c331fb930638420ac6d236db68e9f4c28dab9c03164c3cd0e7967e15/simplejson-3.20.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:30e590e133b06773f0dc9c3f82e567463df40598b660b5adf53eb1c488202544", size = 154367, upload-time = "2025-09-26T16:28:14.393Z" }, + { url = "https://files.pythonhosted.org/packages/09/1d/171009bd35c7099d72ef6afd4bb13527bab469965c968a17d69a203d62a6/simplejson-3.20.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:552f55745044a24c3cb7ec67e54234be56d5d6d0e054f2e4cf4fb3e297429be5", size = 150579, upload-time = "2025-09-26T16:28:21.337Z" }, + { url = "https://files.pythonhosted.org/packages/43/f1/b392952200f3393bb06fbc4dd975fc63a6843261705839355560b7264eb2/simplejson-3.20.2-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:133ae2098a8e162c71da97cdab1f383afdd91373b7ff5fe65169b04167da976b", size = 152598, upload-time = "2025-09-26T16:28:24.962Z" }, + { url = "https://files.pythonhosted.org/packages/f4/b4/d6b7279e52a3e9c0fa8c032ce6164e593e8d9cf390698ee981ed0864291b/simplejson-3.20.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7977640af7b7d5e6a852d26622057d428706a550f7f5083e7c4dd010a84d941f", size = 150498, upload-time = "2025-09-26T16:28:26.114Z" }, + { url = "https://files.pythonhosted.org/packages/99/21/603709455827cdf5b9d83abe726343f542491ca8dc6a2528eb08de0cf034/simplejson-3.20.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f28ee755fadb426ba2e464d6fcf25d3f152a05eb6b38e0b4f790352f5540c769", size = 154717, upload-time = "2025-09-26T16:28:30.288Z" }, + { url = "https://files.pythonhosted.org/packages/05/5b/83e1ff87eb60ca706972f7e02e15c0b33396e7bdbd080069a5d1b53cf0d8/simplejson-3.20.2-py3-none-any.whl", hash = "sha256:3b6bb7fb96efd673eac2e4235200bfffdc2353ad12c54117e1e4e2fc485ac017", size = 57309, upload-time = "2025-09-26T16:29:35.312Z" }, +] + +[[package]] +name = "tabulate" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/fe/802052aecb21e3797b8f7902564ab6ea0d60ff8ca23952079064155d1ae1/tabulate-0.9.0.tar.gz", hash = "sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c", size = 81090, upload-time = "2022-10-06T17:21:48.54Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" }, +] + +[[package]] +name = "tensorboardx" +version = "2.6.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "protobuf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2b/c5/d4cc6e293fb837aaf9f76dd7745476aeba8ef7ef5146c3b3f9ee375fe7a5/tensorboardx-2.6.4.tar.gz", hash = "sha256:b163ccb7798b31100b9f5fa4d6bc22dad362d7065c2f24b51e50731adde86828", size = 4769801, upload-time = "2025-06-10T22:37:07.419Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/1d/b5d63f1a6b824282b57f7b581810d20b7a28ca951f2d5b59f1eb0782c12b/tensorboardx-2.6.4-py3-none-any.whl", hash = "sha256:5970cf3a1f0a6a6e8b180ccf46f3fe832b8a25a70b86e5a237048a7c0beb18e2", size = 87201, upload-time = "2025-06-10T22:37:05.44Z" }, +] + +[[package]] +name = "tensorstore" +version = "0.1.80" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/18/7b91daa9cf29dbb6bfdd603154f355c9069a9cd8c757038fe52b0f613611/tensorstore-0.1.80.tar.gz", hash = "sha256:4158fe76b96f62d12a37d7868150d836e089b5280b2bdd363c43c5d651f10e26", size = 7090032, upload-time = "2025-12-10T21:35:10.941Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/15/0734521f8b648e2c43a00f1bc99a7195646c9e4e31f64ab22a15ac84e75c/tensorstore-0.1.80-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:acb8d52fadcefafef4ef8ecca3fc99b1d0e3c5c5a888766484c3e39f050be7f5", size = 19013402, upload-time = "2025-12-10T21:34:24.961Z" }, + { url = "https://files.pythonhosted.org/packages/48/85/55addd16896343ea2731388028945576060139dda3c68a15d6b00158ef6f/tensorstore-0.1.80-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bc28a58c580253a526a4b6d239d18181ef96f1e285a502dbb03ff15eeec07a5b", size = 21007488, upload-time = "2025-12-10T21:34:28.093Z" }, + { url = "https://files.pythonhosted.org/packages/f7/fa/d5de3f1b711773e33a329b5fe11de1265b77a13f2a2447fe685ee5d0c1bc/tensorstore-0.1.80-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de63843706fdfe9565a45567238c5b1e55a0b28bbde6524200b31d29043a9a16", size = 19013246, upload-time = "2025-12-10T21:34:38.507Z" }, + { url = "https://files.pythonhosted.org/packages/87/ee/e874b5a495a7aa14817772a91095971f3a965a4cef5b52ad06a8e15c924f/tensorstore-0.1.80-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c8dbbdd31cbb28eccfb23dbbd4218fe67bfc32e9cb452875a485b81031c949d", size = 21008391, upload-time = "2025-12-10T21:34:41.332Z" }, + { url = "https://files.pythonhosted.org/packages/b6/df/71f317633a0cd5270b85d185ac5ce91a749930fc076205d3fae4f1f043ed/tensorstore-0.1.80-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:07e4a84bacf70b78305831897068a9b5ad30326e63bbeb92c4bf7e565fcf5e9e", size = 19020675, upload-time = "2025-12-10T21:34:51.168Z" }, + { url = "https://files.pythonhosted.org/packages/2b/35/f03cdb5edf8e009ff73e48c0c3d0f692a70a7ffc5e393f2ea1761eff89b5/tensorstore-0.1.80-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2b353b0bd53fedd77fc5a12a1c1a91cacc3cf59e3dd785529c5a54b31d1c7b1", size = 21009171, upload-time = "2025-12-10T21:34:53.979Z" }, + { url = "https://files.pythonhosted.org/packages/42/47/5c63024ced48e3f440c131babedef2f5398f48ab81c1aeee6c6193491d1c/tensorstore-0.1.80-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b7c5dd434bba4ee08fe46bbbdb25c60dd3d47ccb4b8561a9751cf1526da52b8", size = 19024739, upload-time = "2025-12-10T21:35:04.324Z" }, + { url = "https://files.pythonhosted.org/packages/6e/16/d08ade819949e0622f27e949c15b09f7b86ac18f8ac7c4d8bdfb4a711076/tensorstore-0.1.80-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e93df6d34ff5f0f6be245f4d29b99a7c1eef8ad91b50686adf57a5eeea99cb74", size = 21024449, upload-time = "2025-12-10T21:35:08.149Z" }, +] + +[[package]] +name = "tokamax" +version = "0.0.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "einshape", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "immutabledict", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jaxtyping", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "qwix", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tensorboardx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typeguard", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/77/ef/262ad37c150d98bcf24ce9bed8ff39dd247ba079b85a5522d423475be49a/tokamax-0.0.11.tar.gz", hash = "sha256:bba14f9ec86df573c7993e7e1e809f924417598249afd9d3a459e0ee3135749f", size = 424560, upload-time = "2026-03-09T19:47:17.293Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/61/6b2f5819373b989c975b82c591ca91e7e584e84d17b4487e47c76915b907/tokamax-0.0.11-py3-none-any.whl", hash = "sha256:ba90451944d4badc63a2110b7f183e5cb584375ed68364fc7cd6a98e27d004c3", size = 600218, upload-time = "2026-03-09T19:47:15.731Z" }, +] + +[[package]] +name = "toolz" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, +] + +[[package]] +name = "tqdm" +version = "4.67.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, +] + +[[package]] +name = "treescope" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f0/2a/d13d3c38862632742d2fe2f7ae307c431db06538fd05ca03020d207b5dcc/treescope-0.1.10.tar.gz", hash = "sha256:20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95", size = 138870, upload-time = "2025-08-08T05:43:48.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/2b/36e984399089c026a6499ac8f7401d38487cf0183839a4aa78140d373771/treescope-0.1.10-py3-none-any.whl", hash = "sha256:dde52f5314f4c29d22157a6fe4d3bd103f9cae02791c9e672eefa32c9aa1da51", size = 182255, upload-time = "2025-08-08T05:43:46.673Z" }, +] + +[[package]] +name = "typeguard" +version = "2.13.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/38/c61bfcf62a7b572b5e9363a802ff92559cb427ee963048e1442e3aef7490/typeguard-2.13.3.tar.gz", hash = "sha256:00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4", size = 40604, upload-time = "2021-12-10T21:09:39.158Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/bb/d43e5c75054e53efce310e79d63df0ac3f25e34c926be5dffb7d283fb2a8/typeguard-2.13.3-py3-none-any.whl", hash = "sha256:5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1", size = 17605, upload-time = "2021-12-10T21:09:37.844Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + +[[package]] +name = "zipp" +version = "3.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, +] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/6d/db/ddb11011826ed7db9d0e485d13df79b58586bfdec56e5c84a928a9a78c1c/zstandard-0.25.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bfc4e20784722098822e3eee42b8e576b379ed72cca4a7cb856ae733e62192ea", size = 5063001, upload-time = "2025-09-14T22:17:31.044Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/e3678b4e776db00f9f7b2fe58e547e8928ef32727d7a1ff01dea010f3f13/zstandard-0.25.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8e735494da3db08694d26480f1493ad2cf86e99bdd53e8e9771b2752a5c0246a", size = 5547173, upload-time = "2025-09-14T22:17:36.084Z" }, + { url = "https://files.pythonhosted.org/packages/4e/d5/ba05ed95c6b8ec30bd468dfeab20589f2cf709b5c940483e31d991f2ca58/zstandard-0.25.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3a39c94ad7866160a4a46d772e43311a743c316942037671beb264e395bdd611", size = 5046736, upload-time = "2025-09-14T22:17:37.891Z" }, + { url = "https://files.pythonhosted.org/packages/50/d5/870aa06b3a76c73eced65c044b92286a3c4e00554005ff51962deef28e28/zstandard-0.25.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:172de1f06947577d3a3005416977cce6168f2261284c02080e7ad0185faeced3", size = 5576368, upload-time = "2025-09-14T22:17:40.206Z" }, + { url = "https://files.pythonhosted.org/packages/5d/35/398dc2ffc89d304d59bc12f0fdd931b4ce455bddf7038a0a67733a25f550/zstandard-0.25.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3c83b0188c852a47cd13ef3bf9209fb0a77fa5374958b8c53aaa699398c6bd7b", size = 4954022, upload-time = "2025-09-14T22:17:41.879Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e5/fbd822d5c6f427cf158316d012c5a12f233473c2f9c5fe5ab1ae5d21f3d8/zstandard-0.25.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4f187a0bb61b35119d1926aee039524d1f93aaf38a9916b8c4b78ac8514a0aaf", size = 5360113, upload-time = "2025-09-14T22:17:48.893Z" }, + { url = "https://files.pythonhosted.org/packages/a4/86/a48e56320d0a17189ab7a42645387334fba2200e904ee47fc5a26c1fd8ca/zstandard-0.25.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223415140608d0f0da010499eaa8ccdb9af210a543fac54bce15babbcfc78439", size = 5058095, upload-time = "2025-09-14T22:17:59.498Z" }, + { url = "https://files.pythonhosted.org/packages/31/dc/cc50210e11e465c975462439a492516a73300ab8caa8f5e0902544fd748b/zstandard-0.25.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e09bb6252b6476d8d56100e8147b803befa9a12cea144bbe629dd508800d1ad0", size = 5560402, upload-time = "2025-09-14T22:18:05.954Z" }, + { url = "https://files.pythonhosted.org/packages/c9/ae/56523ae9c142f0c08efd5e868a6da613ae76614eca1305259c3bf6a0ed43/zstandard-0.25.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a9ec8c642d1ec73287ae3e726792dd86c96f5681eb8df274a757bf62b750eae7", size = 4955108, upload-time = "2025-09-14T22:18:07.68Z" }, + { url = "https://files.pythonhosted.org/packages/f6/49/fac46df5ad353d50535e118d6983069df68ca5908d4d65b8c466150a4ff1/zstandard-0.25.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4203ce3b31aec23012d3a4cf4a2ed64d12fea5269c49aed5e4c3611b938e4088", size = 5359591, upload-time = "2025-09-14T22:18:16.465Z" }, +]