TensorFlowのKeras関連モジュールの補完が効かない

起こったこと

概ねタイトル通りだが、詳しく状況を記載していく。
筆者は普段、VSCodeのPylanceのanalyzerおよび補完機能を使用してPythonコードを書いている。ところが、TensorFlowのバージョン更新を境にtf.keras名前空間に属するモジュールの補完が効かなくなってしまった（その他のライブラリや関数では補完が効くので、VSCodeの破損ではない）。
tf.kerasにはlayersやmixed_precisionなども属しているため、影響はかなり致命的である。
ちなみに、詳しくは後述するが、TensorFlowのバージョンやコードの書き方によって補完の有無は大きく左右されるようなので、「同じように悩んでいた」という方もいれば、「TensorFlow使ってますが事もありませんよ？何かの間違いでは？」という方もいらっしゃるかもしれない。

原因

どうやら、アップデートで導入されたlazyloaderという仕組みが原因らしい。
これは簡単に説明すると、あるモジュールを使用する直前にimportする仕組みのことだ。このような動作にすることで、プログラムが利用するリソース（メモリ等）を極力節約しようというのが狙いのようだ。一瞬ピンときませんでしたが、要はPyTorchにtensorflow-datasetsを使うとtfdsのimport直後にVRAM全部持ってかれてPyTorchが死ぬやつですね、わかりますわかります。
そして、このlazyloaderで読み込みが保留されたライブラリは、VSCodeでも"なかったこと"にされて補完がなされないらしい。別にVSCodeのほうで補完のために読み込む分にはリソース面の問題などは特に発生しないはずなので、あまりいい仕様ではないように筆者には思えるが…。しかしながら、残念なことに2023年12月現在のVSCodeはこのような挙動となっている。

ということで、この厄介な現象が発生するのはlazyloaderが導入されたtf2.6以降である（コードのチェックおよび実際のインストールで確認済み）。
ただし、不思議なことにlazyloader導入以降のTensorFlowでも補完のされ方が異なる。具体的には、tf2.6/2.7では少し補完が弱くなる程度なのに対し、tf2.8~2.10ではlayers関連の補完が全く効かない。
さらには、書き方によっても補完の結果が異なる。from tensorflow.keras import layersでは補完がなされないものの、import tensorflow.keras.layers as layersでは問題なく補完が行われる。
残念ながら筆者はコード解析に詳しくないので、なぜこのような挙動の違いが生まれるのかは理解していない。

対策

残念ながら根本的な解決策はない…と思われる。が、筆者が最も実用的と感じたものから順に3つほど紹介する。かなり悩んでいるので、根本的な解決策をご存じの方がいらっしゃれば是非ともコメントで指摘して欲しい。

1. python.analysis.extraPathsの指定

VSCodeには現在開かれていない外部の別のコードを補完のために参照する機能がある。これは本来、自作のライブラリなどを利用するためにある機能なのだと思われるが…これを使うとなんといい感じに補完を復活させることができる。
具体的にはまず、tf2.5がインストールされたpythonの仮想環境を用意する。筆者が使用しているAnacondaではGUIでポチポチするだけで（本記事の趣旨とは逸れるので手順は割愛）簡単に仮想環境を作成できる。注意点として、新しすぎるPythonはtf2.5にサポートされていないので、仮想環境のpythonのバージョンは3.8や3.9あたりが良いだろう。
そして、VSCodeの設定ファイルの"python.analysis.extraPaths"でこの仮想環境のsite-packages（インストールしたpythonライブラリが入っているところ）を指定する。
文字で書くとやや難解なのでより具体的な手順を示すと

①VSCodeで「ctrl」と「,」を同時押しした後に右上の「設定（JSON）を開く」を押し
②"python.analysis.extraPaths": ["C:\\Users\\XXX\\anaconda3\\envs\\tf\\Lib\\site-packages"],という一行を設定ファイルに付け加える
（これはAnacondaを使用している私の場合はこうなったいう話であり、仮想環境の作成方法によって当然仮想環境のパスも異なってくる…はず）

このように設定すると、なんと仮想環境のtf2.5の補完をtf2.10でうまいこと利用できてしまう。
原理的に、この方法ならば昔と同様の補完が完全に復活する。もっとも、VSCodeの機能を本来想定されていないであろう使い方をしているというのが少々気がかりではあるが…

2. TensorFlowのバージョンを下げる

2番目の方法はシンプルで、TensorFlowのバージョンアップが原因でトラブルが発生してしまうのならば、TensorFlowのバージョンを上げなければよい。
TensorFlowのバージョンをきちんと補完が効くv2.5.3に固定してしまおう。
ただ、tf2.5は2021年ごろと結構昔のリリースであるというのがシンプルなこの案の最大の欠点であろう。たとえば実際、筆者はtf2.7や2.10では問題なく動く2入力モデルがtf2.5では動かないというトラブルに遭遇した。

3. 補完してもらえる書き方を毎回心がける

これは筆者としてはもっとも気が進まない解決策だ。
import tensorflow.keras.layers as layersといったようなきちんと補完が効く書き方を必ず心がければよい。
ただ、本来は何ら問題がないであろう書き方が使えないというのはなかなかにストレスがたまるし、何よりも原理が謎なので全ての場面をこれで乗り切れるのかは正直よくわからない。
コーディング規約等に厳しくないプロジェクトではもしかしたら一考の余地があるかもしれない。

まとめ

3つほど解決法を考えたが、どれも欠点があって完璧ではないという印象だ。
筆者は1番や2番の方法を用いることで、今のところは大きな問題なくコーディングできているが、やはり根本的な解決がなされることを祈るばかりである。

環境

Windows11 Pro
VSCode+Pylance拡張
TensorFlow 2.10.1