# tensorplay.ao.pruning API

Source: https://www.tensorplay.cn/docs/api/tensorplay.ao.pruning.html

## Functions 11

[#](#api-tensorplay.ao.pruning.compute_nparams_to_prune)

### compute_nparams_to_prune

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.compute_nparams_to_prune.html)

```python
tensorplay.ao.pruning.compute_nparams_to_prune(amount: int | float, tensor_size: int) → int
```

Convert a pruning amount into an absolute unit count.

An integer amount is returned unchanged; a float amount is read
as a fraction of tensor_size and rounded to the nearest integer.

Parameters:

- amount – pruning quantity (absolute int, or fractional float in [0, 1]).

- tensor_size – number of units available for pruning.

Returns:

The number of units to prune.

[#](#api-tensorplay.ao.pruning.custom_from_mask)

### custom_from_mask

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.custom_from_mask.html)

```python
tensorplay.ao.pruning.custom_from_mask(module: Module, name: str, mask: Tensor) → Module
```

Prune module[name] with a pre-computed binary mask.

Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding mask;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- mask – binary mask to be applied to the parameter.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = custom_from_mask(
...     nn.Linear(5, 3), name="bias", mask=tensorplay.tensor([0, 1, 0])
... )
>>> print(m.bias_mask)
tensor([0., 1., 0.])
```

[#](#api-tensorplay.ao.pruning.global_unstructured)

### global_unstructured

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.global_unstructured.html)

```python
tensorplay.ao.pruning.global_unstructured(parameters: Iterable[tuple[Module, str]], pruning_method: type[BasePruningMethod], importance_scores: dict[tuple[Module, str], Tensor] | None = None, **kwargs: Any) → None
```

Prune several tensors jointly under a single unstructured budget.

Aggregates the importance scores of every listed parameter into one
vector, computes a single mask under the shared amount budget, and
slices that mask back onto each parameter. Modifies the modules in place
by:

- adding a named buffer called name + '_mask' for every listed parameter;

- replacing each parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- parameters – iterable of (module, name) tuples identifying the parameters to prune globally, i.e. by aggregating all values before deciding which units to remove.

- pruning_method – a pruning method class from this package (or a user-defined subclass of [BasePruningMethod](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html#tensorplay.ao.pruning.BasePruningMethod)) whose PRUNING_TYPE is 'unstructured'.

- importance_scores – mapping from (module, name) tuples to the corresponding importance-scores tensor (same shape as the parameter). Parameters absent from the mapping use their own values as importance scores.

- kwargs – keyword arguments forwarded to pruning_method, typically amount: the quantity of units to prune across all listed parameters (a float fraction in [0, 1] or an absolute int).

Raises:

[TypeError](https://docs.python.org/3/builtins/exceptions.html#TypeError) – if parameters is not an iterable, if
    importance_scores is not a dict, or if the PRUNING_TYPE
    of pruning_method is not 'unstructured'.

> **Note**
>
> Global pruning is restricted to unstructured methods: a structured norm is only comparable across channels of equal size, which cannot be guaranteed across heterogeneous parameters.

Examples

```
>>> # xdoctest: +SKIP
>>> net = nn.Sequential(nn.Linear(10, 4), nn.Linear(4, 1))
>>> parameters_to_prune = (
...     (net[0], "weight"),
...     (net[1], "weight"),
... )
>>> global_unstructured(
...     parameters_to_prune,
...     pruning_method=L1Unstructured,
...     amount=10,
... )
```

[#](#api-tensorplay.ao.pruning.identity)

### identity

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.identity.html)

```python
tensorplay.ao.pruning.identity(module: Module, name: str) → Module
```

Attach the pruning reparameterization to module[name] without
pruning any unit.

Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

> **Note**
>
> The mask is a tensor of ones.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = identity(nn.Linear(2, 3), "bias")
>>> print(m.bias_mask)
tensor([1., 1., 1.])
```

[#](#api-tensorplay.ao.pruning.is_pruned)

### is_pruned

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.is_pruned.html)

