def exec(self): x = self.inputs["x"] y_shape_dict = AxisKeyDict() for axis, index in self.indices.items(): if isinstance(index, slice): index = normalize_slice(index, x.shape_dict[axis]) y_shape_dict[axis] = ((abs(index.stop - index.start) - 1) // abs(index.step)) + 1 elif isinstance(index, int): pass # Remove axis elif index is None: y_shape_dict[axis] = 1 # Insert axis y = Variable(list(y_shape_dict.values()), Order(list(y_shape_dict.keys()))) self.append_output("y", y) return y,
def exec(self): A = self.inputs["A"] B = self.inputs["B"] c_shape_dict = AxisKeyDict() for axis in A.order.axes: if axis not in self.axes[0]: c_shape_dict[axis] = A.shape_dict[axis] for axis in B.order.axes: if axis not in self.axes[1]: c_shape_dict[axis] = B.shape_dict[axis] C = Variable(list(c_shape_dict.values()), Order(list(c_shape_dict.keys()))) self.append_output("C", C) for axis in C.order.axes: self.attributes.add(Tensorwise(self, axis=axis)) return C,
def __call__(self, x: Variable): # assert index is valid for axis, index in self.indices.items(): if axis in x.order.axes: if isinstance(index, slice): index = normalize_slice(index, x.shape_dict[axis]) valid_start = -x.shape_dict[ axis] <= index.start <= x.shape_dict[axis] valid_stop = -x.shape_dict[ axis] <= index.stop <= x.shape_dict[axis] if not valid_start or not valid_stop: raise ValueError(f""" [Slice] Index {index} in {axis} is out of range: (x.order) = {x.order} (x.shape) = {x.shape} (indices) = {self.indices} (indices[{axis.name}]) = {index} """) if ((abs(index.stop - index.start) - 1) // abs(index.step)) + 1 < 0: raise ValueError(f""" [Slice] Slice operator doesn't support 0-size output: (x.order) = {x.order} (x.shape) = {x.shape} (indices) = {self.indices} (indices[{axis.name}]) = {index} """) elif isinstance(index, int): if not -x.shape_dict[axis] <= index < x.shape_dict[axis]: raise ValueError(f""" [Slice] Index {index} in {axis} is out of range: (x.order) = {x.order} (x.shape) = {x.shape} (indices) = {self.indices} (indices[{axis.name}]) = {index} (valid range) = [{-x.shape_dict[axis]}, {x.shape_dict[axis]}) """) elif index is None: raise ValueError(f""" [Slice] Axis {axis} is already exist: (x.order) = {x.order} (x.shape) = {x.shape} (indices) = {self.indices} (indices[{axis.name}]) = {index} """) else: if index is not None: raise ValueError(f""" [Slice] Axis {axis} is not exist in input variable. In this case, index must be "None" (=insert new axis): (x.order) = {x.order} (x.shape) = {x.shape} (indices) = {self.indices} (indices[{axis.name}]) = {index} """) if all(isinstance(index, int) for index in self.indices.values()): raise NotImplementedError(f""" [Slice] Accessing to one element is not supported: (indices) = {self.indices} """) y_shape_dict = AxisKeyDict() for axis, index in self.indices.items(): if isinstance(index, slice): index = normalize_slice(index, x.shape_dict[axis]) y_shape_dict[axis] = ( (abs(index.stop - index.start) - 1) // abs(index.step)) + 1 elif isinstance(index, int): pass # Remove axis elif index is None: y_shape_dict[axis] = 1 # Insert axis y = Variable(list(y_shape_dict.values()), Order(list(y_shape_dict.keys()))) for axis in x.order.axes: if axis in self.indices: index = self.indices[axis] if isinstance( index, slice ) and index.start is None and index.stop is None and index.step is None: # This axis is not sliced. self.attributes.add(Tensorwise(axis)) else: # This axis is not sliced. self.attributes.add(Tensorwise(axis)) self.append_input("x", x) self.append_output("y", y) return y,
def __call__(self, A: Variable, B: Variable): for axis in self.axes[0]: assert axis in A.order.axes, f""" [Tensordot] Input variable "A" must have axes "{axis}": (op) = {self} (op.axes[0]) = {self.axes[0]} (A) = {A}""" for axis in A.order.axes: if axis not in self.axes[0]: assert axis in self.axes[1] or axis not in B.order.axes, f""" [Tensordot] Axes of "A" which are not reduced must not be contained in "B": (op) = {self} (A.order.axes) = {A.order.axes} (B.order.axes) = {B.order.axes} (op.axes) = {self.axes}""" for axis in self.axes[1]: assert axis in B.order.axes, f""" [Tensordot] Input variable "B" must have axes "{axis}": (op) = {self} (op.axes[1]) = {self.axes[1]} (B) = {B}""" for axis in B.order.axes: if axis not in self.axes[1]: assert axis in self.axes[0] or axis not in A.order.axes, f""" [Tensordot] Axes of "B" which are not reduced must not be contained in "A": (op) = {self} (A.order.axes) = {A.order.axes} (B.order.axes) = {B.order.axes} (op.axes) = {self.axes}""" reduction_size_a = mul(A.shape_dict[a] for a in self.axes[0]) reduction_size_b = mul(B.shape_dict[a] for a in self.axes[1]) assert reduction_size_a == reduction_size_b, f""" [Tensordot] Reduction size of "A" and "B" must be same: (A) = {A} (B) = {B} (axes) = {self.axes} (reduction size of A) = {reduction_size_a} (reduction size of B) = {reduction_size_b} """ c_shape_dict = AxisKeyDict() for axis in A.order.axes: if axis not in self.axes[0]: c_shape_dict[axis] = A.shape_dict[axis] for axis in B.order.axes: if axis not in self.axes[1]: c_shape_dict[axis] = B.shape_dict[axis] C = Variable(list(c_shape_dict.values()), Order(list(c_shape_dict.keys()))) for axis in C.order.axes: self.attributes.add(Tensorwise(axis)) self.append_input("A", A) self.append_input("B", B) self.append_output("C", C) return C,