Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/ops/stateless_random_ops.py: 36%
234 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Stateless random ops which take seed as a tensor input."""
16import enum
17import numpy as np
19from tensorflow.python.framework import constant_op
20from tensorflow.python.framework import dtypes
21from tensorflow.python.framework import ops
22from tensorflow.python.ops import array_ops
23from tensorflow.python.ops import array_ops_stack
24from tensorflow.python.ops import bitwise_ops
25from tensorflow.python.ops import gen_random_index_shuffle_ops
26from tensorflow.python.ops import gen_stateless_random_ops
27from tensorflow.python.ops import gen_stateless_random_ops_v2
28from tensorflow.python.ops import math_ops
29from tensorflow.python.ops import shape_util
30from tensorflow.python.util import deprecation
31from tensorflow.python.util import dispatch
32from tensorflow.python.util.tf_export import tf_export
35ops.NotDifferentiable("StatelessMultinomial")
36ops.NotDifferentiable("StatelessRandomBinomial")
37ops.NotDifferentiable("StatelessRandomNormal")
38ops.NotDifferentiable("StatelessRandomPoisson")
39ops.NotDifferentiable("StatelessRandomUniform")
40ops.NotDifferentiable("StatelessRandomUniformInt")
41ops.NotDifferentiable("StatelessRandomUniformFullInt")
42ops.NotDifferentiable("StatelessTruncatedNormal")
43ops.NotDifferentiable("StatelessRandomNormalV2")
44ops.NotDifferentiable("StatelessRandomUniformV2")
45ops.NotDifferentiable("StatelessRandomUniformIntV2")
46ops.NotDifferentiable("StatelessRandomUniformFullIntV2")
47ops.NotDifferentiable("StatelessTruncatedNormalV2")
48ops.NotDifferentiable("StatelessRandomShuffle")
49ops.NotDifferentiable("RandomIndexShuffle")
52@tf_export("random.Algorithm", "random.experimental.Algorithm")
53class Algorithm(enum.Enum):
54 """A random-number-generation (RNG) algorithm.
56 Many random-number generators (e.g. the `alg` argument of
57 `tf.random.Generator` and `tf.random.stateless_uniform`) in TF allow
58 you to choose the algorithm used to generate the (pseudo-)random
59 numbers. You can set the algorithm to be one of the options below.
61 * `PHILOX`: The Philox algorithm introduced in the paper ["Parallel
62 Random Numbers: As Easy as 1, 2,
63 3"](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf).
64 * `THREEFRY`: The ThreeFry algorithm introduced in the paper
65 ["Parallel Random Numbers: As Easy as 1, 2,
66 3"](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf).
67 * `AUTO_SELECT`: Allow TF to automatically select the algorithm
68 depending on the accelerator device. Note that with this option,
69 running the same TF program on different devices may result in
70 different random numbers. Also note that TF may select an
71 algorithm that is different from `PHILOX` and `THREEFRY`.
72 """
73 # The numbers here must match framework/rng_alg.h
74 PHILOX = 1
75 THREEFRY = 2
76 AUTO_SELECT = 3
79def unsupported_alg_error_msg(alg):
80 """Produces the unsupported-algorithm error message."""
81 if isinstance(alg, int):
82 philox = Algorithm.PHILOX.value
83 threefry = Algorithm.THREEFRY.value
84 auto_select = Algorithm.AUTO_SELECT.value
85 elif isinstance(alg, str):
86 philox = "philox"
87 threefry = "threefry"
88 auto_select = "auto_select"
89 else:
90 philox = Algorithm.PHILOX
91 threefry = Algorithm.THREEFRY
92 auto_select = Algorithm.AUTO_SELECT
93 return (f"Argument `alg` got unsupported value {alg}. Supported values are "
94 f"{philox} for the Philox algorithm, "
95 f"{threefry} for the ThreeFry algorithm, and "
96 f"{auto_select} for auto-selection.")
99def convert_alg_to_int(alg):
100 """Converts algorithm to an integer.
102 Args:
103 alg: can be one of these types: integer, Algorithm, Tensor, string. Allowed
104 strings are "philox" and "threefry".
106 Returns:
107 An integer, unless the input is a Tensor in which case a Tensor is returned.
108 """
109 if isinstance(alg, int):
110 return alg
111 if isinstance(alg, Algorithm):
112 return alg.value
113 if isinstance(alg, ops.Tensor):
114 return alg
115 if isinstance(alg, str):
116 # canonicalized alg
117 canon_alg = alg.strip().lower().replace("-", "").replace("_", "")
118 if canon_alg == "philox":
119 return Algorithm.PHILOX.value
120 elif canon_alg == "threefry":
121 return Algorithm.THREEFRY.value
122 elif canon_alg == "autoselect":
123 return Algorithm.AUTO_SELECT.value
124 else:
125 raise ValueError(unsupported_alg_error_msg(alg))
126 else:
127 raise TypeError(
128 f"Can't convert argument `alg` (of value {alg} and type {type(alg)}) "
129 f"to int.")
132def _get_key_counter(seed, alg):
133 """Calculates the key and counter to pass to raw RNG ops.
135 This function calculates the key and counter that will be passed to
136 the raw RNG ops like `StatelessRandomUniformV2`. Depending on the
137 input `alg`, the key and counter may be scrambled or copied from
138 `seed`. If `alg` is `"auto_select"`, the key and counter will be
139 determined at runtime based on device type.
141 Args:
142 seed: An integer tensor of shape [2]. The seed to calculate the
143 key and counter from.
144 alg: The RNG algorithm. See `tf.random.stateless_uniform` for an
145 explanation.
147 Returns:
148 A pair (key, counter) suitable for V2 stateless RNG ops like
149 `StatelessRandomUniformV2`.
150 """
151 if alg == Algorithm.AUTO_SELECT.value:
152 key, counter = gen_stateless_random_ops_v2.stateless_random_get_key_counter(
153 seed)
154 elif alg == Algorithm.PHILOX.value:
155 key, counter = _philox_scramble_seed(seed)
156 elif alg == Algorithm.THREEFRY.value:
157 key = array_ops.reshape(
158 uint32s_to_uint64(math_ops.cast(seed, dtypes.uint32)), [1])
159 counter = array_ops.zeros([1], dtypes.uint64)
160 else:
161 raise ValueError(unsupported_alg_error_msg(alg))
162 return key, counter
165def _get_key_counter_alg(seed, alg):
166 if alg is None:
167 alg = Algorithm.AUTO_SELECT.value
168 alg = convert_alg_to_int(alg)
169 key, counter = _get_key_counter(seed, alg)
170 return key, counter, alg
173def _philox_scramble_seed(seed):
174 # the same scrambling procedure as core/kernels/stateless_random_ops.cc
175 key = constant_op.constant([0x02461e293ec8f720], dtypes.uint64)
176 counter = math_ops.cast(seed, dtypes.uint64)
177 mix = gen_stateless_random_ops_v2.stateless_random_uniform_full_int_v2(
178 [4], key=key, counter=counter, dtype=dtypes.uint32,
179 alg=Algorithm.PHILOX.value)
180 key = array_ops.reshape(uint32s_to_uint64(mix[:2]), [1])
181 counter = array_ops_stack.stack([0, uint32s_to_uint64(mix[2:])], axis=0)
182 return key, counter
185def uint32s_to_uint64(x):
186 return bitwise_ops.bitwise_or(
187 math_ops.cast(x[0], dtypes.uint64),
188 bitwise_ops.left_shift(math_ops.cast(x[1], dtypes.uint64),
189 constant_op.constant(32, dtypes.uint64)))
192@tf_export("random.split", "random.experimental.stateless_split")
193@dispatch.add_dispatch_support
194def split(seed, num=2, alg="auto_select"):
195 """Splits an RNG seed into `num` new seeds by adding a leading axis.
197 Example:
199 >>> seed = [1, 2]
200 >>> new_seeds = tf.random.split(seed, num=3)
201 >>> print(new_seeds)
202 tf.Tensor(
203 [[1105988140 1738052849]
204 [-335576002 370444179]
205 [ 10670227 -246211131]], shape=(3, 2), dtype=int32)
206 >>> tf.random.stateless_normal(shape=[3], seed=new_seeds[0, :])
207 <tf.Tensor: shape=(3,), dtype=float32, numpy=array([-0.59835213, -0.9578608 ,
208 0.9002807 ], dtype=float32)>
210 Args:
211 seed: an RNG seed (a tensor with shape [2] and dtype `int32` or
212 `int64`). (When using XLA, only `int32` is allowed.)
213 num: optional, a positive integer or scalar tensor indicating the number of
214 seeds to produce (default 2).
215 alg: The RNG algorithm used to generate the random numbers. See
216 `tf.random.stateless_uniform` for a detailed explanation.
218 Returns:
219 A tensor with shape [num, 2] representing `num` new seeds. It will have the
220 same dtype as `seed` (if `seed` doesn't have an explict dtype, the dtype
221 will be determined by `tf.convert_to_tensor`).
222 """
223 seed = ops.convert_to_tensor(seed)
224 return stateless_random_uniform(shape=[num, 2], seed=seed, dtype=seed.dtype,
225 minval=None, maxval=None, alg=alg)
228@tf_export("random.fold_in", "random.experimental.stateless_fold_in")
229@dispatch.add_dispatch_support
230def fold_in(seed, data, alg="auto_select"):
231 """Folds in data to an RNG seed to form a new RNG seed.
233 For example, in a distributed-training setting, suppose we have a master seed
234 and a replica ID. We want to fold the replica ID into the master seed to
235 form a "replica seed" to be used by that replica later on, so that different
236 replicas will generate different random numbers but the reproducibility of the
237 whole system can still be controlled by the master seed:
239 >>> master_seed = [1, 2]
240 >>> replica_id = 3
241 >>> replica_seed = tf.random.experimental.stateless_fold_in(
242 ... master_seed, replica_id)
243 >>> print(replica_seed)
244 tf.Tensor([1105988140 3], shape=(2,), dtype=int32)
245 >>> tf.random.stateless_normal(shape=[3], seed=replica_seed)
246 <tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.03197195, 0.8979765 ,
247 0.13253039], dtype=float32)>
249 Args:
250 seed: an RNG seed (a tensor with shape [2] and dtype `int32` or
251 `int64`). (When using XLA, only `int32` is allowed.)
252 data: an `int32` or `int64` scalar representing data to be folded in to the
253 seed.
254 alg: The RNG algorithm used to generate the random numbers. See
255 `tf.random.stateless_uniform` for a detailed explanation.
257 Returns:
258 A new RNG seed that is a deterministic function of the inputs and is
259 statistically safe for producing a stream of new pseudo-random values. It
260 will have the same dtype as `data` (if `data` doesn't have an explict dtype,
261 the dtype will be determined by `tf.convert_to_tensor`).
262 """
263 data = ops.convert_to_tensor(data)
264 seed1 = stateless_random_uniform(shape=[], seed=seed, dtype=data.dtype,
265 minval=None, maxval=None, alg=alg)
266 return array_ops_stack.stack([seed1, data])
269@tf_export("random.experimental.index_shuffle")
270@dispatch.add_dispatch_support
271def index_shuffle(index, seed, max_index):
272 """Outputs the position of `index` in a permutation of `[0, ..., max_index]`.
274 For each possible `seed` and `max_index` there is one pseudorandom
275 permutation of the sequence `S=[0, ..., max_index]`. Instead of
276 materializing the full array we can compute the new position of any
277 integer `i` (`0 <= i <= max_index`) in `S`. This can be useful for
278 very large `max_index`s by avoiding allocating large chunks of
279 memory.
281 In the simplest case, `index` and `max_index` are scalars, and
282 `seed` is a length-2 vector (as typical for stateless RNGs). But
283 you can add a leading batch dimension to all of them. If some of
284 them don't have the batch dimension while others do, `index_shuffle`
285 will add a batch dimension to the former by broadcasting.
287 The input `index` and output can be used as indices to shuffle a
288 vector. For example:
290 >>> vector = tf.constant(['e0', 'e1', 'e2', 'e3'])
291 >>> indices = tf.random.experimental.index_shuffle(
292 ... index=tf.range(4), seed=[5, 9], max_index=3)
293 >>> print(indices)
294 tf.Tensor([2 0 1 3], shape=(4,), dtype=int32)
295 >>> shuffled_vector = tf.gather(vector, indices)
296 >>> print(shuffled_vector)
297 tf.Tensor([b'e2' b'e0' b'e1' b'e3'], shape=(4,), dtype=string)
299 More usefully, it can be used in a streaming (aka online) scenario such as
300 `tf.data`, where each element of `vector` is processed individually and the
301 whole `vector` is never materialized in memory.
303 >>> dataset = tf.data.Dataset.range(10)
304 >>> dataset = dataset.map(
305 ... lambda idx: tf.random.experimental.index_shuffle(idx, [5, 8], 9))
306 >>> print(list(dataset.as_numpy_iterator()))
307 [3, 8, 0, 1, 2, 7, 6, 9, 4, 5]
309 This operation is stateless (like the `tf.random.stateless_*`
310 functions), meaning the output is fully determined by the `seed`
311 (other inputs being equal). Each `seed` choice corresponds to one
312 permutation, so when calling this function multiple times for the
313 same shuffling, please make sure to use the same `seed`. For
314 example:
316 >>> seed = [5, 9]
317 >>> idx0 = tf.random.experimental.index_shuffle(0, seed, 3)
318 >>> idx1 = tf.random.experimental.index_shuffle(1, seed, 3)
319 >>> idx2 = tf.random.experimental.index_shuffle(2, seed, 3)
320 >>> idx3 = tf.random.experimental.index_shuffle(3, seed, 3)
321 >>> shuffled_vector = tf.gather(vector, [idx0, idx1, idx2, idx3])
322 >>> print(shuffled_vector)
323 tf.Tensor([b'e2' b'e0' b'e1' b'e3'], shape=(4,), dtype=string)
325 Args:
326 index: An integer scalar tensor or vector with values in `[0,
327 max_index]`. It can be seen as either a value `v` in the
328 sequence `S=[0, ..., max_index]` to be permutated, or as an
329 index of an element `e` in a shuffled vector.
330 seed: A tensor of shape [2] or [n, 2] with dtype `int32`,
331 `uint32`, `int64` or `uint64`. The RNG seed. If the rank is
332 unknown during graph-building time it must be 1 at runtime.
333 max_index: A non-negative tensor with the same shape and dtype as
334 `index`. The upper bound (inclusive).
336 Returns:
337 If all inputs were scalar (shape [2] for `seed`), the output will
338 be a scalar with the same dtype as `index`. The output can be seen
339 as the new position of `v` in `S`, or as the index of `e` in the
340 vector before shuffling. If one or multiple inputs were vectors
341 (shape [n, 2] for `seed`), then the output will be a vector of the
342 same size which each element shuffled independently. Scalar values
343 are broadcasted in this case.
344 """
345 # We expect users to pass a seed with shape [2] to be consistent with other
346 # stateless_* ops, but the raw op expects shape [3].
347 seed = ops.convert_to_tensor(seed)
348 # Pad the first dimension with an arbitrary number since our raw op expects
349 # shape [3].
350 if seed.shape.rank is None:
351 paddings = [[1, 0]]
352 else:
353 paddings = [[1, 0]] + (seed.shape.rank - 1) * [[0, 0]]
354 seed = array_ops.pad(seed, paddings, constant_values=498247692)
355 return gen_random_index_shuffle_ops.random_index_shuffle(
356 index, seed=seed, max_index=max_index, rounds=4)
359@tf_export("random.experimental.stateless_shuffle")
360@dispatch.add_dispatch_support
361def stateless_shuffle(value, seed, alg="auto_select", name=None):
362 """Randomly and deterministically shuffles a tensor along its first dimension.
364 The tensor is shuffled along dimension 0, such that each `value[j]` is mapped
365 to one and only one `output[i]`. For example, a mapping that might occur for a
366 3x2 tensor is:
368 ```python
369 [[1, 2], [[5, 6],
370 [3, 4], ==> [1, 2],
371 [5, 6]] [3, 4]]
372 ```
374 >>> v = tf.constant([[1, 2], [3, 4], [5, 6]])
375 >>> shuffled = tf.random.experimental.stateless_shuffle(v, seed=[8, 9])
376 >>> print(shuffled)
377 tf.Tensor(
378 [[5 6]
379 [1 2]
380 [3 4]], shape=(3, 2), dtype=int32)
382 This is a stateless version of `tf.random.shuffle`: if run twice with the
383 same `value` and `seed`, it will produce the same result. The
384 output is consistent across multiple runs on the same hardware (and between
385 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
386 hardware.
388 Args:
389 value: A Tensor to be shuffled.
390 seed: A shape [2] Tensor. The seed to the random number generator. Must have
391 dtype `int32` or `int64`.
392 alg: The RNG algorithm used to generate the random numbers. See
393 `tf.random.stateless_uniform` for a detailed explanation.
394 name: A name for the operation.
396 Returns:
397 A tensor of same shape and type as `value`, shuffled along its first
398 dimension.
399 """
400 with ops.name_scope(name, "stateless_shuffle", [value, seed]) as name:
401 key, counter, alg = _get_key_counter_alg(seed, alg)
402 return gen_stateless_random_ops_v2.stateless_shuffle(
403 value, key=key, counter=counter, alg=alg)
406@tf_export("random.stateless_uniform")
407@dispatch.add_dispatch_support
408def stateless_random_uniform(shape,
409 seed,
410 minval=0,
411 maxval=None,
412 dtype=dtypes.float32,
413 name=None,
414 alg="auto_select"):
415 """Outputs deterministic pseudorandom values from a uniform distribution.
417 This is a stateless version of `tf.random.uniform`: if run twice with the
418 same seeds and shapes, it will produce the same pseudorandom numbers. The
419 output is consistent across multiple runs on the same hardware (and between
420 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
421 hardware.
423 The generated values follow a uniform distribution in the range
424 `[minval, maxval)`. The lower bound `minval` is included in the range, while
425 the upper bound `maxval` is excluded.
427 For floats, the default range is `[0, 1)`. For ints, at least `maxval` must
428 be specified explicitly.
430 In the integer case, the random integers are slightly biased unless
431 `maxval - minval` is an exact power of two. The bias is small for values of
432 `maxval - minval` significantly smaller than the range of the output (either
433 `2**32` or `2**64`).
435 For full-range (i.e. inclusive of both max and min) random integers, pass
436 `minval=None` and `maxval=None` with an integer `dtype`. For an integer dtype
437 either both `minval` and `maxval` must be `None` or neither may be `None`. For
438 example:
439 ```python
440 ints = tf.random.stateless_uniform(
441 [10], seed=(2, 3), minval=None, maxval=None, dtype=tf.int32)
442 ```
444 Args:
445 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
446 seed: A shape [2] Tensor, the seed to the random number generator. Must have
447 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
448 minval: A Tensor or Python value of type `dtype`, broadcastable with
449 `shape` (for integer types, broadcasting is not supported, so it needs to
450 be a scalar). The lower bound on the range of random values to
451 generate. Pass `None` for full-range integers. Defaults to 0.
452 maxval: A Tensor or Python value of type `dtype`, broadcastable with
453 `shape` (for integer types, broadcasting is not supported, so it needs to
454 be a scalar). The upper bound on the range of random values to generate.
455 Defaults to 1 if `dtype` is floating point. Pass `None` for full-range
456 integers.
457 dtype: The type of the output: `float16`, `bfloat16`, `float32`, `float64`,
458 `int32`, or `int64`. For unbounded uniform ints (`minval`, `maxval` both
459 `None`), `uint32` and `uint64` may be used. Defaults to `float32`.
460 name: A name for the operation (optional).
461 alg: The RNG algorithm used to generate the random numbers. Valid
462 choices are `"philox"` for [the Philox
463 algorithm](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf),
464 `"threefry"` for [the ThreeFry
465 algorithm](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf),
466 and `"auto_select"` (default) for the system to automatically
467 select an algorithm based the device type. Values of
468 `tf.random.Algorithm` can also be used. Note that with
469 `"auto_select"`, the outputs of this function may change when
470 it is running on a different device.
472 Returns:
473 A tensor of the specified shape filled with random uniform values.
475 Raises:
476 ValueError: If `dtype` is integral and only one of `minval` or `maxval` is
477 specified.
478 """
479 dtype = dtypes.as_dtype(dtype)
480 accepted_dtypes = (dtypes.float16, dtypes.bfloat16, dtypes.float32,
481 dtypes.float64, dtypes.int32, dtypes.int64, dtypes.uint32,
482 dtypes.uint64)
483 if dtype not in accepted_dtypes:
484 raise ValueError(
485 f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are "
486 f"{accepted_dtypes}.")
487 if dtype.is_integer:
488 if (minval is None) != (maxval is None):
489 raise ValueError(
490 f"For integer `dtype` argument {dtype}, argument `minval` and "
491 f"`maxval` must be both None or not None. Got `minval`={minval} and "
492 f"`maxval`={maxval}.")
493 if minval is not None and dtype in (dtypes.uint32, dtypes.uint64):
494 raise ValueError(
495 f"Argument `dtype` got invalid value {dtype} when argument `minval` "
496 f"is not None. Please don't use unsigned integers in this case.")
497 elif maxval is None:
498 maxval = 1
499 with ops.name_scope(name, "stateless_random_uniform",
500 [shape, seed, minval, maxval]) as name:
501 shape = shape_util.shape_tensor(shape)
502 if dtype.is_integer and minval is None:
503 key, counter, alg = _get_key_counter_alg(seed, alg)
504 result = (
505 gen_stateless_random_ops_v2.stateless_random_uniform_full_int_v2(
506 shape, key=key, counter=counter, dtype=dtype, alg=alg, name=name))
507 else:
508 minval = ops.convert_to_tensor(minval, dtype=dtype, name="min")
509 maxval = ops.convert_to_tensor(maxval, dtype=dtype, name="max")
510 if dtype.is_integer:
511 key, counter, alg = _get_key_counter_alg(seed, alg)
512 result = gen_stateless_random_ops_v2.stateless_random_uniform_int_v2(
513 shape,
514 key=key,
515 counter=counter,
516 minval=minval,
517 maxval=maxval,
518 alg=alg,
519 name=name)
520 else:
521 key, counter, alg = _get_key_counter_alg(seed, alg)
522 rnd = gen_stateless_random_ops_v2.stateless_random_uniform_v2(
523 shape, key=key, counter=counter, dtype=dtype, alg=alg)
524 result = math_ops.add(rnd * (maxval - minval), minval, name=name)
525 shape_util.maybe_set_static_shape(result, shape)
526 return result
529@tf_export("random.stateless_binomial")
530@dispatch.add_dispatch_support
531def stateless_random_binomial(shape,
532 seed,
533 counts,
534 probs,
535 output_dtype=dtypes.int32,
536 name=None):
537 """Outputs deterministic pseudorandom values from a binomial distribution.
539 The generated values follow a binomial distribution with specified count and
540 probability of success parameters.
542 This is a stateless version of `tf.random.Generator.binomial`: if run twice
543 with the same seeds and shapes, it will produce the same pseudorandom numbers.
544 The output is consistent across multiple runs on the same hardware (and
545 between CPU and GPU), but may change between versions of TensorFlow or on
546 non-CPU/GPU hardware.
548 Example:
550 ```python
551 counts = [10., 20.]
552 # Probability of success.
553 probs = [0.8]
555 binomial_samples = tf.random.stateless_binomial(
556 shape=[2], seed=[123, 456], counts=counts, probs=probs)
558 counts = ... # Shape [3, 1, 2]
559 probs = ... # Shape [1, 4, 2]
560 shape = [3, 4, 3, 4, 2]
561 # Sample shape will be [3, 4, 3, 4, 2]
562 binomial_samples = tf.random.stateless_binomial(
563 shape=shape, seed=[123, 456], counts=counts, probs=probs)
564 ```
566 Args:
567 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
568 seed: A shape [2] Tensor, the seed to the random number generator. Must have
569 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
570 counts: Tensor. The counts of the binomial distribution. Must be
571 broadcastable with `probs`, and broadcastable with the rightmost
572 dimensions of `shape`.
573 probs: Tensor. The probability of success for the binomial distribution.
574 Must be broadcastable with `counts` and broadcastable with the rightmost
575 dimensions of `shape`.
576 output_dtype: The type of the output. Default: tf.int32
577 name: A name for the operation (optional).
579 Returns:
580 samples: A Tensor of the specified shape filled with random binomial
581 values. For each i, each samples[..., i] is an independent draw from
582 the binomial distribution on counts[i] trials with probability of
583 success probs[i].
585 """
586 with ops.name_scope(name, "stateless_random_binomial",
587 [shape, seed, counts, probs]) as name:
588 shape = shape_util.shape_tensor(shape)
589 probs = ops.convert_to_tensor(
590 probs, dtype_hint=dtypes.float32, name="probs")
591 counts = ops.convert_to_tensor(
592 counts, dtype_hint=probs.dtype, name="counts")
593 result = gen_stateless_random_ops.stateless_random_binomial(
594 shape=shape, seed=seed, counts=counts, probs=probs, dtype=output_dtype)
595 shape_util.maybe_set_static_shape(result, shape)
596 return result
599@tf_export("random.stateless_gamma")
600@dispatch.add_dispatch_support
601def stateless_random_gamma(shape,
602 seed,
603 alpha,
604 beta=None,
605 dtype=dtypes.float32,
606 name=None):
607 """Outputs deterministic pseudorandom values from a gamma distribution.
609 The generated values follow a gamma distribution with specified concentration
610 (`alpha`) and inverse scale (`beta`) parameters.
612 This is a stateless version of `tf.random.gamma`: if run twice with the same
613 seeds and shapes, it will produce the same pseudorandom numbers. The output is
614 consistent across multiple runs on the same hardware (and between CPU and
615 GPU),
616 but may change between versions of TensorFlow or on non-CPU/GPU hardware.
618 A slight difference exists in the interpretation of the `shape` parameter
619 between `stateless_gamma` and `gamma`: in `gamma`, the `shape` is always
620 prepended to the shape of the broadcast of `alpha` with `beta`; whereas in
621 `stateless_gamma` the `shape` parameter must always encompass the shapes of
622 each of `alpha` and `beta` (which must broadcast together to match the
623 trailing dimensions of `shape`).
625 Note: Because internal calculations are done using `float64` and casting has
626 `floor` semantics, we must manually map zero outcomes to the smallest
627 possible positive floating-point value, i.e., `np.finfo(dtype).tiny`. This
628 means that `np.finfo(dtype).tiny` occurs more frequently than it otherwise
629 should. This bias can only happen for small values of `alpha`, i.e.,
630 `alpha << 1` or large values of `beta`, i.e., `beta >> 1`.
632 The samples are differentiable w.r.t. alpha and beta.
633 The derivatives are computed using the approach described in
634 (Figurnov et al., 2018).
636 Example:
638 ```python
639 samples = tf.random.stateless_gamma([10, 2], seed=[12, 34], alpha=[0.5, 1.5])
640 # samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
641 # the samples drawn from each distribution
643 samples = tf.random.stateless_gamma([7, 5, 2], seed=[12, 34], alpha=[.5, 1.5])
644 # samples has shape [7, 5, 2], where each slice [:, :, 0] and [:, :, 1]
645 # represents the 7x5 samples drawn from each of the two distributions
647 alpha = tf.constant([[1.], [3.], [5.]])
648 beta = tf.constant([[3., 4.]])
649 samples = tf.random.stateless_gamma(
650 [30, 3, 2], seed=[12, 34], alpha=alpha, beta=beta)
651 # samples has shape [30, 3, 2], with 30 samples each of 3x2 distributions.
653 with tf.GradientTape() as tape:
654 tape.watch([alpha, beta])
655 loss = tf.reduce_mean(tf.square(tf.random.stateless_gamma(
656 [30, 3, 2], seed=[12, 34], alpha=alpha, beta=beta)))
657 dloss_dalpha, dloss_dbeta = tape.gradient(loss, [alpha, beta])
658 # unbiased stochastic derivatives of the loss function
659 alpha.shape == dloss_dalpha.shape # True
660 beta.shape == dloss_dbeta.shape # True
661 ```
663 Args:
664 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
665 seed: A shape [2] Tensor, the seed to the random number generator. Must have
666 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
667 alpha: Tensor. The concentration parameter of the gamma distribution. Must
668 be broadcastable with `beta`, and broadcastable with the rightmost
669 dimensions of `shape`.
670 beta: Tensor. The inverse scale parameter of the gamma distribution. Must be
671 broadcastable with `alpha` and broadcastable with the rightmost dimensions
672 of `shape`.
673 dtype: Floating point dtype of `alpha`, `beta`, and the output.
674 name: A name for the operation (optional).
676 Returns:
677 samples: A Tensor of the specified shape filled with random gamma values.
678 For each i, each `samples[..., i] is an independent draw from the gamma
679 distribution with concentration alpha[i] and scale beta[i].
681 """
682 with ops.name_scope(name, "stateless_random_gamma",
683 [shape, seed, alpha, beta]) as name:
684 shape = shape_util.shape_tensor(shape)
685 alpha = ops.convert_to_tensor(alpha, dtype=dtype, name="alpha")
686 beta = ops.convert_to_tensor(
687 beta if beta is not None else 1, name="beta", dtype=dtype)
688 broadcast_shape = array_ops.broadcast_dynamic_shape(
689 array_ops.shape(alpha), array_ops.shape(beta))
690 alpha_broadcast = array_ops.broadcast_to(alpha, broadcast_shape)
691 alg = "auto_select"
692 key, counter, alg = _get_key_counter_alg(seed, alg)
693 rnd = gen_stateless_random_ops_v2.stateless_random_gamma_v3(
694 shape, key=key, counter=counter, alg=alg, alpha=alpha_broadcast)
695 result = math_ops.maximum(
696 np.finfo(alpha.dtype.as_numpy_dtype).tiny, rnd / beta)
697 shape_util.maybe_set_static_shape(result, shape)
698 return result
701@tf_export("random.stateless_poisson")
702@dispatch.add_dispatch_support
703def stateless_random_poisson(shape,
704 seed,
705 lam,
706 dtype=dtypes.int32,
707 name=None):
708 """Outputs deterministic pseudorandom values from a Poisson distribution.
710 The generated values follow a Poisson distribution with specified rate
711 parameter.
713 This is a stateless version of `tf.random.poisson`: if run twice with the same
714 seeds and shapes, it will produce the same pseudorandom numbers. The output is
715 consistent across multiple runs on the same hardware, but may change between
716 versions of TensorFlow or on non-CPU/GPU hardware.
718 A slight difference exists in the interpretation of the `shape` parameter
719 between `stateless_poisson` and `poisson`: in `poisson`, the `shape` is always
720 prepended to the shape of `lam`; whereas in `stateless_poisson` the shape of
721 `lam` must match the trailing dimensions of `shape`.
723 Example:
725 ```python
726 samples = tf.random.stateless_poisson([10, 2], seed=[12, 34], lam=[5, 15])
727 # samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
728 # the samples drawn from each distribution
730 samples = tf.random.stateless_poisson([7, 5, 2], seed=[12, 34], lam=[5, 15])
731 # samples has shape [7, 5, 2], where each slice [:, :, 0] and [:, :, 1]
732 # represents the 7x5 samples drawn from each of the two distributions
734 rate = tf.constant([[1.], [3.], [5.]])
735 samples = tf.random.stateless_poisson([30, 3, 1], seed=[12, 34], lam=rate)
736 # samples has shape [30, 3, 1], with 30 samples each of 3x1 distributions.
737 ```
739 Args:
740 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
741 seed: A shape [2] Tensor, the seed to the random number generator. Must have
742 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
743 lam: Tensor. The rate parameter "lambda" of the Poisson distribution. Shape
744 must match the rightmost dimensions of `shape`.
745 dtype: Dtype of the samples (int or float dtypes are permissible, as samples
746 are discrete). Default: int32.
747 name: A name for the operation (optional).
749 Returns:
750 samples: A Tensor of the specified shape filled with random Poisson values.
751 For each i, each `samples[..., i]` is an independent draw from the Poisson
752 distribution with rate `lam[i]`.
754 """
755 with ops.name_scope(name, "stateless_random_poisson",
756 [shape, seed, lam]) as name:
757 shape = shape_util.shape_tensor(shape)
758 result = gen_stateless_random_ops.stateless_random_poisson(
759 shape, seed=seed, lam=lam, dtype=dtype)
760 shape_util.maybe_set_static_shape(result, shape)
761 return result
764@tf_export("random.stateless_normal")
765@dispatch.add_dispatch_support
766def stateless_random_normal(shape,
767 seed,
768 mean=0.0,
769 stddev=1.0,
770 dtype=dtypes.float32,
771 name=None,
772 alg="auto_select"):
773 """Outputs deterministic pseudorandom values from a normal distribution.
775 This is a stateless version of `tf.random.normal`: if run twice with the
776 same seeds and shapes, it will produce the same pseudorandom numbers. The
777 output is consistent across multiple runs on the same hardware (and between
778 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
779 hardware.
781 Args:
782 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
783 seed: A shape [2] Tensor, the seed to the random number generator. Must have
784 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
785 mean: A 0-D Tensor or Python value of type `dtype`. The mean of the normal
786 distribution.
787 stddev: A 0-D Tensor or Python value of type `dtype`. The standard deviation
788 of the normal distribution.
789 dtype: The float type of the output: `float16`, `bfloat16`, `float32`,
790 `float64`. Defaults to `float32`.
791 name: A name for the operation (optional).
792 alg: The RNG algorithm used to generate the random numbers. See
793 `tf.random.stateless_uniform` for a detailed explanation.
795 Returns:
796 A tensor of the specified shape filled with random normal values.
797 """
798 with ops.name_scope(name, "stateless_random_normal",
799 [shape, seed, mean, stddev]) as name:
800 shape = shape_util.shape_tensor(shape)
801 mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
802 stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
803 key, counter, alg = _get_key_counter_alg(seed, alg)
804 rnd = gen_stateless_random_ops_v2.stateless_random_normal_v2(
805 shape, key=key, counter=counter, dtype=dtype, alg=alg)
806 result = math_ops.add(rnd * stddev, mean, name=name)
807 shape_util.maybe_set_static_shape(result, shape)
808 return result
811@tf_export("random.stateless_truncated_normal")
812@dispatch.add_dispatch_support
813def stateless_truncated_normal(shape,
814 seed,
815 mean=0.0,
816 stddev=1.0,
817 dtype=dtypes.float32,
818 name=None,
819 alg="auto_select"):
820 """Outputs deterministic pseudorandom values, truncated normally distributed.
822 This is a stateless version of `tf.random.truncated_normal`: if run twice with
823 the same seeds and shapes, it will produce the same pseudorandom numbers. The
824 output is consistent across multiple runs on the same hardware (and between
825 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
826 hardware.
828 The generated values follow a normal distribution with specified mean and
829 standard deviation, except that values whose magnitude is more than 2 standard
830 deviations from the mean are dropped and re-picked.
832 Args:
833 shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
834 seed: A shape [2] Tensor, the seed to the random number generator. Must have
835 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
836 mean: A 0-D Tensor or Python value of type `dtype`. The mean of the
837 truncated normal distribution.
838 stddev: A 0-D Tensor or Python value of type `dtype`. The standard deviation
839 of the normal distribution, before truncation.
840 dtype: The type of the output.
841 name: A name for the operation (optional).
842 alg: The RNG algorithm used to generate the random numbers. See
843 `tf.random.stateless_uniform` for a detailed explanation.
845 Returns:
846 A tensor of the specified shape filled with random truncated normal values.
847 """
848 with ops.name_scope(name, "stateless_truncated_normal",
849 [shape, seed, mean, stddev]) as name:
850 shape = shape_util.shape_tensor(shape)
851 mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
852 stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
853 key, counter, alg = _get_key_counter_alg(seed, alg)
854 rnd = gen_stateless_random_ops_v2.stateless_truncated_normal_v2(
855 shape, key=key, counter=counter, dtype=dtype, alg=alg)
856 result = math_ops.add(rnd * stddev, mean, name=name)
857 shape_util.maybe_set_static_shape(result, shape)
858 return result
861@tf_export(v1=["random.stateless_multinomial"])
862@dispatch.add_dispatch_support
863@deprecation.deprecated(
864 date=None, instructions="Use `tf.random.stateless_categorical` instead.")
865def stateless_multinomial(logits,
866 num_samples,
867 seed,
868 output_dtype=dtypes.int64,
869 name=None):
870 """Draws deterministic pseudorandom samples from a multinomial distribution.
872 This is a stateless version of `tf.random.categorical`: if run twice with the
873 same seeds and shapes, it will produce the same pseudorandom numbers. The
874 output is consistent across multiple runs on the same hardware (and between
875 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
876 hardware.
878 Example:
880 ```python
881 # samples has shape [1, 5], where each value is either 0 or 1 with equal
882 # probability.
883 samples = tf.random.stateless_categorical(
884 tf.math.log([[0.5, 0.5]]), 5, seed=[7, 17])
885 ```
887 Args:
888 logits: 2-D Tensor with shape `[batch_size, num_classes]`. Each slice
889 `[i, :]` represents the unnormalized log-probabilities for all classes.
890 num_samples: 0-D. Number of independent samples to draw for each row slice.
891 seed: A shape [2] Tensor, the seed to the random number generator. Must have
892 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
893 output_dtype: The integer type of the output: `int32` or `int64`. Defaults
894 to `int64`.
895 name: Optional name for the operation.
897 Returns:
898 The drawn samples of shape `[batch_size, num_samples]`.
899 """
900 with ops.name_scope(name, "stateless_multinomial", [logits, seed]):
901 return stateless_multinomial_categorical_impl(logits, num_samples,
902 output_dtype, seed)
905@tf_export("random.stateless_categorical")
906@dispatch.add_dispatch_support
907def stateless_categorical(logits,
908 num_samples,
909 seed,
910 dtype=dtypes.int64,
911 name=None):
912 """Draws deterministic pseudorandom samples from a categorical distribution.
914 This is a stateless version of `tf.categorical`: if run twice with the
915 same seeds and shapes, it will produce the same pseudorandom numbers. The
916 output is consistent across multiple runs on the same hardware (and between
917 CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU
918 hardware.
921 Example:
923 ```python
924 # samples has shape [1, 5], where each value is either 0 or 1 with equal
925 # probability.
926 samples = tf.random.stateless_categorical(
927 tf.math.log([[0.5, 0.5]]), 5, seed=[7, 17])
928 ```
930 Args:
931 logits: 2-D Tensor with shape `[batch_size, num_classes]`. Each slice
932 `[i, :]` represents the unnormalized log-probabilities for all classes.
933 num_samples: 0-D. Number of independent samples to draw for each row slice.
934 seed: A shape [2] Tensor, the seed to the random number generator. Must have
935 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
936 dtype: The integer type of the output: `int32` or `int64`. Defaults to
937 `int64`.
938 name: Optional name for the operation.
940 Returns:
941 The drawn samples of shape `[batch_size, num_samples]`.
942 """
943 with ops.name_scope(name, "stateless_categorical", [logits, seed]):
944 return stateless_multinomial_categorical_impl(logits, num_samples, dtype,
945 seed)
948def stateless_multinomial_categorical_impl(logits, num_samples, dtype, seed):
949 """Implementation for stateless multinomial/categorical ops (v1/v2)."""
950 logits = ops.convert_to_tensor(logits, name="logits")
951 dtype = dtypes.as_dtype(dtype) if dtype else dtypes.int64
952 accepted_dtypes = (dtypes.int32, dtypes.int64)
953 if dtype not in accepted_dtypes:
954 raise ValueError(
955 f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are "
956 f"{accepted_dtypes}.")
957 return gen_stateless_random_ops.stateless_multinomial(
958 logits, num_samples, seed, output_dtype=dtype)
961@dispatch.add_dispatch_support
962@tf_export("random.stateless_parameterized_truncated_normal")
963def stateless_parameterized_truncated_normal(shape,
964 seed,
965 means=0.0,
966 stddevs=1.0,
967 minvals=-2.0,
968 maxvals=2.0,
969 name=None):
970 """Outputs random values from a truncated normal distribution.
972 The generated values follow a normal distribution with specified mean and
973 standard deviation, except that values whose magnitude is more than 2 standard
974 deviations from the mean are dropped and re-picked.
977 Examples:
979 Sample from a Truncated normal, with deferring shape parameters that
980 broadcast.
982 >>> means = 0.
983 >>> stddevs = tf.math.exp(tf.random.uniform(shape=[2, 3]))
984 >>> minvals = [-1., -2., -1000.]
985 >>> maxvals = [[10000.], [1.]]
986 >>> y = tf.random.stateless_parameterized_truncated_normal(
987 ... shape=[10, 2, 3], seed=[7, 17],
988 ... means=means, stddevs=stddevs, minvals=minvals, maxvals=maxvals)
989 >>> y.shape
990 TensorShape([10, 2, 3])
992 Args:
993 shape: A 1-D integer `Tensor` or Python array. The shape of the output
994 tensor.
995 seed: A shape [2] Tensor, the seed to the random number generator. Must have
996 dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
997 means: A `Tensor` or Python value of type `dtype`. The mean of the truncated
998 normal distribution. This must broadcast with `stddevs`, `minvals` and
999 `maxvals`, and the broadcasted shape must be dominated by `shape`.
1000 stddevs: A `Tensor` or Python value of type `dtype`. The standard deviation
1001 of the truncated normal distribution. This must broadcast with `means`,
1002 `minvals` and `maxvals`, and the broadcasted shape must be dominated by
1003 `shape`.
1004 minvals: A `Tensor` or Python value of type `dtype`. The minimum value of
1005 the truncated normal distribution. This must broadcast with `means`,
1006 `stddevs` and `maxvals`, and the broadcasted shape must be dominated by
1007 `shape`.
1008 maxvals: A `Tensor` or Python value of type `dtype`. The maximum value of
1009 the truncated normal distribution. This must broadcast with `means`,
1010 `stddevs` and `minvals`, and the broadcasted shape must be dominated by
1011 `shape`.
1012 name: A name for the operation (optional).
1014 Returns:
1015 A tensor of the specified shape filled with random truncated normal values.
1016 """
1017 with ops.name_scope(name, "stateless_parameterized_truncated_normal",
1018 [shape, means, stddevs, minvals, maxvals]) as name:
1019 shape_tensor = shape_util.shape_tensor(shape)
1020 means_tensor = ops.convert_to_tensor(means, name="means")
1021 stddevs_tensor = ops.convert_to_tensor(stddevs, name="stddevs")
1022 minvals_tensor = ops.convert_to_tensor(minvals, name="minvals")
1023 maxvals_tensor = ops.convert_to_tensor(maxvals, name="maxvals")
1024 rnd = gen_stateless_random_ops.stateless_parameterized_truncated_normal(
1025 shape_tensor, seed, means_tensor, stddevs_tensor, minvals_tensor,
1026 maxvals_tensor)
1027 shape_util.maybe_set_static_shape(rnd, shape)
1028 return rnd