def _iota(*, dtype, shape, dimension):
    dtype = to_tf_dtype(dtype)
    shape_tf = _eval_shape(shape)
    vec = tf.range(tf.cast(shape_tf[dimension], tf.int32), dtype=tf.int32)
    vec_shape = [((- 1) if (i == dimension) else 1) for i in range(len(shape))]
    return tf.cast(tf.broadcast_to(tf.reshape(vec, vec_shape), shape_tf), dtype)
def _iota(*, dtype, shape, dimension):
    dtype = to_tf_dtype(dtype)
    shape_tf = _eval_shape(shape)
    vec = tf.range(tf.cast(shape_tf[dimension], tf.int32), dtype=tf.int32)
    vec_shape = [((- 1) if (i == dimension) else 1) for i in range(len(shape))]
    return tf.cast(tf.broadcast_to(tf.reshape(vec, vec_shape), shape_tf), dtype)