Add pfor converters for:

  Fill
  StatelessMultinomial
  StatelessRandomBinomial
  StatelessRandomGammaV2
  StatelessRandomNormal
  StatelessRandomPoisson
  StatelessRandomUniform
  StatelessRandomUniformInt
  StatelessTruncatedNormal
  UnsortedSegmentMax
  UnsortedSegmentMin
  UnsortedSegmentProd

PiperOrigin-RevId: 294771643
Change-Id: Iba575475be98cc1bfb7e740f752275ab72b18ee4

tensorflow/python/ops/parallel_for/array_test.py

@@ -121,6 +121,13 @@ def loop_fn(i):
 
     self._test_loop_fn(loop_fn, 3)
 
+  def test_fill(self):
+
+    def loop_fn(i):
+      return array_ops.fill((2, 3), i)
+
+    self._test_loop_fn(loop_fn, 3)
+
   def test_broadcast_to(self):
     x = random_ops.random_uniform([3, 2, 1, 3])
 

tensorflow/python/ops/parallel_for/control_flow_ops_test.py

@@ -50,6 +50,7 @@
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import rnn
 from tensorflow.python.ops import rnn_cell
+from tensorflow.python.ops import stateless_random_ops
 from tensorflow.python.ops import tensor_array_grad  # pylint: disable=unused-import
 from tensorflow.python.ops import tensor_array_ops
 from tensorflow.python.ops import variables
@@ -636,6 +637,41 @@ def loop_fn(i):
     self._test_loop_fn(loop_fn, 5)
 
 
+class StatelessRandomTest(PForTestCase):
+
+  # This test currently only tests that the vectorized and non-vectorized
+  # outputs have same shapes. This is needed since under XLA compilation,
+  # stateless random numbers can generate different random numbers.
+  # TODO(agarwal): switch to checking for actual values matching once
+  # b/149402339 is resolved.
+  def run_and_assert_equal(self, targets1, targets2):
+    outputs = self._run_targets(targets1, targets2)
+    n = len(outputs) // 2
+    for i in range(n):
+      self.assertAllEqual(outputs[i].shape, outputs[i + n].shape)
+
+  # TODO(agarwal): add tests for other random functions
+  def test_multinomial(self):
+    seeds = [[1, 2], [3, 4]]
+    logits = random_ops.random_uniform([2, 3, 4])
+
+    def loop_fn(i):
+      logits_0 = array_ops.gather(logits, 0)
+      logits_i = array_ops.gather(logits, i)
+      seeds_0 = array_ops.gather(seeds, 0)
+      seeds_i = array_ops.gather(seeds, i)
+      return (stateless_random_ops.stateless_categorical(
+          logits=logits_i, num_samples=3, seed=seeds_i),
+              stateless_random_ops.stateless_categorical(
+                  logits=logits_i, num_samples=3, seed=seeds_0),
+              stateless_random_ops.stateless_categorical(
+                  logits=logits_0, num_samples=3, seed=seeds_i),
+              stateless_random_ops.stateless_categorical(
+                  logits=logits_0, num_samples=3, seed=seeds_0))
+
+    self._test_loop_fn(loop_fn, 2)
+
+
 class LoggingTest(PForTestCase):
 
   @test_util.run_v1_only("b/122612051")

tensorflow/python/ops/parallel_for/math_test.py

@@ -448,7 +448,10 @@ def loop_fn(i):
             out_dtypes = out_dtypes + [dtypes.int32]
           self._test_loop_fn(loop_fn, 2)
 
-  def test_unsorted_segment_sum(self):
+  @parameterized.parameters(
+      (math_ops.unsorted_segment_sum,), (math_ops.unsorted_segment_min,),
+      (math_ops.unsorted_segment_max,), (math_ops.unsorted_segment_prod,))
+  def test_unsorted_segment_reduction(self, reduction_op):
     t = random_ops.random_uniform([3, 3, 2])
     for segment_ids_dtype in (dtypes.int32, dtypes.int64):
       for num_segments_dtype in (dtypes.int32, dtypes.int64):
@@ -462,9 +465,9 @@ def loop_fn(i):
           data_0 = array_ops.gather(t, 0)
           seg_ids = array_ops.gather(segment_ids, i)
           seg_ids_0 = array_ops.gather(segment_ids, 0)
