Exposing TFLite entry points with Python

Corbin_Foucart · October 10, 2023, 8:21am

I’m interested in hooking into some of the simple TFLite reference functions from Python directly in order to explore their behavior. For example, if I’d like to call the function

MultiplyByQuantizedMultiplier

defined here:

github.com

tensorflow/tensorflow/blob/115b2dbfbe97457c116ecd188972e00e484a1d00/tensorflow/lite/kernels/internal/common.h#L262


      
              for (int i = 0; i < bias_size; i++) {
                array_data[array_offset + i] = ActivationFunctionWithMinMax(
                    array_data[array_offset + i] + bias_data[i], clamp_min, clamp_max);
              }
            }
          #endif
          }
          
          // Single-rounding MultiplyByQuantizedMultiplier
          #if TFLITE_SINGLE_ROUNDING
          inline int32_t MultiplyByQuantizedMultiplier(int32_t x,
                                                       int32_t quantized_multiplier,
                                                       int shift) {
            TFLITE_DCHECK(quantized_multiplier >= 0);
            TFLITE_DCHECK(shift >= -31 && shift <= 30);
          
            const int64_t total_shift = 31 - shift;
            const int64_t round = static_cast<int64_t>(1) << (total_shift - 1);
            int64_t result = x * static_cast<int64_t>(quantized_multiplier) + round;
            result = result >> total_shift;

from Python, is there a relatively straightforward way to add a new cmake target to expose this function directly with pybind11? I’m building tensorflow from source and compiling it using cmake

tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh native

Any guidance would be appreciated!