src/transformers/models/yoso/modeling_yoso.py/YosoLSHCumulation/backward
class YosoLSHCumulation:    def backward(ctx, grad):
        grad = to_contiguous(grad)

        query_mask, key_mask, query_hash_code, key_hash_code, query, key, value = ctx.saved_tensors
        config = ctx.config

        use_cuda = grad.is_cuda
        hash_code_len = config["hash_code_len"]
        hashtable_capacity = int(2**hash_code_len)

        if config["lsh_backward"]:
            grad_value = lsh_cumulation.lsh_cumulation(
                key_mask, key_hash_code, query_mask, query_hash_code, grad, hashtable_capacity, use_cuda, 1
            )
            grad_query = lsh_cumulation.lsh_weighted_cumulation(
                query_mask,
                query_hash_code,
                grad,
                key_mask,
                key_hash_code,
                value,
                (hash_code_len / 2) * key,
                hashtable_capacity,
                use_cuda,
                4,
            )
            grad_key = lsh_cumulation.lsh_weighted_cumulation(
                key_mask,
                key_hash_code,
                value,
                query_mask,
                query_hash_code,
                grad,
                (hash_code_len / 2) * query,
                hashtable_capacity,
                use_cuda,
                4,
            )
        else:
            expectation = (1 - torch.acos(torch.matmul(query, key.transpose(-1, -2))) / math.pi) ** hash_code_len
            expectation = expectation * query_mask[:, :, None] * key_mask[:, None, :]
            weighted_exp = torch.matmul(grad, value.transpose(-1, -2)) * expectation
            grad_query = torch.matmul(weighted_exp, (hash_code_len / 2) * key)
            grad_key = torch.matmul(weighted_exp.transpose(-1, -2), (hash_code_len / 2) * query)
            grad_value = torch.matmul(expectation.transpose(-1, -2), grad)

        return None, None, grad_query, grad_key, grad_value, None