Add a batch-1 test case for batch_to_space_nd

PiperOrigin-RevId: 252686434
This commit is contained in:
Jared Duke 2019-06-11 13:29:21 -07:00 committed by TensorFlower Gardener
parent 2b4245af99
commit da3013c1fc
2 changed files with 18 additions and 1 deletions

View File

@ -121,6 +121,14 @@ TEST(BatchToSpaceNDOpTest, SimpleConstTestInt8) {
{1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16}));
}
TEST(BatchToSpaceNDOpTest, BatchOneConstTest) {
BatchToSpaceNDOpConstModel m({1, 2, 2, 1}, {1, 1}, {0, 0, 0, 0});
m.SetInput<float>({1, 2, 3, 4});
m.Invoke();
EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 2, 1}));
EXPECT_THAT(m.GetOutput<float>(), ElementsAreArray({1, 2, 3, 4}));
}
TEST(BatchToSpaceNDOpTest, SimpleDynamicTest) {
BatchToSpaceNDOpDynamicModel m({4, 2, 2, 1});
m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});

View File

@ -2875,7 +2875,16 @@ def make_batch_to_space_nd_tests(options):
"constant_block_shape": [True, False],
"constant_crops": [True, False],
},
# Non-4D use case: 1 bath dimension, 3 spatial dimensions, 2 others.
# Single batch (no-op)
{
"dtype": [tf.float32],
"input_shape": [[1, 3, 3, 1]],
"block_shape": [[1, 1]],
"crops": [[[0, 0], [0, 0]], [[1, 1], [1, 1]]],
"constant_block_shape": [True],
"constant_crops": [True],
},
# Non-4D use case: 1 batch dimension, 3 spatial dimensions, 2 others.
{
"dtype": [tf.float32],
"input_shape": [[8, 2, 2, 2, 1, 1]],