29void OneHot(
const int32_t depth,
const T on_value,
const T off_value, int32_t axis,
30 const Shape &indices_shape,
const TI *indices_data,
const Shape &, T *output_data)
38 int prefix_dim_size = 1;
39 for (
int i = 0; i < axis; ++i)
41 prefix_dim_size *= indices_shape.
Dims(i);
43 const int suffix_dim_size = indices_shape.
FlatSize() / prefix_dim_size;
51 for (
int i = 0; i < prefix_dim_size; ++i)
53 for (
int j = 0; j < depth; ++j)
55 for (
int k = 0; k < suffix_dim_size; ++k, ++output_data)
58 static_cast<int>(indices_data[i * suffix_dim_size + k]) == j ? on_value : off_value;
void OneHot(const int32_t depth, const T on_value, const T off_value, int32_t axis, const Shape &indices_shape, const TI *indices_data, const Shape &, T *output_data)