diff options
| author |  Jared Duke <jdduke@google.com> | 2018-07-26 10:53:21 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-26 10:56:43 -0700 |
| commit | 6e658c0a5ca77677a954a34fb98f241c592c970d (patch) | |
| tree | b645103887539af5232b3f70d80a2eb9b77ed63a /tensorflow/contrib/lite/kernels/register.cc | |
| parent | 0a3155f7fbf56df5e81c7cbf35afd45173359635 (diff) | |
Add one_hot op support to TFLite
PiperOrigin-RevId: 206185190
Diffstat (limited to 'tensorflow/contrib/lite/kernels/register.cc')
| -rw-r--r-- | tensorflow/contrib/lite/kernels/register.cc | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc index 0b70bed308..da69b85041 100644 --- a/tensorflow/contrib/lite/kernels/register.cc +++ b/tensorflow/contrib/lite/kernels/register.cc @@ -107,6 +107,7 @@ TfLiteRegistration* Register_SHAPE(); TfLiteRegistration* Register_POW(); TfLiteRegistration* Register_FAKE_QUANT(); TfLiteRegistration* Register_PACK(); +TfLiteRegistration* Register_ONE_HOT(); BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_RELU, Register_RELU()); @@ -197,6 +198,7 @@ BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_POW, Register_POW()); AddBuiltin(BuiltinOperator_FAKE_QUANT, Register_FAKE_QUANT(), 1, 2); AddBuiltin(BuiltinOperator_PACK, Register_PACK()); + AddBuiltin(BuiltinOperator_ONE_HOT, Register_ONE_HOT()); // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that // custom ops aren't always included by default. |