manipulate.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef TVM_RELAX_ATTRS_MANIPULATE_H_
#define TVM_RELAX_ATTRS_MANIPULATE_H_
 
#include <tvm/relax/expr.h>
#include <tvm/tir/index_map.h>
 
namespace tvm {
namespace relax {
 
struct ConcatAttrs : public tvm::AttrsNode<ConcatAttrs> {
  Optional<Integer> axis;
 
  TVM_DECLARE_ATTRS(ConcatAttrs, "relax.attrs.ConcatAttrs") {
    TVM_ATTR_FIELD(axis).describe(
        "The axis at which the input arrays are concatenated."
        "Should lie in range `[-ndim, ndim)`.");
  }
};  // struct ConcatAttrs
 
struct ExpandDimsAttrs : public tvm::AttrsNode<ExpandDimsAttrs> {
  Array<Integer> axis;
 
  TVM_DECLARE_ATTRS(ExpandDimsAttrs, "relax.attrs.ExpandDimsAttrs") {
    TVM_ATTR_FIELD(axis).describe(
        "The axes at which the input array are expanded. "
        "All values are required to lie in range `[-data.ndim - 1, data.ndim]`, "
        "with the convention of negative indexing.");
  }
};  // struct ExpandDimsAttrs
 
struct LayoutTransformAttrs : public tvm::AttrsNode<LayoutTransformAttrs> {
  tir::IndexMap index_map;
  // pad_value is chosen to be of PrimValue type, as it represents constant TIR POD expression. This
  // needs to be revisited in case PrimValue is evolved to represent symbolic expression in future.
  Optional<PrimValue> pad_value;
  Optional<Array<IntImm>> axis_separators;
  Optional<Array<IntImm>> input_axis_separators;
 
  TVM_DECLARE_ATTRS(LayoutTransformAttrs, "relax.attrs.LayoutTransformAttrs") {
    TVM_ATTR_FIELD(index_map).describe("The layout transformation to apply.");
    TVM_ATTR_FIELD(pad_value).describe(
        "The specific value to be used to pad if the layout transform would result in implicit "
        "padding. If not specified, the compiler is free to choose any value.");
    TVM_ATTR_FIELD(axis_separators)
        .describe("The separators between input axes when generating flat output axes");
    TVM_ATTR_FIELD(input_axis_separators)
        .describe("The separators between axes to regenerate output");
  }
};  // struct LayoutTransformAttrs
 
struct PermuteDimsAttrs : public tvm::AttrsNode<PermuteDimsAttrs> {
  Optional<Array<Integer>> axes;
 
  TVM_DECLARE_ATTRS(PermuteDimsAttrs, "relax.attrs.PermuteDimsAttrs") {
    TVM_ATTR_FIELD(axes).describe("The target axes order, reverse order if not specified.");
  }
};  // struct PermuteDimsAttrs
 
struct SplitAttrs : public tvm::AttrsNode<SplitAttrs> {
  ObjectRef indices_or_sections;
  int axis;
 
  TVM_DECLARE_ATTRS(SplitAttrs, "relax.attrs.SplitAttrs") {
    TVM_ATTR_FIELD(indices_or_sections)
        .describe("The input array of indices or the number of split sections.");
    TVM_ATTR_FIELD(axis).describe("The axis to be splitted");
  }
};  // struct SplitAttrs
 
struct SqueezeAttrs : public tvm::AttrsNode<SqueezeAttrs> {
  Optional<Array<Integer>> axis;
 
  TVM_DECLARE_ATTRS(SqueezeAttrs, "relax.attrs.SqueezeAttrs") {
    TVM_ATTR_FIELD(axis).describe(
        "The axis to squeeze in the input tensor."
        "If `axis = None`, all axis of dimension 1 get squeezed;"
        "Else, the dimension in axes get squeezed."
        "It is an error if an axis does not has dimension 1.");
  }
};  // struct SqueezeAttrs
 
struct RepeatAttrs : public tvm::AttrsNode<RepeatAttrs> {
  int repeats;
  Optional<Integer> axis;
 
