Skip to content

Commit

Permalink
[tmva][sofie] Fix Tile operator
Browse files Browse the repository at this point in the history
Fix the casting of the input repeats vector data of the Tile operator.
The type of data should be int64_t and not size_t
The previous code could have some issue when the size of size_t is 32 bits
  • Loading branch information
lmoneta committed Dec 20, 2024
1 parent 8720c03 commit a4a9b0e
Showing 1 changed file with 19 additions and 12 deletions.
31 changes: 19 additions & 12 deletions tmva/sofie/inc/TMVA/ROperator_Tile.hxx
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,17 @@ public:
}
fShapeInput=model.GetTensorShape(fNInput);

// Retrieve the data pointer for the repeats tensor
// if repeats vector is not initialized we cannot deduce shape of output
// not support for time being this case
if (!model.IsInitializedTensor(fNRepeats)) {
throw std::runtime_error("TMVA SOFIE Tile Op: non-initialized repeats input is not supported");
}

// Retrieve the data pointer for the repeats tensor
auto repptr = model.GetInitializedTensorData(fNRepeats);
// Cast the raw pointer to the appropriate type (size_t*)
auto repeat_shape = static_cast<size_t*>(repptr.get());

if (repeat_shape == nullptr) {
auto repeats_data = static_cast<int64_t*>(repptr.get());
if (repeats_data == nullptr) {
throw std::runtime_error("Failed to retrieve the data for the repeats tensor.");
}
// Get the shape of the repeats tensor to determine the number of elements
Expand All @@ -66,12 +71,18 @@ public:
throw std::runtime_error("Repeats tensor is not 1D.");
}
size_t num_elements = repeats_shape[0];
// Convert the data to a vector
std::vector<size_t> repeats_vector(repeat_shape, repeat_shape + num_elements);
// Convert the data to a vector of size_t
std::vector<size_t> repeats_vector(num_elements);
std::copy(repeats_data, repeats_data + num_elements, repeats_vector.begin());


fShapeY = ShapeInference({fShapeInput,repeats_vector})[0];

model.AddIntermediateTensor(fNY, model.GetTensorType(fNInput), fShapeY);

if (model.Verbose())
std::cout << "Tile: " << fNInput << " " << ConvertShapeToString(fShapeInput) << " -> " << fNY << " with shape " << ConvertShapeToString(fShapeY)
<< " given repeats " << ConvertShapeToString(repeats_vector) << std::endl;
}

std::string Generate(std::string OpName){
Expand All @@ -89,17 +100,13 @@ public:
std::string output = "tensor_" + fNY;
out << "///-------- Tile operator\n";
out << "{\n"; // add scope to re-use same names
out << "std::vector<int> input_shape = " << ConvertShapeToString(fShapeInput) << ";\n";
std::vector<size_t> repeats = fShapeY;
for (size_t i = 0; i < repeats.size(); i++)
repeats[i] /= fShapeInput[i];
out << "const int input_shape[" << fShapeInput.size() << "] = " << ConvertShapeToString(fShapeInput) << ";\n";

out << "std::vector<int> repeats = " << ConvertShapeToString(repeats) << ";\n";
out << "int inputLength = " << ConvertShapeToLength(fShapeInput) << ";\n";
out << "int s = 1;\n";
// loop from inverse dim order
out << "for (int i = " << fShapeInput.size()-1 << "; i >=0; i--) {\n";
out << SP << "int r = repeats[i];\n";
out << SP << "int r = tensor_" << fNRepeats << "[i];\n";
// we cannot exclude case where repeats=1 since we need offset
//out << SP << "if (r == 1 && i < " << fShapeInput.size()-1 << ") continue;\n";
out << SP << "int i_offset = 0, o_offset = 0;\n";
Expand Down

0 comments on commit a4a9b0e

Please sign in to comment.