From 4d546e2d70d27589c228d4179fb7b9eee576299d Mon Sep 17 00:00:00 2001 From: Dmitriy Suponitskiy Date: Thu, 10 Sep 2026 20:43:35 -0400 Subject: [PATCH 1/2] Expose automorphism key indices as a Python list - Add GetExistingEvalAutomorphismKeyIndices to query generated keys by tag. - Document its behavior, demonstrate bootstrapping key counts, and add tests for rotation and bootstrapping keys. --- examples/pke/simple-ckks-bootstrapping.py | 3 ++ src/include/docstrings/cryptocontext_docs.h | 14 +++++ src/lib/bindings.cpp | 6 +++ tests/test_ckks.py | 60 +++++++++++++++++++++ 4 files changed, 83 insertions(+) diff --git a/examples/pke/simple-ckks-bootstrapping.py b/examples/pke/simple-ckks-bootstrapping.py index 612eb58..56ceba2 100644 --- a/examples/pke/simple-ckks-bootstrapping.py +++ b/examples/pke/simple-ckks-bootstrapping.py @@ -52,6 +52,9 @@ def simple_bootstrap_example(): cryptocontext.EvalMultKeyGen(key_pair.secretKey) cryptocontext.EvalBootstrapKeyGen(key_pair.secretKey, num_slots) + key_indices = cryptocontext.GetExistingEvalAutomorphismKeyIndices(key_pair.secretKey.GetKeyTag()) + print(f"Number of bootstrapping automorphism keys: {len(key_indices)}") + x = [0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 4.0, 5.0] encoded_length = len(x) diff --git a/src/include/docstrings/cryptocontext_docs.h b/src/include/docstrings/cryptocontext_docs.h index f9fa759..c1c7943 100644 --- a/src/include/docstrings/cryptocontext_docs.h +++ b/src/include/docstrings/cryptocontext_docs.h @@ -957,6 +957,20 @@ const char* cc_GetEvalAutomorphismKeyMap_docs = R"pbdoc( :rtype: EvalKeyMap )pbdoc"; +const char* cc_GetExistingEvalAutomorphismKeyIndices_docs = R"pbdoc( + Get the automorphism indices of all evaluation keys held for a secret key tag, including + the conjugation keys generated by EvalBootstrapKeyGen. + + These are automorphism indices, not the slot offsets passed to EvalRotateKeyGen. Map an + offset with FindAutomorphismIndex() to look it up here; that function takes non-negative + offsets only. + + :param keyTag: secret key identifier, obtained with secretKey.GetKeyTag() + :type keyTag: str + :return: sorted list of unique automorphism indices; empty if no keys exist for the tag + :rtype: list[int] +)pbdoc"; + const char* cc_GetEvalSumKeyMap_docs = R"pbdoc( Get a map of summation keys (each is composed of several automorphism keys) for a specific secret key tag :return: EvalKeyMap: key map diff --git a/src/lib/bindings.cpp b/src/lib/bindings.cpp index 051bff4..ff56ff1 100644 --- a/src/lib/bindings.cpp +++ b/src/lib/bindings.cpp @@ -1148,6 +1148,12 @@ void bind_crypto_context(py::module &m) { .def_static("GetEvalAutomorphismKeyMap", &CryptoContextImpl::GetEvalAutomorphismKeyMapPtr, py::arg("keyTag") = "", py::doc(cc_GetEvalAutomorphismKeyMap_docs)) + .def_static("GetExistingEvalAutomorphismKeyIndices", [](const std::string& keyTag) { + const auto indices = CryptoContextImpl::GetExistingEvalAutomorphismKeyIndices(keyTag); + return std::vector(indices.begin(), indices.end()); + }, + py::arg("keyTag") = "", + py::doc(cc_GetExistingEvalAutomorphismKeyIndices_docs)) .def_static("SerializeEvalMultKey", [](const std::string &filename, const SerType::SERBINARY &sertype, std::string keyTag = "") { std::ofstream outfile(filename, std::ios::out | std::ios::binary); bool res = CryptoContextImpl::SerializeEvalMultKey(outfile, sertype, keyTag); diff --git a/tests/test_ckks.py b/tests/test_ckks.py index 48342c7..59d5b30 100644 --- a/tests/test_ckks.py +++ b/tests/test_ckks.py @@ -55,3 +55,63 @@ def test_add_two_numbers(ckks_context): raw_added = [a + b for (a, b) in zip(*raw)] total = sum(abs(a - b) for (a, b) in zip(raw_added, final_added)) assert total < 1e-3 + + +def test_existing_eval_automorphism_key_indices(ckks_context): + _, cc, existing_keys = ckks_context + keys = cc.KeyGen() + key_tag = keys.secretKey.GetKeyTag() + assert cc.GetExistingEvalAutomorphismKeyIndices(keyTag=key_tag) == [] + + cc.EvalRotateKeyGen(keys.secretKey, [1, -2]) + # CKKS rotation offsets map to powers of 5 modulo the cyclotomic order. + cyclotomic_order = 2 * cc.GetRingDimension() + expected = sorted(pow(5, offset, cyclotomic_order) for offset in [1, -2]) + indices = cc.GetExistingEvalAutomorphismKeyIndices(key_tag) + assert isinstance(indices, list) + assert indices == expected + assert fhe.CryptoContext.GetExistingEvalAutomorphismKeyIndices(key_tag) == expected + # The documented way to go from a slot offset to an entry in this list. This also + # cross-checks the 5^offset formula above. FindAutomorphismIndex rejects negative + # offsets, so only the positive one round-trips. + assert cc.FindAutomorphismIndex(1) == pow(5, 1, cyclotomic_order) + assert cc.FindAutomorphismIndex(1) in indices + with pytest.raises(TypeError): + cc.FindAutomorphismIndex(-2) + # keyTag defaults to "", which matches no key map. + assert cc.GetExistingEvalAutomorphismKeyIndices() == [] + + # Repeated key generation must not duplicate indices or affect another tag. + cc.EvalRotateKeyGen(keys.secretKey, [1, 3]) + updated = sorted(expected + [pow(5, 3, cyclotomic_order)]) + assert cc.GetExistingEvalAutomorphismKeyIndices(key_tag) == updated + assert cc.GetExistingEvalAutomorphismKeyIndices(existing_keys.secretKey.GetKeyTag()) == expected + assert indices == expected + + +def test_existing_bootstrap_key_indices(): + parameters = fhe.CCParamsCKKSRNS() + parameters.SetSecurityLevel(fhe.HEStd_NotSet) + parameters.SetRingDim(512) + parameters.SetSecretKeyDist(fhe.UNIFORM_TERNARY) + level_budget = [2, 2] + depth = fhe.FHECKKSRNS.GetBootstrapDepth(level_budget, fhe.UNIFORM_TERNARY) + parameters.SetMultiplicativeDepth(depth + 2) + parameters.SetScalingTechnique(fhe.FIXEDAUTO) + parameters.SetScalingModSize(78 if fhe.get_native_int() == 128 else 59) + parameters.SetFirstModSize(89 if fhe.get_native_int() == 128 else 60) + cc = fhe.GenCryptoContext(parameters) + for feature in [fhe.PKE, fhe.KEYSWITCH, fhe.LEVELEDSHE, fhe.ADVANCEDSHE, fhe.FHE]: + cc.Enable(feature) + + slots = 8 + cc.EvalBootstrapSetup(level_budget, [0, 0], slots) + keys = cc.KeyGen() + key_tag = keys.secretKey.GetKeyTag() + assert cc.GetExistingEvalAutomorphismKeyIndices(key_tag) == [] + cc.EvalBootstrapKeyGen(keys.secretKey, slots) + indices = cc.GetExistingEvalAutomorphismKeyIndices(key_tag) + assert isinstance(indices, list) + assert indices == sorted(set(indices)) + assert len(indices) > 1 + assert 2 * cc.GetRingDimension() - 1 in indices # Conjugation key. From e262e30f6974ef05c470b0ec98da932e8d873b80 Mon Sep 17 00:00:00 2001 From: Dmitriy Suponitskiy Date: Thu, 10 Sep 2026 22:17:14 -0400 Subject: [PATCH 2/2] Fix automorphism index bindings to accept signed rotation offsets - Accept signed 32-bit offsets in FindAutomorphismIndex and FindAutomorphismIndices, converting them to the C++ API's unsigned form. - Update docstrings and add CKKS coverage for negative offsets, empty lists, and out-of-range inputs. --- src/include/docstrings/cryptocontext_docs.h | 9 ++++---- src/lib/bindings.cpp | 9 ++++++-- tests/test_ckks.py | 25 ++++++++++++++++----- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/src/include/docstrings/cryptocontext_docs.h b/src/include/docstrings/cryptocontext_docs.h index c1c7943..b7f1e71 100644 --- a/src/include/docstrings/cryptocontext_docs.h +++ b/src/include/docstrings/cryptocontext_docs.h @@ -961,9 +961,8 @@ const char* cc_GetExistingEvalAutomorphismKeyIndices_docs = R"pbdoc( Get the automorphism indices of all evaluation keys held for a secret key tag, including the conjugation keys generated by EvalBootstrapKeyGen. - These are automorphism indices, not the slot offsets passed to EvalRotateKeyGen. Map an - offset with FindAutomorphismIndex() to look it up here; that function takes non-negative - offsets only. + These are automorphism indices, not the slot offsets passed to EvalRotateKeyGen. Map a + signed offset with FindAutomorphismIndex() to look it up here. :param keyTag: secret key identifier, obtained with secretKey.GetKeyTag() :type keyTag: str @@ -1590,7 +1589,7 @@ const char* cc_EvalAutomorphismKeyGen_docs = R"pbdoc( const char* cc_FindAutomorphismIndex_docs = R"pbdoc( Finds an automorphism index for a given vector index using a scheme-specific algorithm - :param idx: regular vector index + :param idx: signed 32-bit rotation offset, as accepted by EvalRotateKeyGen :type idx: int :return: the automorphism index :rtype: int @@ -1599,7 +1598,7 @@ const char* cc_FindAutomorphismIndex_docs = R"pbdoc( const char* cc_FindAutomorphismIndices_docs = R"pbdoc( Finds automorphism indices for a given list of vector indices using a scheme-specific algorithm - :param idxList: list of indices + :param idxList: list of signed 32-bit rotation offsets, as accepted by EvalRotateKeyGen :type idxList: List[int] :return: a list of automorphism indices :rtype: List[int] diff --git a/src/lib/bindings.cpp b/src/lib/bindings.cpp index ff56ff1..b5727e7 100644 --- a/src/lib/bindings.cpp +++ b/src/lib/bindings.cpp @@ -1110,10 +1110,15 @@ void bind_crypto_context(py::module &m) { py::arg("ciphertextVec")) .def("EvalAddManyInPlace", &CryptoContextImpl::EvalAddManyInPlace, py::arg("ciphertextVec")) - .def("FindAutomorphismIndex", &CryptoContextImpl::FindAutomorphismIndex, + .def("FindAutomorphismIndex", [](const CryptoContextImpl& self, int32_t idx) { + return self.FindAutomorphismIndex(static_cast(idx)); + }, py::arg("idx"), py::doc(cc_FindAutomorphismIndex_docs)) - .def("FindAutomorphismIndices", &CryptoContextImpl::FindAutomorphismIndices, + .def("FindAutomorphismIndices", [](const CryptoContextImpl& self, const std::vector& idxList) { + // The C++ API carries signed rotation offsets in uint32_t values. + return self.FindAutomorphismIndices(std::vector(idxList.begin(), idxList.end())); + }, py::arg("idxList"), py::doc(cc_FindAutomorphismIndices_docs)) .def("GetEvalSumKeyMap", diff --git a/tests/test_ckks.py b/tests/test_ckks.py index 59d5b30..26fd747 100644 --- a/tests/test_ckks.py +++ b/tests/test_ckks.py @@ -72,12 +72,10 @@ def test_existing_eval_automorphism_key_indices(ckks_context): assert indices == expected assert fhe.CryptoContext.GetExistingEvalAutomorphismKeyIndices(key_tag) == expected # The documented way to go from a slot offset to an entry in this list. This also - # cross-checks the 5^offset formula above. FindAutomorphismIndex rejects negative - # offsets, so only the positive one round-trips. - assert cc.FindAutomorphismIndex(1) == pow(5, 1, cyclotomic_order) - assert cc.FindAutomorphismIndex(1) in indices - with pytest.raises(TypeError): - cc.FindAutomorphismIndex(-2) + # cross-checks the 5^offset formula above for both rotation directions. + for offset in [1, -2]: + assert cc.FindAutomorphismIndex(offset) == pow(5, offset, cyclotomic_order) + assert cc.FindAutomorphismIndex(offset) in indices # keyTag defaults to "", which matches no key map. assert cc.GetExistingEvalAutomorphismKeyIndices() == [] @@ -89,6 +87,21 @@ def test_existing_eval_automorphism_key_indices(ckks_context): assert indices == expected +def test_find_automorphism_indices_signed_offsets(ckks_context): + _, cc, _ = ckks_context + cyclotomic_order = 2 * cc.GetRingDimension() + offsets = [0, 1, -2, 3, -1, -2] + expected = [pow(5, offset, cyclotomic_order) for offset in offsets] + assert cc.FindAutomorphismIndices(idxList=offsets) == expected + assert [cc.FindAutomorphismIndex(idx=offset) for offset in offsets] == expected + assert cc.FindAutomorphismIndices([]) == [] + for offset in [-(2**31) - 1, 2**31]: + with pytest.raises(TypeError): + cc.FindAutomorphismIndex(offset) + with pytest.raises(TypeError): + cc.FindAutomorphismIndices([1, offset]) + + def test_existing_bootstrap_key_indices(): parameters = fhe.CCParamsCKKSRNS() parameters.SetSecurityLevel(fhe.HEStd_NotSet)