|
|
|
@@ -21,6 +21,36 @@ namespace lite { |
|
|
|
#ifdef PRIMITIVE_WRITEABLE |
|
|
|
int SoftMax::GetAxis() const { return this->primitive_->value.AsSoftMax()->axis; } |
|
|
|
|
|
|
|
int SoftMax::UnPackAttr(const Primitive &prim, const std::vector<AnfNodePtr> &inputs) { |
|
|
|
if (this->primitive_ == nullptr) { |
|
|
|
this->primitive_ = new (std::nothrow) schema::PrimitiveT; |
|
|
|
if (this->primitive_ == nullptr) { |
|
|
|
MS_LOG(ERROR) << "new primitiveT failed"; |
|
|
|
return RET_ERROR; |
|
|
|
} |
|
|
|
this->primitive_->value.type = schema::PrimitiveType_SoftMax; |
|
|
|
} |
|
|
|
if (this->primitive_->value.type != schema::PrimitiveType_SoftMax) { |
|
|
|
MS_LOG(ERROR) << "Primitive type is error :" << this->primitive_->value.type; |
|
|
|
return RET_ERROR; |
|
|
|
} |
|
|
|
if (this->primitive_->value.value == nullptr) { |
|
|
|
auto attr = new (std::nothrow) schema::SoftMaxT(); |
|
|
|
if (attr == nullptr) { |
|
|
|
MS_LOG(ERROR) << "new primitiveT value failed"; |
|
|
|
return RET_ERROR; |
|
|
|
} |
|
|
|
auto prim_axis = GetValue<int>(prim.GetAttr("axis")); |
|
|
|
attr->axis = prim_axis; |
|
|
|
this->primitive_->value.value = attr; |
|
|
|
if (this->primitive_->value.value == nullptr) { |
|
|
|
MS_LOG(ERROR) << "primitive value is nullptr"; |
|
|
|
return RET_ERROR; |
|
|
|
} |
|
|
|
} |
|
|
|
return RET_OK; |
|
|
|
} |
|
|
|
|
|
|
|
void SoftMax::SetAxis(int axis) { this->primitive_->value.AsSoftMax()->axis = axis; } |
|
|
|
|
|
|
|
#else |
|
|
|
|