import tensorflow as tf

    # Calculate the output size.
    x_shape = tf.shape(x)
    out_shape = calc_out_size_4d(x_shape, ksize, strides, padding)

    # Pad input.
    x_ = _pad_input(
        x, ksize, strides, padding, bsize=[1, blk_shape[1], blk_shape[2], 1], bstrides=bstrides)

    # In matrix multiplication mode, the block patch should be the same as the kernel size.
    assert_shape = tf.assert_equal(
        tf.stack([blk_shape[1], blk_shape[2]]),
        tf.stack([ksize[0], ksize[1]]),
        message='Expect blk_indices.shape[1] == w.shape[0] and blk_indices.shape[2] == w.shape[1].')

    # Currently we do not support strides > 1 in this matrix multiplication mode. Could be supported
    # in the future.
    assert_strides = tf.assert_equal(
        tf.cast(tf.stack([strides[1], strides[2]]), tf.int64),
        tf.constant([1, 1], dtype=tf.int64),
        message='Strides > 1 not supported.')
