1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "absl/types/span.h" 17 #include "tensorflow/compiler/xla/client/xla_builder.h" 18 #include "tensorflow/compiler/xla/types.h" 19 20 #ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ 21 #define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ 22 23 namespace xla { 24 25 // Slices input starting from the base_indices and within the window_sizes, 26 // using the supplied strides. This is the equivalent of the Python slicing op 27 // [base_indices : base_indices+window_sizes : stride]. 28 XlaOp DynamicStridedSlice(XlaOp input, absl::Span<const XlaOp> base_indices, 29 absl::Span<const int64> window_sizes, 30 absl::Span<const int64> strides); 31 32 // Updates a slice of 'x', i.e., 33 // x[start[0], ..., start[n]] = update 34 XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span<const int64> start); 35 36 // Performs a slice in the minor dimensions of a tensor. 37 // x[..., start[0]:end[0], ..., start[n]:end[n]] 38 XlaOp SliceInMinorDims(XlaOp x, absl::Span<const int64> start, 39 absl::Span<const int64> end); 40 41 // Updates a slice of 'x', where 'start' contains a list of minor dimensions: 42 // x[..., start[0]:..., ..., start[n]:...] = update 43 XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, 44 absl::Span<const int64> start); 45 46 // Performs a dynamic slice in the minor dimensions of a tensor. 47 XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span<const XlaOp> starts, 48 absl::Span<const int64> sizes); 49 50 XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, 51 absl::Span<const XlaOp> starts); 52 53 // Gathers values along an axis specified by dim. 54 // 55 // For a 3-D tensor the output is specified by: 56 // 57 // out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 58 // out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 59 // out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 60 // 61 // If `input` is an n-dimensional tensor with size 62 // [X0,X1,X2,..XN] and dim = i `index` must be an n-dimensional tensor with size 63 // [X0,X1,...Y,Xi+1,...,X[N] where y >= 1 and `out` will have the same sizes as 64 // `index`. 65 XlaOp TorchGather(XlaOp input, XlaOp index, int64 dim, bool sparse = true); 66 67 // idx = index[i][j][k] 68 // output[idx][j][k] = combiner(input[idx][j][k], src[i][j][k]) # if dim == 0 69 // output[i][idx][k] = combiner(input[i][idx][k], src[i][j][k]) # if dim == 1 70 // output[i][j][idx] = combiner(input[i][j][idx], src[i][j][k]) # if dim == 2 71 XlaOp TorchScatterDense(XlaOp input, XlaOp index, XlaOp src, int64 dim, 72 const std::function<XlaOp(XlaOp, XlaOp)>& combiner); 73 74 // Returns a new tensor which indexes the input tensor along dimension dim using 75 // the entries in index. 76 // 77 // The returned tensor has the same number of dimensions as the original tensor 78 // (input). The dimth dimension has the same size as the length of index; other 79 // dimensions have the same size as in the original tensor. 80 // 81 // This operation supports 0 or more major batch dimensions that act like a 82 // multidimensional loop over both the input and the index. 83 XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim, 84 int64 batch_dims = 0); 85 86 } // namespace xla 87 88 #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ 89