Home
last modified time | relevance | path

Searched refs:block_rank (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dbatchtospace_op.cc31 const int block_rank = block_shape.size(); in BatchToSpace() local
34 ctx, input_rank >= 1 + block_rank, in BatchToSpace()
35 errors::InvalidArgument("input rank should be >= ", 1 + block_rank, in BatchToSpace()
38 remainder_shape.remove_prefix(1 + block_rank); in BatchToSpace()
43 block_rank == xla::ShapeUtil::GetDimension(crops.shape(), 0) && in BatchToSpace()
45 errors::InvalidArgument("crops should have shape [", block_rank, in BatchToSpace()
53 for (int i = 0; i < block_rank; ++i) { in BatchToSpace()
70 std::vector<int64> reshaped_shape(input_rank + block_rank); in BatchToSpace()
72 reshaped_shape[block_rank] = batch_size / block_num_elems; in BatchToSpace()
74 reshaped_shape.begin() + block_rank + 1); in BatchToSpace()
[all …]
Dspacetobatch_op.cc31 const int block_rank = block_shape.size(); in SpaceToBatch() local
34 ctx, input_rank >= 1 + block_rank, in SpaceToBatch()
35 errors::InvalidArgument("input rank should be >= ", 1 + block_rank, in SpaceToBatch()
38 remainder_shape.remove_prefix(1 + block_rank); in SpaceToBatch()
43 block_rank == xla::ShapeUtil::GetDimension(paddings.shape(), 0) && in SpaceToBatch()
45 errors::InvalidArgument("paddings should have shape [", block_rank, in SpaceToBatch()
57 for (int i = 0; i < block_rank; ++i) { in SpaceToBatch()
89 std::vector<int64> reshaped_padded_shape(input_rank + block_rank); in SpaceToBatch()
91 for (int i = 0; i < block_rank; ++i) { in SpaceToBatch()
102 reshaped_padded_shape.begin() + 1 + 2 * block_rank); in SpaceToBatch()
[all …]