try ints

maneeshs · May 2, 2020 · 4dd711d · 4dd711d
1 parent 5c383e6
commit 4dd711d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc
@@ -257,7 +257,7 @@ Status GetGradients(OpKernelContext* context,
 
 tensorflow::Status GetNumSamples(
     tensorflow::OpKernelContext* context,
-    std::vector<std::vector<unsigned int>>* parsed_num_samples) {
+    std::vector<std::vector<int>>* parsed_num_samples) {
   const Tensor* input_num_samples;
   Status status = context->input("num_samples", &input_num_samples);
   if (!status.ok()) {
@@ -270,13 +270,14 @@ tensorflow::Status GetNumSamples(
                                input_num_samples->dims(), "."));
   }
 
-  const auto matrix_num_samples = input_num_samples->matrix<unsigned int>();
+  const auto matrix_num_samples = input_num_samples->matrix<int>();
   parsed_num_samples->reserve(matrix_num_samples.dimension(0));
   for (unsigned int i = 0; i < matrix_num_samples.dimension(0); i++) {
-    std::vector<unsigned int> sub_parsed_num_samples;
+    std::vector<int> sub_parsed_num_samples;
     sub_parsed_num_samples.reserve(matrix_num_samples.dimension(1));
-    for (int j = 0; j < matrix_num_samples.dimension(1); j++) {
-      sub_parsed_num_samples.push_back(matrix_num_samples(i, j));
+    for (unsigned int j = 0; j < matrix_num_samples.dimension(1); j++) {
+      const int num_samples = matrix_num_samples(i, j);
+      sub_parsed_num_samples.push_back(num_samples);
     }
     parsed_num_samples->push_back(sub_parsed_num_samples);
   }

diff --git a/tensorflow_quantum/core/ops/parse_context.h b/tensorflow_quantum/core/ops/parse_context.h
@@ -78,7 +78,7 @@ tensorflow::Status GetGradients(tensorflow::OpKernelContext* context,
 // Parses the number of samples from the 'num_samples' input tensor.
 tensorflow::Status GetNumSamples(
     tensorflow::OpKernelContext* context,
-    std::vector<std::vector<unsigned int>>* parsed_num_samples);
+    std::vector<std::vector<int>>* parsed_num_samples);
 
 }  // namespace tfq
 

diff --git a/tensorflow_quantum/core/ops/tfq_simulate_sampled_expectation_op.cc b/tensorflow_quantum/core/ops/tfq_simulate_sampled_expectation_op.cc
@@ -79,7 +79,7 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
                     programs.size(), " circuits and ", pauli_sums.size(),
                     " paulisums.")));
 
-    std::vector<std::vector<unsigned int>> num_samples;
+    std::vector<std::vector<int>> num_samples;
     OP_REQUIRES_OK(context, GetNumSamples(context, &num_samples));
 
     OP_REQUIRES(context, num_samples.size() == pauli_sums.size(),
@@ -172,7 +172,7 @@ REGISTER_OP("TfqSimulateSampledExpectation")
     .Input("symbol_names: string")
     .Input("symbol_values: float")
     .Input("pauli_sums: string")
-    .Input("num_samples: uint32")
+    .Input("num_samples: int")
     .Output("expectations: float")
     .SetShapeFn([](tensorflow::shape_inference::InferenceContext *c) {
       tensorflow::shape_inference::ShapeHandle programs_shape;