arrscope 0.1.0__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (32) hide show
  1. arrscope-0.3.0/PKG-INFO +158 -0
  2. arrscope-0.3.0/README.md +131 -0
  3. arrscope-0.3.0/examples/renders.py +71 -0
  4. {arrscope-0.1.0 → arrscope-0.3.0}/main.py +26 -2
  5. {arrscope-0.1.0 → arrscope-0.3.0}/pyproject.toml +1 -1
  6. arrscope-0.3.0/src/arrscope/_api.py +157 -0
  7. arrscope-0.3.0/src/arrscope/_config.py +10 -0
  8. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/_types.py +11 -0
  9. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/adapters/_core.py +4 -1
  10. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/renderers/html.py +33 -3
  11. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/renderers/terminal.py +51 -27
  12. {arrscope-0.1.0 → arrscope-0.3.0}/uv.lock +1 -1
  13. arrscope-0.1.0/PKG-INFO +0 -185
  14. arrscope-0.1.0/README.md +0 -158
  15. arrscope-0.1.0/src/arrscope/_api.py +0 -131
  16. arrscope-0.1.0/src/arrscope/_config.py +0 -14
  17. {arrscope-0.1.0 → arrscope-0.3.0}/.gitignore +0 -0
  18. {arrscope-0.1.0 → arrscope-0.3.0}/docs/ADR-001-visual-grammar.md +0 -0
  19. {arrscope-0.1.0 → arrscope-0.3.0}/docs/ADR-002-architecture.md +0 -0
  20. {arrscope-0.1.0 → arrscope-0.3.0}/docs/ADR-003-api-design.md +0 -0
  21. {arrscope-0.1.0 → arrscope-0.3.0}/docs/GLOSSARY.md +0 -0
  22. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/__init__.py +0 -0
  23. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/__main__.py +0 -0
  24. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/_format.py +0 -0
  25. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/_layout.py +0 -0
  26. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/adapters/__init__.py +0 -0
  27. {arrscope-0.1.0 → arrscope-0.3.0}/src/arrscope/renderers/__init__.py +0 -0
  28. {arrscope-0.1.0 → arrscope-0.3.0}/test.ipynb +0 -0
  29. {arrscope-0.1.0 → arrscope-0.3.0}/tests/test_adapter.py +0 -0
  30. {arrscope-0.1.0 → arrscope-0.3.0}/tests/test_format.py +0 -0
  31. {arrscope-0.1.0 → arrscope-0.3.0}/tests/test_layout.py +0 -0
  32. {arrscope-0.1.0 → arrscope-0.3.0}/tests/test_stats.py +0 -0
