Restructure Python wheels following TF[Lite] importer changes

[ScottTodd]

Following the work in #13037 and other issues/PRs, our TensorFlow and TensorFlow Lite importers are now pure Python packages/scripts, with no direct C++ dependencies on TensorFlow. As a result, our iree_tools_tf and iree_tools_tflite-*.whl Python wheels are now significantly simpler to build and smaller in size (~50MB before -> ~30KB after). We also no longer have a iree_tools_xla-*.whl package.

I'd like to take this chance to see if we can further simplify or restructure our released Python packages.

I was also looking at adding badges to our repo README like these:

for those, it would be nice to have a smaller list (ideally just one package).

Some ideas:

A: Merge all iree_tools_* packages into iree_compiler

Include the Python sources for the TF/TFLite tools in the main package, but make them optional by adding checks around code like this:
https://github.com/openxla/iree/blob/060616cea34d062434097db155c458edb814e235/integrations/tensorflow/python_projects/iree_tf/iree/tools/tf/scripts/iree_import_tf/__main__.py#L60-L72

try:
  import tensorflow as tf
except ImportError as e:
  # print error / instructions for installing

Pros:

• Just one package to install for the IREE compiler and import tools

Cons:

• Adding TensorFlow deps (even optional ones) to the core compiler is bad layering
• Optional imports can't enforce dependency version requirements (tf-nightly > some version)

B: Merge just iree_tools_tf and iree_tools_tflite

These packages are small and similar enough that they could be merged. If they have the same requirements (tensorflow/tf-nightly), is there any harm in merging them?

[stellaraccident]

Option B, merging both into iree-tools-tf, seems the right thing. We can't introduce TF dependencies to the compiler.

[hcindyl]

Sorry just caught up with this since the IREE nightly release finally has this change.

We need to update https://openxla.github.io/iree/getting-started/tflite/#using-command-line-tools if it is assumed tensorflow python package is installed. The alternative is to add tensorflow in setup.py so it is included when the wheels are installed.

The iree-import-tflite command line behavior at iree-tools-tflite-20230509.514

Traceback (most recent call last):
  File ".../iree-import-tflite", line 5, in <module>
    from iree.tools.tflite.scripts.iree_import_tflite.__main__ import main
  File ".../iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 14, in <module>
    from tensorflow.python.pywrap_mlir import experimental_tflite_to_tosa_bytecode
ModuleNotFoundError: No module named 'tensorflow'

Please advise on the proper flow to set the wheels up.

[ScottTodd]

The current source of truth for what dependencies is required is https://github.com/openxla/iree/blob/main/integrations/tensorflow/test/requirements.txt (as of May 9 2023, that is tf-nightly==2.13.0.dev20230501). Once the Python APIs (experimental_tflite_to_tosa_bytecode and experimental_convert_saved_model_to_mlir) are stable / available in non-nightly TF releases, setup should stabilize.

[hcindyl]

The current source of truth for what dependencies is required is https://github.com/openxla/iree/blob/main/integrations/tensorflow/test/requirements.txt (as of May 9 2023, that is tf-nightly==2.13.0.dev20230501). Once the Python APIs (experimental_tflite_to_tosa_bytecode and experimental_convert_saved_model_to_mlir) are stable / available in non-nightly TF releases, setup should stabilize.

Understood. In that case, in the webdoc https://openxla.github.io/iree/getting-started/tflite/#prerequisites should also include tensorflow (or point to the requirments.txt). Is that correct?

[ScottTodd]

Yep. Maybe we should just include tf-nightly in the recommended pip install instructions.

[ScottTodd]

I'm wondering if we'll want to keep the iree-tools-tf and iree-tools-tflite packages separate now that TensorFlow Lite got renamed to LiteRT back in September: https://developers.googleblog.com/en/tensorflow-lite-is-now-litert/. The TFLite python packages got replaced with https://pypi.org/project/ai-edge-litert/ and code got moved from TF to https://github.com/google-ai-edge/litert.

I haven't seen or heard that much about either framework lately, but TFLite/LiteRT is still in use in a few spots (like #19402).

[ScottTodd]

I'm now thinking option A ("Merge all iree_tools_* packages into iree_compiler") could work well, using extras like [tf] and [litert]. We have some precedent with the [onnx] extra and the iree-import-onnx tool. The deps at this point are all optional Python dependencies, not C++ source dependencies.