pub_soft/alphafold3
3
0
Fork 0
mirror of https://github.com/google-deepmind/alphafold3.git synced 2026-06-02 11:54:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Explicitly set backend='triton' in pallas_call invocations.

Browse Source 
This change is being made because the default backend for `pallas_call` is changing. Explicitly setting `backend='triton'` ensures that these calls continue to use the Triton backend, maintaining current behavior.

PiperOrigin-RevId: 822019610
Change-Id: I7ba83894fcd960d424502335c8b904d76c88a733
This commit is contained in:
Benjamin Chetioui
2025-10-21 02:25:57 -07:00
committed by Copybara-Service
parent c9121bf646
commit cf404610b5
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/alphafold3/jax/gated_linear_unit/matmul_ext.py
View File

@@ -221,6 +221,7 @@ def _gated_linear_unit(
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)