@@ -0,0 +1,158 @@
1
+ Metadata-Version: 2.4
2
+ Name: arrscope
3
+ Version: 0.3.0
4
+ Summary: Beautiful n-dimensional array visualization for Python
5
+ Project-URL: Homepage, https://github.com/vizarray/arrscope
6
+ Project-URL: Repository, https://github.com/vizarray/arrscope
7
+ Author: arrscope contributors
8
+ License: MIT
9
+ Keywords: arrays,data-science,machine-learning,numpy,visualization
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Topic :: Scientific/Engineering :: Visualization
17
+ Requires-Python: >=3.11
18
+ Requires-Dist: numpy>=1.24
19
+ Requires-Dist: rich>=13.0
20
+ Provides-Extra: jax
21
+ Requires-Dist: jax>=0.4; extra == 'jax'
22
+ Provides-Extra: tf
23
+ Requires-Dist: tensorflow>=2.12; extra == 'tf'
24
+ Provides-Extra: torch
25
+ Requires-Dist: torch>=2.0; extra == 'torch'
26
+ Description-Content-Type: text/markdown
27
+
28
+ # arrscope
29
+
30
+ Visualize n-dimensional arrays in the terminal and Jupyter — with structural trees, tiled mosaics, and color-coded value maps.
31
+
32
+ ```python
33
+ from arrscope import scope
34
+ import numpy as np
35
+
36
+ scope(np.random.rand(3, 4, 5))
37
+ ```
38
+
39
+ ## Features
40
+
41
+ - **1D → 6D+**: Tiered visual grammar — lists, grids, trees, nested layers
42
+ - **Named axes**: Attach semantics (`batch`, `heads`, `h`, `w`)
43
+ - **Three color modes**:
44
+ - `dtype` — semantic colors by data type (blue=float, green=int, …)
45
+ - `heatmap` — diverging colormap (red → light → blue) by value
46
+ - `sparsity` — zeros as `·`, non-zeros in bold
47
+ - **Two render styles**:
48
+ - `tree` (default) — hierarchical branch view with colored guide lines
49
+ - `mosaic` — all 2D sub-slices tiled side by side as numeric tables
50
+ - **Method chaining**: `scope(arr, mode="heatmap").tree().mosaic()` — switch styles without re-specifying
51
+ - **Stats overlay**: min, max, mean, std, zero%, NaN count (always shown)
52
+ - **Head/tail truncation**: large dims show first/last N slices with `…` (default 20)
53
+ - **Smart precision**: auto-detects significant figures for floats
54
+ - **Terminal + Jupyter**: Rich ANSI + static HTML/CSS with dark mode auto-detect
55
+ - **Multi-framework**: NumPy, PyTorch, TensorFlow, JAX, tinygrad (lazy imports)
56
+ - **Zero config**: color always on, stats always on, auto-detect terminal vs notebook
57
+
58
+ ## Install
59
+
60
+ ```bash
61
+ pip install arrscope
62
+ ```
63
+
64
+ Only requires `numpy` + `rich`. Torch/TF/JAX/tinygrad are optional.
65
+
66
+ ## Quick start
67
+
68
+ ```python
69
+ from arrscope import scope
70
+ import numpy as np
71
+
72
+ # Auto-detect — last 2 dims form the grid
73
+ scope(np.random.rand(3, 4, 5))
74
+
75
+ # Named axes
76
+ scope(
77
+ np.random.rand(2, 8, 32, 32),
78
+ axes=['batch', 'heads', 'h', 'w'],
79
+ grid=['h', 'w'],
80
+ title='Attention heads',
81
+ )
82
+
83
+ # Custom grid dims
84
+ scope(data, axes=['a', 'b', 'c', 'd'], grid=['a', 'b'])
85
+ ```
86
+
87
+ ## Color modes
88
+
89
+ ```python
90
+ scope(arr, mode='dtype') # default
91
+ scope(arr, mode='heatmap') # diverging colormap
92
+ scope(arr, mode='sparsity') # · for zeros, bold for non-zeros
93
+ ```
94
+
95
+ ## Render styles
96
+
97
+ ```python
98
+ scope(arr) # tree (default) — click-to-expand layers
99
+ scope(arr, render_style='mosaic') # all sub-slices tiled side by side
100
+
101
+ # Method chaining for switching styles
102
+ r = scope(arr, mode='heatmap')
103
+ print(r.tree())
104
+ print(r.mosaic())
105
+ ```
106
+
107
+ ## CLI
108
+
109
+ ```bash
110
+ arrscope 3x4x5
111
+ arrscope 2x3x32x32 --axes batch heads h w --grid h w --mode heatmap
112
+ arrscope 20x4x5 --max-height 6
113
+ ```
114
+
115
+ ## Framework support
116
+
117
+ ```python
118
+ import torch
119
+ scope(torch.randn(2, 3, 4))
120
+
121
+ import tensorflow as tf
122
+ scope(tf.random.uniform((2, 3, 4)))
123
+
124
+ import jax.numpy as jnp
125
+ scope(jnp.array([[1, 2], [3, 4]]))
126
+
127
+ from tinygrad import Tensor
128
+ scope(Tensor.randn(3, 4))
129
+ ```
130
+
131
+ ## API
132
+
133
+ ```python
134
+ scope(
135
+ arr,
136
+ axes=None, # list[str] — name each dimension
137
+ grid=None, # list[str | int] — which dims form the leaf grid
138
+ title=None, # str — heading above the visualization
139
+ max_height=20, # int | None — rows before truncation, None disables
140
+ fmt=None, # str — format spec like '.4f'
141
+ mode='dtype', # 'dtype' | 'heatmap' | 'sparsity'
142
+ render_style='tree', # 'tree' | 'mosaic'
143
+ )
144
+ ```
145
+
146
+ ## Development
147
+
148
+ ```bash
149
+ git clone https://github.com/vizarray/arrscope
150
+ cd arrscope
151
+ uv sync
152
+ uv run pytest
153
+ uv run python main.py
154
+ ```
155
+
156
+ ## License
157
+
158
+ MIT
@@ -0,0 +1,131 @@
1
+ # arrscope
2
+
3
+ Visualize n-dimensional arrays in the terminal and Jupyter — with structural trees, tiled mosaics, and color-coded value maps.
4
+
5
+ ```python
6
+ from arrscope import scope
7
+ import numpy as np
8
+
9
+ scope(np.random.rand(3, 4, 5))
10
+ ```
11
+
12
+ ## Features
13
+
14
+ - **1D → 6D+**: Tiered visual grammar — lists, grids, trees, nested layers
15
+ - **Named axes**: Attach semantics (`batch`, `heads`, `h`, `w`)
16
+ - **Three color modes**:
17
+ - `dtype` — semantic colors by data type (blue=float, green=int, …)
18
+ - `heatmap` — diverging colormap (red → light → blue) by value
19
+ - `sparsity` — zeros as `·`, non-zeros in bold
20
+ - **Two render styles**:
21
+ - `tree` (default) — hierarchical branch view with colored guide lines
22
+ - `mosaic` — all 2D sub-slices tiled side by side as numeric tables
23
+ - **Method chaining**: `scope(arr, mode="heatmap").tree().mosaic()` — switch styles without re-specifying
24
+ - **Stats overlay**: min, max, mean, std, zero%, NaN count (always shown)
25
+ - **Head/tail truncation**: large dims show first/last N slices with `…` (default 20)
26
+ - **Smart precision**: auto-detects significant figures for floats
27
+ - **Terminal + Jupyter**: Rich ANSI + static HTML/CSS with dark mode auto-detect
28
+ - **Multi-framework**: NumPy, PyTorch, TensorFlow, JAX, tinygrad (lazy imports)
29
+ - **Zero config**: color always on, stats always on, auto-detect terminal vs notebook
30
+
31
+ ## Install
32
+
33
+ ```bash
34
+ pip install arrscope
35
+ ```
36
+
37
+ Only requires `numpy` + `rich`. Torch/TF/JAX/tinygrad are optional.
38
+
39
+ ## Quick start
40
+
41
+ ```python
42
+ from arrscope import scope
43
+ import numpy as np
44
+
45
+ # Auto-detect — last 2 dims form the grid
46
+ scope(np.random.rand(3, 4, 5))
47
+
48
+ # Named axes
49
+ scope(
50
+ np.random.rand(2, 8, 32, 32),
51
+ axes=['batch', 'heads', 'h', 'w'],
52
+ grid=['h', 'w'],
53
+ title='Attention heads',
54
+ )
55
+
56
+ # Custom grid dims
57
+ scope(data, axes=['a', 'b', 'c', 'd'], grid=['a', 'b'])
58
+ ```
59
+
60
+ ## Color modes
61
+
62
+ ```python
63
+ scope(arr, mode='dtype') # default
64
+ scope(arr, mode='heatmap') # diverging colormap
65
+ scope(arr, mode='sparsity') # · for zeros, bold for non-zeros
66
+ ```
67
+
68
+ ## Render styles
69
+
70
+ ```python
71
+ scope(arr) # tree (default) — click-to-expand layers
72
+ scope(arr, render_style='mosaic') # all sub-slices tiled side by side
73
+
74
+ # Method chaining for switching styles
75
+ r = scope(arr, mode='heatmap')
76
+ print(r.tree())
77
+ print(r.mosaic())
78
+ ```
79
+
80
+ ## CLI
81
+
82
+ ```bash
83
+ arrscope 3x4x5
84
+ arrscope 2x3x32x32 --axes batch heads h w --grid h w --mode heatmap
85
+ arrscope 20x4x5 --max-height 6
86
+ ```
87
+
88
+ ## Framework support
89
+
90
+ ```python
91
+ import torch
92
+ scope(torch.randn(2, 3, 4))
93
+
94
+ import tensorflow as tf
95
+ scope(tf.random.uniform((2, 3, 4)))
96
+
97
+ import jax.numpy as jnp
98
+ scope(jnp.array([[1, 2], [3, 4]]))
99
+
100
+ from tinygrad import Tensor
101
+ scope(Tensor.randn(3, 4))
102
+ ```
103
+
104
+ ## API
105
+
106
+ ```python
107
+ scope(
108
+ arr,
109
+ axes=None, # list[str] — name each dimension
110
+ grid=None, # list[str | int] — which dims form the leaf grid
111
+ title=None, # str — heading above the visualization
112
+ max_height=20, # int | None — rows before truncation, None disables
113
+ fmt=None, # str — format spec like '.4f'
114
+ mode='dtype', # 'dtype' | 'heatmap' | 'sparsity'
115
+ render_style='tree', # 'tree' | 'mosaic'
116
+ )
117
+ ```
118
+
119
+ ## Development
120
+
121
+ ```bash
122
+ git clone https://github.com/vizarray/arrscope
123
+ cd arrscope
124
+ uv sync
125
+ uv run pytest
126
+ uv run python main.py
127
+ ```
128
+
129
+ ## License
130
+
131
+ MIT
@@ -0,0 +1,71 @@
1
+ """arrscope render style examples.
2
+
3
+ Run this file to see both render styles in action.
4
+ Use method chaining to switch styles without re-specifying params.
5
+
6
+ Usage:
7
+ uv run python examples/renders.py
8
+ """
9
+
10
+ import numpy as np
11
+
12
+ from arrscope import scope
13
+
14
+ # ── data ──────────────────────────────────────────────────────────────
15
+ arr_3d = np.arange(60).reshape(3, 4, 5).astype(float)
16
+ eye = np.eye(8)
17
+ sparse = np.zeros((6, 8), dtype=float)
18
+ sparse[0, 0] = 1.5
19
+ sparse[2, 3] = -2.3
20
+ sparse[4, 5] = 42.0
21
+ sparse[5, 7] = 0.7
22
+
23
+
24
+ # ── individual calls ──────────────────────────────────────────────────
25
+ print("=" * 56)
26
+ print("1 tree — hierarchical branch view (default)")
27
+ print("=" * 56)
28
+ print(scope(arr_3d, render_style="tree"))
29
+
30
+ print("=" * 56)
31
+ print("2 mosaic — numeric tables tiled side by side")
32
+ print("=" * 56)
33
+ print(scope(arr_3d, render_style="mosaic"))
34
+
35
+
36
+ # ── method chaining ───────────────────────────────────────────────────
37
+ print("=" * 56)
38
+ print("3 method chaining — scope().tree() / .mosaic()")
39
+ print("=" * 56)
40
+
41
+ r = scope(eye, title="Identity", mode="heatmap")
42
+
43
+ print("--- .tree() ---")
44
+ print(r.tree())
45
+
46
+ print("--- .mosaic() ---")
47
+ print(r.mosaic())
48
+
49
+
50
+ # ── all three modes with mosaic ──────────────────────────────────────
51
+ print("=" * 56)
52
+ print("4 mosaic + all three color modes")
53
+ print("=" * 56)
54
+
55
+ print("--- dtype ---")
56
+ print(scope(sparse, render_style="mosaic", mode="dtype"))
57
+
58
+ print("--- heatmap ---")
59
+ print(scope(sparse, render_style="mosaic", mode="heatmap"))
60
+
61
+ print("--- sparsity ---")
62
+ print(scope(sparse, render_style="mosaic", mode="sparsity"))
63
+
64
+
65
+ # ── large array in mosaic ────────────────────────────────────────────
66
+ print("=" * 56)
67
+ print("5 large array — (8, 8, 8) as mosaic + heatmap")
68
+ print("=" * 56)
69
+ print(scope(np.random.rand(8, 8, 8), render_style="mosaic", mode="heatmap"))
70
+
71
+ print("\nDone!")
@@ -62,7 +62,31 @@ print(scope(sparse, mode="sparsity"))
62
62
  print("\n1️⃣2️⃣ HEATMAP + 3D + global stats")
