[u-simha]

What:

• Add support for palettizing weights that are already sparse (e.g. from a pruned model) jointly with their sparsity, instead of applying sparsity during palettization
• Support per-tensor, per-grouped-channel, per-channel-scale, quantized-LUT and (CoreAI) vector palettization for joint sparsity; see the tables below
• Document the option in the PalettizationSpec docstring and docs/src/palettization/config.md, and add a changelog fragment

How:

• Add a public sparsity_detection_threshold option to PalettizationSpec (a float in [0, 1]). None (default) turns detection off. Otherwise, on the first forward pass, a weight with at least that fraction of exact zeros (and at least one zero) is treated as sparse; its nonzero mask is kept fixed afterwards. No pruning is applied.
• For detected weights, k-means fits the palette on the nonzero weights only, and the fake-palettized weight is lut[indices] * mask, so the zeros stay exactly zero:
  • Per-tensor and balanced grouped layers use all 2**n_bits centroids.
  • Grouped layers whose channel groups hold different numbers of nonzeros use 2**n_bits - 1 centroids and reserve a zero LUT entry, so the zeros the CoreAI export stores for padding look up exactly 0.0.
  • With vector palettization (cluster_dim = d > 1, per-tensor or grouped along axis 0), the stored values of each group are clustered in vectors of d consecutive values, so there are nnz // d indices per group and the LUT is [num_groups, 2**n_bits, d]. If a group's nnz % d != 0, its smallest nnz % d nonzeros are zeroed out of the mask once, where the mask is detected, so training and export see the same mask. Groups with fewer vectors than the fullest group also store some of their zeros, chosen at that point too; a stored zero can share a vector with a nonzero and come back as a small nonzero.
• Update prepare_for_export for palettization:
  • CoreAI: lut_to_dense on the nonzero indices, followed by sparse_to_dense (and constexpr_blockwise_shift_scale for a per-channel scale and/or a quantized LUT). Grouped LUTs are passed as [num_groups, num_clusters, 1]; lut_to_dense maps the g-th contiguous block of the flat nonzero indices to group g, so every group must store the same number of elements.
  • CoreML: add PRUNING to the compression types, giving constexpr_lut_to_sparse -> constexpr_sparse_to_dense (and constexpr_blockwise_shift_scale for a per-channel scale). CoreML recovers each element's group from the mask.
• Export errors name the layer.
• Add end-to-end tests on pre-sparsified MNIST and ResNet-50 models for both backends, logic tests of prepare() and finalize() on a small conv/conv/linear model (including a by-hand round trip of the stored indices), and tests for the padding helper, the centroid counts and the rejections.

Supported configurations:

Palettization without sparsity:

Option	CoreML	CoreAI	Notes
Per-tensor	Yes	Yes
Per-grouped-channel	Yes	Yes
Vector (cluster_dim > 1)	Yes	Yes
Per-channel scale	Yes	Yes
Quantized LUT (lut_qspec)	Yes	Yes
Combinations of vector, per-channel scale and quantized LUT	No	Yes	For CoreML, any two of them raise CoreMLExportError. For CoreAI there is no export guard, and per-channel scale + quantized LUT is tested.

Joint sparsity + palettization of an already-sparse weight:

Option	CoreML	CoreAI	Notes
Per-tensor	Yes	Yes
Grouped along axis 0	Yes	Yes	For CoreML, any number of nonzeros per group works. For CoreAI, groups with fewer nonzeros than the fullest one also store some zeros, which look up exactly 0.0, so every group stores the same number of elements; the LUT reserves an entry for 0.0. Groups that already hold equal numbers of nonzeros pay nothing.
Grouped along another axis	Yes	No	For CoreAI, a grouped LUT maps contiguous blocks of the flat nonzero list to groups, which only holds along axis 0.
Per-channel scale	Yes	Yes
Quantized LUT (lut_qspec)	No	Yes	For CoreML, coremltools cannot convert the quantization step after its sparse LUT op. For CoreAI, the zero point must be 0 (a symmetric lut_qspec), so masked-out weights dequantize to exactly 0, and it is not supported together with vector palettization. Combined with a per-channel scale, the two are fused into one scale.
Vector (cluster_dim > 1)	No	Yes	For CoreML, constexpr_lut_to_sparse takes its mask on the grid of vectors, so a vector is kept or dropped as a whole and cannot express nonzero values grouped into vectors. For CoreAI, per-tensor or grouped along axis 0 (grouping along another axis is rejected when the sparsity is detected); it works with a per-channel scale but not with a quantized LUT.

Example config:

config = KMeansPalettizerConfig(
    global_config=ModuleKMeansPalettizerConfig(
        op_state_spec={
            "weight": PalettizationSpec(
                n_bits=4,
                granularity=PerGroupedChannelGranularity(axis=0, group_size=16),
                sparsity_detection_threshold=0.4,
            )
        }
    )
)

