# images = (?, 28, 28, 1)  # MNIST images

# logits = (?, 28, 28, 256)

import tensorflow as tf

images = tf.placeholder(tf.float32, shape=[None, 28, 28, 1])
# ...network goes here ...
logits = tf.placeholder(tf.float32, shape=[None, 28, 28, 256])

# Softmax Cross Entropy.
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
    logits=tf.reshape(logits, (-1, 256)),  # [BatchHeightWidthChannel, Distribution]
    labels=tf.to_int32(tf.reshape(images, shape=(-1,))  # [TruePixelValues]
))