63
63
  print(scope(np.random.rand(2, 3, 4), mode="heatmap"))
64
64
 
65
- print("\n1️⃣3️⃣ STATS OFF")
66
- print(scope(np.eye(5), stats=False))
65
+ print("\n1️⃣3️⃣ TRUNCATION OFF (max_height=None)")
66
+ print(scope(np.arange(60).reshape(3, 4, 5).astype(float), max_height=None))
67
+
68
+ print("\n1️⃣4️⃣ MOSAIC — 3×4×5 tiled (no tree)")
69
+ arr = np.arange(60).reshape(3, 4, 5).astype(float)
70
+ print(scope(arr, render_style="mosaic"))
71
+
72
+ print("\n1️⃣5️⃣ MOSAIC — 4D with named axes tiled side by side")
73
+ print(
74
+ scope(
75
+ np.random.rand(2, 3, 6, 6),
76
+ axes=["batch", "heads", "h", "w"],
77
+ grid=["h", "w"],
78
+ title="Mosaic: 2×3 attention heads",
79
+ render_style="mosaic",
80
+ )
81
+ )
82
+
83
+ print("\n1️⃣6️⃣ MOSAIC + HEATMAP — tiled grids with diverging colors")
84
+ print(
85
+ scope(
86
+ np.random.rand(2, 4, 5, 5),
87
+ render_style="mosaic",
88
+ mode="heatmap",
89
+ )
90
+ )
67
91
 
