Python tensorflow.ensure_shape() Examples

def sparse_add_self_loops(indices, N=None):
    """
    Given the indices of a square SparseTensor, adds the diagonal entries (i, i)
    and returns the reordered indices.
    :param indices: Tensor of rank 2, the indices to a SparseTensor.
    :param N: the size of the N x N SparseTensor indexed by the indices. If `None`,
    N is calculated as the maximum entry in the indices plus 1.
    :return: Tensor of rank 2, the indices to a SparseTensor.
    """
    N = tf.reduce_max(indices) + 1 if N is None else N
    row, col = indices[..., 0], indices[..., 1]
    mask = tf.ensure_shape(row != col, row.shape)
    sl_indices = tf.range(N, dtype=row.dtype)[:, None]
    sl_indices = tf.repeat(sl_indices, 2, -1)
    indices = tf.concat((indices[mask], sl_indices), 0)
    dummy_values = tf.ones_like(indices[:, 0])
    indices, _ = gen_sparse_ops.sparse_reorder(indices, dummy_values, (N, N))
    return indices 

def parse_fn(self, file):
        config = self.config
        image_size = config.IMAGE_SIZE
        dmap_size = config.MAP_SIZE
        label_size = 1

        def _parse_function(_file):
            _file = _file.decode('UTF-8')
            image_bytes = image_size * image_size * 3
            dmap_bytes = dmap_size * dmap_size
            bin = np.fromfile(_file, dtype='uint8')
            image = np.transpose(bin[0:image_bytes].reshape((3, image_size, image_size)) / 255, (1, 2, 0))
            dmap  = np.transpose(bin[image_bytes:image_bytes+dmap_bytes].reshape((1, dmap_size, dmap_size)) / 255, (1, 2, 0))
            label = bin[image_bytes+dmap_bytes:image_bytes+dmap_bytes+label_size] / 1
            dmap1 = dmap * (1-label)
            dmap2 = np.ones_like(dmap) * label
            dmap = np.concatenate([dmap1, dmap2], axis=2)

            return image.astype(np.float32), dmap.astype(np.float32), label.astype(np.float32)

        image_ts, dmap_ts, label_ts = tf.numpy_function(_parse_function, [file], [tf.float32, tf.float32, tf.float32])
        image_ts = tf.ensure_shape(image_ts, [config.IMAGE_SIZE, config.IMAGE_SIZE, 3])
        dmap_ts  = tf.ensure_shape(dmap_ts,  [config.MAP_SIZE, config.MAP_SIZE, 2])
        label_ts = tf.ensure_shape(label_ts, [1])
        return image_ts, dmap_ts, label_ts 

def gaussianMI(x, y, constellation, M, dtype=tf.float64):
    """
        Computes mutual information with Gaussian auxiliary channel assumption and constellation with uniform porbability distribution

        x: (1, N), N normalized complex samples at the transmitter, where N is the batchSize/sampleSize
        y: (1, N), N normalized complex observations at the receiver, where N is the batchSize/sampleSize
        constellation: (1, M), normalized complex constellation of order M
        
        Transcribed from Dr. Tobias Fehenberger MATLAB code.
        See: https://www.fehenberger.de/#sourcecode
    """
    if len(constellation.shape) == 1:
        constellation = tf.expand_dims(constellation, axis=0)
    if len(y.shape) == 1:
        y = tf.expand_dims(y, axis=0)
    if len(x.shape) == 1:
        x = tf.expand_dims(x, axis=0)
    if y.shape[0] != 1:
        y = tf.linalg.matrix_transpose(y)
    if x.shape[0] != 1:
        x = tf.linalg.matrix_transpose(x)
    if constellation.shape[0] == 1:
        constellation = tf.linalg.matrix_transpose(constellation)

    N = tf.cast( tf.shape(x)[1], dtype )

    PI = tf.constant( np.pi, dtype=dtype )
    REALMIN = tf.constant( np.finfo(float).tiny, dtype=dtype )

    xint = tf.math.argmin(tf.square(tf.abs(x - constellation)), axis=0, output_type=tf.int32)
    x_count = tf.math.bincount(xint)
    x_count = tf.ensure_shape(x_count, (M,))
    P_X = tf.cast(x_count, dtype) / N
        
    N0 = tf.reduce_mean( tf.square( tf.abs(x-y) ) )
    
    qYonX = 1 / ( PI*N0 ) * tf.exp( ( -tf.square(tf.math.real(y)-tf.math.real(x)) -tf.square(tf.math.imag(y)-tf.math.imag(x)) ) / N0 )
    
    qY = []
    for ii in np.arange(M):
        temp = P_X[ii] * (1 / (PI * N0) * tf.exp( ( -tf.square(tf.math.real(y)-tf.math.real(constellation[ii,0])) -tf.square(tf.math.imag(y)-tf.math.imag(constellation[ii,0])) ) / N0) )
        qY.append(temp)
    qY = tf.reduce_sum( tf.concat(qY, axis=0), axis=0)
            
    MI = 1 / N * tf.reduce_sum( log2( tf.math.maximum(qYonX, REALMIN) / tf.math.maximum(qY, REALMIN) ) )
    
    return MI 

