def __init__(self, transformer, op):
super(DimShuffleKernel, self).__init__(transformer)
if isinstance(op, DimshuffleOp):
out = TensorDescriptionWrapper(self.transformer,
op.tensor_description())
(arg, ) = (_ for _ in op.call_info())
in_tensor = TensorDescriptionWrapper(self.transformer,
arg,
ignore_layout=True)
# Reshape the tensors in place with dimshuffle views
in_tensor.shape = tuple(op.in_view.shape)
in_tensor.strides = tuple(op.in_view.strides)
out.shape = tuple(op.out_view.shape)
out.strides = tuple(op.out_view.strides)
dtype = out.dtype
shape = in_tensor.shape
axes = op.axis_order
elif isinstance(op, AssignOp):
(larg, rarg) = (_ for _ in op.call_info())
out = TensorDescriptionWrapper(self.transformer, larg)
in_tensor = TensorDescriptionWrapper(self.transformer, rarg)
dtype = out.dtype
shape = in_tensor.shape
axes = tuple(range(len(shape)))
self.kernel, self.params = get_dimshuffle(dtype, shape, axes,
in_tensor, out)
def __init__(self, transformer, op):
super(DimShuffleKernel, self).__init__(transformer)
out = TensorDescriptionWrapper(op.tensor_description())
(arg, ) = (_ for _ in op.call_info())
in_tensor = TensorDescriptionWrapper(arg, ignore_layout=True)
# Reshape the tensors in place with dimshuffle views
in_tensor.shape = tuple(op.in_view.shape)
in_tensor.strides = tuple(op.in_view.strides)
out.shape = tuple(op.out_view.shape)
out.strides = tuple(op.out_view.strides)
dtype = out.dtype
shape = in_tensor.shape
axes = op.axis_order
self.kernel, self.params = get_dimshuffle(dtype, shape, axes,
in_tensor, out)