```python
tensorplay.ao.pruning.is_pruned(module: Module) → bool
```

Check whether module carries an active pruning reparameterization.

Scans every submodule for forward pre-hooks that are instances of
[BasePruningMethod](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html#tensorplay.ao.pruning.BasePruningMethod).

Parameters:

module – module that is either pruned or unpruned.

Returns:

True when at least one submodule is pruned, False otherwise.

Examples

```
>>> # xdoctest: +SKIP
>>> m = nn.Linear(5, 7)
>>> print(is_pruned(m))
False
>>> random_unstructured(m, name="weight", amount=0.2)
>>> print(is_pruned(m))
True
```

[#](#api-tensorplay.ao.pruning.l1_unstructured)

### l1_unstructured

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.l1_unstructured.html)

```python
tensorplay.ao.pruning.l1_unstructured(module: Module, name: str, amount: int | float, importance_scores: Tensor | None = None) → Module
```

Prune module[name] by removing the units with the smallest
magnitudes.

Removes amount (currently unpruned) units ranked by absolute value.
Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

- importance_scores – tensor of importance scores with the same shape as the parameter; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = l1_unstructured(nn.Linear(2, 3), "weight", amount=0.2)
>>> list(m.state_dict().keys())
['bias', 'weight_orig', 'weight_mask']
```

[#](#api-tensorplay.ao.pruning.ln_structured)

### ln_structured

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.ln_structured.html)

```python
tensorplay.ao.pruning.ln_structured(module: Module, name: str, amount: int | float, n: int | float | str, dim: int, importance_scores: Tensor | None = None) → Module
```

Prune module[name] by removing the channels with the smallest
L``n``-norm along dim.

Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- n – norm order; accepts the orders valid for [tensorplay.linalg.vector_norm()](/docs/generated/tensorplay.linalg.vector_norm.html#tensorplay.linalg.vector_norm) and [tensorplay.linalg.matrix_norm()](/docs/generated/tensorplay.linalg.matrix_norm.html#tensorplay.linalg.matrix_norm).

- dim – axis along which channels are defined.

- importance_scores – tensor of importance scores with the same shape as the parameter; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = ln_structured(
...     nn.Conv2d(5, 3, 2), "weight", amount=0.3, dim=1, n=float("-inf")
... )
```

[#](#api-tensorplay.ao.pruning.random_structured)

### random_structured

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.random_structured.html)

```python
tensorplay.ao.pruning.random_structured(module: Module, name: str, amount: int | float, dim: int) → Module
```

Prune module[name] by removing random channels along dim.

Removes amount (currently unpruned) channels chosen uniformly at
random. Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- dim – axis along which channels are defined.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = random_structured(nn.Linear(5, 3), "weight", amount=3, dim=1)
>>> columns_pruned = int(sum(tensorplay.sum(m.weight, dim=0) == 0))
>>> print(columns_pruned)
3
```

[#](#api-tensorplay.ao.pruning.random_unstructured)

### random_unstructured

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.random_unstructured.html)

```python
tensorplay.ao.pruning.random_unstructured(module: Module, name: str, amount: int | float) → Module
```

Prune module[name] by removing a random subset of its units.

Removes amount (currently unpruned) units chosen uniformly at random.
Modifies the module in place (and also returns it) by:

- adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

- replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

Returns:

The modified (i.e. pruned) module.

Examples

```
>>> # xdoctest: +SKIP
>>> m = random_unstructured(nn.Linear(2, 3), "weight", amount=1)
>>> int(tensorplay.sum(m.weight_mask == 0))
1
```

[#](#api-tensorplay.ao.pruning.remove)

### remove

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.remove.html)

```python
tensorplay.ao.pruning.remove(module: Module, name: str) → Module
```

Make the pruning of module[name] permanent and drop the
reparameterization.

The pruned parameter name remains permanently pruned, the parameter
name + '_orig' is removed from the parameter list, and the buffer
name + '_mask' is removed from the buffers. The pruning hook is also
detached from the module.

> **Note**
>
> Pruning itself is NOT undone or reversed!

Parameters:

- module – module containing the pruned tensor.

- name – parameter name within module whose pruning is removed.

Returns:

The modified module.

Raises:

[ValueError](https://docs.python.org/3/builtins/exceptions.html#ValueError) – if name is not currently pruned on module.

[#](#api-tensorplay.ao.pruning.validate_pruning_amount)

### validate_pruning_amount

function[Full reference ↗](/docs/generated/tensorplay.ao.pruning.validate_pruning_amount.html)

```python
tensorplay.ao.pruning.validate_pruning_amount(amount: int, tensor_size: int) → None
```

Check that an absolute pruning count fits the data being pruned.

Parameters:

- amount – absolute number of units to prune.

- tensor_size – number of units available for pruning.

Raises:

[ValueError](https://docs.python.org/3/builtins/exceptions.html#ValueError) – if amount exceeds tensor_size.

## Classes 8

[#](#api-tensorplay.ao.pruning.BasePruningMethod)

### BasePruningMethod

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html)

```python
class tensorplay.ao.pruning.BasePruningMethod
```

Abstract base class for pruning techniques.

Subclasses must override [compute_mask()](#tensorplay.ao.pruning.BasePruningMethod.compute_mask) and declare a
PRUNING_TYPE class attribute (one of 'unstructured',
'structured' or 'global'). The classmethod [apply()](#tensorplay.ao.pruning.BasePruningMethod.apply) installs
the reparameterization on a module and registers an instance of the
subclass as a forward pre-hook.

```python
classmethod apply(module: Module, name: str, *args: Any, importance_scores: Tensor | None = None, **kwargs: Any) → BasePruningMethod
```

Install the pruning reparameterization for module[name].

Moves the original parameter to name + '_orig', registers the
computed mask as the buffer name + '_mask', stores the masked
values under name and adds a forward pre-hook that re-applies the
mask on every forward pass. If the tensor is already pruned by
another method, the new method is composed with the existing one
through a [PruningContainer](/docs/generated/tensorplay.ao.pruning.PruningContainer.html#tensorplay.ao.pruning.PruningContainer).

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- args – positional arguments forwarded to the subclass constructor.

- importance_scores – tensor of importance scores with the same shape as module[name]; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used as its own importance scores.

- kwargs – keyword arguments forwarded to the subclass constructor.

Returns:

The pruning method (or container of methods) now attached to the
module.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
abstractmethod compute_mask(t: Tensor, default_mask: Tensor) → Tensor
```

Compute the pruning mask for the input tensor t.

Starting from default_mask (a mask of ones when t has never
been pruned), derive the new mask according to the recipe of the
concrete method. Entries already zeroed by default_mask must stay
zero.

Parameters:

- t – tensor whose entries or channels are ranked for pruning, typically the importance scores of the parameter.

- default_mask – mask accumulated from previous pruning iterations; must be respected by the new mask. Same shape as t.

Returns:

The mask to apply to t, with the same shape as t.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by [compute_mask()](#tensorplay.ao.pruning.BasePruningMethod.compute_mask) without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.CustomFromMask)

### CustomFromMask

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.CustomFromMask.html)

```python
class tensorplay.ao.pruning.CustomFromMask(mask: Tensor)
```

Zero out exactly the units designated by a caller-supplied mask.

Parameters:

mask – binary mask whose zeros mark the units to prune.

```python
classmethod apply(module: Module, name: str, mask: Tensor) → BasePruningMethod
```

Install a user-provided mask for module[name].

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- mask – binary mask to be applied to the parameter.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by compute_mask() without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.Identity)

