def METHOD_NAME(dtype):
    batch_size = 3
    image = [
        [[0], [0], [0], [0], [0]],
        [[0], [1], [1], [1], [0]],
        [[0], [1], [1], [1], [0]],
        [[0], [1], [1], [1], [0]],
        [[0], [0], [0], [0], [0]],
    ]
    expected_output = np.array(
        [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 2, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0]
        * batch_size
    )
    images = tf.constant([image] * batch_size, dtype=tf.uint8)
    output = dist_ops.euclidean_dist_transform(images, dtype=dtype)
    output_flat = tf.reshape(output, [-1])
    assert output.shape == [batch_size, 5, 5, 1]
    test_utils.assert_allclose_according_to_type(output_flat, expected_output)