1# Config options to enable/disable C++ kernel for nn.functional.MHA 2# and nn.TransformerEncoder 3import torch 4 5 6_is_fastpath_enabled: bool = True 7 8 9def get_fastpath_enabled() -> bool: 10 """Returns whether fast path for TransformerEncoder and MultiHeadAttention 11 is enabled, or ``True`` if jit is scripting. 12 13 ..note: 14 The fastpath might not be run even if ``get_fastpath_enabled`` returns 15 ``True`` unless all conditions on inputs are met. 16 """ 17 if not torch.jit.is_scripting(): 18 return _is_fastpath_enabled 19 return True 20 21 22def set_fastpath_enabled(value: bool) -> None: 23 """Sets whether fast path is enabled""" 24 global _is_fastpath_enabled 25 _is_fastpath_enabled = value 26