kd.inspect.get_colab_model_overview

kauldron.inspect.get_colab_model_overview(

*,

model: nn.Module,

train_ds: data.Pipeline,

ds_sharding: sharding.ShardingTree,

model_config: konfig.ConfigDict | None = None,

rngs: dict[str, kd_random.PRNGKey],

) → pd.DataFrame

Return pd.DataFrame for displaying the model params, inputs,…