# Conditional Decompositions

**URL:** <https://dev-discuss.pytorch.org/t/conditional-decompositions/2958>\
**Category:** compiler\
**Created:** [April 30, 2025, 9:46am UTC](https://dev-discuss.pytorch.org/t/conditional-decompositions/2958 "2025-04-30T09:46:24Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Max1996](https://avatars.discourse-cdn.com/v4/letter/m/3da27b/32.png) [@Max1996](https://dev-discuss.pytorch.org/u/Max1996)\
**Post date:** [April 30, 2025, 9:46am UTC](https://dev-discuss.pytorch.org/t/conditional-decompositions/2958/1 "2025-04-30T09:46:24Z")

</div>

Hi,

I am trying to implement a conditional decomposition for the torch.aten.mul operator:

In most cases, it should not change, and keep being the torch.aten.mul operator.  
However, in specific cases, I want to decompose the mul to other ops that my flow supports.

I tried to implement it as a custom decomposition, but I am not able to find a solution to leave the operator untouched, if my condition does not apply.

```python
from typing import Sequence, Union, Dict, Callable

import torch
from torch._ops import OperatorBase, OpOverloadPacket
from torch._decomp import register_decomposition
from torch._prims_common.wrappers import out_wrapper

@register_decomposition(aten_op=aten.mul)
@out_wrapper()
def replace_mul(input_a : Union[torch.Tensor, bool, int, float, complex], input_b : Union[torch.Tensor, bool, int, float, complex]) -> torch.Tensor:
    
  if <my condition>:
    return <my custom decomposition>

  # else: we just want to keep the standard mul operator
  ... 
  # return None ## did not work
  return torch.mul(input_a, input_b) ## ends in endless recursive function calls

```

I tried returning torch.mul, but this results in endless recursive calls of the decomposition.  
I also tried returning None, but this fails with an AssertionError.

Is there a mechanism to enable conditional decompositions or an alternative approach that can be used to solve this issue?