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")
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
"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