Updates toy sources to use stateless random number generator.

PiperOrigin-RevId: 464610112
Change-Id: I686969628aea8e68973f35689270d3976c5d6b9d

Author: Johannes Ballé

models/toy_sources/ramp.py

@@ -76,7 +76,7 @@ def _event_shape(self):
   def _sample_n(self, n, seed=None):
     ind = self.index_points
     if self.phase is None:
-      phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+      phase = tf.random.stateless_uniform((n, 1), seed=seed, dtype=self.dtype)
     else:
       phase = tf.fill((n, 1), tf.constant(self.phase, dtype=self.dtype))
     return (ind + phase) % 1 - .5

models/toy_sources/sawbridge.py

@@ -96,15 +96,15 @@ def _event_shape(self):
 
   def _sample_n(self, n, seed=None):
     if self.drop is None:
-      uniform = tf.random.uniform(
+      uniform = tf.random.stateless_uniform(
           (self.order, n, 1), seed=seed, dtype=self.dtype)
     else:
       uniform = tf.fill(
           (self.order, n, 1), tf.constant(self.drop, dtype=self.dtype))
     ind = self.index_points
     if self.stationary:
       if self.phase is None:
-        phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+        phase = tf.random.stateless_uniform((n, 1), seed=seed, dtype=self.dtype)
       else:
         phase = tf.constant(self.phase, dtype=self.dtype)
       ind += phase

models/toy_sources/sinusoid.py

@@ -80,7 +80,7 @@ def _event_shape(self):
   def _sample_n(self, n, seed=None):
     ind = self.index_points
     if self.phase is None:
-      phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+      phase = tf.random.stateless_uniform((n, 1), seed=seed, dtype=self.dtype)
     else:
       phase = tf.fill((n, 1), tf.constant(self.phase, dtype=self.dtype))
     return tf.sin((2 * np.pi) * (ind + phase))

models/toy_sources/sphere.py

@@ -73,10 +73,11 @@ def _event_shape(self):
     return tf.TensorShape([self.order])
 
   def _sample_n(self, n, seed=None):
-    samples = tf.random.normal((n, self.order), seed=seed, dtype=self.dtype)
+    samples = tf.random.stateless_normal(
+        (n, self.order), seed=seed, dtype=self.dtype)
     radius = tf.math.sqrt(tf.reduce_sum(tf.square(samples), -1, keepdims=True))
     if self.width:
-      radius *= tf.random.uniform(
+      radius *= tf.random.stateless_uniform(
           (n, 1), minval=1. - self.width / 2., maxval=1. + self.width / 2.,
           seed=seed, dtype=self.dtype)
     return samples / radius