From 8e43f76b04e5e011938effd59168e12730010a3f Mon Sep 17 00:00:00 2001 From: Raul Torres Date: Mon, 14 Apr 2025 18:05:42 -0500 Subject: [PATCH] Remove redundant flatten/unflatten of PyTrees --- frontend/catalyst/tracing/type_signatures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/tracing/type_signatures.py b/frontend/catalyst/tracing/type_signatures.py index 84a8098b8e..d0d2f3b289 100644 --- a/frontend/catalyst/tracing/type_signatures.py +++ b/frontend/catalyst/tracing/type_signatures.py @@ -227,7 +227,7 @@ def get_decomposed_signature(args, static_argnums): """ dynamic_args, static_args = split_static_args(args, static_argnums) flat_dynamic_args, treedef = tree_flatten(dynamic_args) - flat_signature = get_abstract_signature(flat_dynamic_args) + flat_signature = [shaped_abstractify(arg) for arg in flat_dynamic_args] return flat_signature, treedef, static_args