def rotate_train_nhwc(x, pi):
    sym = tf.random_uniform(
        [],
        minval=0,
        maxval=len(SYMMETRIES),
        dtype=tf.int32,
        seed=123)

    def rotate(tensor):
        # flipLeftRight
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 1) > 0,
            tf.reverse(tensor, axis=[0]),
            tensor)
        # flipUpDown
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 2) > 0,
            tf.reverse(tensor, axis=[1]),
            tensor)
        # flipDiagonal
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 4) > 0,
            tf.transpose(tensor, perm=[1, 0, 2]),
            tensor)
        return tensor

    # TODO(tommadams): use tf.ensure_shape instead of tf.assert_equal.
    squares = go.N * go.N
    assert_shape_pi = tf.assert_equal(pi.shape.as_list(), [squares + 1])

    x_shape = x.shape.as_list()
    assert_shape_x = tf.assert_equal(x_shape, [go.N, go.N, x_shape[2]])

    pi_move = tf.slice(pi, [0], [squares], name="slice_moves")
    pi_pass = tf.slice(pi, [squares], [1], name="slice_pass")
    # Add a final dim so that x and pi have same shape: [N,N,num_features].
    pi_n_by_n = tf.reshape(pi_move, [go.N, go.N, 1])

    with tf.control_dependencies([assert_shape_x, assert_shape_pi]):
        pi_rot = tf.concat(
            [tf.reshape(rotate(pi_n_by_n), [squares]), pi_pass],
            axis=0)

    return rotate(x), pi_rot 

def rotate_train_nchw(x, pi):
    sym = tf.random_uniform(
        [],
        minval=0,
        maxval=len(SYMMETRIES),
        dtype=tf.int32,
        seed=123)

    def rotate(tensor):
        # flipLeftRight
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 1) > 0,
            tf.reverse(tensor, axis=[1]),
            tensor)
        # flipUpDown
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 2) > 0,
            tf.reverse(tensor, axis=[2]),
            tensor)
        # flipDiagonal
        tensor = tf.where(
            tf.bitwise.bitwise_and(sym, 4) > 0,
            tf.transpose(tensor, perm=[0, 2, 1]),
            tensor)
        return tensor

    # TODO(tommadams): use tf.ensure_shape instead of tf.assert_equal.
    squares = go.N * go.N
    assert_shape_pi = tf.assert_equal(pi.shape.as_list(), [squares + 1])

    x_shape = x.shape.as_list()
    assert_shape_x = tf.assert_equal(x_shape, [x_shape[0], go.N, go.N])

    pi_move = tf.slice(pi, [0], [squares], name="slice_moves")
    pi_pass = tf.slice(pi, [squares], [1], name="slice_pass")
    # Add a dim so that x and pi have same shape: [num_features,N,N].
    pi_n_by_n = tf.reshape(pi_move, [1, go.N, go.N])

    with tf.control_dependencies([assert_shape_x, assert_shape_pi]):
        pi_rot = tf.concat(
            [tf.reshape(rotate(pi_n_by_n), [squares]), pi_pass],
            axis=0)

    return rotate(x), pi_rot