Refactored inputs for rand_uniform and refactored multinomial_dyn_test.

src/include/migraphx/op/multinomial.hpp

@@ -27,6 +27,8 @@
  *         each category, or bucket.  This does not require the standard multinomial
  *         distribution but instead takes a probability distribution as an input.
  *
+ *          In the large number limit, the fractional counts approach the multinomial distribution.
+ *
  *      Inputs:   args[0] - a vector of probabilities for each category.  Values are running totals
  as provided by op prefix_scan_sum.
  *                          Values are log normalized (i.e. start with any set of numbers > 0, then
@@ -87,7 +89,7 @@ struct multinomial
         // return a static shape
         if((not inputs.front().dynamic()) or (inputs.front().dyn_dims().front().is_fixed()))
         {
-            if((not inputs.back().dynamic()) or (inputs.back().dyn_dims().front().is_fixed()))
+            if((not inputs.back().dynamic()) or (inputs.back().dyn_dims().back().is_fixed()))
             {
                 size_t batch = {inputs.front().max_lens().front()};
                 size_t sample_size{inputs.back().max_lens().back()};
@@ -96,7 +98,7 @@ struct multinomial
         }
         return {dtype,
                 {inputs.front().to_dynamic().dyn_dims().front(),
-                 inputs.back().to_dynamic().dyn_dims().front()}};
+                 inputs.back().to_dynamic().dyn_dims().back()}};
     }
 
     argument compute(const dyn_output& dyn_out, std::vector<argument> args) const

src/include/migraphx/op/rand_uniform.hpp

@@ -53,14 +53,11 @@ struct rand_uniform
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        return pack(f(self.dtype, "dtype"), f(self.sample_size, "sample_size"), f(self.seed, "seed"));
-    }
-
-    value attributes() const
-    {
-        return {{"sample_size", sample_size}, {"seed", seed}};
+        return pack(
+            f(self.dtype, "dtype"), f(self.sample_size, "sample_size"), f(self.seed, "seed"));
     }
 
+    value attributes() const { return {{"sample_size", sample_size}, {"seed", seed}}; }
 
     std::string name() const { return "rand_uniform"; }
     shape normalize_compute_shape(std::vector<shape> inputs) const
@@ -81,21 +78,20 @@ struct rand_uniform
         }
     }
 
-
     argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
+        (void)args; // suppress compiler warning
         argument result{dyn_out.computed_shape};
 
         std::mt19937 gen(seed);
         std::uniform_real_distribution<> dis(0.0, 1.0);