-          return (math_ops.unsorted_segment_sum(data, seg_ids, num_segments),
-                  math_ops.unsorted_segment_sum(data_0, seg_ids, num_segments),
-                  math_ops.unsorted_segment_sum(data, seg_ids_0, num_segments))
+          return (reduction_op(data, seg_ids, num_segments),
+                  reduction_op(data_0, seg_ids, num_segments),
+                  reduction_op(data, seg_ids_0, num_segments))
 
         # pylint: enable=cell-var-from-loop
 

tensorflow/python/ops/parallel_for/pfor.py

@@ -1845,6 +1845,21 @@ def _convert_reshape(pfor_input):
   return wrap(array_ops.reshape(t, new_shape), True)
 
 
+@RegisterPFor("Fill")
+def _convert_fill(pfor_input):
+  dims = pfor_input.unstacked_input(0)
+  value = pfor_input.stacked_input(1)
+  # Expand the rank of `value`
+  new_shape = array_ops.concat(
+      [[-1], array_ops.ones([array_ops.size(dims)], dtype=dtypes.int32)],
+      axis=0)
+  value = array_ops.reshape(value, new_shape)
+  # Compute the new output shape
+  new_dims = array_ops.concat([pfor_input.pfor.loop_len_vector, dims], axis=0)
+  # Broadcast
+  return wrap(array_ops.broadcast_to(value, new_dims), True)
+
+
 @RegisterPFor("BroadcastTo")
 def _convert_broadcast_to(pfor_input):
   t = pfor_input.stacked_input(0)
@@ -2399,8 +2414,11 @@ def _convert_biasadd(pfor_input):
     return wrap(nn_ops.bias_add(t, bias, data_format=data_format), True)
 
 
-@RegisterPFor("UnsortedSegmentSum")
-def _convert_unsortedsegmentsum(pfor_input):
+@RegisterPForWithArgs("UnsortedSegmentSum", math_ops.unsorted_segment_sum)
+@RegisterPForWithArgs("UnsortedSegmentMax", math_ops.unsorted_segment_max)
+@RegisterPForWithArgs("UnsortedSegmentMin", math_ops.unsorted_segment_min)
+@RegisterPForWithArgs("UnsortedSegmentProd", math_ops.unsorted_segment_prod)
+def _convert_unsortedsegmentsum(pfor_input, _, op_func):
   pfor_input.stack_inputs([0, 1])
   data = pfor_input.stacked_input(0)
   segment_ids = pfor_input.stacked_input(1)
@@ -2419,7 +2437,7 @@ def _convert_unsortedsegmentsum(pfor_input):
   segment_ids += segment_offset
   num_segments = math_ops.cast(num_segments, dtypes.int64) * math_ops.cast(
       n, dtypes.int64)
-  output = math_ops.unsorted_segment_sum(data, segment_ids, num_segments)
+  output = op_func(data, segment_ids, num_segments)
   new_output_shape = array_ops.concat(
       [[n, -1], array_ops.shape(output)[1:]], axis=0)
   output = array_ops.reshape(output, new_output_shape)
@@ -2859,6 +2877,24 @@ def _convert_multinomial(pfor_input):
   return wrap(stacked_samples, True)
 
 
+@RegisterPFor("StatelessMultinomial")
+@RegisterPFor("StatelessRandomBinomial")
+@RegisterPFor("StatelessRandomGammaV2")
+@RegisterPFor("StatelessRandomNormal")
+@RegisterPFor("StatelessRandomPoisson")
+@RegisterPFor("StatelessRandomUniform")
+@RegisterPFor("StatelessRandomUniformInt")
+@RegisterPFor("StatelessTruncatedNormal")
+def _convert_stateless_multinomial(pfor_input):
+  # Unlike stateful random ops, for stateless ones we want better
+  # reproducibility based on seed. Hence we don't want to use a similar strategy
+  # as used for stateful ones where we generate a possibly different set of
+  # random numbers under vectorization.
+  # Unfortunately, the kernels currently are not necessarily setup to do this
+  # efficiently and hence we fallback to a sequential loop for vectorization.
+  return _fallback_converter(pfor_input)
+
+
 # linalg_ops