From 1e9f9376d0644e23d47ca38a122afd1fea8ecbc0 Mon Sep 17 00:00:00 2001 From: Akshay <59740725+Killua7362@users.noreply.github.com> Date: Wed, 6 Sep 2023 07:40:55 +0300 Subject: [PATCH 1/2] fix: nested maps on torch.compile for torch tensors --- ivy/functional/ivy/nest.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/ivy/functional/ivy/nest.py b/ivy/functional/ivy/nest.py index 5965c48416e08..042826a6087de 100644 --- a/ivy/functional/ivy/nest.py +++ b/ivy/functional/ivy/nest.py @@ -1160,12 +1160,14 @@ def nested_map( to_ignore = ivy.default(to_ignore, ()) extra_nest_types = ivy.default(extra_nest_types, ()) if include_derived is True: - include_derived = {tuple: True, list: True, dict: True} + include_derived = {str(tuple): True, str(list): True, str(dict): True} elif not include_derived: include_derived = {} for t in (tuple, list, dict): if t not in include_derived: - include_derived[t] = False + include_derived[str(t)] = False + # to ensure all keys are strings + include_derived = {str(key): value for key, value in include_derived.items()} if ivy.exists(max_depth) and _depth > max_depth: return x class_instance = type(x) @@ -1182,7 +1184,7 @@ def nested_map( _tuple_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[tuple] + if include_derived[str(tuple)] else (lambda x_, t_: type(x_) is t_) ), ) @@ -1190,7 +1192,7 @@ def nested_map( _list_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[list] + if include_derived[str(list)] else (lambda x_, t_: type(x_) is t_) ), ) @@ -1198,11 +1200,10 @@ def nested_map( _dict_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[dict] + if include_derived[str(dict)] else (lambda x_, t_: type(x_) is t_) ), ) - if tuple_check_fn(x, tuple) and not isinstance(x, to_ignore): ret_list = [ nested_map( From b17aa3b8e8b51e0dadc433ab5aac0d5515676f9d Mon Sep 17 00:00:00 2001 From: Akshay <59740725+Killua7362@users.noreply.github.com> Date: Wed, 6 Sep 2023 09:23:56 +0300 Subject: [PATCH 2/2] refractored types to strings --- ivy/data_classes/array/conversions.py | 8 ++++---- ivy/data_classes/container/conversions.py | 8 ++++---- ivy/func_wrapper.py | 6 +++--- ivy/functional/frontends/jax/func_wrapper.py | 8 ++++---- ivy/functional/frontends/mxnet/func_wrapper.py | 6 +++--- ivy/functional/frontends/numpy/func_wrapper.py | 10 +++++----- ivy/functional/frontends/onnx/func_wrapper.py | 6 +++--- ivy/functional/frontends/paddle/func_wrapper.py | 6 +++--- .../frontends/tensorflow/func_wrapper.py | 2 +- ivy/functional/frontends/torch/func_wrapper.py | 6 +++--- ivy/functional/ivy/nest.py | 15 +++++++-------- 11 files changed, 40 insertions(+), 41 deletions(-) diff --git a/ivy/data_classes/array/conversions.py b/ivy/data_classes/array/conversions.py index 30ac4445c0192..8da2954f31249 100644 --- a/ivy/data_classes/array/conversions.py +++ b/ivy/data_classes/array/conversions.py @@ -53,7 +53,7 @@ def _to_ivy(x: Any) -> Any: def to_ivy( x: Union[ivy.Array, ivy.NativeArray, Iterable], nested: bool = False, - include_derived: Optional[Dict[type, bool]] = None, + include_derived: Optional[Dict[str, bool]] = None, ) -> Union[ivy.Array, ivy.NativeArray, Iterable]: """ Return the input array converted to an ivy.Array instance if it is a native array @@ -84,7 +84,7 @@ def to_ivy( def args_to_ivy( *args: Iterable[Any], - include_derived: Optional[Dict[type, bool]] = None, + include_derived: Optional[Dict[str, bool]] = None, **kwargs: Dict[str, Any], ) -> Tuple[Iterable[Any], Dict[str, Any]]: """ @@ -115,7 +115,7 @@ def args_to_ivy( def to_native( x: Union[ivy.Array, ivy.NativeArray, Iterable], nested: bool = False, - include_derived: Optional[Dict[type, bool]] = None, + include_derived: Optional[Dict[str, bool]] = None, cont_inplace: bool = False, to_ignore: Optional[Union[type, Tuple[type]]] = None, ) -> Union[ivy.Array, ivy.NativeArray, Iterable]: @@ -157,7 +157,7 @@ def to_native( def args_to_native( *args: Iterable[Any], - include_derived: Dict[type, bool] = None, + include_derived: Dict[str, bool] = None, cont_inplace: bool = False, to_ignore: Optional[Union[type, Tuple[type]]] = None, **kwargs: Dict[str, Any], diff --git a/ivy/data_classes/container/conversions.py b/ivy/data_classes/container/conversions.py index 910729cc613fe..00cdf6cdcfef8 100644 --- a/ivy/data_classes/container/conversions.py +++ b/ivy/data_classes/container/conversions.py @@ -18,7 +18,7 @@ class _ContainerWithConversions(ContainerBase): def _static_to_native( x: Union[ivy.Array, ivy.NativeArray, ivy.Container], nested: Union[bool, ivy.Container] = False, - include_derived: Optional[Union[Dict[type, bool], ivy.Container]] = None, + include_derived: Optional[Union[Dict[str, bool], ivy.Container]] = None, key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, to_apply: Union[bool, ivy.Container] = True, prune_unapplied: Union[bool, ivy.Container] = False, @@ -78,7 +78,7 @@ def _static_to_native( def to_native( self: ivy.Container, nested: Union[bool, ivy.Container] = False, - include_derived: Optional[Union[Dict[type, bool], ivy.Container]] = None, + include_derived: Optional[Union[Dict[str, bool], ivy.Container]] = None, key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, to_apply: Union[bool, ivy.Container] = True, prune_unapplied: Union[bool, ivy.Container] = False, @@ -138,7 +138,7 @@ def to_native( def _static_to_ivy( x: Union[ivy.Array, ivy.NativeArray, ivy.Container], nested: Union[bool, ivy.Container] = False, - include_derived: Optional[Union[Dict[type, bool], ivy.Container]] = None, + include_derived: Optional[Union[Dict[str, bool], ivy.Container]] = None, key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, to_apply: Union[bool, ivy.Container] = True, prune_unapplied: Union[bool, ivy.Container] = False, @@ -199,7 +199,7 @@ def _static_to_ivy( def to_ivy( self: ivy.Container, nested: Union[bool, ivy.Container] = False, - include_derived: Optional[Union[Dict[type, bool], ivy.Container]] = None, + include_derived: Optional[Union[Dict[str, bool], ivy.Container]] = None, key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, to_apply: Union[bool, ivy.Container] = True, prune_unapplied: Union[bool, ivy.Container] = False, diff --git a/ivy/func_wrapper.py b/ivy/func_wrapper.py index 6ff998368f2e6..e7fa2ca9acce7 100644 --- a/ivy/func_wrapper.py +++ b/ivy/func_wrapper.py @@ -494,7 +494,7 @@ def _inputs_to_ivy_arrays(*args, **kwargs): has_out = True # convert all arrays in the inputs to ivy.Array instances ivy_args, ivy_kwargs = ivy.args_to_ivy( - *args, **kwargs, include_derived={tuple: True} + *args, **kwargs, include_derived={"tuple": True} ) if has_out: ivy_kwargs["out"] = out @@ -564,7 +564,7 @@ def _outputs_to_ivy_arrays(*args, **kwargs): ret = fn(*args, **kwargs) # convert all arrays in the return to `ivy.Array` instances return ( - ivy.to_ivy(ret, nested=True, include_derived={tuple: True}) + ivy.to_ivy(ret, nested=True, include_derived={"tuple": True}) if ivy.array_mode else ret ) @@ -594,7 +594,7 @@ def output_to_native_arrays(fn: Callable) -> Callable: @functools.wraps(fn) def _output_to_native_arrays(*args, **kwargs): ret = fn(*args, **kwargs) - return ivy.to_native(ret, nested=True, include_derived={tuple: True}) + return ivy.to_native(ret, nested=True, include_derived={"tuple": True}) _output_to_native_arrays.outputs_to_native_arrays = True return _output_to_native_arrays diff --git a/ivy/functional/frontends/jax/func_wrapper.py b/ivy/functional/frontends/jax/func_wrapper.py index cace90dd4fe33..ba89504f7243c 100644 --- a/ivy/functional/frontends/jax/func_wrapper.py +++ b/ivy/functional/frontends/jax/func_wrapper.py @@ -111,10 +111,10 @@ def _inputs_to_ivy_arrays_jax(*args, **kwargs): has_out = True # convert all arrays in the inputs to ivy.Array instances new_args = ivy.nested_map( - args, _to_ivy_array, include_derived={tuple: True}, shallow=False + args, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) new_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True}, shallow=False + kwargs, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) # add the original out argument back to the keyword arguments if has_out: @@ -153,10 +153,10 @@ def _outputs_to_frontend_arrays_jax(*args, **kwargs): return _from_ivy_array_to_jax_frontend_array_weak_type( ret, nested=True, - include_derived={tuple: True}, + include_derived={"tuple": True}, ) return _from_ivy_array_to_jax_frontend_array( - ret, nested=True, include_derived={tuple: True} + ret, nested=True, include_derived={"tuple": True} ) return _outputs_to_frontend_arrays_jax diff --git a/ivy/functional/frontends/mxnet/func_wrapper.py b/ivy/functional/frontends/mxnet/func_wrapper.py index 20b358488ea89..9beab9f42365d 100644 --- a/ivy/functional/frontends/mxnet/func_wrapper.py +++ b/ivy/functional/frontends/mxnet/func_wrapper.py @@ -81,10 +81,10 @@ def _inputs_to_ivy_arrays_mxnet(*args, **kwargs): """ # convert all arrays in the inputs to ivy.Array instances new_args = ivy.nested_map( - args, _to_ivy_array, include_derived={tuple: True}, shallow=False + args, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) new_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True}, shallow=False + kwargs, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) return fn(*new_args, **new_kwargs) @@ -105,7 +105,7 @@ def _outputs_to_frontend_arrays_mxnet(*args, **kwargs): ret = fn(*args, **kwargs) # convert all arrays in the return to `frontend.Tensorflow.tensor` instances - return ivy.nested_map(ret, _ivy_array_to_mxnet, include_derived={tuple: True}) + return ivy.nested_map(ret, _ivy_array_to_mxnet, include_derived={"tuple": True}) _outputs_to_frontend_arrays_mxnet.outputs_to_frontend_arrays = True return _outputs_to_frontend_arrays_mxnet diff --git a/ivy/functional/frontends/numpy/func_wrapper.py b/ivy/functional/frontends/numpy/func_wrapper.py index 495820f649e0f..a001f245ec3d5 100644 --- a/ivy/functional/frontends/numpy/func_wrapper.py +++ b/ivy/functional/frontends/numpy/func_wrapper.py @@ -194,7 +194,7 @@ def _set_order(args, order): ) if order in ["K", "A", None]: check_order = ivy.nested_map( - args, _check_C_order, include_derived={tuple: True}, shallow=False + args, _check_C_order, include_derived={"tuple": True}, shallow=False ) if all(v is None for v in check_order) or any( ivy.multi_index_nest(check_order, ivy.all_nested_indices(check_order)) @@ -447,9 +447,9 @@ def _inputs_to_ivy_arrays_np(*args, **kwargs): The return of the function, with ivy arrays passed in the arguments. """ # convert all arrays in the inputs to ivy.Array instances - ivy_args = ivy.nested_map(args, _to_ivy_array, include_derived={tuple: True}) + ivy_args = ivy.nested_map(args, _to_ivy_array, include_derived={"tuple": True}) ivy_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True} + kwargs, _to_ivy_array, include_derived={"tuple": True} ) return fn(*ivy_args, **ivy_kwargs) @@ -509,10 +509,10 @@ def _outputs_to_frontend_arrays(*args, order="K", **kwargs): # convert all returned arrays to `ndarray` instances if order == "F": return ivy.nested_map( - ret, _ivy_to_numpy_order_F, include_derived={tuple: True} + ret, _ivy_to_numpy_order_F, include_derived={"tuple": True} ) else: - return ivy.nested_map(ret, _ivy_to_numpy, include_derived={tuple: True}) + return ivy.nested_map(ret, _ivy_to_numpy, include_derived={"tuple": True}) if "order" in list(inspect.signature(fn).parameters.keys()): contains_order = True diff --git a/ivy/functional/frontends/onnx/func_wrapper.py b/ivy/functional/frontends/onnx/func_wrapper.py index ced490f7f96a1..d1e784c7a197e 100644 --- a/ivy/functional/frontends/onnx/func_wrapper.py +++ b/ivy/functional/frontends/onnx/func_wrapper.py @@ -58,10 +58,10 @@ def _inputs_to_ivy_arrays_onnx(*args, **kwargs): """ # convert all arrays in the inputs to ivy.Array instances new_args = ivy.nested_map( - args, _to_ivy_array, include_derived={tuple: True}, shallow=False + args, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) new_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True}, shallow=False + kwargs, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) return fn(*new_args, **new_kwargs) @@ -82,7 +82,7 @@ def _outputs_to_frontend_arrays_onnx(*args, **kwargs): # convert all arrays in the return to `frontend.onnx.Tensor` instances return _from_ivy_array_to_onnx_frontend_tensor( - ret, nested=True, include_derived={tuple: True} + ret, nested=True, include_derived={"tuple": True} ) return _outputs_to_frontend_arrays_onnx diff --git a/ivy/functional/frontends/paddle/func_wrapper.py b/ivy/functional/frontends/paddle/func_wrapper.py index 1af63572cc413..b0ddb992e05e9 100644 --- a/ivy/functional/frontends/paddle/func_wrapper.py +++ b/ivy/functional/frontends/paddle/func_wrapper.py @@ -50,10 +50,10 @@ def new_fn(*args, **kwargs): """ # convert all input arrays to ivy.Array instances new_args = ivy.nested_map( - args, _to_ivy_array, include_derived={tuple: True}, shallow=False + args, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) new_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True}, shallow=False + kwargs, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) return fn(*new_args, **new_kwargs) @@ -82,7 +82,7 @@ def new_fn(*args, **kwargs): ivy.unset_default_float_dtype() # convert all arrays in the return to `paddle_frontend.Tensor` instances return _from_ivy_array_to_paddle_frontend_tensor( - ret, nested=True, include_derived={tuple: True} + ret, nested=True, include_derived={"tuple": True} ) return new_fn diff --git a/ivy/functional/frontends/tensorflow/func_wrapper.py b/ivy/functional/frontends/tensorflow/func_wrapper.py index ca14f5fc1845a..06fee643e81e0 100644 --- a/ivy/functional/frontends/tensorflow/func_wrapper.py +++ b/ivy/functional/frontends/tensorflow/func_wrapper.py @@ -217,7 +217,7 @@ def _outputs_to_frontend_arrays_tf(*args, **kwargs): # convert all arrays in the return to `frontend.Tensorflow.tensor` instances return ivy.nested_map( - ret, _ivy_array_to_tensorflow, include_derived={tuple: True} + ret, _ivy_array_to_tensorflow, include_derived={"tuple": True} ) _outputs_to_frontend_arrays_tf.outputs_to_frontend_arrays = True diff --git a/ivy/functional/frontends/torch/func_wrapper.py b/ivy/functional/frontends/torch/func_wrapper.py index 84e16bb9d797d..da2fad9b18517 100644 --- a/ivy/functional/frontends/torch/func_wrapper.py +++ b/ivy/functional/frontends/torch/func_wrapper.py @@ -145,10 +145,10 @@ def _inputs_to_ivy_arrays_torch(*args, **kwargs): ) # convert all input arrays to ivy.Array instances new_args = ivy.nested_map( - args, _to_ivy_array, include_derived={tuple: True}, shallow=False + args, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) new_kwargs = ivy.nested_map( - kwargs, _to_ivy_array, include_derived={tuple: True}, shallow=False + kwargs, _to_ivy_array, include_derived={"tuple": True}, shallow=False ) return fn(*new_args, **new_kwargs) @@ -202,7 +202,7 @@ def outputs_to_frontend_arrays_torch(*args, **kwargs): ret = _from_ivy_array_to_torch_frontend_tensor( ret, nested=True, - include_derived={tuple: True}, + include_derived={"tuple": True}, requires_grad=kwargs.get( "requires_grad", any( diff --git a/ivy/functional/ivy/nest.py b/ivy/functional/ivy/nest.py index 042826a6087de..c0826e5e8d653 100644 --- a/ivy/functional/ivy/nest.py +++ b/ivy/functional/ivy/nest.py @@ -1036,7 +1036,7 @@ def nested_map( x: Union[ivy.Array, ivy.NativeArray, Iterable], /, fn: Callable, - include_derived: Optional[Union[Dict[type, bool], bool]] = None, + include_derived: Optional[Union[Dict[str, bool], bool]] = None, to_ignore: Optional[Union[type, Tuple[type]]] = None, to_mutable: bool = False, max_depth: Optional[int] = None, @@ -1160,14 +1160,13 @@ def nested_map( to_ignore = ivy.default(to_ignore, ()) extra_nest_types = ivy.default(extra_nest_types, ()) if include_derived is True: - include_derived = {str(tuple): True, str(list): True, str(dict): True} + include_derived = {"tuple": True, "list": True, "dict": True} elif not include_derived: include_derived = {} - for t in (tuple, list, dict): + for t in ("tuple", "list", "dict"): if t not in include_derived: - include_derived[str(t)] = False + include_derived[t] = False # to ensure all keys are strings - include_derived = {str(key): value for key, value in include_derived.items()} if ivy.exists(max_depth) and _depth > max_depth: return x class_instance = type(x) @@ -1184,7 +1183,7 @@ def nested_map( _tuple_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[str(tuple)] + if include_derived["tuple"] else (lambda x_, t_: type(x_) is t_) ), ) @@ -1192,7 +1191,7 @@ def nested_map( _list_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[str(list)] + if include_derived["list"] else (lambda x_, t_: type(x_) is t_) ), ) @@ -1200,7 +1199,7 @@ def nested_map( _dict_check_fn, ( (lambda x_, t_: isinstance(x_, t_)) - if include_derived[str(dict)] + if include_derived["dict"] else (lambda x_, t_: type(x_) is t_) ), )