fix multinomial (#664)

Author: Samantha Andow

functorch/csrc/BatchRulesRandomness.cpp

@@ -207,6 +207,39 @@ std::tuple<Tensor,Tensor> native_dropout_batching_rule(const Tensor& tensor, dou
   return std::make_tuple(output, mask);
 }
 
+Tensor multinomial_batching_rule(const Tensor& self, const int64_t num_samples, const bool replacement, const c10::optional<Generator> generator) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  const auto cur_level = maybe_layer->layerId();
+
+  Tensor self_value;
+  optional<int64_t> self_bdim;
+  std::tie(self_value, self_bdim) = unwrapTensorAtLevel(self, cur_level);
+  self_value = moveBatchDimToFront(self_value, self_bdim);
+
+  RandomnessType randomness = maybe_layer->randomness();
+  check_randomness(randomness, self_bdim.has_value());
+
+  if (randomness == RandomnessType::Different && !self_bdim) {
+    auto shape = self_value.sizes();
+    VmapDimVector shapeVec(1, maybe_layer->batchSize());
+    shapeVec.reserve(shape.size() + 1);
+    shapeVec.insert(shapeVec.end(), shape.begin(), shape.end());
+    self_value = self_value.expand(shapeVec);
+  }
+  if (self_value.dim() == 3 && (self_bdim || randomness == RandomnessType::Different)) {
+    self_value = reshape_dim_into(1, 0, self_value);
+  }
+  auto out = multinomial(self_value, num_samples, replacement, generator);
+  if (randomness == RandomnessType::Same && !self_bdim) {
+    return out;
+  }
+  if(self_value.dim() == 3 && self_bdim) {
+    out = out.reshape(self.sizes());
+  }
+  return makeBatched(out, 0, cur_level);
+}
+
 template <typename A, A a, typename C>
 struct RandomBatchRuleHelper;
 
@@ -420,7 +453,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
 
   UNARY_POINTWISE_RANDOM(_standard_gamma);
   UNARY_POINTWISE_RANDOM(_sample_dirichlet);
-  UNARY_POINTWISE_RANDOM(multinomial);
+  m.impl("multinomial", multinomial_batching_rule);
   UNARY_POINTWISE_RANDOM(poisson);
   UNARY_POINTWISE_RANDOM(bernoulli);
 

test/test_vmap.py

@@ -3970,20 +3970,13 @@ def test_random_unary_out_of_place(self, device, use_generator, randomness, batc
             lambda t, _: torch.poisson(t, **kwargs),
         ]
 
-        # TODO(samdow): fix multinomial and readd
-        def flatten_op(t, ignored):
-            return torch.multinomial(t, 10, **kwargs)
-
         B0 = 4
         seed = 1234567
         in_dims = self._in_dims(batched_input)
 
         for op in ops:
             always_batched = torch.randn(B0, device=device)
             passed = self._get_image(batched_input, B0, device)
-            if op == flatten_op:
-                passed = passed.flatten(1, -1)
-
             if randomness == 'error':
                 self._assert_throws_in_error_mode(op, (passed, always_batched), in_dims=in_dims)
                 return
@@ -4009,6 +4002,72 @@ def flatten_op(t, ignored):
                 for i in range(B0):
                     self.assertEqual(vmap_result[i], expected)
 
+    @parametrize('use_generator', [True, False])
+    @parametrize('randomness', ['error', 'same', 'different'])
+    @parametrize('batched_call', [True, False])
+    @parametrize('batched_input', ["first", "last", "none"])
+    def test_multinomial(self, device, use_generator, randomness, batched_call, batched_input):
+        def flatten_input(input, batch_call, batch_location):
+            if batch_call and batch_location != "none":
+                final_size = 3  # [B0, B, N]
+            elif not batch_call and batch_location == "none":
+                final_size = 1  # [N]
+            else:
+                final_size = 2  # [B0, N] or [B, N]
+
+            start_idx = final_size - 1
+            end_idx = -1
+            if batch_location == "last":
+                start_idx -= 1
+                end_idx -= 1   # gets to correct final size because using negative indices
+
+            ret = input.flatten(start_idx, end_idx)
+            assert ret.dim() == final_size
+            return ret
+
+        def op(input, _):
+            return torch.multinomial(input, 10, **kwargs)
+
+        generator = torch.Generator(device=device)
+        orig_state = generator.get_state()
+        kwargs = {'generator': generator} if use_generator else {}
+
+        B0 = 4
+        seed = 1234567
+        in_dims = self._in_dims(batched_input)
+
+        always_batched = torch.randn(B0, device=device)
+        passed = self._get_image(batched_input, B0, device)
+        passed = flatten_input(passed, batched_call, batched_input)
+        if randomness == 'error':
+            self._assert_throws_in_error_mode(op, (passed, always_batched), in_dims=in_dims)
+            return
+        if randomness == 'same' and batched_input != "none":
+            self._assert_throws_in_same_mode_batched(op, (passed, always_batched), in_dims=in_dims)
+            return
+
+        generator = self._reset_random(generator, orig_state, use_generator, seed)
+        vmap_result = vmap(op, in_dims=in_dims, randomness=randomness)(passed, always_batched)
+
+        generator = self._reset_random(generator, orig_state, use_generator, seed)
+
+        if randomness == "different":
+            if batched_input == "none":
+                passed = passed.expand(B0, *passed.shape)
+            if batched_input == "last":
+                passed = passed.movedim(-1, 0)
+            orig_passed_size = passed.shape[:2] if batched_call else passed.shape[:1]
+            passed = passed.flatten(0, 1) if batched_call else passed
+            expected = op(passed, always_batched)
+            expected.reshape(*orig_passed_size, 10)
+            self._assert_all_slices_unique(vmap_result)
+            self.assertEqual(vmap_result, expected)
+        else:
+            expected = op(passed, always_batched)
+            self._assert_all_slices_equal(vmap_result)
+            for i in range(B0):
+                self.assertEqual(vmap_result[i], expected)
+
     def test_unsupported_random(self, device):
         x = torch.randn(3, device=device)
         y = x.abs()