Python fold_in示例

示例#1

文件： sharded.py 项目： beoy/probability

 def _sample_n(self, n, seed, **kwargs):
     seed = samplers.sanitize_seed(seed, salt='sharded_sample')
     seed = samplers.fold_in(seed, tf.cast(self.replica_id, tf.int32))
     return self.distribution.sample(sample_shape=n, seed=seed, **kwargs)

示例#2

 def _sample_n(self, n, seed, **kwargs):
   seed = samplers.sanitize_seed(seed, salt='sharded_independent_sample')
   seed = samplers.fold_in(seed, tf.cast(self.replica_id, tf.int32))
   return super(ShardedIndependent, self)._sample_n(n, seed, **kwargs)

示例#3

文件： distribute_lib.py 项目： stjordanis/probability

 def fold_in(seed, axes):
     for name in axes:
         axis_index = get_axis_index(name)
         seed = samplers.fold_in(seed, tf.cast(axis_index, tf.int32))
     return seed

示例#4

 def _sample_n(self, n, seed, **kwargs):
   seed = samplers.sanitize_seed(seed, salt='sharded_sample')
   for axis_name in self.experimental_shard_axis_names:
     axis_index = distribute_lib.get_axis_index(axis_name)
     seed = samplers.fold_in(seed, tf.cast(axis_index, tf.int32))
   return self.distribution.sample(sample_shape=n, seed=seed, **kwargs)