Skip to content

Commit c9ac82c

Browse files
loisloGoogle-ML-Automation
authored andcommitted
[XLA:GPU] Add missing BF16_BF16_F32_X9 matmul option in config.py
Extend the list of possible default algorithms that dot could use. PiperOrigin-RevId: 736879149
1 parent 5944c9e commit c9ac82c

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

jax/_src/config.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1404,8 +1404,8 @@ def _update_disable_jit_thread_local(val):
14041404
'ANY_F8_ANY_F8_F32', 'ANY_F8_ANY_F8_F32_FAST_ACCUM', 'ANY_F8_ANY_F8_ANY',
14051405
'ANY_F8_ANY_F8_ANY_FAST_ACCUM', 'F16_F16_F16', 'F16_F16_F32',
14061406
'BF16_BF16_BF16', 'BF16_BF16_F32', 'BF16_BF16_F32_X3',
1407-
'BF16_BF16_F32_X6', 'TF32_TF32_F32', 'TF32_TF32_F32_X3', 'F32_F32_F32',
1408-
'F64_F64_F64',
1407+
'BF16_BF16_F32_X6', 'BF16_BF16_F32_X9', 'TF32_TF32_F32',
1408+
'TF32_TF32_F32_X3', 'F32_F32_F32', 'F64_F64_F64',
14091409
],
14101410
default=None,
14111411
help=('Control the default matmul and conv precision for 32bit inputs.\n\n'

0 commit comments

Comments
 (0)