  TVM_DECLARE_ATTRS(RepeatAttrs, "relax.attrs.RepeatAttrs") {
    TVM_ATTR_FIELD(repeats).describe("The number of repetitions.");
    TVM_ATTR_FIELD(axis).describe(
        "The axis along which to repeat values. The negative numbers are interpreted "
        "counting from the backward. By default, use the flattened input array, and "
        "return a flat output array.");
  }
};  // struct RepeatAttrs
 
struct TileAttrs : public tvm::AttrsNode<TileAttrs> {
  Array<Integer> repeats;
 
  TVM_DECLARE_ATTRS(TileAttrs, "relax.attrs.TileAttrs") {
    TVM_ATTR_FIELD(repeats).describe("The number of repetitions of data along each axis.");
  }
};  // struct TileAttrs
 
struct FlipAttrs : public tvm::AttrsNode<FlipAttrs> {
  Integer axis;
  TVM_DECLARE_ATTRS(FlipAttrs, "relax.attrs.FlipAttrs") {
    TVM_ATTR_FIELD(axis)
        .set_default(NullValue<Integer>())
        .describe("The axis along which to flip over.");
  }
};  // struct FlipAttrs
 
struct GatherElementsAttrs : public tvm::AttrsNode<GatherElementsAttrs> {
  Integer axis;
 
  TVM_DECLARE_ATTRS(GatherElementsAttrs, "relax.attrs.GatherElementsAttrs") {
    TVM_ATTR_FIELD(axis).set_default(0).describe("The axis along which to index.");
  }
};  // struct GatherElementsAttrs
 
struct GatherNDAttrs : public tvm::AttrsNode<GatherNDAttrs> {
  Integer batch_dims;
  TVM_DECLARE_ATTRS(GatherNDAttrs, "relax.attrs.GatherNDAttrs") {
    TVM_ATTR_FIELD(batch_dims).set_default(Integer(0)).describe("The number of batch dims.");
  }
};  // struct GatherNDAttrs
 
struct ScatterElementsAttrs : public tvm::AttrsNode<ScatterElementsAttrs> {
  Integer axis;
  String reduction;
 
  TVM_DECLARE_ATTRS(ScatterElementsAttrs, "relax.attrs.ScatterElementsAttrs") {
    TVM_ATTR_FIELD(axis).set_default(0).describe("The axis over which to select values.");
    TVM_ATTR_FIELD(reduction).set_default("update").describe(
        "Reduction mode of the scatter elements, "
        "either \"update\", \"add\", \"mul\", \"mean\", \"min\" or \"max\".");
  }
};  // struct ScatterElementsAttrs
 
struct ScatterNDAttrs : public tvm::AttrsNode<ScatterNDAttrs> {
  String reduction;
 
  TVM_DECLARE_ATTRS(ScatterNDAttrs, "relax.attrs.ScatterNDAttrs") {
    TVM_ATTR_FIELD(reduction).set_default("update").describe(
        "Accumulation mode of the ScatterND, "
        "either \"update\", \"add\", \"mul\", \"min\" or \"max\".");
  }
};  // struct ScatterNDAttrs
 
struct OneHotAttrs : public tvm::AttrsNode<OneHotAttrs> {
  int depth;
  int axis;
 
  TVM_DECLARE_ATTRS(OneHotAttrs, "relax.attrs.OneHotAttrs") {
    TVM_ATTR_FIELD(depth).describe("Depth of the one hot dimension.");
    TVM_ATTR_FIELD(axis).set_default(-1).describe("Axis to fill.");
  }
};  // struct OneHotAttrs
 
}  // namespace relax
}  // namespace tvm
 
#endif  // TVM_RELAX_ATTRS_MANIPULATE_H_