### Identity

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.Identity.html)

```python
class tensorplay.ao.pruning.Identity
```

Prune nothing and only install the mask reparameterization.

The generated mask is a tensor of ones, which is useful to prepare a
module for later iterative pruning without removing any unit yet.

```python
classmethod apply(module: Module, name: str) → BasePruningMethod
```

Install the identity (all-ones mask) reparameterization.

Parameters:

- module – module containing the tensor to reparameterize.

- name – parameter name within module on which pruning acts.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by compute_mask() without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.L1Unstructured)

### L1Unstructured

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.L1Unstructured.html)

```python
class tensorplay.ao.pruning.L1Unstructured(amount: int | float)
```

Zero out the units with the smallest magnitudes.

Ranks all currently unpruned units by absolute value and removes the
amount smallest ones.

Parameters:

amount – quantity of units to prune. A float in [0, 1] denotes
the fraction of units to prune; an int denotes the absolute
number of units to prune.

```python
classmethod apply(module: Module, name: str, amount: int | float, importance_scores: Tensor | None = None) → BasePruningMethod
```

Install magnitude-based unstructured pruning for module[name].

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

- importance_scores – tensor of importance scores with the same shape as the parameter; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by compute_mask() without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.LnStructured)

### LnStructured

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.LnStructured.html)

