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 // Updates a slice of 'x', i.e., 26 // x[start[0], ..., start[n]] = update 27 XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span<const int64> start); 28 29 // Performs a slice in the minor dimensions of a tensor. 30 // x[..., start[0]:end[0], ..., start[n]:end[n]] 31 XlaOp SliceInMinorDims(XlaOp x, absl::Span<const int64> start, 32 absl::Span<const int64> end); 33 34 // Updates a slice of 'x', where 'start' contains a list of minor dimensions: 35 // x[..., start[0]:..., ..., start[n]:...] = update 36 XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, 37 absl::Span<const int64> start); 38 39 // Performs a dynamic slice in the minor dimensions of a tensor. 40 XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span<const XlaOp> starts, 41 absl::Span<const int64> sizes); 42 43 XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, 44 absl::Span<const XlaOp> starts); 45 46 // Gathers values along an axis specified by dim. 47 // 48 // For a 3-D tensor the output is specified by: 49 // 50 // out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 51 // out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 52 // out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 53 // 54 // If `input` is an n-dimensional tensor with size 55 // [X0,X1,X2,..XN] and dim = i `index` must be an n-dimensional tensor with size 56 // [X0,X1,...Y,Xi+1,...,X[N] where y >= 1 and `out` will have the same sizes as 57 // `index`. 58 XlaOp TorchGather(XlaOp input, XlaOp index, int64 dim); 59 60 // Returns a new tensor which indexes the input tensor along dimension dim using 61 // the entries in index. 62 // 63 // The returned tensor has the same number of dimensions as the original tensor 64 // (input). The dimth dimension has the same size as the length of index; other 65 // dimensions have the same size as in the original tensor. 66 XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim); 67 68 } // namespace xla 69 70 #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ 71