wehub-resource-sync
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
# 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.
|
|
"""SliceScatter operator"""
|
|
|
|
from tvm import topi
|
|
|
|
from . import utils
|
|
|
|
|
|
def slice_scatter(input_tensor, src, start, end, step, axis):
|
|
"""
|
|
Scatters a slice of src into input along the given axis (SSA form).
|
|
|
|
Args:
|
|
input_tensor (te.Tensor): The input tensor to scatter into.
|
|
src (te.Tensor): The source tensor to scatter from.
|
|
start (int): The starting index of the slice.
|
|
end (int): The ending index of the slice.
|
|
step (int): The step size of the slice.
|
|
axis (int): The axis to scatter along.
|
|
|
|
Returns:
|
|
list[te.Tensor]: A list containing the output tensor with the slice scattered.
|
|
"""
|
|
|
|
dim_size_expr = input_tensor.shape[axis] # Expression for dimension size
|
|
dim_size = utils.get_const_int(dim_size_expr) # Dimension size (as constant int)
|
|
|
|
if start == 0 and end == dim_size and step == 1:
|
|
return topi.identity(src)
|
|
|
|
mask = topi.full((dim_size,), "bool", True)
|
|
idx = topi.arange(start=0, stop=dim_size, step=1, dtype="int64")
|
|
|
|
if start != 0:
|
|
mask = topi.logical_and(mask, topi.greater_equal(idx, start))
|
|
|
|
if end != dim_size:
|
|
mask = topi.logical_and(mask, topi.less(idx, end))
|
|
|
|
if step != 1:
|
|
step_mask = topi.equal(topi.floor_mod(idx - start, step), 0)
|
|
mask = topi.logical_and(mask, step_mask)
|
|
|
|
mask_shape_base = [1] * len(input_tensor.shape)
|
|
mask_shape_base[axis] = dim_size
|
|
mask_shape = tuple(mask_shape_base)
|
|
|
|
mask_reshaped = topi.reshape(mask, mask_shape)
|
|
|
|
idx_new_pre = idx - start + (step - 1)
|
|
idx_new_div = topi.floor_divide(idx_new_pre, step)
|
|
idx_new = topi.clip(idx_new_div, 0, dim_size - 1)
|
|
|
|
temp = topi.take(src, idx_new, axis=axis)
|
|
|
|
mask_shape_expanded_base = list(input_tensor.shape)
|
|
mask_shape_expanded = tuple(mask_shape_expanded_base)
|
|
|
|
mask_expanded = topi.broadcast_to(mask_reshaped, mask_shape_expanded)
|
|
|
|
output = topi.where(mask_expanded, temp, input_tensor)
|
|
|
|
return [output]
|