import tensorflow as tf

    # 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].')

