Skip to content
Snippets Groups Projects

Update vit operators

Merged Houssem ROUIS requested to merge hrouis/aidge_core:update_vit_operators into dev
Files
7
@@ -27,25 +27,26 @@
@@ -27,25 +27,26 @@
#include "aidge/utils/Types.h"
#include "aidge/utils/Types.h"
namespace Aidge {
namespace Aidge {
enum class GatherAttr { Axis };
enum class GatherAttr { Indices, GatheredShape, Axis };
class Gather_Op : public OperatorTensor,
class Gather_Op : public OperatorTensor,
public Registrable<Gather_Op,
public Registrable<Gather_Op,
std::string,
std::string,
std::unique_ptr<OperatorImpl>(const Gather_Op&)>,
std::unique_ptr<OperatorImpl>(const Gather_Op&)>,
public StaticAttributes<GatherAttr, int> {
public StaticAttributes<GatherAttr, std::vector<std::int64_t>, std::vector<DimSize_t>, std::int64_t> {
public:
public:
static const std::string Type;
static const std::string Type;
Gather_Op() = delete;
Gather_Op() = delete;
using Attributes_ = StaticAttributes<GatherAttr, std::vector<std::int64_t>, std::vector<DimSize_t>, std::int64_t>;
using Attributes_ = StaticAttributes<GatherAttr, int>;
template <GatherAttr e> using attr = typename Attributes_::template attr<e>;
template <GatherAttr e> using attr = typename Attributes_::template attr<e>;
Gather_Op(int axis)
Gather_Op(const std::vector<std::int64_t>& indices, const std::vector<DimSize_t>& gatheredShape, std::int64_t axis)
: OperatorTensor(Type, 2, 0, 1),
: OperatorTensor(Type, 1, 0, 1),
Attributes_(
Attributes_(
 
attr<GatherAttr::Indices>(indices),
 
attr<GatherAttr::GatheredShape>(gatheredShape),
attr<GatherAttr::Axis>(axis))
attr<GatherAttr::Axis>(axis))
{}
{}
@@ -76,21 +77,21 @@ public:
@@ -76,21 +77,21 @@ public:
}
}
static const std::vector<std::string> getInputsName(){
static const std::vector<std::string> getInputsName(){
return {"data_input", "indexes"};
return {"data_input"};
}
}
static const std::vector<std::string> getOutputsName(){
static const std::vector<std::string> getOutputsName(){
return {"data_output"};
return {"data_output"};
}
}
};
};
inline std::shared_ptr<Node> Gather(int axis = 0, const std::string& name = "") {
inline std::shared_ptr<Node> Gather( const std::vector<std::int64_t>& indices, const std::vector<DimSize_t>& gatheredShape, std::int64_t axis = 0, const std::string& name = "") {
return std::make_shared<Node>(std::make_shared<Gather_Op>(axis), name);
return std::make_shared<Node>(std::make_shared<Gather_Op>(indices, gatheredShape, axis), name);
}
}
} // namespace Aidge
} // namespace Aidge
namespace {
namespace {
template <>
template <>
const char *const EnumStrings<Aidge::GatherAttr>::data[] = {"Axis"};
const char *const EnumStrings<Aidge::GatherAttr>::data[] = {"Indices", "GatheredShape", "Axis"};
}
}
#endif /* AIDGE_CORE_OPERATOR_GATHER_H_ */
#endif /* AIDGE_CORE_OPERATOR_GATHER_H_ */
Loading