import tensorflow as tf from .. import config from ..distribution import Distribution @Distribution def Uniform(name=None): X = tf.placeholder(config.dtype, name=name) Distribution.logp = tf.fill(tf.shape(X), config.dtype(0)) def integral(lower, upper): return tf.cond( tf.logical_or( tf.is_inf(tf.cast(lower, config.dtype)), tf.is_inf(tf.cast(upper, config.dtype)) ), lambda: tf.constant(1, dtype=config.dtype), lambda: tf.cast(upper, config.dtype) - tf.cast(lower, config.dtype), ) Distribution.integral = integral return X @Distribution def UniformInt(name=None): X = tf.placeholder(config.int_dtype, name=name) Distribution.logp = tf.fill(tf.shape(X), config.dtype(0)) def integral(lower, upper): val = tf.cond( tf.logical_or( tf.is_inf(tf.ceil(tf.cast(lower, config.dtype))), tf.is_inf(tf.floor(tf.cast(upper, config.dtype))) ), lambda: tf.constant(1, dtype=config.dtype), lambda: tf.cast(upper, config.dtype) - tf.cast(lower, config.dtype), ) return val Distribution.integral = integral return X