68
92
  print("\n Done!")
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "arrscope"
7
- version = "0.1.0"
7
+ version = "0.3.0"
8
8
  description = "Beautiful n-dimensional array visualization for Python"
9
9
  readme = "README.md"
10
10
  license = { text = "MIT" }
@@ -0,0 +1,157 @@
1
+ from typing import Literal
2
+
3
+ from arrscope._format import format_value
4
+ from arrscope._layout import build_layout, compute_stats
5
+ from arrscope._types import VizOutput
6
+ from arrscope.adapters import to_numpy
7
+ from arrscope.renderers import render_html, render_terminal
8
+
9
+ Mode = Literal["dtype", "heatmap", "sparsity"]
10
+ RenderStyle = Literal["tree", "mosaic"]
11
+
12
+
13
+ def scope(
14
+ arr,
15
+ *,
16
+ axes: list[str] | None = None,
17
+ grid: list[str | int] | None = None,
18
+ title: str | None = None,
19
+ max_height: int | None = 20,
20
+ fmt: str | None = None,
21
+ mode: Mode | None = None,
22
+ render_style: RenderStyle | None = None,
23
+ ) -> VizOutput:
24
+ """Visualize an n-dimensional array in the terminal and/or Jupyter.
25
+
26
+ Renders a structural tree view with 2D value grids at the leaves.
27
+ Works with NumPy, PyTorch, TensorFlow, JAX, and tinygrad arrays.
28
+
29
+ Args:
30
+ arr: The array to visualize. Accepts np.ndarray, torch.Tensor,
31
+ tf.Tensor, jax.Array, tinygrad.Tensor, or any object
32
+ implementing ``__array__``.
33
+ axes: Names for each dimension, e.g. ``["batch", "heads", "h", "w"]``.
34
+ Used in headers and for ``grid`` axis selection by name.
35
+ grid: Which dimensions form the leaf 2D grid. Given as axis indices
36
+ (``[0, 1]``), axis names (``["h", "w"]``), or omitted to
37
+ default to the last two dimensions.
38
+ title: Optional heading displayed above the visualization.
39
+ max_height: Max tree rows before truncation. The first half and
40
+ last half of children are shown with ``... (N more)`` in between.
41
+ Defaults to 20. Set to ``None`` to disable truncation.
42
+ fmt: Format spec for float values, e.g. ``".4f"``, ``".2e"``.
43
+ Auto-detected from value range when not set.
44
+ mode: Colorization strategy:
45
+
46
+ - ``"dtype"`` (default) — Semantic colors by data type
47
+ (blue=float, green=int, magenta=bool, yellow=str).
48
+ Negative values tinted red, positive tinted blue.
49
+ - ``"heatmap"`` — Diverging colormap (red → light → blue)
50
+ scaled to the array's global min/max range.
51
+ - ``"sparsity"`` — Zeros rendered as dim ``·``, non-zero
52
+ values in bold.
53
+
54
+ render_style: Visual layout:
55
+
56
+ - ``"tree"`` (default) — Hierarchical branch view with
57
+ click-to-expand nested layers.
58
+ - ``"mosaic"`` — All 2D sub-slices tiled side by side as
59
+ numeric tables.
60
+
61
+ Returns:
62
+ A ``VizOutput`` with ``.ansi`` and ``.html`` attributes.
63
+ Use ``.tree()`` or ``.mosaic()`` to switch render styles
64
+ without re-specifying parameters.
65
+
66
+ Examples:
67
+ >>> import numpy as np
68
+ >>> from arrscope import scope
69
+ >>> scope(np.random.rand(3, 4, 5))
70
+
71
+ >>> scope(np.eye(5), title="Identity", mode="sparsity")
72
+
73
+ >>> scope(
74
+ ... np.random.rand(2, 8, 32, 32),
75
+ ... axes=["batch", "heads", "h", "w"],
76
+ ... mode="heatmap",
77
+ ... )
78
+ """
79
+ arr_np = to_numpy(arr)
80
+
81
+ mode = mode or "dtype"
82
+ rs = render_style or "tree"
83
+
84
+ if mode not in ("dtype", "heatmap", "sparsity"):
85
+ raise ValueError(f"Unknown mode '{mode}'. Expected one of: dtype, heatmap, sparsity")
86
+
87
+ if rs not in ("tree", "mosaic"):
88
+ raise ValueError(f"Unknown render_style '{rs}'. Expected one of: tree, mosaic")
89
+
90
+ grid_dims = _resolve_grid_dims(arr_np, axes, grid)
91
+ use_html = _in_jupyter()
92
+
93
+ def _render(style: str) -> VizOutput:
94
+ node = build_layout(arr_np, grid_dims=grid_dims, max_height=max_height)
95
+ gs = compute_stats(arr_np)
96
+ if gs:
97
+ node.stats = gs
98
+
99
+ out = VizOutput(_rerender=_render)
100
+ if use_html:
101
+ out.html = _wrap_html(render_html(node, mode=mode, global_stats=gs, render_style=style), title)
102
+ out.ansi = _wrap_ansi(render_terminal(node, mode=mode, global_stats=gs, render_style=style), title)
103
+ return out
104
+
105
+ return _render(rs)
106
+
107
+
108
+ def _resolve_grid_dims(
109
+ arr,
110
+ axes: list[str] | None,
111
+ grid: list[str | int] | None,
112
+ ) -> list[int] | None:
113
+ if grid is None:
114
+ return None
115
+
116
+ if axes is None:
117
+ return [g if isinstance(g, int) else arr.ndim - 2 + (list(axes or []).index(g))
118
+ for g in grid]
119
+
120
+ name_to_idx = {name: i for i, name in enumerate(axes)}
121
+ resolved: list[int] = []
122
+ for g in grid:
123
+ if isinstance(g, str):
124
+ if g not in name_to_idx:
125
+ raise ValueError(
126
+ f"Axis '{g}' not found in axes={axes}. "
127
+ f"Available axes: {list(name_to_idx.keys())}"
128
+ )
129
+ resolved.append(name_to_idx[g])
130
+ else:
131
+ resolved.append(g)
132
+ return resolved
133
+
134
+
135
+ def _in_jupyter() -> bool:
136
+ try:
137
+ from IPython import get_ipython
138
+
139
+ shell = get_ipython()
140
+ return shell is not None and "IPKernelApp" in shell.config
141
+ except (ImportError, AttributeError):
142
+ return False
143
+
144
+
145
+ def _wrap_html(body: str, title: str | None) -> str:
146
+ if title:
147
+ return f'<h3 style="font-family: sans-serif; margin: 4px 0;">{title}</h3>\n{body}'
148
+ return body
149
+
150
+
151
+ def _wrap_ansi(body: str, title: str | None) -> str:
152
+ if title:
153
+ from rich.text import Text
154
+
155
+ t = Text(title, style="bold")
156
+ return f"{t}\n{body}"
157
+ return body
@@ -0,0 +1,10 @@
1
+ from dataclasses import dataclass
2
+
3
+
4
+ @dataclass
5
+ class Config:
6
+ mode: str = "dtype"
7
+ render_style: str = "tree" # tree | mosaic
8
+
9
+
10
+ config: Config = Config()
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ from collections.abc import Callable
3
4
  from dataclasses import dataclass, field