```python
class tensorplay.ao.pruning.LnStructured(amount: int | float, n: int | float | str, dim: int = -1)
```

Zero out the channels with the smallest L``n``-norm.

Reduces the tensor to one L``n``-norm per channel along dim and
keeps the amount-largest channels, zeroing the rest.

Parameters:

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- n – norm order; accepts the orders valid for [tensorplay.linalg.vector_norm()](/docs/generated/tensorplay.linalg.vector_norm.html#tensorplay.linalg.vector_norm) and [tensorplay.linalg.matrix_norm()](/docs/generated/tensorplay.linalg.matrix_norm.html#tensorplay.linalg.matrix_norm) (e.g. a positive number, float('inf'), float('-inf'), 'fro', 'nuc').

- dim – axis along which channels are defined. Default: -1.

```python
classmethod apply(module: Module, name: str, amount: int | float, n: int | float | str, dim: int, importance_scores: Tensor | None = None) → BasePruningMethod
```

Install norm-based structured pruning for module[name].

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- n – norm order; accepts the orders valid for [tensorplay.linalg.vector_norm()](/docs/generated/tensorplay.linalg.vector_norm.html#tensorplay.linalg.vector_norm) and [tensorplay.linalg.matrix_norm()](/docs/generated/tensorplay.linalg.matrix_norm.html#tensorplay.linalg.matrix_norm).

- dim – axis along which channels are defined.

- importance_scores – tensor of importance scores with the same shape as the parameter; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
compute_mask(t: Tensor, default_mask: Tensor) → Tensor
```

Compute a channel mask for t by per-channel L``n``-norm.

Parameters:

- t – tensor whose channels are pruned.

- default_mask – mask accumulated from previous pruning iterations; must be respected by the new mask. Same shape as t.

Returns:

The mask to apply to t, with the same shape as t.

Raises:

[IndexError](https://docs.python.org/3/builtins/exceptions.html#IndexError) – if self.dim is not a valid axis of t.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by [compute_mask()](#tensorplay.ao.pruning.LnStructured.compute_mask) without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.PruningContainer)

### PruningContainer

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.PruningContainer.html)

```python
class tensorplay.ao.pruning.PruningContainer(*args: BasePruningMethod)
```

Sequence of pruning methods applied iteratively to the same tensor.

Tracks the order in which methods were added and combines successive
pruning calls: each new method only ranks the entries or channels that
the previous masks left unpruned.

Accepts as argument an instance of a [BasePruningMethod](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html#tensorplay.ao.pruning.BasePruningMethod) or an
iterable of them.

```python
add_pruning_method(method: BasePruningMethod | None) → None
```

Add a child pruning method to the container.

Parameters:

method – child pruning method to be added to the container.

Raises:

- [TypeError](https://docs.python.org/3/builtins/exceptions.html#TypeError) – if method is neither None nor a [BasePruningMethod](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html#tensorplay.ao.pruning.BasePruningMethod) instance.

- [ValueError](https://docs.python.org/3/builtins/exceptions.html#ValueError) – if method acts on a different tensor name than the methods already held by the container.

```python
classmethod apply(module: Module, name: str, *args: Any, importance_scores: Tensor | None = None, **kwargs: Any) → BasePruningMethod
```

Install the pruning reparameterization for module[name].

Moves the original parameter to name + '_orig', registers the
computed mask as the buffer name + '_mask', stores the masked
values under name and adds a forward pre-hook that re-applies the
mask on every forward pass. If the tensor is already pruned by
another method, the new method is composed with the existing one
through a [PruningContainer](#tensorplay.ao.pruning.PruningContainer).

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- args – positional arguments forwarded to the subclass constructor.

- importance_scores – tensor of importance scores with the same shape as module[name]; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used as its own importance scores.

- kwargs – keyword arguments forwarded to the subclass constructor.

Returns:

The pruning method (or container of methods) now attached to the
module.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
compute_mask(t: Tensor, default_mask: Tensor) → Tensor
```

Apply the latest method and merge its mask into default_mask.

The new partial mask is computed on the entries or channels that
default_mask has not zeroed out. Which portion of t the new
mask is derived from depends on the PRUNING_TYPE of the last
method:

- 'unstructured': the mask is computed from the flattened list of entries not yet masked;

- 'structured': the mask is computed from the channels that still hold at least one unmasked entry;

- 'global': the mask is computed across all entries.

Parameters:

- t – tensor representing the parameter to prune (same shape as default_mask).

- default_mask – mask accumulated from previous pruning iterations.

Returns:

The mask combining the effects of default_mask and of the
latest method, with the same shape as default_mask and t.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by [compute_mask()](#tensorplay.ao.pruning.PruningContainer.compute_mask) without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.RandomStructured)

### RandomStructured

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.RandomStructured.html)

```python
class tensorplay.ao.pruning.RandomStructured(amount: int | float, dim: int = -1)
```

Zero out entire randomly selected channels of a tensor.

Parameters:

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- dim – axis along which channels are defined. Default: -1.

```python
classmethod apply(module: Module, name: str, amount: int | float, dim: int = -1) → BasePruningMethod
```

Install random structured pruning for module[name].

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

- dim – axis along which channels are defined. Default: -1.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
compute_mask(t: Tensor, default_mask: Tensor) → Tensor
```

Compute a channel mask for t by drawing channels at random.

Parameters:

- t – tensor whose channels are pruned.

- default_mask – mask accumulated from previous pruning iterations; must be respected by the new mask. Same shape as t.

Returns:

The mask to apply to t, with the same shape as t.

Raises:

[IndexError](https://docs.python.org/3/builtins/exceptions.html#IndexError) – if self.dim is not a valid axis of t.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by [compute_mask()](#tensorplay.ao.pruning.RandomStructured.compute_mask) without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!

[#](#api-tensorplay.ao.pruning.RandomUnstructured)

### RandomUnstructured

class[Full reference ↗](/docs/generated/tensorplay.ao.pruning.RandomUnstructured.html)

```python
class tensorplay.ao.pruning.RandomUnstructured(amount: int | float)
```

Zero out a uniformly random subset of the currently unpruned units.

Parameters:

amount – quantity of units to prune. A float in [0, 1] denotes
the fraction of units to prune; an int denotes the absolute
number of units to prune.

```python
classmethod apply(module: Module, name: str, amount: int | float) → BasePruningMethod
```

Install random unstructured pruning for module[name].

Parameters:

- module – module containing the tensor to prune.

- name – parameter name within module on which pruning acts.

- amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

```python
apply_mask(module: Module) → Tensor
```

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns
their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

```python
prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor
```

Return a pruned copy of the input tensor t.

Applies the rule implemented by compute_mask() without any
module-side reparameterization.

Parameters:

- t – tensor to prune (same shape as default_mask).

- importance_scores – tensor of importance scores with the same shape as t; each entry ranks the corresponding element of t. When unspecified, t itself is used.

- default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

```python
remove(module: Module) → None
```

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original
values is written back into the parameter name, and the auxiliary
parameter name + '_orig' and buffer name + '_mask' are
dropped.

> **Note**
>
> Pruning itself is NOT undone or reversed!