def replace_slice(
dest_node, # type: Node
src_node, # type: Node
lower_bounds, # type: List[int]
upper_bounds, # type: List[int]
strides=None, # type: List[int]
name=None, # type: str
):
# type: (...) -> Node
"""Return a copy of `dest_node` with the specified slice overwritten by the `src_node` data.
:param dest_node: The node providing data to be overwritten by the specified slice.
:param src_node: The node providing data for overwriting.
:param lower_bounds: The (inclusive) lower-bound coordinates for the replaced slice.
:param upper_bounds: The (exclusive) upper-bound coordinates for the replaced slice.
:param strides: The strides for the replaced slice.
:param name: The optional name for the output new node.
:return: The new node with copy of `dest_node` with the specified slice overwritten
by the `src_node`.
"""
if strides is None:
return ReplaceSlice(dest_node, src_node, Coordinate(lower_bounds),
Coordinate(upper_bounds))
else:
return ReplaceSlice(dest_node, src_node, Coordinate(lower_bounds),
Coordinate(upper_bounds), Strides(strides))
def test_replace_slice():
element_type = Type.f32
A = Parameter(element_type, Shape([6, 4]))
B = Parameter(element_type, Shape([3, 2]))
parameter_list = [A, B]
input_arr_a = np.zeros(24, dtype=np.float32).reshape(6, 4)
input_arr_b = np.ones(6, dtype=np.float32).reshape(3, 2)
lower_bounds = [0, 1]
upper_bounds = [3, 3]
function = Function(
NodeVector([
ReplaceSlice(A, B, Coordinate(lower_bounds),
Coordinate(upper_bounds))
]), parameter_list, 'test')
backend = Backend.create(test.BACKEND_NAME)
a = backend.create_tensor(element_type, Shape([6, 4]))
b = backend.create_tensor(element_type, Shape([3, 2]))
result = backend.create_tensor(element_type, Shape([6, 4]))
a.write(util.numpy_to_c(input_arr_a), 0, 24 * 4)
b.write(util.numpy_to_c(input_arr_b), 0, 6 * 4)
result_arr = np.zeros(24, dtype=np.float32).reshape(6, 4)
result.write(util.numpy_to_c(result_arr), 0, 24 * 4)
handle = backend.compile(function)
handle.call([result], [a, b])
result.read(util.numpy_to_c(result_arr), 0, 24 * 4)
result_arr_ref = np.copy(input_arr_a)
result_arr_ref[lower_bounds[0]:upper_bounds[0],
lower_bounds[1]:upper_bounds[1]] = input_arr_b
assert np.allclose(result_arr, result_arr_ref)
#test with strides
lower_bounds = [0, 0]
upper_bounds = [5, 3]
strides = [2, 2]
function = Function(
NodeVector([
ReplaceSlice(A, B, Coordinate(lower_bounds),
Coordinate(upper_bounds), Strides(strides))
]), parameter_list, 'test')
backend = Backend.create(test.BACKEND_NAME)
handle = backend.compile(function)
handle.call([result], [a, b])
result.read(util.numpy_to_c(result_arr), 0, 24 * 4)
result_arr_ref = np.copy(input_arr_a)
result_arr_ref[::strides[0], ::strides[1]] = input_arr_b
assert np.allclose(result_arr, result_arr_ref)
def test_slice():
element_type = Type.f32
shape = Shape([6, 6])
A = Parameter(element_type, shape)
parameter_list = [A]
input_arr = np.arange(36, dtype=np.float32).reshape(6, 6)
lower_bounds = [1, 1]
upper_bounds = [5, 5]
function = Function(
NodeVector(
[Slice(A, Coordinate(lower_bounds), Coordinate(upper_bounds))]),
parameter_list, 'test')
backend = Backend.create(test.BACKEND_NAME)
a = backend.create_tensor(element_type, shape)
result = backend.create_tensor(element_type, Shape([4, 4]))
a.write(util.numpy_to_c(input_arr), 0, 36 * 4)
result_arr = np.zeros(16, dtype=np.float32).reshape(4, 4)
result.write(util.numpy_to_c(result_arr), 0, 16 * 4)
handle = backend.compile(function)
handle.call([result], [a])
result.read(util.numpy_to_c(result_arr), 0, 64)
result_arr_ref = input_arr[lower_bounds[0]:upper_bounds[0],
lower_bounds[1]:upper_bounds[1]]
assert np.allclose(result_arr, result_arr_ref)
#test with strides
strides = [1, 2]
function = Function(
NodeVector([
Slice(A, Coordinate(lower_bounds), Coordinate(upper_bounds),
Strides(strides))
]), parameter_list, 'test')
backend = Backend.create(test.BACKEND_NAME)
result = backend.create_tensor(element_type, Shape([4, 2]))
result_arr = np.zeros(8, dtype=np.float32).reshape(4, 2)
result.write(util.numpy_to_c(result_arr), 0, 8 * 4)
handle = backend.compile(function)
handle.call([result], [a])
result.read(util.numpy_to_c(result_arr), 0, 32)
result_arr_ref = result_arr_ref[::strides[0], ::strides[1]]
assert np.allclose(result_arr, result_arr_ref)
def slice(node, lower_bounds, upper_bounds, strides=None, name=None):
# type: (Node, List[int], List[int], List[int], str) -> Node
"""Take a slice of an input tensor, (sub-tensor) that resides within a bounding box.
Optionally this function may be provided with stride along each axis.
:param node: The tensor we want to slice.
:param lower_bounds: The (inclusive) lower-bound coordinates for the tensor slice.
:param upper_bounds: The (exclusive) upper-bound coordinates for the tensor slice.
:param strides: The strides for the tensor slice.
:param name: Optional name for the output node.
:return: Return node that represents a slice of input nodes data.
"""
if strides is None:
return Slice(node, Coordinate(lower_bounds), Coordinate(upper_bounds))
else:
return Slice(node, Coordinate(lower_bounds), Coordinate(upper_bounds), Strides(strides))
def visit(self, op, x):
self.computation.set_op_rank(op)
# op.args[0] : x
# op.slices
lowers = []
uppers = []
strides = []
axes_to_remove = []
for axis, s in zip(x.axes, op.slices):
if isinstance(s, int):
lowers.append(s)
uppers.append(s + 1)
strides.append(1)
axes_to_remove.append(axis)
else:
if s.start is None:
lowers.append(0)
else:
lowers.append(s.start)
if s.step is None:
strides.append(1)
else:
strides.append(s.step)
if s.stop is None:
uppers.append(axis.length)
else:
uppers.append(s.stop)
op_element_type = self.computation.lookup_cpp_op(x)
ngraph_sliced = PyngSlice(op_element_type, Coordinate(lowers),
Coordinate(uppers), Strides(strides))
if axes_to_remove:
ngraph_sliced = PyngReshape(
ngraph_sliced, AxisVector(list(range(0, len(x.axes)))),
Shape(list(op.axes.lengths)))
self.computation.register_cpp_op(op, ngraph_sliced)