Some models are using "TF2" but also using Session and passing
a ConfigProto. This is TF1 code running on the TF2 MLIR bridge.
The TF2 MLIR bridge assumed that this was not possible. This cl
updates the TF2 version of the MLIR bridge to support a ConfigProto
passed in via Session.
Disable MLIR bridge presubmit testing of saved_model_test.py because this fix reveals that the test is actually broken. It is using TF2 control flow but loading a model with Session.
PiperOrigin-RevId: 335915669
Change-Id: Ib50bef389449ce0011878dd50b73856e9c520289
for TensorFlow V1 even when the user overrides Session's default ConfigProto.
In TensorFlow V1, when the user doesn't override Session's default ConfigProto, the ConfigProto
uses the value from TF_XLA_FLAGS (via context.config). However, when the user passes in their
own ConfigProto, enable_mlir_bridge is defaulted to false.
PiperOrigin-RevId: 320661288
Change-Id: I5deacb4e12b6551a57c2496c35830ac3af25fe28
On Windows, Bazel populates this path with `/`s only making proper path
management impossible without sanitizing the path up front. This changes to
accessing the env var through an indirection layer which will fix path problems
on Windows when the codebase is ready to switch over.
PiperOrigin-RevId: 296083765
Change-Id: I26bbaf83ba5e3fafd3ab0a0de08f6cb597b94477
This pass is similar to the MLIR bridge pass but is specific to TensorFlow V1 Graphs. In addition, this pass will only trigger on non function graphs as they are already handled by the MLIR bridge pass.
PiperOrigin-RevId: 294521035
Change-Id: If80bb613b51ff9fd8b2dad61b4357a55452a4093
This reworks the MLIR bridge pass from an GraphOptimizationPass to a FunctionOptimizationPass for the function library runtime. Control ret node names are also passed into and updated from the registry pass run (MLIR to Graph).
PiperOrigin-RevId: 293703916
Change-Id: I644765f4c5294962450cacb9c748582a579989c5
* This fixes a correctness problem where we were relying on the variable's "name" attribute, which is merely advisory.
* Should match the load heuristics in Python by reading the TrackableObjectGraph from the checkpoint, re-associating it with the SavedObjectGraph and using that to restore the variable.
* Has a side-effect of eliminating the dependency of the importer on the CPU runtime and kernels, which should reduce necessary dependencies to compile by tens of megabytes.
Open questions:
* I created a new bundle_v2.h because V2 is so substantially different from V1. Also, it has a different dependency surface area and I therefore wanted to use it as a different (lighter) library. Advise if you would like this organized differently. I tried to factor it so that a future C++ SavedModel loader/runner could be built on this as well (with some incremental work).
* There didn't seem to be a facility for generating the testdata SavedModels. I ended up just writing a small standalone python script that I can invoke from the command line to generate the new ones. It isn't wired up in any way but it should be possible to do so at some point (I'm not familiar enough with the test infra on this side of the tree, so please advise).
* I could see factoring the SavedModelV2Bundle::RestoreObjects method into a dedicated facility for loading arbitrary checkpoints. Not sure what the roadmap here is so just kept is simple for now.
PiperOrigin-RevId: 282439157
Change-Id: I685df47ab621917eed270319cde220498f191b74
This moves a function from tpu_rewrite_device_util for getting devices, and a function for adding devices from mlir_bridge_pass, into its own util.
PiperOrigin-RevId: 273610226
This takes every device in the device set of the GraphOptimizationPassOptions and converts their parsed name into a string. The device strings are then added to the module op as an ArrayAttr<StringAttr> attribute. Devices are necessary for certain passes (e.g. TPURewritePass).
PiperOrigin-RevId: 272029573
This pass is implemented as early as possible in the pipeline and allows to
implement MLIR-based graph transformation. At the moment it is enabled by
an experimental session config flag only, and won't do anything useful
but round-tripping the Graph to MLIR.
PiperOrigin-RevId: 266036600
