No autocompletion for return value of `tensorflow_datasets.load(...)`

/!\ PLEASE INCLUDE THE FULL STACKTRACE AND CODE SNIPPET

Short description
As title.

Environment information

• Operating System: macOS 12.2.1
• Python version: 3.8.12
• tensorflow related versions:

tensorflow-datasets      4.5.2
tensorflow-macos         2.8.0
tensorflow-metadata      1.7.0
tensorflow-metal         0.4.0

• Does the issue still exists with the last tfds-nightly package (pip install --upgrade tfds-nightly) ?

Sorry, I don't want to try this in case something cannot be reverted.

Reproduction instructions

import tensorflow as tf
import tensorflow_datasets as tfds

def _test_perf():

    (ds_train, ds_test), ds_info = tfds.load(
        "mnist",
        split=["train", "test"],
        shuffle_files=True,
        as_supervised=True,
        with_info=True,
    )

    ds_train.

If you share a colab, make sure to update the permissions to share it.

I didn't use colab.

Link to logs
no.

Expected behavior
I expect to see these functions in the completion menu when ds_train. typed:

ds_train = ds_train.map(normalize_img, num_parallel_calls=AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits["train"].num_examples)
ds_train = ds_train.batch(BATCH_SIZE)
ds_train = ds_train.prefetch(AUTOTUNE)

Additional context
no. Ping me if you need some specific info.

The entry point is tensorflow_datasets.load in the provided Python reproduction; start by reproducing the missing completion on ds_train and inspect the returned-value typing. Done means the returned dataset exposes the expected map, cache, shuffle, batch, and prefetch methods in autocomplete.