Decoupling tensorflow from keras in stubs

[Avasam]

Whilst working on #11696, I was having issues referencing the source code. After reading the doc and a bit of investigation, I realized why:

The runtime of https://github.com/python/typeshed/tree/main/stubs/tensorflow/tensorflow/keras actually are just references to the keras source code at https://github.com/keras-team/keras/tree/master/keras (import keras), and is generated / injected in import tensorflow.keras by https://github.com/keras-team/keras/blob/6454a4888a494c20ab0ea1dc6912cbcc13c5f940/pip_build.py#L62

Even tensorflow's own doc links back to karas source code https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/Callback

>>> from keras import callbacks
2024-04-03 23:28:35.387416: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-04-03 23:28:36.959410: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
>>> callbacks.__file__
'...\\typeshed\\.venv\\lib\\site-packages\\keras\\callbacks\\__init__.py'
>>> from tensorflow.keras import callbacks 
>>> callbacks.__file__
'...\\typeshed\\.venv\\lib\\site-packages\\keras\\callbacks\\__init__.py'

In other words, keras "pollutes" the tensorflow namespace! I don't know if the type stubs specifications allow two distributions to write to the same stub-only package, which would allow us to perfectly reflect what's truly happening. Second best thing would be to have stubs for keras and keep tensorflow.keras in the tensorflow stubs (since tensorflow is dependent on keras anyway)

CC @hoel-bagard & @hmc-cs-mdrissi

[hmc-cs-mdrissi]

The actual way source code looks for tensorflow is not goal of stubs. Tensorflow does a lot of dynamic import behaviors and re-exports symbols using decorators (tf_export) in other modules. The exact placement is unstable and changes regularly in very breaking ways. Keras source code has changed rather drastically in past several versions with folders being added to wheel that may not even exist in GitHub repository either. The wheel vs repository being inconsistent was done I think to hide implementation details of source organization even further.

Instead the goal of the stubs is to model the documentation and public api surface. tf.keras.optimizers.Optimizer is publicly documented as the type and works at runtime. Source may have it in keras package as keras.optimizers.Optimizer. Or maybe it’s keras.src.optimizers.Optimizer in the wheel but not repository. And tensorflow/keras were separate sources several years ago, got merged into one repo, and then got separated into two repos again years later. I do not want stubs to capture organizational implementation choices not captured by documentation/stability expectations of the library.

If you would like to add keras stubs folder for publicly documented members with keras root and then re-export them in tensorflow other publicly documented aliases I would be comfortable and supportive of that. I push back strongly on trying to mimic source code implementation though and any details not part of documentation.

[Avasam]

Thank you, that's exactly the kind of comments and historical insight I was looking for.

If you would like to add keras stubs folder for publicly documented members with keras root and then re-export them in tensorflow other publicly documented aliases I would be comfortable and supportive of that.

I'd like to do that yes. Given how tightly coupled keras seems to be to tensorflow, I wouldn't mind if the package was exposed as a top-level under tensorflow stubs. (With only the types re-exported by tensorflow)

Since the source code actually exists on disk (not a c-extension or dynamically generated module), I do believe there's value added for both devs using the stubs and typeshed contributors, to be able to access the source from within their working environment / IDE

[hmc-cs-mdrissi]

On Keras/Tensorflow coupling, originally Keras was high level wrapper of tensorflow and theano. Theano is an older ml library before tensorflow. Tensorflow was mostly upgrade over theano so several years ago theano stop being maintained and recommended its users move to other ml libraries. After this for years Keras only supported tensorflow.

Recently (last couple months), Keras 3 came out and now Keras is trying to be high level wrapper of multiple major ml libraries again. This time Keras supports tensorflow, Jax, and PyTorch.

My own personal experience is mostly focused on tensorflow though. I have moderate experience doing projects in PyTorch and have never used jax.

[PabloLION]

I've looked for available stub files for Keras3.3 and TF2.16, but I had no luck. In my VSCode with pylance(pyright), there are always many warnings and errors on the keras typing. I no option but to ignore them. I tried to use my own stub file but the decoupling seems complicated.
Hope someone can fix this soon

[hmc-cs-mdrissi]

The current status is pip install types-tensorflow. That will give you all that is available today.

There are many apis not covered today and they continue to gradually be added at pace of open source. Tensorflow apis tend to be better than keras specific ones and I know keras 3 is major change from keras 2. My codebase crashes if I try to upgrade and various dependencies of mine conflict. So the codebase I work on is will probably stay Keras 2 for a while.

I’ll happily review keras 3 prs just not an area I’m likely to work on soon.