| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=invalid-name,consider-using-enumerate,redefined-outer-name |
| """Injective transformation operators""" |
| from __future__ import absolute_import as _abs |
| import tvm |
| from tvm import te |
| from tvm import topi |
| from . import cpp |
| from . import tag |
| from .util import within_index, make_idx |
| |
| |
| def expand_dims(a, axis, num_newaxis=1): |
| """Expand the shape of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be expanded. |
| |
| num_newaxis: int, optional |
| Number of newaxis to be inserted on axis |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.expand_dims(a, axis, num_newaxis) |
| |
| |
| def expand_like(a, shape_like, axis): |
| """Expand an input array with the shape of second array. |
| This operation can always be composed of unsqueezing and |
| expanding dims on those unsqueezed axes. |
| |
| Examples |
| -------- |
| .. code-block:: |
| |
| input = [ 12. 19. 27.] |
| input.shape = (3,) |
| |
| new_shape_array = [[[1,2],[2,3],[1,3]], |
| [[1,4],[4,3],[5,2]], |
| [[7,1],[7,2],[7,3]]] |
| new_shape_array.shape = (3, 3, 2) |
| |
| expand_like(input, [1,2], new_shape_array) = |
| [[[12,12],[12,12],[12,12]], |
| [[19,19],[19,19],[19,19]], |
| [[27,27],[27,27],[27,27]]] |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be expanded. |
| shape_like : tvm.te.Tensor |
| The tensor to with target shape. |
| axis: list of int |
| axis to be expanded on |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| odim = len(axis) + len(a.shape) |
| if odim != len(shape_like.shape): |
| if len(a.shape) == 1 and len(axis) == len(shape_like.shape): |
| # A special case: `a` is a scalar represented as a 1-dim tensor |
| return te.compute(shape_like.shape, lambda *idxs: a(0)) |
| raise ValueError( |
| "shape inconsistent when expand_like ({}, {}, {})".format( |
| len(axis), len(a.shape), len(shape_like.shape) |
| ) |
| ) |
| |
| real_axis = topi.reduction._get_real_axis(len(shape_like.shape), axis) |
| real_axis = sorted(real_axis) |
| |
| def _compute(*idxs): |
| indices = [] |
| axis_index = 0 |
| for i in range(0, len(idxs)): |
| if i not in real_axis: |
| indices.append(idxs[i]) |
| axis_index += 1 |
| return a(*indices) |
| |
| return te.compute(shape_like.shape, _compute) |
| |
| |
| def transpose(a, axes=None): |
| """Permute the dimensions of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be expanded. |
| |
| axes: tuple of ints, optional |
| By default, reverse the dimensions. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.transpose(a, axes) |
| |
| |
| def flip(a, axis=0): |
| """Flip/reverse elements of an array in a particular axis. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be expanded. |
| |
| axis : int, optional |
| The axis along which the tensors will be reveresed. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.flip(a, axis) |
| |
| |
| def reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0): |
| """Reverse the tensor for variable length slices. |
| Input is first sliced along batch axis and then elements are reversed along seq axis. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be reversed. |
| |
| seq_lengths : tvm.te.Tensor |
| A 1D Tensor with length a.dims[batch_axis] |
| Must be one of the following types: int32, int64 |
| if seq_lengths[i] > a.dims[seq_axis], it is rounded to a.dims[seq_axis] |
| if seq_lengths[i] < 1, it is rounded to 1 |
| |
| seq_axis : int, optional |
| The axis along which the elements will be reversed. Default is 1. |
| |
| batch_axis : int, optional |
| The axis along which the tensor will be sliced. Default is 0. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| The computed result of same shape and type as of input. |
| |
| """ |
| return cpp.reverse_sequence(a, seq_lengths, seq_axis, batch_axis) |
| |
| |
| def strided_slice(a, begin, end, strides=None, slice_mode="end"): |
| """Slice of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be sliced. |
| |
| begin : list of int |
| The indices to begin with in the slicing. |
| |
| end : list of int |
| Indicies indicating end of the slice. |
| |
| strides : list of int, optional |
| Specifies the stride values, it can be negative |
| in that case, the input tensor will be reversed |
| in that particular axis. |
| |
| slice_mode : str, optional |
| The slice mode [end, size]. |
| end - The ending indices for the slice [default]. |
| size - The input strides will be ignored, input end in this mode indicates |
| the sizeof a slice starting at the location specified by begin. If end[i] |
| is -1, all remaining elements in that dimension are included in the slice. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| if strides is None: |
| strides = [] |
| return cpp.strided_slice(a, begin, end, strides, slice_mode) |
| |
| |
| @tvm.te.tag_scope(tag=tag.INJECTIVE + ",strided_set") |
| def strided_set(a, v, begin, end, strides=None): |
| """Set slice of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be sliced. |
| |
| v : tvm.te.Tensor |
| The values to set |
| |
| begin: tvm.te.Tensor |
| The indices to begin with in the slicing. |
| |
| end: tvm.te.Tensor |
| Indicies indicating end of the slice. |
| |
| strides: tvm.te.Tensor, optional |
| Specifies the stride values, it can be negative |
| in that case, the input tensor will be reversed |
| in that particular axis. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| n = len(a.shape) |
| |
| if len(begin.shape) != 1: |
| raise ValueError("begin should be a vector") |
| if not begin.dtype == "int32": |
| raise TypeError("begin should be int32") |
| if len(end.shape) != 1: |
| raise ValueError("end should be a vector") |
| if not end.dtype == "int32": |
| raise TypeError("end should be int32") |
| if strides is not None: |
| if len(strides.shape) != 1: |
| raise ValueError("strides should be a vector") |
| if not strides.dtype == "int32": |
| raise TypeError("strides should be int32") |
| |
| def _max(a, b): |
| return tvm.tir.Select(a > b, a, b) |
| |
| if strides is None: |
| strides = [tvm.tir.const(1, "int32")] * n |
| else: |
| strides = [ |
| tvm.tir.if_then_else(strides.shape[0] > i, strides[i], tvm.tir.const(1, "int32")) |
| for i in range(n) |
| ] |
| |
| begin = [ |
| tvm.tir.if_then_else( |
| begin.shape[0] > i, |
| begin[i], |
| tvm.tir.Select(strides[i] > 0, tvm.tir.const(0, "int32"), a.shape[i]), |
| ) |
| for i in range(n) |
| ] |
| end = [ |
| tvm.tir.if_then_else( |
| end.shape[0] > i, |
| end[i], |
| tvm.tir.Select(strides[i] > 0, a.shape[i] + 1, -(a.shape[i] + 1)), |
| ) |
| for i in range(n) |
| ] |
| |
| # Convert negative indexes |
| for i in range(n): |
| begin[i] = tvm.tir.if_then_else(begin[i] < 0, begin[i] + a.shape[i], begin[i]) |
| end[i] = tvm.tir.if_then_else(end[i] < 0, end[i] + a.shape[i], end[i]) |
| |
| def _select(*indices): |
| from_val = [] |
| index_tuple = [] |
| for i in range(n): |
| from_val.append(within_index(begin[i], end[i], strides[i], indices[i])) |
| index_tuple.append(make_idx(begin[i], end[i], strides[i], a.shape[i], indices[i])) |
| return tvm.tir.if_then_else(tvm.tir.all(*from_val), v(*index_tuple), a(*indices)) |
| |
| return te.compute(a.shape, _select, name="strided_set") |
| |
| |
| def reshape(a, newshape): |
| """Reshape the array |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be reshaped |
| newshape : tuple of ints |
| The new shape |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.reshape(a, newshape) |
| |
| |
| def squeeze(a, axis=None): |
| """Remove single-dimensional entries from the shape of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| |
| axis : None or int or tuple of ints, optional |
| Selects a subset of the single-dimensional entries in the shape. |
| If an axis is selected with shape entry greater than one, an error is raised. |
| |
| Returns |
| ------- |
| squeezed : tvm.te.Tensor |
| """ |
| return cpp.squeeze(a, axis) |
| |
| |
| def concatenate(a_tuple, axis=0): |
| """Join a sequence of arrays along an existing axis. |
| |
| Parameters |
| ---------- |
| a_tuple : tuple of tvm.te.Tensor |
| The arrays to concatenate |
| |
| axis : int, optional |
| The axis along which the arrays will be joined. Default is 0. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.concatenate(a_tuple, axis) |
| |
| |
| def stack(a, axis): |
| """Repeats the whole array multiple times. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be stacked. |
| |
| axis : int, optional |
| The axis in the result array along which the input arrays are stacked. |
| |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.stack(a, axis) |
| |
| |
| def split(ary, indices_or_sections, axis=0): |
| """Split an array into multiple sub-arrays. |
| |
| Parameters |
| ---------- |
| ary : tvm.te.Tensor |
| |
| indices_or_sections : int or 1-D array |
| |
| axis : int |
| |
| Returns |
| ------- |
| ret : tuple of tvm.te.Tensor |
| """ |
| return cpp.split(ary, indices_or_sections, axis) |
| |
| |
| def take(a, indices, axis=None, mode="clip"): |
| """Take elements from an array along an axis. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The source array. |
| |
| indices : tvm.te.Tensor |
| The indices of the values to extract. |
| |
| axis : int, optional |
| The axis over which to select values. By default, |
| the flattened input array is used. |
| |
| mode : str, optional |
| Specifies how out-of-bound indices will behave. |
| clip - clip to the range (default) |
| wrap - wrap around the indices |
| fast - no clip or wrap around (user must make sure indices are in-bound) |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| if axis is None: |
| return cpp.take(a, indices, mode) |
| return cpp.take(a, indices, int(axis), mode) |
| |
| |
| def gather(data, axis, indices): |
| """Gather values along given axis from given indices. |
| |
| E.g. for a 3D tensor, output is computed as: |
| |
| .. code-block:: python |
| |
| out[i][j][k] = data[indices[i][j][k]][j][k] # if axis == 0 |
| out[i][j][k] = data[i][indices[i][j][k]][k] # if axis == 1 |
| out[i][j][k] = data[i][j][indices[i][j][k]] # if axis == 2 |
| |
| ``indices`` must have same shape as ``data``, except at dimension ``axis`` |
| which must just be not null. Output will have same shape as ``indices``. |
| |
| Parameters |
| ---------- |
| data : tvm.te.Tensor |
| The input data to the operator. |
| |
| axis: int |
| The axis along which to index. |
| |
| indices : tvm.te.Tensor |
| The indices of the values to extract. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.gather(data, axis, indices) |
| |
| |
| def gather_nd(a, indices): |
| """Gather elements from a n-dimension array.. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The source array. |
| |
| indices : tvm.te.Tensor |
| The indices of the values to extract. |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.gather_nd(a, indices) |
| |
| |
| def matmul(a, b, transp_a=False, transp_b=False): |
| """ |
| Creates an operation that calculates a matrix multiplication (row-major notation): |
| A(i, k) * B(k, j) |
| if trans_a == trans_b, the usual transposed combinations, otherwise |
| |
| Parameters |
| ---------- |
| a : The matrix A |
| b : The matrix B |
| trans_a : Is A's layout transposed? |
| trans_b : Is B's layout transposed? |
| |
| Returns |
| ------- |
| A Tensor whose op member is the matmul operation |
| """ |
| return cpp.matmul(a, b, transp_a, transp_b) |
| |
| |
| def tensordot(a, b, axes): |
| """A generalization of matrix multiplication to tensor. |
| |
| Parameters |
| ---------- |
| a : The tensor A |
| b : The tensor B |
| axes : The number of dimensions to reduce over |
| |
| Returns |
| ------- |
| A Tensor computing the result |
| """ |
| if isinstance(axes, int): |
| return cpp.tensordot(a, b, axes) |
| if isinstance(axes[0], int): |
| return cpp.tensordot(a, b, (axes[0],), (axes[1],)) |
| return cpp.tensordot(a, b, axes[0], axes[1]) |
| |
| |
| def arange(start, stop=None, step=1, dtype="float32"): |
| """Creates a tensor with evenly spaced values within a given interval. |
| |
| Parameters |
| ---------- |
| start : tvm.Expr, optional |
| Start of interval. The interval includes this value. The default start |
| value is 0. |
| |
| stop : tvm.Expr |
| Stop of interval. The interval does not include this value. |
| |
| step : tvm.Expr, optional |
| Spacing between values. The default step size is 1. |
| |
| dtype : str, optional |
| The target data type. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| The resulting tensor. |
| """ |
| if stop is None: |
| stop = start |
| start = 0 |
| return cpp.arange(start, stop, step, dtype) |
| |
| |
| def meshgrid(a_tuple, indexing): |
| """Create coordinate matrices from coordinate vectors. |
| |
| Parameters |
| ---------- |
| a_tuple : tuple of tvm.te.Tensor |
| The coordinate vectors or scalars. |
| |
| indexing : str |
| Indexing mode, either "ij" or "xy". |
| |
| Returns |
| ------- |
| result : tuple of tvm.te.Tensor |
| The resulting grids for each axis. |
| """ |
| return cpp.meshgrid(a_tuple, indexing) |
| |
| |
| def repeat(a, repeats, axis): |
| """Repeats elements of an array. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be repeated. |
| |
| repeats: int, required |
| Number of repetitions for each element |
| |
| axis: int, optional |
| The axis along which to repeat values |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.repeat(a, repeats, axis) |
| |
| |
| def tile(a, reps): |
| """Repeats the whole array multiple times. |
| |
| Parameters |
| ---------- |
| a : tvm.te.Tensor |
| The tensor to be tiled. |
| |
| reps: tuple of ints, required |
| The number of times for repeating the tensor |
| |
| Returns |
| ------- |
| ret : tvm.te.Tensor |
| """ |
| return cpp.tile(a, reps) |
| |
| |
| def layout_transform(array, src_layout, dst_layout): |
| """Transform the layout according to src_layout and dst_layout |
| |
| Parameters |
| ---------- |
| array : tvm.te.Tensor |
| The source array. |
| |
| src_layout : str |
| the source layout. |
| |
| dst_layout : str |
| the destination layout. |
| """ |
| return cpp.layout_transform(array, src_layout, dst_layout) |
| |
| |
| def shape(array, dtype="int32"): |
| """Get the shape of input array |
| |
| Parameters |
| ---------- |
| array : tvm.te.Tensor |
| The source tensor. |
| |
| dtype : str, optional |
| The target data type. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| The resulting tensor. |
| """ |
| return cpp.shape(array, dtype) |
| |
| |
| def sequence_mask(data, valid_length, mask_value=0, axis=0): |
| """Sets all elements outside the expected length of the sequence to a constant value. |
| |
| This function takes an n-dimensional input array of the form [MAX_LENGTH, batch_size, ...] or |
| [batch_size, MAX_LENGTH, ...] and returns an array of the same shape. |
| |
| `axis` means the axis of the length dimension and can only be 0 or 1. If `axis` is 0, |
| the data must have shape [MAX_LENGTH, batch_size, ...]. Otherwise (axis=1), the data must have |
| shape [batch_size, MAX_LENGTH, ...]. |
| |
| `valid_length` gives the length of each sequence. `valid_length` should be |
| a 1D int array with positive ints and has dimension [batch_size,]. |
| |
| Parameters |
| ---------- |
| data : tvm.te.Tensor |
| N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...] |
| depending on the value of `axis`. |
| |
| valid_length : tvm.te.Tensor |
| 1-D with shape [batch_size,] |
| |
| mask_value : float, optional |
| The masking value, default 0 |
| |
| axis : int, optional |
| axis of the length dimension, must be 0 or 1, default 0 |
| |
| Returns |
| ------- |
| output : tvm.te.Tensor |
| N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...] |
| depending on the value of `axis`. |
| """ |
| |
| assert len(data.shape) >= 2, "only support data.ndim >= 2, received data.shape = {}".format( |
| data.shape |
| ) |
| assert axis in (0, 1), "only support axis = 0, 1, received axis = {}".format(axis) |
| return cpp.sequence_mask(data, valid_length, mask_value, axis) |
| |
| |
| def ndarray_size(array, dtype="int32"): |
| """Get the number of elements of input array |
| |
| Parameters |
| ---------- |
| array : tvm.te.Tensor |
| The source tensor. |
| |
| dtype : str, optional |
| The target data type. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| The resulting tensor. |
| """ |
| return cpp.ndarray_size(array, dtype) |
| |
| |
| def where(condition, x, y): |
| """Get the elements, either from x or y, depending on the condition. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| The condition array. |
| |
| x : tvm.te.Tensor |
| First array to be selected. |
| |
| y : tvm.te.Tensor |
| Second array to be selected. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| A Tensor selected from x or y depending on condition. |
| """ |
| return cpp.where(condition, x, y) |
| |
| |
| def one_hot(indices, on_value, off_value, depth, axis, dtype): |
| """ |
| Returns a one-hot tensor where the locations repsented by indices take value on_value, |
| other locations take value off_value. |
| Final dimension is <indices outer dimensions> x depth x <indices inner dimensions>. |
| |
| Parameters |
| ---------- |
| indices : tvm.te.Tensor |
| Locations to set to on_value. |
| |
| on_value : tvm.te.Tensor |
| Value to fill at indices. |
| |
| off_value : tvm.te.Tensor |
| Value to fill at all other positions besides indices. |
| |
| depth : int |
| Depth of the one-hot dimension. |
| |
| axis : int |
| Axis to fill. |
| |
| dtype : relay.DataType |
| Data type of the output tensor. |
| |
| Returns |
| ------- |
| ret : relay.Expr |
| The one-hot tensor. |
| |
| Examples |
| -------- |
| .. code-block:: python |
| |
| indices = [0, 1, 2] |
| |
| relay.one_hot(indices, 3) = |
| [[1, 0, 0], |
| [0, 1, 0], |
| [0, 0, 1]] |
| """ |
| return cpp.one_hot(indices, on_value, off_value, depth, axis, dtype) |
| |
| |
| def unravel_index(indices, shape): |
| """Convert a flat index or array of flat indices into a tuple of coordinate arrays. |
| |
| Example:: |
| - unravel_index([22, 41, 37], [7, 6]) = [[3, 6, 6], [4, 5, 1]] |
| |
| Parameters |
| ---------- |
| indices : relay.Expr |
| An integer array containing indices. |
| |
| shape : relay.Expr |
| The shape of the array. |
| |
| Returns |
| ------- |
| result : relay.Expr |
| The tuple of coordinate arrays. |
| """ |
| |
| return cpp.unravel_index(indices, shape) |
| |
| |
| def sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=0): |
| """Converts a sparse representation into a dense tensor. |
| |
| Example:: |
| - sparse_to_dense([[0, 0], [1, 1]], [2, 2], [3, 3], 0) = [[3, 0], [0, 3]] |
| |
| Parameters |
| ---------- |
| sparse_indices : tvm.te.Tensor |
| A 0-D, 1-D, or 2-D tensor of integers containing location of sparse values. |
| |
| output_shape : A list of integers |
| Shape of the dense output tensor. |
| |
| sparse_values : tvm.te.Tensor |
| A 0-D or 1-D tensor containing the sparse values for the sparse indices. |
| |
| default_value : tvm.te.Tensor |
| A 0-D tensor containing the default value for the remaining locations. |
| Defaults to 0. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| Dense tensor of shape output_shape. Has the same type as sparse_values. |
| """ |
| |
| return cpp.sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value) |
| |
| |
| def matrix_set_diag(data, diagonal, k=0, align="RIGHT_LEFT"): |
| """ |
| Returns a tensor with the diagonals of input tensor replaced with the provided diagonal values. |
| |
| Parameters |
| ---------- |
| data : relay.Expr |
| Input Tensor. |
| |
| diagonal : relay.Expr |
| Values to be filled in the diagonal. |
| |
| k : int or tuple of int, optional |
| Diagonal Offset(s). The diagonal or range of diagonals to set. (0 by default) |
| Positive value means superdiagonal, 0 refers to the main diagonal, and |
| negative value means subdiagonals. k can be a single integer (for a single diagonal) |
| or a pair of integers specifying the low and high ends of a matrix band. |
| k[0] must not be larger than k[1]. |
| |
| align : string, optional |
| Some diagonals are shorter than max_diag_len and need to be padded. |
| align is a string specifying how superdiagonals and subdiagonals should be aligned, |
| respectively. There are four possible alignments: "RIGHT_LEFT" (default), "LEFT_RIGHT", |
| "LEFT_LEFT", and "RIGHT_RIGHT". "RIGHT_LEFT" aligns superdiagonals to the right |
| (left-pads the row) and subdiagonals to the left (right-pads the row). It is the packing |
| format LAPACK uses. cuSPARSE uses "LEFT_RIGHT", which is the opposite alignment. |
| |
| Returns |
| ------- |
| result : relay.Expr |
| New tensor with given diagonal values. |
| |
| Examples |
| -------- |
| .. code-block:: python |
| |
| data = [[[7, 7, 7, 7], |
| [7, 7, 7, 7], |
| [7, 7, 7, 7]], |
| [[7, 7, 7, 7], |
| [7, 7, 7, 7], |
| [7, 7, 7, 7]]] |
| |
| diagonal = [[1, 2, 3], |
| [4, 5, 6]] |
| |
| topi.matrix_set_diag(input, diagonal) = |
| [[[1, 7, 7, 7], |
| [7, 2, 7, 7], |
| [7, 7, 3, 7]], |
| [[4, 7, 7, 7], |
| [7, 5, 7, 7], |
| [7, 7, 6, 7]]] |
| """ |
| if isinstance(k, (tuple, list)): |
| k_one = k[0] |
| if len(k) >= 2: |
| k_two = k[1] |
| else: |
| k_two = k[0] |
| else: |
| k_one = k |
| k_two = k |
| |
| super_diag_right_align = align[:5] == "RIGHT" |
| sub_diag_right_align = align[-5:] == "RIGHT" |
| |
| return cpp.matrix_set_diag( |
| data, diagonal, k_one, k_two, super_diag_right_align, sub_diag_right_align |
| ) |
| |
| |
| def adv_index(data, indices): |
| """Numpy style indexing with tensors. |
| |
| Parameters |
| ---------- |
| data : tvm.te.Tensor |
| Input data. |
| |
| indices : A list of tvm.te.Tensor |
| Tensor index. |
| |
| Returns |
| ------- |
| result : tvm.te.Tensor |
| Output tensor |
| """ |
| return cpp.adv_index(data, indices) |
