jax.experimental.pallas.tpu.prng_seed

Contents

jax.experimental.pallas.tpu.prng_seed#

jax.experimental.pallas.tpu.prng_seed(*seeds)[source]#

Sets the seed for PRNG.

Parameters:: seeds (int | jax.Array) – One or more integer seeds for setting the PRNG seed. If more than one seed is passed in, the seed material will be mixed before setting the internal PRNG state.
Return type:: None