# Questions about handling non-scalar kwargs in pattern matcher and kwargs capture in replacements

**URL:** <https://dev-discuss.pytorch.org/t/questions-about-handling-non-scalar-kwargs-in-pattern-matcher-and-kwargs-capture-in-replacements/2904>\
**Category:** compiler\
**Created:** [April 11, 2025, 11:48am UTC](https://dev-discuss.pytorch.org/t/questions-about-handling-non-scalar-kwargs-in-pattern-matcher-and-kwargs-capture-in-replacements/2904 "2025-04-11T11:48:23Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![mriganks-amd](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/mriganks-amd/32/2662_2.png) [@mriganks-amd](https://dev-discuss.pytorch.org/u/mriganks-amd)\
**Post date:** [April 11, 2025, 11:48am UTC](https://dev-discuss.pytorch.org/t/questions-about-handling-non-scalar-kwargs-in-pattern-matcher-and-kwargs-capture-in-replacements/2904/1 "2025-04-11T11:48:23Z")

</div>

I’m working with the pattern matcher in PyTorch, specifically using ‘register\_replacement’ (not ‘gen\_register\_replacement’). I have two related questions about handling kwargs:

1. How can we pass non-scalar kwargs to pattern matcher candidates?

Here’s a simplified example of what I’m trying to achieve:

```auto
# Pattern with a non-scalar kwarg (e.g., torch.dtype or bool)
def _my_pattern(arg_0, arg_1, kwarg_0):
    mm_0 = torch.ops.aten.mm(arg_0, arg_1)
    res = my_custom_op(mm_0, kw=kwarg_0)
    return res

# Replacement
def _my_replacement(arg_0, arg_1, kwarg_0):
    res = my_fused_op(arg_0, arg_1, kw=kwarg_0)
    return res

# Setting up candidates
arg_1 = functools.partial(torch.empty, (256, 32), device="cpu", requires_grad=True, dtype=torch.float)
arg_2 = functools.partial(torch.empty, (32, 512), device="cpu", requires_grad=True, dtype=torch.float)

# How do we handle non-scalar kwargs here?
# The scalar_workaround in register_replacement only works for values like alpha=1.3

```

1. Alternative approach: Is it possible to capture kwargs from the pattern in the replacement?

If we don’t specify the kwargs in the pattern, can we capture them from the FX graph? For example:

```auto
# Pattern without explicit kwargs
def _my_pattern(arg_0, arg_1):
    mm_0 = torch.ops.aten.mm(arg_0, arg_1)
    res = my_custom_op(mm_0) # has kwargs in the actual graph
    return res

# Replacement - can we capture kwargs from pattern?
def _my_replacement(arg_0, arg_1):
    # How to get kwargs from my_custom_op in the pattern?
    res = my_fused_op(arg_0, arg_1)
    return res

```

I’ve been looking at `torch/_inductor/fx_passes/fuse_attention.py` for implementation guidance but haven’t found a clear solution for either approach. Any suggestions would be appreciated.

---

<div class="post-metadata">

**Author:** ![lgovedic](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/lgovedic/32/2583_2.png) [@lgovedic](https://dev-discuss.pytorch.org/u/lgovedic)\
**Post date:** [April 22, 2025, 4:31am UTC](https://dev-discuss.pytorch.org/t/questions-about-handling-non-scalar-kwargs-in-pattern-matcher-and-kwargs-capture-in-replacements/2904/2 "2025-04-22T04:31:06Z")

</div>

In the second case, you can capture the `kwargs` inside the `extra_check` function (callback that does the final check to validate a match). That callback gets the match object passed as a parameter, and if it returns `True`, the replacement occurs. You could either save the kwarg, let the replacement happen, and post process, or you could return `False` from the check and manually process the match and execute the replacement either later or inside the `extra_check`.

---

<div class="post-metadata">

**Author:** ![mriganks-amd](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/mriganks-amd/32/2662_2.png) [@mriganks-amd](https://dev-discuss.pytorch.org/u/mriganks-amd)\
**Post date:** [May 21, 2025, 5:54am UTC](https://dev-discuss.pytorch.org/t/questions-about-handling-non-scalar-kwargs-in-pattern-matcher-and-kwargs-capture-in-replacements/2904/3 "2025-05-21T05:54:51Z")

</div>

Thanks @lgovedic for the suggestion. However, in [post-grad passes](https://github.com/pytorch/pytorch/blob/134179474539648ba7dee1317959529fbd0e7f89/torch/_inductor/fx_passes/post_grad.py#L437) there seems to be a better way to pass any kwarg to a pattern with the help of ‘register\_graph\_pattern’ decorator as opposed to ‘register\_replacement’, this works best for my present use case.