Added dynamic input handling to multinomial.hpp. Tests don't pass because it...

Added dynamic input handling to multinomial.hpp.  Tests don't pass because it still requires both a random uniform, and dynamic-friendly prefix_scap operations to accept dynamic inputs

src/include/migraphx/op/multinomial.hpp

 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -21,6 +21,34 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
+
+/**
+ * Multinomial or categorical distribution.  Performs a sampling and returns a count of
+ *         each category, or bucket.  This does not require the standard multinomial
+ *         distribution but instead takes a probability distribution as an input.
+ *
+ *      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
+ *                          val[i] = log(val[i]) / sum (log(val[0]) + log(val[1])+ ...) )
+ *                          This input has Rank 2.  Dimension 0 is batch #.  The size of dimension
+ *                          1 is the number of categories.
+ *                args[1] - a vector of random numbers.  Can be 1-, 2- or more dimensions but all
+ *                          except the last dim must be 1.
+ *
+ *                          Values as created by a std::mt19937 like this:
+ *
+                            size_t sample_size = 100000;
+                            float seed         = 0.0f;
+                            std::mt19937 gen(seed);
+                            std::uniform_real_distribution<> dis(0.0, 1.0);
+                            std::vector<float> rand_samples(sample_size);
+                            std::generate(rand_samples.begin(), rand_samples.end(), [&]() { return
+ dis(gen); });
+ *
+        Output:   A vector of category each input.
+ *
+*/
 #ifndef MIGRAPHX_GUARD_OPERATORS_MULTINOMIAL_HPP
 #define MIGRAPHX_GUARD_OPERATORS_MULTINOMIAL_HPP
 
@@ -48,11 +76,27 @@ struct multinomial
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this, true}.has(2).only_dims(2);
+
+        if(inputs.back().ndim() < 1)
+            MIGRAPHX_THROW("Multinomial: Second input shape (sample) has no dimensions");
         if(not contains({shape::int32_type, shape::int64_type}, dtype))
             MIGRAPHX_THROW(
                 "Multinomial: Invalid output type. Valid types are int32_type and int64_type.");
 
-        return inputs.front().normalize_standard();
+        // Output takes one dimension from each of the two input shapes.  If they are both fixed,
+        // 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()))
+            {
+                size_t batch = {inputs.front().max_lens().front()};
+                size_t sample_size{inputs.back().max_lens().back()};
+                return {dtype, {batch, sample_size}};
+            }
+        }
+        return {dtype,
+                {inputs.front().to_dynamic().dyn_dims().front(),
+                 inputs.back().to_dynamic().dyn_dims().front()}};
     }
 
     argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
@@ -68,8 +112,12 @@ struct multinomial
                     auto idx       = args[1].get_shape().multi(i);
                     auto cdf_begin = cdf.begin() + (idx[0] * class_size);
                     auto cdf_end   = cdf_begin + class_size;
+
+                    // std::upper_bound returns an iterator to the bucket the value belongs in,
+                    // when normalized by the probability distribution dist
                     auto sample_iter =
                         std::upper_bound(cdf_begin, cdf_end, dist[i] * *(std::prev(cdf_end)));
+                    // convert iterator to an integer index
                     output[i] = std::distance(cdf_begin, sample_iter);
                 });
             });

test/ref_ops_test.cpp

@@ -4905,9 +4905,10 @@ TEST_CASE(multinomial_test)
     auto rs_lit = mm->add_literal(migraphx::literal{rs, rand_samples});
 
     migraphx::shape s{migraphx::shape::float_type, {1, 5}};
-    std::vector<int> dist{15, 25, 15, 25, 20};
+    std::vector<int> dist{85, 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);
@@ -4921,13 +4922,23 @@ TEST_CASE(multinomial_test)
     mm->add_instruction(migraphx::make_op("multinomial"), cdf, rs_lit);
     p.compile(migraphx::make_target("ref"));
     auto result = p.eval({}).back();
+    // result_vec is a list of indices, or category labels, for each slot
     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.
     std::vector<int> res_dist(5, 0);
     for(const auto& r : result_vec)
         res_dist[r]++;
+
+    // To check the result, normalize the original probability distribution dist
+    // and the sampling result res_dist; they should be close
+
+    // Total the unnormalized probabilities
     auto dist_sum     = std::accumulate(dist.begin(), dist.end(), 0);
+
+    // Total the number of values returned
     auto res_dist_sum = std::accumulate(res_dist.begin(), res_dist.end(), 0);
     std::vector<float> norm(5);
     std::vector<float> res_norm(5);
@@ -4937,6 +4948,7 @@ TEST_CASE(multinomial_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 result ="); for(auto aa:res_norm)printf(", %f", aa);printf("\n");
     EXPECT(migraphx::verify_range(norm, res_norm, 100000));
 }