Add a batch-1 test case for batch_to_space_nd
PiperOrigin-RevId: 252686434
This commit is contained in:
parent
2b4245af99
commit
da3013c1fc
@ -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});
|
||||
|
@ -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]],
|
||||
|
Loading…
Reference in New Issue
Block a user