Skip to content

Commit

Permalink
full_like to full decomposition moving to decomposition.py for dynami…
Browse files Browse the repository at this point in the history
…c case
  • Loading branch information
apbose committed Nov 19, 2024
1 parent 3e8d735 commit b73bb83
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 65 deletions.
12 changes: 12 additions & 0 deletions py/torch_tensorrt/dynamo/lowering/_decompositions.py
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,18 @@ def instance_norm_decomposition(
)


@register_torch_trt_decomposition(
torch.ops.aten.full_like, registry=TORCH_TRT_DECOMPOSITIONS
)
def full_like_decomposition(*args, **kwargs) -> torch.Tensor:
input = args[0]
shape = args[0].shape
fill_value = args[1]
kwargs["dtype"] = input.dtype
kwargs["device"] = to_torch_device(default_device())
return torch.full(shape, fill_value, dtype=kwargs["dtype"], device=kwargs["device"])


def get_decompositions(
enable_experimental_decompositions: bool = False,
) -> Dict[OpOverload, Callable[[Any], Any]]:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
from .remove_detach import remove_detach
from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones
from .repair_input_as_output import repair_input_as_output
from .replace_full_like_with_full import replace_full_like_with_full
from .replace_max_pool_with_indices import replace_max_pool_with_indices
from .view_to_reshape import view_to_reshape

Expand All @@ -27,7 +26,6 @@
lower_linear,
fuse_prims_broadcast,
replace_max_pool_with_indices,
replace_full_like_with_full,
view_to_reshape,
remove_assert_scalar,
accumulate_fp32_matmul,
Expand Down

This file was deleted.

0 comments on commit b73bb83

Please sign in to comment.