[danieldk]

This change adds a function to get all the kernelized layers and functions in a model, along with the repository the kernel was retrieved from. In order to do so, we add _kernel_repo to kernelized layers.

One annoyance is that to keep layers pickable, we need to deal with the Rust data structures that are used inside the repo information. To do so, we use the existing JSON (de)serialization in pickling.

Example use:

from kernels import get_kernelized_layers

kernelized = get_kernelized_layers(model)
kernelized.layers
{'layers.0.mlp.act_fn': KernelizedLayer(module=..., repo=LayerRepository(...))}
kernelized.funcs
{('layers.0.attn', 'apply_rotary_pos_emb'): KernelizedFunc(module=..., repo=FuncRepository(...))}

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

Coverage report — kernels/

Measured on: Python 3.10 / Torch 2.13.0.
Other CI configurations are not included in this number.
Hardware-gated code paths (ROCm/XPU/NPU/Darwin/Windows) are excluded or unreachable on the Linux+CUDA runner.

Total coverage: 88.5% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	14	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	130	14	89%	53, 59-60, 63-64, 102, 165-170, 199, 219
src/kernels/archs.py	56	1	98%	94
src/kernels/backends.py	215	63	71%	42, 46, 50-53, 70, 92, 110, 119, 123, 127-129, 150, 159, 163, 167-169, 190, 201, 203, 210-213, 222, 226, 230-255, 263, 286-306
src/kernels/compat.py	9	1	89%	5
src/kernels/deps.py	70	1	99%	56
src/kernels/hf_hub.py	63	2	97%	21, 23
src/kernels/importer.py	50	5	90%	95, 99, 102, 116-117
src/kernels/install.py	21	7	67%	76-100
src/kernels/layer/__init__.py	7	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	12	75%	42, 47-49, 91, 96-98, 101, 149, 152, 156
src/kernels/layer/func.py	93	5	95%	94, 119, 195, 315, 390
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/inspect.py	27	0	100%
src/kernels/layer/kernelize.py	82	8	90%	259, 297, 305-306, 312, 316, 332-334
src/kernels/layer/layer.py	320	14	96%	212, 259, 435, 575-576, 597, 605, 616, 657, 661, 674, 729, 770, 845
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	144	36	75%	33, 36-43, 74, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257
src/kernels/load.py	71	2	97%	338, 378
src/kernels/locking.py	89	64	28%	35-83, 91-98, 102-125, 137, 152-159, 165-175, 179-186
src/kernels/python_deps.py	58	6	90%	59-60, 64-65, 101, 104
src/kernels/resolver.py	162	1	99%	248
src/kernels/status.py	50	2	96%	25, 79
src/kernels/validate.py	88	5	94%	9, 100, 167, 190-191
src/kernels/variants.py	278	17	94%	65, 96, 117, 147, 256-257, 299-302, 304, 388-394, 400-406, 455-461
src/kernels/verify.py	127	6	95%	46, 202-204, 318-319
TOTAL	2400	277	88%

Updated by the Test kernels workflow on commit 84fced9ecc54233e7d4ec9de05067fe3d697abff.

[drbh]

lgtm!