| 117 | |
| 118 | template <KernelType kernel_type> |
| 119 | TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { |
| 120 | BatchToSpaceNDContext op_context(context, node); |
| 121 | |
| 122 | // Resize the output tensor if the output tensor is dynamic. |
| 123 | if (IsDynamicTensor(op_context.output)) { |
| 124 | TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); |
| 125 | } |
| 126 | |
| 127 | #define TF_LITE_BATCH_TO_SPACE_ND(type, scalar) \ |
| 128 | type::BatchToSpaceND(GetTensorShape(op_context.input), \ |
| 129 | GetTensorData<scalar>(op_context.input), \ |
| 130 | GetTensorShape(op_context.block_shape), \ |
| 131 | GetTensorData<int32_t>(op_context.block_shape), \ |
| 132 | GetTensorShape(op_context.crops), \ |
| 133 | GetTensorData<int32_t>(op_context.crops), \ |
| 134 | GetTensorShape(op_context.output), \ |
| 135 | GetTensorData<scalar>(op_context.output)) |
| 136 | switch (op_context.input->type) { // Already know in/out types are same. |
| 137 | case kTfLiteFloat32: |
| 138 | if (kernel_type == kReference) { |
| 139 | TF_LITE_BATCH_TO_SPACE_ND(reference_ops, float); |
| 140 | } else { |
| 141 | TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, float); |
| 142 | } |
| 143 | break; |
| 144 | case kTfLiteUInt8: |
| 145 | if (kernel_type == kReference) { |
| 146 | TF_LITE_BATCH_TO_SPACE_ND(reference_ops, uint8_t); |
| 147 | } else { |
| 148 | TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, uint8_t); |
| 149 | } |
| 150 | break; |
| 151 | case kTfLiteInt8: |
| 152 | if (kernel_type == kReference) { |
| 153 | TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int8_t); |
| 154 | } else { |
| 155 | TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int8_t); |
| 156 | } |
| 157 | break; |
| 158 | case kTfLiteInt32: |
| 159 | if (kernel_type == kReference) { |
| 160 | TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int32_t); |
| 161 | } else { |
| 162 | TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int32_t); |
| 163 | } |
| 164 | break; |
| 165 | case kTfLiteInt64: |
| 166 | if (kernel_type == kReference) { |
| 167 | TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int64_t); |
| 168 | } else { |
| 169 | TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int64_t); |
| 170 | } |
| 171 | break; |
| 172 | default: |
| 173 | context->ReportError( |
| 174 | context, "Type %d is currently not supported by BatchToSpace.", |
| 175 | op_context.input->type); |
| 176 | return kTfLiteError; |
nothing calls this directly
no test coverage detected