4
5
 
5
6
  import numpy as np
@@ -18,11 +19,15 @@ class VizNode:
18
19
  stats: dict | None = None
19
20
 
20
21
 
22
+ _RerenderFn = Callable[[str], "VizOutput"]
23
+
24
+
21
25
  @dataclass
22
26
  class VizOutput:
23
27
  html: str = ""
24
28
  ansi: str = ""
25
29
  metadata: dict = field(default_factory=dict)
30
+ _rerender: _RerenderFn | None = None
26
31
 
27
32
  def _repr_html_(self) -> str:
28
33
  return self.html
@@ -37,3 +42,9 @@ class VizOutput:
37
42
 
38
43
  def __repr__(self) -> str:
39
44
  return self.ansi
45
+
46
+ def tree(self) -> VizOutput:
47
+ return self._rerender("tree") if self._rerender else self
48
+
49
+ def mosaic(self) -> VizOutput:
50
+ return self._rerender("mosaic") if self._rerender else self
@@ -23,11 +23,14 @@ def to_numpy(arr) -> np.ndarray:
23
23
  if "jax" in module:
24
24
  return np.asarray(arr)
25
25
 
26
+ if "tinygrad" in module:
27
+ return arr.numpy()
28
+
26
29
  try:
27
30
  return np.asarray(arr)
