import tensorflow as tf


def categorical(probs, size=1):
    logits = tf.math.log(probs)[None, :]
    return tf.squeeze(tf.random.categorical(logits, num_samples=size), axis=0)


def randint(high, size=(), dtype=tf.int64):
    return tf.random.uniform(size, maxval=high, dtype=dtype)
