layout_transform takes an input tensor x and an attribute index_map of tir::IndexMap type. It also has an optional pad_value attribute.
It transforms x as per the index_map attribute and returns the transformed tensor.
This operator for now allows layout transformations that introduce implicit padding. For example, transforming a tensor of shape (10,) using the lambda i: (i//4, i%4). The output shape will be (3, 4) with two elements padded. The optional pad_value is used to pad if specified, otherwise the compiler is free to choose any value to pad.
Adds
layout_transform
operator in Relax as part of Relax layout planning.layout_transform
takes an input tensorx
and an attributeindex_map
oftir::IndexMap
type. It also has an optionalpad_value
attribute. It transformsx
as per theindex_map
attribute and returns the transformed tensor.This operator for now allows layout transformations that introduce implicit padding. For example, transforming a tensor of shape
(10,)
using thelambda i: (i//4, i%4)
. The output shape will be(3, 4)
with two elements padded. The optionalpad_value
is used to pad if specified, otherwise the compiler is free to choose any value to pad.