Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[New Feature]: Add support for transposing a tensor along a specified axes #443

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

silencrown
Copy link

Description

Hello, while using TenSEAL, I noticed that the current version does not support transposing a tensor along a specified dimension. This may make it more complicated for users to implement functions like multi-head attention. Based on xtensor, I have modified the existing code to add support for the aforementioned operation in TenSEAL. In this code submission, users can utilize the Python method transpose(0, 2, 1). This operation is consistent with NumPy and can be adapted with minimal changes to support operations similar to PyTorch's transpose(-2, -1).

Affected Dependencies

No

How has this been tested?

The test code is included in this pull request.

Checklist

@silencrown
Copy link
Author

Here's a simple example that might better describe what this PR does. In the following example, the three operators np.transpose(data, [0, 2, 1] torch_tensor.transpose(-2, -1) ckks_tensor.transpose([0, 2, 1] are equivalent.

shape = [2, 3, 4]
expected_shape = [2, 4, 3]

data = np.arange(24).reshape(shape)

# numpy
np_trans = np.transpose(data, [0, 2, 1])
        
# torch
pt_trans = np.array(torch.tensor(data).transpose(-2, -1).tolist())

# CKKSTensor
enc_tensor = cxt_man.encrypt(data)
enc_trans = enc_tensor.transpose([0, 2, 1])
dec_trans = np.array(enc_trans.decrypt().tolist())

assert list(dec_trans.shape) == expected_shape
assert np.allclose(pt_trans, np_trans, rtol=0, atol=0.01)
assert np.allclose(dec_trans, np_trans, rtol=0, atol=0.01)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant