45 lines
1.2 KiB
C
45 lines
1.2 KiB
C
//
|
|
// Created by xiangzhou on 11/9/23.
|
|
//
|
|
|
|
|
|
#pragma once
|
|
#include "C_Tensor.h"
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
//----------------------------------------------------------------------------
|
|
// Gather
|
|
//------------------------------------------------------------------------
|
|
struct GatherAttribute {
|
|
// user
|
|
const char* name_;
|
|
int axis_; // axis to gather on
|
|
};
|
|
typedef struct GatherAttribute GatherAttribute;
|
|
typedef struct GatherAttribute GatherElementsAttribute;
|
|
|
|
//----------------------------------------------------------------------------
|
|
// GatherND
|
|
//------------------------------------------------------------------------
|
|
|
|
struct GatherNDAttribute {
|
|
// user
|
|
const char* name_;
|
|
int batch_dim_; // axis to gather on
|
|
};
|
|
typedef struct GatherNDAttribute GatherNDAttribute;
|
|
|
|
//----------------------------------------------------------------------------
|
|
// ScatterND
|
|
//------------------------------------------------------------------------
|
|
typedef struct Attribute ScatterNDAttribute;
|
|
|
|
CREATE_DECL_RUN_FUN(Gather)
|
|
CREATE_DECL_RUN_FUN(GatherElements)
|
|
CREATE_DECL_RUN_FUN(GatherND)
|
|
CREATE_DECL_RUN_FUN(ScatterND)
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif |