Runtime error on reshape operation to higher order tensor

[tsunami935]

Summary

HELayers is unable to encode an ONNX model with a reshape operation to a higher order, ie (BxTxC) -> (BxTxHxD) where H x D = C. Reshape operations to an equal or lower order work without issue.

Running the code below produces the following error:

terminate called after throwing an instance of 'std::_Nested_exception<std::runtime_error>'
  what():  Neural network internal simulation failed: Neural network forward pass of [3] Reshape (node_view) layer failed: /data/helayers/src/helayers/math/TensorDimensionMapping.cpp:690: selectOriginalDims: Assertion failed: dims.size() <= getOrigOrder()
terminate called recursively
Aborted (core dumped)

Minimal Reproducible Code

# Make model and export to ONNX

import torch
import torch.nn as nn

class LinearReshapeModuleMH(nn.Module):
    def __init__(self, dim_in, h):
        super().__init__()
        assert dim_in % h == 0
        self.dim_in = dim_in
        self.h = h
        self.l1 = nn.Linear(dim_in, dim_in)

    def forward(self, x: torch.Tensor):
        b, t, c = x.size()
        x = self.l1(x)

        # If we include an unsqueeze(-1) operation here then the reshape works when loaded into pyhelayers.

        x = x.reshape(b, t, self.h, c // self.h)
        return x
    
model = LinearReshapeModuleMH(256, 4).to(dtype=torch.float32)
model.eval()
output_path = "models/reshape-mh.onnx"
torch.onnx.export(
    model,  # model being run
    torch.randn((1,32,256)).to(dtype=torch.float32), 
    output_path, 
    opset_version=17,
    export_params=True, 
    external_data=False,
    input_names=["input"],  
    output_names=["output"],
)

def force_allowzero_zero(onnx_path_in, onnx_path_out):
    import onnx
    """
    Torch 2.11+ has allowzero=1 for all reshape nodes by default.
    HELayers requires allowzero=0. 
    Since we know that there are no 0 dimensions, we set allowzero=0 for all reshape nodes.
    """
    model = onnx.load(onnx_path_in)
    for node in model.graph.node:
        if node.op_type == "Reshape":
            node.attribute.pop(0)
            node.attribute.append(onnx.helper.make_attribute("allowzero", 0))

    onnx.save(model, onnx_path_out)
force_allowzero_zero(output_path, output_path)

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

# Load and run with pyhelayers

import numpy as np
import os
import psutil
import time
import pyhelayers

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

# Reshape Module
input_file = ["models/reshape-mh.onnx"]
input_sample = [(np.random.rand(1, 32, 256).astype(np.float64) - 0.5) * 1.0]

# The error shows up regardless of context
context = pyhelayers.SealCkksContext()
he_run_req = pyhelayers.HeRunRequirements()
he_run_req.set_he_context_options([context])
he_run_req.set_model_encrypted(False)
he_run_req.set_lazy_mode(pyhelayers.LazyMode.LAZY_ENCODING)
print("Created HE context and run requirements.")

nn = pyhelayers.NeuralNet()
print("Initialized network.")
nn.encode(input_file, he_run_req, hyper_params) # The program fails here!
print("Encoded model.")
context = nn.get_created_he_context()

model_io_encoder = pyhelayers.ModelIoEncoder(nn)
samples = pyhelayers.EncryptedData(context)
model_io_encoder.encode_encrypt(samples, input_sample)
print("Encoded samples.")

predictions = pyhelayers.EncryptedData(context)
nn.predict(predictions, samples)
print("Predictions ready.")

plain_predictions = model_io_encoder.decrypt_decode_output(predictions)

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

The model runs without issue when loaded with the ONNX runtime.
I also found that if an unsqeeze(-1) operation is applied before the reshape operation,
then the model is successfully encoded, and the program runs without issues.
This suggests to me that there is an issue with how the reshape operation is encoded in HELayers.

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 will check it soon and update.