Skip to content

Commit

Permalink
Rich Summary: uses rich to style the summary console output (#167)
Browse files Browse the repository at this point in the history 
* use rich

* simplify header

* update notebook

* final version

* black .

* random change to fix ci
  • Loading branch information
@cgarciae
cgarciae committed Feb 24, 2021
1 parent 7affe7d commit 690eee1
Show file tree
Hide file tree
Showing 9 changed files with 230 additions and 151 deletions.
63 changes: 38 additions & 25 deletions docs/getting-started/high-level-api.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 3,
"metadata": {
"colab": {},
"colab_type": "code",
Expand Down Expand Up @@ -195,37 +195,50 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓\n",
"┃<span style=\"font-weight: bold\"> Layer </span>┃<span style=\"font-weight: bold\"> Outputs Shape </span>┃<span style=\"font-weight: bold\"> Trainable </span>┃<span style=\"font-weight: bold\"> Non-trainable </span>┃\n",
"┃ ┃ ┃<span style=\"font-weight: bold\"> Parameters </span>┃<span style=\"font-weight: bold\"> Parameters </span>┃\n",
"┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩\n",
"│ Inputs │ (64, 28, 28) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">uint8</span> │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ flatten <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Flatten</span> │ (64, 784) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ linear <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Linear</span> │ (64, 300) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ <span style=\"color: #008000; text-decoration-color: #008000\">235,500</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">942.0 KB</span> │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ relu │ (64, 300) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ linear_1 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Linear</span> │ (64, 100) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ <span style=\"color: #008000; text-decoration-color: #008000\">30,100</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">120.4 KB</span> │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ relu │ (64, 100) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ linear_2 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Linear</span> │ (64, 10) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ <span style=\"color: #008000; text-decoration-color: #008000\">1,010</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">4.0 KB</span> │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│ * <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">MLP</span> │ (64, 10) <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">float32</span> │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼───────────────┤\n",
"│<span style=\"font-weight: bold\"> </span>│<span style=\"font-weight: bold\"> Total </span>│<span style=\"font-weight: bold\"> </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">266,610</span><span style=\"font-weight: bold\"> </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">1.1 MB</span><span style=\"font-weight: bold\"> </span>│<span style=\"font-weight: bold\"> </span>│\n",
"└───────────────────┴───────────────────────┴─────────────────────┴───────────────┘\n",
"<span style=\"font-weight: bold\"> </span>\n",
"<span style=\"font-weight: bold\"> Total Parameters: </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">266,610</span><span style=\"font-weight: bold\"> </span><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; font-weight: bold\">1.1 MB</span><span style=\"font-weight: bold\"> </span>\n",
"</pre>\n"
],
"text/plain": [
"<rich.jupyter.JupyterRenderable at 0x7fd0ac5a21f0>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"╒═══════════════════╤═══════════════════════╤═════════════════════╤═════════════════╕\n",
"│ Layer │ Outputs Shape │ Trainable │ Non-trainable │\n",
"│ │ │ Parameters │ Parameters │\n",
"╞═══════════════════╪═══════════════════════╪═════════════════════╪═════════════════╡\n",
"│ Inputs │ (64, 28, 28) uint8 │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ flatten Flatten │ (64, 784) float32 │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ linear Linear │ (64, 300) float32 │ 235,500 942.0 KB │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ relu │ (64, 300) float32 │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ linear_1 Linear │ (64, 100) float32 │ 30,100 120.4 KB │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ relu │ (64, 100) float32 │ │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ linear_2 Linear │ (64, 10) float32 │ 1,010 4.0 KB │ │\n",
"├───────────────────┼───────────────────────┼─────────────────────┼─────────────────┤\n",
"│ * MLP │ (64, 10) float32 │ │ │\n",
"╘═══════════════════╧═══════════════════════╧═════════════════════╧═════════════════╛\n",
"Total Parameters: 266,610 1.1 MB\n",
"Trainable Parameters: 266,610 1.1 MB\n",
"Non-trainable Parameters: 0\n",
"\n"
]
}
Expand Down
12 changes: 6 additions & 6 deletions elegy/generalized_module/elegy_module_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,11 +99,11 @@ def call(self, x):
assert "2" in lines[11]
assert "8 B" in lines[11]

assert "21" in lines[13]
assert "84 B" in lines[13]
assert "9" in lines[13]
assert "36 B" in lines[13]

assert "9" in lines[14]
assert "36 B" in lines[14]
assert "12" in lines[13]
assert "48 B" in lines[13]

assert "12" in lines[15]
assert "48 B" in lines[15]
assert "21" in lines[16]
assert "84 B" in lines[16]
12 changes: 6 additions & 6 deletions elegy/generalized_module/haiku_module_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,11 +111,11 @@ def f(x):
assert "2" in lines[11]
assert "8 B" in lines[11]

assert "21" in lines[13]
assert "84 B" in lines[13]
assert "9" in lines[13]
assert "36 B" in lines[13]

assert "9" in lines[14]
assert "36 B" in lines[14]
assert "12" in lines[13]
assert "48 B" in lines[13]

assert "12" in lines[15]
assert "48 B" in lines[15]
assert "21" in lines[16]
assert "84 B" in lines[16]
12 changes: 6 additions & 6 deletions elegy/generalized_module/linen_module_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,11 @@ def __call__(self, x):
assert "2" in lines[11]
assert "8 B" in lines[11]

assert "21" in lines[13]
assert "84 B" in lines[13]
assert "9" in lines[13]
assert "36 B" in lines[13]

assert "9" in lines[14]
assert "36 B" in lines[14]
assert "12" in lines[13]
assert "48 B" in lines[13]

assert "12" in lines[15]
assert "48 B" in lines[15]
assert "21" in lines[16]
assert "84 B" in lines[16]

0 comments on commit 690eee1

Please sign in to comment.