Intellij IDEA / PyCharmでkerasのコード補完が効かない

Jetbrains製のIDEであるIntellij IDEAやPyCharmを使っていた時、tensorflow.keras以下がコード補完できなかった。

標準ライブラリや他のライブラリは補完されており、tensorflowもkeras以外は補完されていたので、tf.keras特有の問題ではないかと思った。

tensorflowの__init__.pyを編集することで解決できた。

まずは site-packages/tensorflow/__init__.py を開く。

import tensorflow の tensorflow にカーソルを合わせて ⌘B を押せば飛ぶことが出来る。

387行目辺り。

_keras_module = "keras.api._v2.keras"
keras = _LazyLoader("keras", globals(), _keras_module)
_module_dir = _module_util.get_parent_dir_for_name(_keras_module)
if _module_dir:
  _current_module.__path__ = [_module_dir] + _current_module.__path__
setattr(_current_module, "keras", keras)

以下のように書き換え。

import typing as _typing
if _typing.TYPE_CHECKING:
  from keras.api._v2 import keras
else:
  _keras_module = "keras.api._v2.keras"
  keras = _LazyLoader("keras", globals(), _keras_module)
  _module_dir = _module_util.get_parent_dir_for_name(_keras_module)
  if _module_dir:
    _current_module.__path__ = [_module_dir] + _current_module.__path__
  setattr(_current_module, "keras", keras)

編集後IDEを再起動し、kerasのコード補完が使えるようになった。

参考

PyCharm doesn't resolve anything under tensorflow.keras · Issue #53144 · tensorflow/tensorflow

System information OS Platform and Distribution: Windows 10 TensorFlow installed from: binary TensorFlow version: 2.7.0 ...