{"id":16814499,"url":"https://github.com/findus23/jax-array-info","last_synced_at":"2025-07-15T19:31:47.340Z","repository":{"id":192611644,"uuid":"687063407","full_name":"Findus23/jax-array-info","owner":"Findus23","description":"Debugging tool to print information (especially sharding) about jax arrays","archived":false,"fork":false,"pushed_at":"2024-11-14T16:09:11.000Z","size":29,"stargazers_count":4,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2024-11-14T16:39:53.584Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Findus23.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-09-04T14:23:38.000Z","updated_at":"2024-11-14T16:09:17.000Z","dependencies_parsed_at":null,"dependency_job_id":"3a33070c-18b9-4555-9d1f-58133086af76","html_url":"https://github.com/Findus23/jax-array-info","commit_stats":null,"previous_names":["findus23/jax-array-info"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Findus23%2Fjax-array-info","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Findus23%2Fjax-array-info/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Findus23%2Fjax-array-info/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Findus23%2Fjax-array-info/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Findus23","download_url":"https://codeload.github.com/Findus23/jax-array-info/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":226064568,"owners_count":17568035,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-10-13T10:29:38.832Z","updated_at":"2025-07-15T19:31:47.328Z","avatar_url":"https://github.com/Findus23.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# jax-array-info\n\nThis package contains two functions for debugging jax `Array`s:\n\n```bash\npip install git+https://github.com/Findus23/jax-array-info.git\n```\n\n```python\nfrom jax_array_info import sharding_info, sharding_vis, simple_array_info, print_array_stats, pretty_memory_stats\n```\n\n## `sharding_info(arr)`\n\n`sharding_info(arr)` prints general information about a jax or numpy array with special focus on sharding (\nsupporting `SingleDeviceSharding`, `GSPMDSharding`, `PositionalSharding`, `NamedSharding` and `PmapSharding`)\n\n```python\nsome_array = jax.numpy.zeros(shape=(N, N, N), dtype=jax.numpy.float32)\nsome_array = jax.device_put(some_array, NamedSharding(mesh, P(None, \"gpus\")))\nsharding_info(some_array, \"some_array\")\n```\n\n```text\n╭───────────────── some_array ─────────────────╮\n│ shape: (128, 128, 128)                       │\n│ dtype: float32                               │\n│ size: 8.0 MiB                                │\n│ NamedSharding: P(None, 'gpus')               │\n│ axis 1 is sharded: CPU 0 contains 0:16 (1/8) │\n│                    Total size: 128           │\n╰──────────────────────────────────────────────╯\n```\n\n## `simple_array_info()`\n\n`sharding_info()` uses a jax callback to make sure it can get sharding information in as many situations as possible.\nBut this means it is broken when used in some type of functions (e.g. when [using\n`shard_map`](https://github.com/jax-ml/jax/issues/23936)). `simple_array_info()` gives the same output, with the\nadvantage of working everywhere (it is equivalent to a `print`) and the tradeoff that it is not guaranteed to always\nreport correct sharding information (e.g. inside jitted functions).\n\n## `print_array_stats()`\n\nShows a nice overview over the all currently allocated arrays ordered by their size. To save space, scalar values are\ngrouped by dtype.\n\n**Disclaimer**: This uses `jax.live_arrays()` to get its information. There might be allocated arrays that are missing\nin this view. Also keep in mind that this only reports arrays that are acutally allocated, so calling it inside a jitted\nfunction won't give useful information.\n\n```python\narr = jax.numpy.zeros(shape=(16, 16, 16))\narr2 = jax.device_put(jax.numpy.zeros(shape=(2, 16, 4)), NamedSharding(mesh, P(None, \"gpus\")))\nscalar = jax.numpy.array(42)\n\nprint_array_stats()\n```\n\n```text\n             allocated jax arrays              \n┏━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓\n┃ size     ┃ shape        ┃ dtype   ┃      sharded      ┃\n┡━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩\n│ 16.0 KiB │ (16, 16, 16) │ float32 │                   │\n│ 64.0 B   │ (2, 16, 4)   │ float32 │ ✔ (512.0 B total) │\n├──────────┼──────────────┼─────────┼───────────────────┤\n│ 4.0 B    │ 1×s          │ int32   │                   │\n├──────────┼──────────────┼─────────┼───────────────────┤\n│ 16.1 KiB │              │         │                   │\n└──────────┴──────────────┴─────────┴───────────────────┘\n```\n\n## `sharding_vis(arr)`\n\nA modified version\nof [\n`jax.debug.visualize_array_sharding()`](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.visualize_array_sharding.html)\nthat also supports arrays with more than 2 dimensions (by ignoring non-sharded dimensions in the visualisation until\nreaching 2 dimensions)\n\n```python\narray = jax.numpy.zeros(shape=(N, N, N), dtype=jax.numpy.float32)\narray = jax.device_put(array, NamedSharding(mesh, P(None, \"gpus\")))\nsharding_vis(array)\n```\n\n```text\n─────────── showing dims [0, 1] from original shape (128, 128, 128) ────────────\n┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n│ CPU 0 │ CPU 1 │ CPU 2 │ CPU 3 │ CPU 4 │ CPU 5 │ CPU 6 │ CPU 7 │\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n│       │       │       │       │       │       │       │       │\n└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘\n```\n\n## `pretty_memory_stats(device)`\n\nThis is printing the output of `device.memory_stats()` as a table. The data (and if any data is available) depends on\nthe device\n\n```python\npretty_memory_stats(jax.local_devices()[0])\n```\n\n```text\n               Memory Stats of cuda:0                \n                     222 allocs                      \n┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━┓\n┃ name                     ┃     size ┃  size (raw) ┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━┩\n│ bytes_in_use             │  3.6 GiB │  3897557504 │\n│ peak_bytes_in_use        │ 29.9 GiB │ 32087482368 │\n│ largest_alloc_size       │ 23.5 GiB │ 25232940032 │\n│ bytes_limit              │ 33.3 GiB │ 35714334720 │\n│ bytes_reserved           │    0.0 B │           0 │\n│ peak_bytes_reserved      │    0.0 B │           0 │\n│ largest_free_block_bytes │    0.0 B │           0 │\n│ pool_bytes               │ 33.3 GiB │ 35714334720 │\n│ peak_pool_bytes          │ 33.3 GiB │ 35714334720 │\n└──────────────────────────┴──────────┴─────────────┘\n```\n\n## Examples\n\nYou can find many examples of how arrays can be sharded in jax and how the output would look like in [\n`tests/test_jax.py`](./tests/test_jax.py). For examples of sharding arrays along multiple jax processes check [\n`test_multihost.py`](./tests/test_multihost.py) and [`multihost.py`](./tests/multihost.py)\n\n## Assigning Labels to Jax Arrays\n\nWhile the output of `print_array_stats()` is very useful to see which arrays are currently allocated, it is missing the\nconnection to where the array is currently used. While in many cases the shape allows to guess which array a row is\nrefering to, it would be nice to also know its name.\nUnfortunately we only get a reference to all currently allocated arrays from `jax.live_arrays()` which means we don't\nknow where the array is assigned and what the variable is called.\nWe can do an ugly workaround to still get more useful names in the `print_array_stats()` output:\nWhenever `simple_array_info` is used like `simple_array_info(some_array,\"some_array\")`, we will modify the `some_array`\njax array by setting `some_array._custom_label = \"some_array\"`. Therefore this label becomes a part of the jax `Array`\nobject and we can use this label in our array overview.\n\nThis of course has some limitations (it only works if `simple_array_info` is used outside of jitted code) and is an ugly\nhack, so it possibly will break some things. Therefore this feature can be disabled globally using\n\n```python\nfrom jax_array_info import sharding_info, print_array_stats, config as array_info_config\n\narray_info_config.assign_labels_to_arrays = False\n```\n\n#### Example\n\n```python\nimport jax\nfrom jax_array_info import sharding_info, print_array_stats\n\nsome_array = jax.numpy.zeros(shape=(128, 128, 128))\nsharding_info(some_array, \"some_array\")\n\nprint_array_stats()\n```\n\n```text\n╭────── some_array ──────╮\n│ shape: (128, 128, 128) │\n│ dtype: float32         │\n│ size: 8.0 MiB          │\n│ not sharded            │\n╰────────────────────────╯\n                     allocated jax arrays                     \n┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━┓\n┃ size    ┃ shape           ┃ dtype   ┃ sharded ┃ label      ┃\n┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━┩\n│ 8.0 MiB │ (128, 128, 128) │ float32 │         │ some_array │\n├─────────┼─────────────────┼─────────┼─────────┼────────────┤\n│ 8.0 MiB │                 │         │         │            │\n└─────────┴─────────────────┴─────────┴─────────┴────────────┘\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffindus23%2Fjax-array-info","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffindus23%2Fjax-array-info","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffindus23%2Fjax-array-info/lists"}