[yixingli-apple]

The CoreML path has no check for lut_qspec. Nothing in the PR tests that combination, and the CoreAI path rejects it. I think we should either reject lut_qspec in the CoreML path, or add some tests for it if it's valid?

[yixingli-apple]

In some corner cases, _sparsity = 0 (say for a few layers they end up like this), does the export pipeline still work? Can we add one or two test cases for them?

[aseemw]

From the PR description only (not looked at the code yet):

Support both per-tensor and per-grouped-channel (axis 0) palettization for joint sparsity

Why does it say axis=0? Axis is configurable so it should work with whatever axis has been set, right?
(to be precise, right now in the code base axis is configurable only for scalar mode, for vector it is hardcoded to 0)

Add a private _sparsity_detection_threshold field to PalettizationSpec

Any reason to keep it private for now? I think we make it a public flag, if well tested, for the scenarios supported.

For detected weights, k-means reserves LUT entry 0 for exactly 0.0 and fits the other 2**n_bits - 1 centroids on the nonzero weights only, so the zeros stay exactly zero after palettization.

This should only happen in the per_group_channel case. For per_tensor, we should use all 2^n centroids.

For grouped granularity on the CoreAI backend, every group must hold the same number of nonzeros. Layers with uneven groups log a warning and are exported with plain palettization instead (exact, since the zeros index the 0.0 LUT entry).

I am bit confused about whether we need both the conditions to be true (using 2^n-1 centroids on non zero data + each group having the same number of 0s) for the "lut_to_dense" ---> "spars_to_dense" chaining to work properly together.

For the CoreAI backend, vector palettization, quantized LUTs and per-channel scales are not supported with joint sparsity, nor is grouped granularity on an axis other than 0.
For the CoreML backend, vector palettization and per-channel scales are not supported with joint sparsity.

Whats the reason for not supporting some of these cases? Is it just something to keep this PR simple, and add those cases later, or is it something fundamental in these cases which won't work with combined sparsity?

[aseemw]

Can you please add unit test(s) that takes a simple one layer model (e.g. one linear or one conv layer), and applies the following 4 configs:

1. scalar + per_tensor
2. scalar + per_group_channel
3. vector + per_group_channel
4. vector + per_tensor

and for each backend :Core AI and Core ML, verify the correct pattern of ops, that is the correct chaining / sequence of ops in the graph. Now this could be I guess done after the finalize call itself? Or we may need to convert and then check the pattern in the MLIR/MIL graphs.

(optional: for this unit test, we check that the outputs match between the prepared model and the aimodel/mlpackage outputs )

[aseemw]

We should also add another test (test_grouped_with_uneven_groups can be updated for this purpose) that checks that 2^n-1 centroids are used when expected and 2^n centroids are used when thats expected to happen.

[aseemw]

(BTW, if you find vector kmeans taking too much time, although for a single layer it should be fine, if you apply to mnist or resnet model it could take some time, then you can use max_iter=1, n_init=1 for unit testing purposes by using the VectorKMeansConfig introduced in #137).

[u-simha]

done

[fatih-ilhan]

Regarding per-group palletization + sparsity, lut_to_dense has to know which entries of the flat nonzero list belong to which group LUT. The CoreAI op has no per-element group information unlike CoreML. That's why same number of nonzero are needed. Workaround is to pad groups with fewer nonzero to have the same with the group with max nonzero by padding real zeroes and having one zero in the LUT used for them.

However, this workaround reduces the savings in the sparse export format - maybe 5% depending on how sparsity is distributed (probably more if group size is smaller). Is my understanding correct?

[aseemw]

Regarding per-group palletization + sparsity, lut_to_dense has to know which entries of the flat nonzero list belong to which group LUT. The CoreAI op has no per-element group information unlike CoreML. That's why same number of nonzero are needed. Workaround is to pad groups with fewer nonzero to have the same with the group with max nonzero by padding real zeroes and having one zero in the LUT used for them.

However, this workaround reduces the savings in the sparse export format - maybe 5% depending on how sparsity is distributed (probably more if group size is smaller). Is my understanding correct?

Thats correct and @u-simha is updating this PR to implement that "padding" (currently it doesn't do it).
Yeah it will have a hit, how much of that is (whether its 5% or lower/higher than that) will depend on the distribution of zeros/non-zeros.

[aseemw]

Looks most good.

• Left a few comments on tests, I think some of the can be reduced / simplified.

• Whats the reason to keep the flag as private ?

[aseemw]

I took a look again, the testing coverage is comprehensive and nice. Maybe it could be cleaned up a bit to make it more readable, but not a blocker at all.

Approving the PR. Thanks for covering all the different configurations!

Good to merge from my side once

• we make that a public flag
• add a short update to the documentation : to the PalettizationSpec doc string, explaining what this flag does more cleary.