-        size_t index(dyn_out.computed_shape.elements());
+        size_t elts(dyn_out.computed_shape.elements());
         // Use of our visitor and par_for replaces a call like
         //   std::vector<float> rand_samples(sample_size);
         //   std::generate(rand_samples.begin(), rand_samples.end(), [&]() { return dis(gen); });
 
         result.visit([&](auto output) {
-            par_for(sample_size, [&](auto i)
-            {
+            par_for(elts, [&](auto i) {
                 output[i] = dis(gen);
                 // output[i] = rand_samples[i];
             });

test/ref_ops_test.cpp

@@ -4909,7 +4909,6 @@ TEST_CASE(multinomial_test)
     std::vector<int> dist{15, 25, 15, 25, 20};
     std::vector<float> data(5);
     std::transform(dist.begin(), dist.end(), data.begin(), [&](auto d) { return std::log(d); });
-printf("data="); for(auto aa:data)printf(", %f", aa);printf("\n");
     auto input = mm->add_literal(migraphx::literal(s, data));
 
     auto maxes = mm->add_instruction(migraphx::make_op("reduce_max", {{"axes", {1}}}), input);
@@ -4927,8 +4926,8 @@ printf("data="); for(auto aa:data)printf(", %f", aa);printf("\n");
     std::vector<int32_t> result_vec(sample_size);
     result.visit([&](auto output) { result_vec.assign(output.begin(), output.end()); });
 
-    // res_dist is a count, or histogram, of the number of samples in each category.  This is the sampled
-    // distribution.
+    // res_dist is a count, or histogram, of the number of samples in each category.  This is the
+    // sampled distribution.
     std::vector<int> res_dist(5, 0);
     for(const auto& r : result_vec)
         res_dist[r]++;
@@ -4949,7 +4948,7 @@ printf("data="); for(auto aa:data)printf(", %f", aa);printf("\n");
     std::transform(res_dist.begin(), res_dist.end(), res_norm.begin(), [&](auto n) {
         return static_cast<double>(n) / res_dist_sum;
     });
-printf("cumulative distribution of result ="); for(auto aa:res_norm)printf(", %f", aa);printf("\n");
+
     EXPECT(migraphx::verify_range(norm, res_norm, 100000));
 }
 
@@ -4962,29 +4961,33 @@ TEST_CASE(multinomial_dyn_test)
     float seed         = 0.0f;
 
     //      Shape of the random data
-    migraphx::shape rs{migraphx::shape::float_type, {{1, 2}, {23, sample_size + 1}}};
+    migraphx::shape rs{migraphx::shape::float_type, {{1, 2}, {2, sample_size + 1}}};
     auto input = mm->add_parameter("Input_1", rs);
 
-    // the probability distribution, which also defines the number of categories
-    migraphx::shape s{migraphx::shape::float_type, {{1, 2}, {5, 6}}};
+    // Shape of the probability distribution, which also defines the number of categories
+    migraphx::shape s{migraphx::shape::float_type, {{1, 1}, {5, 6}}};
     std::vector<int> dist{15, 25, 15, 25, 20};
     std::vector<float> data(5);
 
-    // todo:  make this an instruction too
-    std::transform(dist.begin(), dist.end(), data.begin(), [&](auto d) { return std::log(d); });
+    std::transform(dist.begin(), dist.end(), data.begin(), [&](auto d) { return d; });
 
     auto input2 = mm->add_parameter("Input_2", s);
 
-    auto maxes = mm->add_instruction(migraphx::make_op("reduce_max", {{"axes", {1}}}), input2);
-    auto mb_maxes =
-        mm->add_instruction(migraphx::make_op("multibroadcast"), maxes, input2);
-    auto cdf = mm->add_instruction(migraphx::make_op("sub"), input2, mb_maxes);
+    // The next several instructions log-normalize the probability distribution,
+    // as required by the multinomial operation
+    auto logs = mm->add_instruction(migraphx::make_op("log"), input2);
+
+    auto maxes    = mm->add_instruction(migraphx::make_op("reduce_max", {{"axes", {1}}}), logs);
+    auto mb_maxes = mm->add_instruction(migraphx::make_op("multibroadcast"), maxes, input2);
+    auto cdf      = mm->add_instruction(migraphx::make_op("sub"), logs, mb_maxes);
+
     cdf = mm->add_instruction(migraphx::make_op("exp"), cdf);
     cdf = mm->add_instruction(
         migraphx::make_op("prefix_scan_sum", {{"axis", 1}, {"exclusive", false}}), cdf);
 
     auto randoms = mm->add_instruction(migraphx::make_op("rand_uniform", {{"seed", seed}}), input);
     mm->add_instruction(migraphx::make_op("multinomial"), cdf, randoms);
+
     p.compile(migraphx::make_target("ref"));
 
     // Create a dummy input in the shape we want for the random data
@@ -4996,12 +4999,16 @@ TEST_CASE(multinomial_dyn_test)
     params0["Input_2"] = migraphx::argument(input_fixed_shape2, data.data());
     auto result        = p.eval(params0).back();
 
-    std::vector<int32_t> result_vec(input_fixed_shape2.elements());
+    std::vector<float> result_vec(input_fixed_shape2.elements());
     result.visit([&](auto output) { result_vec.assign(output.begin(), output.end()); });
 
+    // Make a categorical histogram of output
     std::vector<int> res_dist(5, 0);
     for(const auto& r : result_vec)
         res_dist[r]++;
+
+    // Rescale or normalize both the input probability distribution and the output
+    // histogram, and compare.  Should be close but not identical.
     auto dist_sum     = std::accumulate(dist.begin(), dist.end(), 0);
     auto res_dist_sum = std::accumulate(res_dist.begin(), res_dist.end(), 0);
     std::vector<float> norm(5);
@@ -5012,8 +5019,6 @@ TEST_CASE(multinomial_dyn_test)
     std::transform(res_dist.begin(), res_dist.end(), res_norm.begin(), [&](auto n) {
         return static_cast<double>(n) / res_dist_sum;
     });
-printf("cumulative distribution of input ="); for(auto aa:norm)printf(", %f", aa);printf("\n");
-printf("cumulative distribution of result ="); for(auto aa:res_norm)printf(", %f", aa);printf("\n");
 
     EXPECT(migraphx::verify_range(norm, res_norm, 100000));
 }