Tensor format for reduce window operations.
| 23 | |
| 24 | // Tensor format for reduce window operations. |
| 25 | class TensorFormat { |
| 26 | public: |
| 27 | TensorFormat(int batch_dimension, int feature_dimension, |
| 28 | absl::Span<const int64> spatial_dimensions) |
| 29 | : batch_dimension_(batch_dimension), |
| 30 | feature_dimension_(feature_dimension), |
| 31 | spatial_dimensions_(spatial_dimensions.begin(), |
| 32 | spatial_dimensions.end()) {} |
| 33 | |
| 34 | int batch_dimension() const { return batch_dimension_; } |
| 35 | |
| 36 | int feature_dimension() const { return feature_dimension_; } |
| 37 | |
| 38 | int spatial_dimension(int dim) const { return spatial_dimensions_[dim]; } |
| 39 | |
| 40 | int num_spatial_dims() const { return spatial_dimensions_.size(); } |
| 41 | |
| 42 | private: |
| 43 | // The number of the dimension that represents the batch. |
| 44 | int batch_dimension_; |
| 45 | // The number of the dimension that represents the features. |
| 46 | int feature_dimension_; |
| 47 | // The dimension numbers for the spatial dimensions. |
| 48 | absl::InlinedVector<int, 4> spatial_dimensions_; |
| 49 | }; |
| 50 | |
| 51 | // Computes the max pool of 'operand'. |
| 52 | XlaOp MaxPool(XlaOp operand, absl::Span<const int64> kernel_size, |
no outgoing calls