Converts the tensor data format to the one required by the XLA pooling library.
| 134 | // Converts the tensor data format to the one required by the XLA pooling |
| 135 | // library. |
| 136 | xla::TensorFormat XlaTensorFormat(tensorflow::TensorFormat data_format, |
| 137 | int num_spatial_dims) { |
| 138 | int num_dims = num_spatial_dims + 2; |
| 139 | int batch_dimension = GetTensorBatchDimIndex(num_dims, data_format); |
| 140 | int feature_dimension = GetTensorFeatureDimIndex(num_dims, data_format); |
| 141 | absl::InlinedVector<int64, 4> spatial_dimensions(num_spatial_dims); |
| 142 | for (int spatial_dim = 0; spatial_dim < num_spatial_dims; ++spatial_dim) { |
| 143 | spatial_dimensions[spatial_dim] = |
| 144 | GetTensorSpatialDimIndex(num_dims, data_format, spatial_dim); |
| 145 | } |
| 146 | return xla::TensorFormat(/*batch_dimension=*/batch_dimension, |
| 147 | /*feature_dimension=*/feature_dimension, |
| 148 | /*spatial_dimensions=*/spatial_dimensions); |
| 149 | } |
| 150 | |
| 151 | class MaxPoolOp : public PoolingOp { |
| 152 | public: |
no test coverage detected