Example #1
0
 def test_format_bytes(self):
     self.assertEqual(utils.format_bytes(0), "0.00 B")
     self.assertEqual(utils.format_bytes(999), "999.00 B")
     self.assertEqual(utils.format_bytes(1234), "1.23 KB")
     self.assertEqual(utils.format_bytes(1235), "1.24 KB")
     self.assertEqual(utils.format_bytes(999010), "999.01 KB")
     self.assertEqual(utils.format_bytes(1e3), "1.00 KB")
     self.assertEqual(utils.format_bytes(2e6), "2.00 MB")
     self.assertEqual(utils.format_bytes(3e9), "3.00 GB")
     self.assertEqual(utils.format_bytes(4e12), "4.00 TB")
     self.assertEqual(utils.format_bytes(5e20), "500000000.00 TB")
Example #2
0
    Column("Module", format_call_stack),
    "config":
    Column("Config", lambda r: repr(r.module_details.module)),
    "owned_params":
    Column("Module params", format_owned_params),
    "input":
    Column("Input", format_input),
    "output":
    Column("Output", format_output),
    "params_size":
    Column("Param count",
           lambda r: "{:,}".format(utils.tree_size(r.module_details.params)),
           "right"),
    "params_bytes":
    Column(
        "Param bytes", lambda r: utils.format_bytes(
            utils.tree_bytes(r.module_details.params)), "right"),
}

DEFAULT_COLUMNS = ("module", "config", "owned_params", "input", "output",
                   "params_size", "params_bytes")
DEFAULT_FILTERS = ("has_output", )


def tabulate(
    f: Union[Callable[..., Any], hk.Transformed, hk.TransformedWithState],
    *,
    columns: Optional[Sequence[str]] = DEFAULT_COLUMNS,
    filters: Optional[Sequence[str]] = DEFAULT_FILTERS,
    tabulate_kwargs={"tablefmt": "grid"},
) -> Callable[..., str]:
    # pylint: disable=line-too-long
Example #3
0
    "has_params": Filter(lambda r: bool(r.module_details.params)),
}

all_columns = {
    "module": Column("Module", format_call_stack),
    "config": Column("Config", lambda r: repr(r.module_details.module)),
    "owned_params": Column("Module params", format_owned_params),
    "input": Column("Input", format_input),
    "output": Column("Output", format_output),
    "params_size": Column(
        "Param count",
        lambda r: "{:,}".format(utils.tree_size(r.module_details.params)),
        "right"),
    "params_bytes": Column(
        "Param bytes",
        lambda r: utils.format_bytes(utils.tree_bytes(r.module_details.params)),
        "right"),
}

DEFAULT_COLUMNS = tuple(all_columns.keys())
DEFAULT_FILTERS = tuple(all_filters.keys())


def tabulate(
    f: Union[Callable[..., Any], hk.Transformed, hk.TransformedWithState],
    *,
    columns: Optional[Sequence[str]] = DEFAULT_COLUMNS,
    filters: Optional[Sequence[str]] = DEFAULT_FILTERS,
    tabulate_kwargs={"tablefmt": "grid"},
) -> Callable[..., str]:
  # pylint: disable=line-too-long