OpenFheCkks backend fails running FFN with polynomial activations of degree >= 6 on 3D inputs

[tsunami935]

Summary

When running a feedforward network (FFN) with 3D input
and polynomial activation of degree >= 6 on OpenFheCkks backend,
there is a recursive terminate call.
This issue does not appear with SealCkks or LattigoCkks backend.

Running python -X faulthandler openfhe_nn_bug.py results in the following error:

Start prediction...
terminate called recursively
terminate called recursively
terminate called after throwing an instance of 'std::runtime_errorFatal Python error: '
Aborted

Thread 0x000070e9ccfff6c0terminate called recursively
 (most recent call first):
  File "terminate called recursively
~/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/lib/python3.11/threading.py", line 331 in wait
  File terminate called recursively
"~/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/lib/python3.11/threading.py", line 629 in wait
  File "~/helayers/.venv/lib/python3.11/site-packages/tqdm/_monitor.py", line 60 in run
  File "~/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/lib/python3.11/threading.py", line 1045 in _bootstrap_inner
  File "~/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/lib/python3.11/threading.py", line 1002 in _bootstrap

Thread 0x000070eb249b9740 (most recent call first):
  File "~/helayers/bugs/openfhe_nn_bug.py", line 78 in <module>
  File "<frozen runpy>", line 88 in _run_code
  File "<frozen runpy>", line 198 in _run_module_as_main
Aborted (core dumped)

Minimal Reproducible Code

import os
import time
import psutil

import torch
import torch.nn as nn
import torch.nn.functional as F

class FFN(nn.Module):
    def __init__(self, d, h, degree):
        super().__init__()
        self.l1 = nn.Linear(d, h, bias=False)
        self.degree = degree
        self.coefs = nn.Parameter(torch.ones(degree) / degree)
        self.l2 = nn.Linear(h, d, bias=False)

    def forward(self, x: torch.Tensor):
        x = self.l1(x)
        x = sum([self.coefs[i] * (x ** i) for i in range(self.degree)])
        x = self.l2(x)
        return x

# If using degree <= 5, then there's no problem
model = FFN(384, 384 * 4, 6).to(dtype=torch.float32)
model.eval()
output_path = "models/openfhe_nn_bug.onnx"
torch.onnx.export(
    model,  # model being run
    torch.randn((1,512,384)).to(dtype=torch.float32), 
    output_path, 
    opset_version=17,
    export_params=True, 
    external_data=False,
    input_names=["input"],  
    output_names=["output"],
)

#################################################################################

import numpy as np
import pyhelayers

# # There's no issues with any other context
# he_context = pyhelayers.LattigoContext()
# he_context = pyhelayers.SealCkksContext()
he_context = pyhelayers.OpenFheCkksContext()
print("Using backend: ", he_context.get_library_name())

hyper_params = pyhelayers.PlainModelHyperParams()
hyper_params.init_random_weights = False
hyper_params.min_rand_value = -0.1
hyper_params.max_rand_value = 0.1
hyper_params.sparse_rate = 0.5
hyper_params.verbose = True

np.random.seed(30)
input_file = [output_path]
input_sample = [(np.random.rand(1, 512, 384).astype(np.float64) - 0.5) * 1.0]

he_run_req = pyhelayers.HeRunRequirements()
he_run_req.set_he_context_options([he_context])
he_run_req.set_model_encrypted(False)
he_run_req.set_lazy_mode(pyhelayers.LazyMode.LAZY_ENCODING)

nn = pyhelayers.NeuralNet()
nn.encode(input_file, he_run_req, hyper_params) # Encode

model_io_encoder = pyhelayers.ModelIoEncoder(nn)
samples = pyhelayers.EncryptedData(he_context)
model_io_encoder.encode_encrypt(samples, input_sample)
print('Encrypted samples ready')

process = psutil.Process(os.getpid())
print("Memory consumption (GB): ", process.memory_info().rss / (1000*1000*1000))

predictions = pyhelayers.EncryptedData(he_context)
print("Start prediction...")
start = time.perf_counter()
nn.predict(predictions, samples) # We crash here.
duration=time.perf_counter() - start

print('Prediction done')
print('predict per sample:', duration, 's')

process = psutil.Process(os.getpid())
print("Memory consumption (GB): ", process.memory_info().rss / (1000*1000*1000))

plain_predictions = model_io_encoder.decrypt_decode_output(predictions)
print('Decrypted predictions ready')

process = psutil.Process(os.getpid())
print("Memory consumption (GB): ", process.memory_info().rss / (1000*1000*1000))

plain_predictions = plain_predictions.astype(np.float32)
print(plain_predictions)

process = psutil.Process(os.getpid())
print("Memory consumption (GB): ", process.memory_info().rss / (1000*1000*1000))

Environment Configuration:

• python 3.11.15
• pyhelayers==1.5.5.3
• torch==2.11
• onnx==1.19.0
• onnx-ir==0.2.1
• onnxscript==0.7.0

[aehud]

Thanks. We'll check this one too. We'll update soon and publish fixes for all three issues together.