random_poisson(rate=+inf) should return +inf, not nan.

Leaving the status quo for integer output dtypes, which is determined by how non-finite floats cast to said integers.

PiperOrigin-RevId: 335966603
Change-Id: I5bfb53a17de32e02a10595d0393f037f982b5194

tensorflow/core/kernels/random_poisson_op.cc

@@ -150,6 +150,16 @@ struct PoissonFunctor<CPUDevice, T, U> {
           }
           continue;
         }
+        if (Eigen::numext::isinf(rate) && rate > CT(0)) {
+          // Fill the rest of the samples for the current rate value.
+          for (int64 sample_idx = output_idx % num_samples;
+               sample_idx < num_samples && output_idx < limit_output;
+               sample_idx++, output_idx++) {
+            U k = Eigen::NumTraits<U>::infinity();
+            samples_rate_output[sample_idx * num_rate] = k;
+          }
+          continue;
+        }
         // Transformed rejection due to Hormann.
         //
         // Given a CDF F(x), and G(x), a dominating distribution chosen such

tensorflow/python/kernel_tests/random/random_poisson_test.py

@@ -171,6 +171,11 @@ class RandomPoissonTest(test.TestCase):
               constant_op.constant([1], dtype=lam_dt), [10],
               dtype=out_dt).eval()
 
+  @test_util.run_deprecated_v1
+  def testInfRate(self):
+    sample = random_ops.random_poisson(shape=[2], lam=np.inf)
+    self.assertAllEqual([np.inf, np.inf], self.evaluate(sample))
+
 
 if __name__ == "__main__":
   test.main()