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)