#
# Copyright 2016 The BigDL Authors.
#
# 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.
#


def _attn_wrapper(origin_attn):
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        attn_output, attn_weights = origin_attn(self,
                                                query=query,
                                                key=key,
                                                value=value,
                                                attention_mask=attention_mask,
                                                head_mask=head_mask)
        if query.device.type == 'xpu' and 1 < query.numel() // query.size(-1) <= 64:
            attn_output = attn_output.clone()
        return attn_output, attn_weights
    return _attn