28
31
  except (TypeError, ValueError) as e:
29
32
  raise TypeError(
30
33
  f"Cannot convert {type(arr).__name__} to numpy array. "
31
34
  f"Supported types: np.ndarray, torch.Tensor, tf.Tensor, jax.Array, "
32
- f"and objects implementing __array__."
35
+ f"tinygrad.Tensor, and objects implementing __array__."
33
36
  ) from e
@@ -10,12 +10,16 @@ from arrscope._types import VizNode
10
10
 
11
11
  def render_html(
12
12
  node: VizNode,
13
- color: bool = True,
14
13
  mode: str = "dtype",
15
14
  global_stats: dict | None = None,
15
+ render_style: str = "tree",
16
16
  ) -> str:
17
17
  css = _css()
18
- body = _render_node_html(node, 0, mode, global_stats)
18
+
19
+ if render_style == "mosaic":
20
+ body = _render_mosaic_html(node, mode, global_stats)
21
+ else:
22
+ body = _render_node_html(node, 0, mode, global_stats)
19
23
 
20
24
  stats_html = ""
21
25
  if global_stats and node.stats:
@@ -28,6 +32,32 @@ def render_html(
28
32
  </div>"""
29
33
 
30
34
 
35
+ def _collect_leaves(node: VizNode) -> list[VizNode]:
36
+ if node.grid_data is not None:
37
+ return [node]
38
+ leaves: list[VizNode] = []
39
+ for child in node.children:
40
+ leaves.extend(_collect_leaves(child))
41
+ return leaves
42
+
43
+
44
+ def _render_mosaic_html(
45
+ node: VizNode,
46
+ mode: str,
47
+ global_stats: dict | None,
48
+ ) -> str:
49
+ leaves = _collect_leaves(node)
50
+ if not leaves:
51
+ return "<span>(empty)</span>"
52
+
53
+ tiles: list[str] = []
54
+ for leaf in leaves:
55
+ if leaf.grid_data is not None:
56
+ tiles.append(_render_grid_html(leaf.grid_data, leaf.label, mode, global_stats))
57
+
58
+ return '<div style="display: flex; flex-wrap: wrap; gap: 16px; align-items: flex-start;">' + "".join(tiles) + "</div>"
59
+
60
+
31
61
  def _css() -> str:
32
62
  return """
33
63
  .arrscope-wrapper {
@@ -125,7 +155,7 @@ def _render_node_html(
125
155
  if not node.children and node.truncated:
126
156
  return f'<div class="arrscope-ellipsis">{html_mod.escape(node.label)}</div>'
127
157
 
128
- is_large_tree = depth > 0 and len(node.children) > 4
158
+ is_large_tree = depth > 0 and len(node.children) > 0
129
159
 
130
160
  parts: list[str] = []
131
161
  if depth == 0: