Repository: google/flax

Issue:
`Module.tabulate` does not support `Enum` inputs
An internal request to add `Enum` input support for `Module.tabulate`
```
class Net(nn.Module):
  @nn.compact
  def __call__(self, inputs):
    x = inputs['x']
    x = nn.Dense(features=2)(x)
    return jnp.sum(x)

inputs = {'x': jnp.ones((1, 1))}
print(Net().tabulate(jax.random.key(0), inputs=inputs)) # this input works fine

class InputEnum(str, enum.Enum):
  x = 'x'

inputs = {InputEnum.x: jnp.ones((1, 1))}
print(Net().tabulate(jax.random.key(0), inputs)) # this Enum input throws an error
```
```
RepresenterError: ('cannot represent an object', <InputEnum.x: 'x'>)
```

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) Report
Attention: `4 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`2f6e5ff`)](https://app.codecov.io/gh/google/flax/commit/2f6e5ff94675acfe88adfd48ec939499aa029874?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 52.10% compared to head [(`094cedf`)](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 52.09%.

| [Files](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) | Patch % | Lines |
|---|---|---|
| [flax/linen/summary.py](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google#diff-ZmxheC9saW5lbi9zdW1tYXJ5LnB5) | 33.33% | [4 Missing :warning: ](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #3485      +/-   ##
==========================================
- Coverage   52.10%   52.09%   -0.02%     
==========================================
  Files          91       91              
  Lines       10968    10973       +5     
==========================================
+ Hits         5715     5716       +1     
- Misses       5253     5257       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/google/flax/pull/3485?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).

chiamp: @cgarciae I moved it to `_normalize_structure` instead

Patch:
diff --git a/docs/guides/data_preprocessing/index.rst b/docs/guides/data_preprocessing/index.rst
index 03087cee..47ca8b56 100644
--- a/docs/guides/data_preprocessing/index.rst
+++ b/docs/guides/data_preprocessing/index.rst
@@ -5,4 +5,3 @@ Data preprocessing
    :maxdepth: 1
 
    full_eval
-   loading_datasets
diff --git a/docs/guides/data_preprocessing/loading_datasets.ipynb b/docs/guides/data_preprocessing/loading_datasets.ipynb
deleted file mode 100644
index dd494f06..00000000
--- a/docs/guides/data_preprocessing/loading_datasets.ipynb
+++ /dev/null
@@ -1,296 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Loading datasets\n",
-    "\n",
-    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/data_preprocessing/loading_datasets.ipynb)\n",
-    "\n",
-    "A neural net written in Jax+Flax expects its input data as `jax.numpy` array instances. Therefore, loading a dataset from any source is as simple as converting it to `jax.numpy` types and reshaping it to the appropriate dimensions for your network.\n",
-    "\n",
-    "As an example, this guide demonstrates how to import [MNIST](http://yann.lecun.com/exdb/mnist/) using the APIs from Torchvision, Tensorflow, and Hugging Face. We'll load the whole dataset into memory. For datasets that don't fit into memory the process is analogous but should be done in a batchwise fashion.\n",
-    "\n",
-    "The MNIST dataset consists of greyscale images of 28x28 pixels of handwritten digits, and has a designated 60k/10k train/test split. The task is to predict the correct class (digit 0, ..., 9) of each image.\n",
-    "\n",
-    "Assuming a CNN-based classifier, the input data should have shape `(B, 28, 28, 1)`, where the trailing singleton dimension denotes the greyscale image channel.\n",
-    "\n",
-    "The labels are simply the integer denoting the digit corresponding to the image. Labels should therefore have shape `(B,)`, to enable loss computation with [`optax.softmax_cross_entropy_with_integer_labels`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "import jax.numpy as jnp"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading from `torchvision.datasets`"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "import torchvision"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "def get_dataset_torch():\n",
-    "    mnist = {\n",
-    "        'train': torchvision.datasets.MNIST('./data', train=True, download=True),\n",
-    "        'test': torchvision.datasets.MNIST('./data', train=False, download=True)\n",
-    "    }\n",
-    "\n",
-    "    ds = {}\n",
-    "\n",
-    "    for split in ['train', 'test']:\n",
-    "        ds[split] = {\n",
-    "            'image': mnist[split].data.numpy(),\n",
-    "            'label': mnist[split].targets.numpy()\n",
-    "        }\n",
-    "\n",
-    "        # cast from np to jnp and rescale the pixel values from [0,255] to [0,1]\n",
-    "        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n",
-    "        ds[split]['label'] = jnp.int16(ds[split]['label'])\n",
-    "\n",
-    "        # torchvision returns shape (B, 28, 28).\n",
-    "        # hence, append the trailing channel dimension.\n",
-    "        ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\n",
-    "\n",
-    "    return ds['train'], ds['test']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "outputId": "be39b756-d13e-4380-b99e-a5cbf61458cc",
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(60000, 28, 28, 1) float32\n",
-      "(60000,) int16\n",
-      "(10000, 28, 28, 1) float32\n",
-      "(10000,) int16\n"
-     ]
-    }
-   ],
-   "source": [
-    "train, test = get_dataset_torch()\n",
-    "print(train['image'].shape, train['image'].dtype)\n",
-    "print(train['label'].shape, train['label'].dtype)\n",
-    "print(test['image'].shape, test['image'].dtype)\n",
-    "print(test['label'].shape, test['label'].dtype)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading from `tensorflow_datasets`"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "import tensorflow_datasets as tfds"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "def get_dataset_tf():\n",
-    "    mnist = tfds.builder('mnist')\n",
-    "    mnist.download_and_prepare()\n",
-    "\n",
-    "    ds = {}\n",
-    "\n",
-    "    for split in ['train', 'test']:\n",
-    "        ds[split] = tfds.as_numpy(mnist.as_dataset(split=split, batch_size=-1))\n",
-    "\n",
-    "        # cast to jnp and rescale pixel values\n",
-    "        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n",
-    "        ds[split]['label'] = jnp.int16(ds[split]['label'])\n",
-    "\n",
-    "    return ds['train'], ds['test']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "outputId": "25d2c468-cbc8-4971-a738-1295ce8c6f16",
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(60000, 28, 28, 1) float32\n",
-      "(60000,) int16\n",
-      "(10000, 28, 28, 1) float32\n",
-      "(10000,) int16\n"
-     ]
-    }
-   ],
-   "source": [
-    "train, test = get_dataset_tf()\n",
-    "print(train['image'].shape, train['image'].dtype)\n",
-    "print(train['label'].shape, train['label'].dtype)\n",
-    "print(test['image'].shape, test['image'].dtype)\n",
-    "print(test['label'].shape, test['label'].dtype)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading from 🤗 Hugging Face `datasets`"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "#!pip install datasets # datasets isn't preinstalled on Colab; uncomment to install\n",
-    "from datasets import load_dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "def get_dataset_hf():\n",
-    "    mnist = load_dataset(\"mnist\")\n",
-    "\n",
-    "    ds = {}\n",
-    "\n",
-    "    for split in ['train', 'test']:\n",
-    "        ds[split] = {\n",
-    "            'image': np.array([np.array(im) for im in mnist[split]['image']]),\n",
-    "            'label': np.array(mnist[split]['label'])\n",
-    "        }\n",
-    "\n",
-    "        # cast to jnp and rescale pixel values\n",
-    "        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n",
-    "        ds[split]['label'] = jnp.int16(ds[split]['label'])\n",
-    "\n",
-    "        # append trailing channel dimension\n",
-    "        ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\n",
-    "\n",
-    "    return ds['train'], ds['test']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {
-    "outputId": "b026b33f-3bdd-4d26-867c-49400fff1c96",
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(60000, 28, 28, 1) float32\n",
-      "(60000,) int16\n",
-      "(10000, 28, 28, 1) float32\n",
-      "(10000,) int16\n"
-     ]
-    }
-   ],
-   "source": [
-    "train, test = get_dataset_hf()\n",
-    "print(train['image'].shape, train['image'].dtype)\n",
-    "print(train['label'].shape, train['label'].dtype)\n",
-    "print(test['image'].shape, test['image'].dtype)\n",
-    "print(test['label'].shape, test['label'].dtype)"
-   ]
-  }
- ],
- "metadata": {
-  "jupytext": {
-   "formats": "ipynb,md:myst"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/docs/guides/data_preprocessing/loading_datasets.md b/docs/guides/data_preprocessing/loading_datasets.md
deleted file mode 100644
index bddbc8cb..00000000
--- a/docs/guides/data_preprocessing/loading_datasets.md
+++ /dev/null
@@ -1,159 +0,0 @@
----
-jupytext:
-  formats: ipynb,md:myst
-  text_representation:
-    extension: .md
-    format_name: myst
-    format_version: 0.13
-    jupytext_version: 1.13.8
----
-
-# Loading datasets
-
-[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/data_preprocessing/loading_datasets.ipynb)
-
-A neural net written in Jax+Flax expects its input data as `jax.numpy` array instances. Therefore, loading a dataset from any source is as simple as converting it to `jax.numpy` types and reshaping it to the appropriate dimensions for your network.
-
-As an example, this guide demonstrates how to import [MNIST](http://yann.lecun.com/exdb/mnist/) using the APIs from Torchvision, Tensorflow, and Hugging Face. We'll load the whole dataset into memory. For datasets that don't fit into memory the process is analogous but should be done in a batchwise fashion.
-
-The MNIST dataset consists of greyscale images of 28x28 pixels of handwritten digits, and has a designated 60k/10k train/test split. The task is to predict the correct class (digit 0, ..., 9) of each image.
-
-Assuming a CNN-based classifier, the input data should have shape `(B, 28, 28, 1)`, where the trailing singleton dimension denotes the greyscale image channel.
-
-The labels are simply the integer denoting the digit corresponding to the image. Labels should therefore have shape `(B,)`, to enable loss computation with [`optax.softmax_cross_entropy_with_integer_labels`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels).
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-import numpy as np
-import jax.numpy as jnp
-```
-
-## Loading from `torchvision.datasets`
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-import torchvision
-```
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-def get_dataset_torch():
-    mnist = {
-        'train': torchvision.datasets.MNIST('./data', train=True, download=True),
-        'test': torchvision.datasets.MNIST('./data', train=False, download=True)
-    }
-
-    ds = {}
-
-    for split in ['train', 'test']:
-        ds[split] = {
-            'image': mnist[split].data.numpy(),
-            'label': mnist[split].targets.numpy()
-        }
-
-        # cast from np to jnp and rescale the pixel values from [0,255] to [0,1]
-        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255
-        ds[split]['label'] = jnp.int16(ds[split]['label'])
-
-        # torchvision returns shape (B, 28, 28).
-        # hence, append the trailing channel dimension.
-        ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)
-
-    return ds['train'], ds['test']
-```
-
-```{code-cell} ipython3
-:outputId: be39b756-d13e-4380-b99e-a5cbf61458cc
-:tags: [skip-execution]
-
-train, test = get_dataset_torch()
-print(train['image'].shape, train['image'].dtype)
-print(train['label'].shape, train['label'].dtype)
-print(test['image'].shape, test['image'].dtype)
-print(test['label'].shape, test['label'].dtype)
-```
-
-## Loading from `tensorflow_datasets`
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-import tensorflow_datasets as tfds
-```
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-def get_dataset_tf():
-    mnist = tfds.builder('mnist')
-    mnist.download_and_prepare()
-
-    ds = {}
-
-    for split in ['train', 'test']:
-        ds[split] = tfds.as_numpy(mnist.as_dataset(split=split, batch_size=-1))
-
-        # cast to jnp and rescale pixel values
-        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255
-        ds[split]['label'] = jnp.int16(ds[split]['label'])
-
-    return ds['train'], ds['test']
-```
-
-```{code-cell} ipython3
-:outputId: 25d2c468-cbc8-4971-a738-1295ce8c6f16
-:tags: [skip-execution]
-
-train, test = get_dataset_tf()
-print(train['image'].shape, train['image'].dtype)
-print(train['label'].shape, train['label'].dtype)
-print(test['image'].shape, test['image'].dtype)
-print(test['label'].shape, test['label'].dtype)
-```
-
-## Loading from 🤗 Hugging Face `datasets`
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-#!pip install datasets # datasets isn't preinstalled on Colab; uncomment to install
-from datasets import load_dataset
-```
-
-```{code-cell} ipython3
-:tags: [skip-execution]
-
-def get_dataset_hf():
-    mnist = load_dataset("mnist")
-
-    ds = {}
-
-    for split in ['train', 'test']:
-        ds[split] = {
-            'image': np.array([np.array(im) for im in mnist[split]['image']]),
-            'label': np.array(mnist[split]['label'])
-        }
-
-        # cast to jnp and rescale pixel values
-        ds[split]['image'] = jnp.float32(ds[split]['image']) / 255
-        ds[split]['label'] = jnp.int16(ds[split]['label'])
-
-        # append trailing channel dimension
-        ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)
-
-    return ds['train'], ds['test']
-```
-
-```{code-cell} ipython3
-:outputId: b026b33f-3bdd-4d26-867c-49400fff1c96
-:tags: [skip-execution]
-
-train, test = get_dataset_hf()
-print(train['image'].shape, train['image'].dtype)
-print(train['label'].shape, train['label'].dtype)
-print(test['image'].shape, test['image'].dtype)
-print(test['label'].shape, test['label'].dtype)
-```
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
index 56fcf205..7fd76624 100644
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -15,6 +15,7 @@
 """Flax Module summary library."""
 
 import dataclasses
+import enum
 import io
 from abc import ABC, abstractmethod
 from types import MappingProxyType
@@ -709,12 +710,17 @@ def _normalize_structure(obj):
   if isinstance(obj, (tuple, list)):
     return tuple(map(_normalize_structure, obj))
   elif isinstance(obj, Mapping):
-    return {k: _normalize_structure(v) for k, v in obj.items()}
+    return {
+      _normalize_structure(k): _normalize_structure(v) for k, v in obj.items()
+    }
   elif dataclasses.is_dataclass(obj):
     return {
       f.name: _normalize_structure(getattr(obj, f.name))
       for f in dataclasses.fields(obj)
     }
+  elif isinstance(obj, enum.Enum):
+    # `yaml.safe_dump` does not support Enum key types so extract the underlying value
+    return obj.value
   else:
     return obj
 


Test patch:
diff --git a/tests/linen/summary_test.py b/tests/linen/summary_test.py
index 6d4be643..b0e651da 100644
--- a/tests/linen/summary_test.py
+++ b/tests/linen/summary_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 from typing import List
 
 import jax
@@ -735,6 +736,26 @@ class SummaryTest(absltest.TestCase):
     lines = rep.splitlines()
     self.assertIn('Total Parameters: 50', lines[-2])
 
+  def test_tabulate_enum(self):
+    class Net(nn.Module):
+      @nn.compact
+      def __call__(self, inputs):
+        x = inputs['x']
+        x = nn.Dense(features=2)(x)
+        return jnp.sum(x)
+
+    class InputEnum(str, enum.Enum):
+      x = 'x'
+
+    inputs = {InputEnum.x: jnp.ones((1, 1))}
+    # test args
+    lines = Net().tabulate(jax.random.key(0), inputs).split('\n')
+    self.assertIn('x: \x1b[2mfloat32\x1b[0m[1,1]', lines[5])
+    # test kwargs
+    lines = Net().tabulate(jax.random.key(0), inputs=inputs).split('\n')
+    self.assertIn('inputs:', lines[5])
+    self.assertIn('x: \x1b[2mfloat32\x1b[0m[1,1]', lines[6])
+
 
 if __name__ == '__main__':
   absltest.main()