Skip to content

Commit 5aca1dd

Browse files
committed
full_like to full decomposition moving to decomposition.py for dynamic case
1 parent 3e8d735 commit 5aca1dd

File tree

3 files changed

+12
-65
lines changed

3 files changed

+12
-65
lines changed

py/torch_tensorrt/dynamo/lowering/_decompositions.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,18 @@ def instance_norm_decomposition(
418418
return torch.nn.functional.batch_norm(
419419
input, running_mean, running_var, weight, bias, False, momentum, eps
420420
)
421+
422+
423+
@register_torch_trt_decomposition(
424+
torch.ops.aten.full_like, registry=TORCH_TRT_DECOMPOSITIONS
425+
)
426+
def full_like_decomposition(*args, **kwargs) -> torch.Tensor:
427+
input = args[0]
428+
shape = args[0].shape
429+
fill_value = args[1]
430+
kwargs["dtype"] = input.dtype
431+
kwargs["device"] = to_torch_device(default_device())
432+
return torch.full(shape, fill_value, dtype=kwargs["dtype"], device=kwargs["device"])
421433

422434

423435
def get_decompositions(

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
from .remove_detach import remove_detach
1515
from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones
1616
from .repair_input_as_output import repair_input_as_output
17-
from .replace_full_like_with_full import replace_full_like_with_full
1817
from .replace_max_pool_with_indices import replace_max_pool_with_indices
1918
from .view_to_reshape import view_to_reshape
2019

@@ -27,7 +26,6 @@
2726
lower_linear,
2827
fuse_prims_broadcast,
2928
replace_max_pool_with_indices,
30-
replace_full_like_with_full,
3129
view_to_reshape,
3230
remove_assert_scalar,
3331
accumulate_fp32_matmul,

py/torch_tensorrt/dynamo/lowering/passes/replace_full_like_with_full.py

Lines changed: 0 additions & 63 deletions
This file was deleted.

0 commit comments

Comments
 (0)