24 #ifndef TVM_TOPI_NN_FLATTEN_H_
25 #define TVM_TOPI_NN_FLATTEN_H_
51 auto ishape = x->shape;
53 for (
size_t i = 1; i < ishape.size(); ++i) {
54 dim = dim * ishape[i];
59 std::vector<PrimExpr> extra_shape;
60 for (
size_t i = 1; i < ishape.size(); ++i) {
63 std::reverse(extra_shape.
begin(), extra_shape.
end());
69 std::vector<PrimExpr> index;
70 for (
auto s : extra_shape) {
75 std::reverse(index.begin(), index.end());
Reference to PrimExprNode.
Definition: expr.h:115
Array, container representing a contiguous sequence of ObjectRefs.
Definition: array.h:289
iterator end() const
Definition: array.h:390
void push_back(const T &item)
push a new item to the back of the list
Definition: array.h:457
iterator begin() const
Definition: array.h:387
Tensor structure representing a possible input, or intermediate computation result.
Definition: tensor.h:102
a named variable in TIR
Definition: var.h:89
Utility functions for handling constants in TVM expressions.
Tensor expression language DSL.
Definition: extracted_task.h:33
Tensor compute(Array< PrimExpr > shape, FCompute fcompute, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={})
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fco...
Tensor flatten(const Tensor &x, std::string name="tensor", std::string tag=kInjective)
Flattens the input tensor into a 2-D tensor by collapsing higher dimensions. This requires the input ...
Definition: flatten.h:50
constexpr auto kInjective
Definition: tags.h:33
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:36
PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b) where a and b are non-negative.
PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder floor(a / b) where a and b are non-negative.
Operation node can generate one or multiple Tensors.