From 6515e43972e40d996653e086ff1c637973966898 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Thu, 31 Aug 2023 16:53:07 +0530 Subject: [PATCH 1/3] refactor(All APIs): Added new pre-commit hook for ordering functions --- .pre-commit-config.yaml | 2 +- ivy/functional/backends/jax/activations.py | 32 +- ivy/functional/backends/jax/creation.py | 138 +- ivy/functional/backends/jax/data_type.py | 168 +- ivy/functional/backends/jax/device.py | 146 +- ivy/functional/backends/jax/elementwise.py | 256 +- .../backends/jax/experimental/activations.py | 40 +- .../backends/jax/experimental/creation.py | 77 +- .../backends/jax/experimental/elementwise.py | 284 +- .../backends/jax/experimental/layers.py | 748 +-- .../jax/experimental/linear_algebra.py | 82 +- .../backends/jax/experimental/manipulation.py | 354 +- .../backends/jax/experimental/random.py | 65 +- .../backends/jax/experimental/statistical.py | 416 +- ivy/functional/backends/jax/general.py | 252 +- ivy/functional/backends/jax/gradients.py | 68 +- ivy/functional/backends/jax/layers.py | 156 +- ivy/functional/backends/jax/linear_algebra.py | 92 +- ivy/functional/backends/jax/manipulation.py | 200 +- ivy/functional/backends/jax/random.py | 98 +- ivy/functional/backends/jax/searching.py | 26 +- ivy/functional/backends/jax/sorting.py | 30 +- ivy/functional/backends/jax/statistical.py | 197 +- ivy/functional/backends/mxnet/activations.py | 16 +- ivy/functional/backends/mxnet/creation.py | 80 +- ivy/functional/backends/mxnet/data_type.py | 127 +- ivy/functional/backends/mxnet/device.py | 64 +- ivy/functional/backends/mxnet/elementwise.py | 240 +- .../mxnet/experimental/activations.py | 12 +- .../backends/mxnet/experimental/creation.py | 40 +- .../mxnet/experimental/elementwise.py | 172 +- .../backends/mxnet/experimental/layers.py | 152 +- .../mxnet/experimental/linear_algebra.py | 100 +- .../mxnet/experimental/manipulation.py | 154 +- .../backends/mxnet/experimental/norms.py | 20 +- .../backends/mxnet/experimental/random.py | 38 +- .../mxnet/experimental/statistical.py | 86 +- ivy/functional/backends/mxnet/general.py | 40 +- ivy/functional/backends/mxnet/gradients.py | 30 +- ivy/functional/backends/mxnet/layers.py | 28 +- .../backends/mxnet/linear_algebra.py | 42 +- ivy/functional/backends/mxnet/manipulation.py | 88 +- ivy/functional/backends/mxnet/random.py | 48 +- ivy/functional/backends/mxnet/searching.py | 18 +- ivy/functional/backends/mxnet/sorting.py | 24 +- ivy/functional/backends/mxnet/statistical.py | 72 +- ivy/functional/backends/numpy/activations.py | 116 +- ivy/functional/backends/numpy/creation.py | 140 +- ivy/functional/backends/numpy/data_type.py | 166 +- ivy/functional/backends/numpy/device.py | 100 +- ivy/functional/backends/numpy/elementwise.py | 574 +- .../numpy/experimental/activations.py | 64 +- .../backends/numpy/experimental/creation.py | 125 +- .../numpy/experimental/elementwise.py | 584 +- .../backends/numpy/experimental/layers.py | 1180 ++-- .../numpy/experimental/linear_algebra.py | 118 +- .../numpy/experimental/manipulation.py | 434 +- .../backends/numpy/experimental/random.py | 60 +- .../backends/numpy/experimental/searching.py | 6 +- .../backends/numpy/experimental/sorting.py | 6 +- .../numpy/experimental/statistical.py | 618 +- ivy/functional/backends/numpy/general.py | 188 +- ivy/functional/backends/numpy/gradients.py | 58 +- ivy/functional/backends/numpy/helpers.py | 50 +- ivy/functional/backends/numpy/layers.py | 120 +- .../backends/numpy/linear_algebra.py | 98 +- ivy/functional/backends/numpy/manipulation.py | 158 +- ivy/functional/backends/numpy/random.py | 73 +- ivy/functional/backends/numpy/searching.py | 26 +- ivy/functional/backends/numpy/sorting.py | 38 +- ivy/functional/backends/numpy/statistical.py | 216 +- ivy/functional/backends/numpy/utility.py | 10 +- ivy/functional/backends/paddle/activations.py | 128 +- ivy/functional/backends/paddle/creation.py | 453 +- ivy/functional/backends/paddle/data_type.py | 139 +- ivy/functional/backends/paddle/device.py | 88 +- ivy/functional/backends/paddle/elementwise.py | 1860 +++--- .../paddle/experimental/activations.py | 72 +- .../backends/paddle/experimental/creation.py | 143 +- .../paddle/experimental/elementwise.py | 590 +- .../backends/paddle/experimental/layers.py | 499 +- .../paddle/experimental/linear_algebra.py | 74 +- .../backends/paddle/experimental/losses.py | 24 +- .../paddle/experimental/manipulation.py | 615 +- .../backends/paddle/experimental/norms.py | 56 +- .../backends/paddle/experimental/random.py | 127 +- .../paddle/experimental/statistical.py | 706 +- ivy/functional/backends/paddle/general.py | 202 +- ivy/functional/backends/paddle/gradients.py | 222 +- ivy/functional/backends/paddle/layers.py | 72 +- .../backends/paddle/linear_algebra.py | 142 +- .../backends/paddle/manipulation.py | 336 +- ivy/functional/backends/paddle/random.py | 111 +- ivy/functional/backends/paddle/searching.py | 50 +- ivy/functional/backends/paddle/sorting.py | 48 +- ivy/functional/backends/paddle/statistical.py | 359 +- .../backends/tensorflow/activations.py | 44 +- .../backends/tensorflow/control_flow_ops.py | 77 +- .../backends/tensorflow/creation.py | 160 +- .../backends/tensorflow/data_type.py | 176 +- ivy/functional/backends/tensorflow/device.py | 126 +- .../backends/tensorflow/elementwise.py | 482 +- .../tensorflow/experimental/activations.py | 44 +- .../tensorflow/experimental/creation.py | 120 +- .../tensorflow/experimental/elementwise.py | 362 +- .../tensorflow/experimental/layers.py | 1494 ++--- .../tensorflow/experimental/linear_algebra.py | 216 +- .../tensorflow/experimental/manipulation.py | 360 +- .../backends/tensorflow/experimental/norms.py | 48 +- .../tensorflow/experimental/random.py | 90 +- .../tensorflow/experimental/statistical.py | 658 +- ivy/functional/backends/tensorflow/general.py | 270 +- .../backends/tensorflow/gradients.py | 180 +- ivy/functional/backends/tensorflow/layers.py | 92 +- .../backends/tensorflow/linear_algebra.py | 60 +- .../backends/tensorflow/manipulation.py | 226 +- ivy/functional/backends/tensorflow/random.py | 82 +- .../backends/tensorflow/searching.py | 46 +- ivy/functional/backends/tensorflow/sorting.py | 46 +- .../backends/tensorflow/statistical.py | 159 +- ivy/functional/backends/torch/activations.py | 114 +- ivy/functional/backends/torch/creation.py | 258 +- ivy/functional/backends/torch/data_type.py | 135 +- ivy/functional/backends/torch/device.py | 107 +- ivy/functional/backends/torch/elementwise.py | 1130 ++-- .../torch/experimental/activations.py | 44 +- .../backends/torch/experimental/creation.py | 149 +- .../torch/experimental/elementwise.py | 340 +- .../backends/torch/experimental/layers.py | 988 +-- .../torch/experimental/linear_algebra.py | 112 +- .../backends/torch/experimental/losses.py | 35 +- .../torch/experimental/manipulation.py | 396 +- .../backends/torch/experimental/norms.py | 116 +- .../backends/torch/experimental/random.py | 104 +- .../backends/torch/experimental/searching.py | 6 +- .../backends/torch/experimental/sorting.py | 6 +- .../torch/experimental/statistical.py | 756 ++- ivy/functional/backends/torch/general.py | 334 +- ivy/functional/backends/torch/gradients.py | 146 +- ivy/functional/backends/torch/layers.py | 124 +- .../backends/torch/linear_algebra.py | 198 +- ivy/functional/backends/torch/manipulation.py | 264 +- ivy/functional/backends/torch/random.py | 97 +- ivy/functional/backends/torch/searching.py | 26 +- ivy/functional/backends/torch/sorting.py | 50 +- ivy/functional/backends/torch/statistical.py | 270 +- ivy/functional/backends/torch/utility.py | 10 +- .../frontends/torch/comparison_ops.py | 2 +- ivy/functional/ivy/activations.py | 368 +- ivy/functional/ivy/constants.py | 81 +- ivy/functional/ivy/control_flow_ops.py | 202 +- ivy/functional/ivy/creation.py | 2441 +++---- ivy/functional/ivy/data_type.py | 2036 +++--- ivy/functional/ivy/device.py | 1391 ++-- ivy/functional/ivy/elementwise.py | 3932 ++++++------ .../ivy/experimental/activations.py | 340 +- ivy/functional/ivy/experimental/creation.py | 990 +-- .../ivy/experimental/elementwise.py | 1264 ++-- ivy/functional/ivy/experimental/general.py | 8 + ivy/functional/ivy/experimental/layers.py | 4041 ++++++------ .../ivy/experimental/linear_algebra.py | 1755 ++--- ivy/functional/ivy/experimental/losses.py | 178 +- .../ivy/experimental/manipulation.py | 3478 +++++----- ivy/functional/ivy/experimental/norms.py | 332 +- ivy/functional/ivy/experimental/random.py | 204 +- .../ivy/experimental/sparse_array.py | 1118 ++-- .../ivy/experimental/statistical.py | 1232 ++-- ivy/functional/ivy/general.py | 5674 ++++++++--------- ivy/functional/ivy/gradients.py | 1686 +++-- ivy/functional/ivy/layers.py | 3223 +++++----- ivy/functional/ivy/linear_algebra.py | 487 +- ivy/functional/ivy/losses.py | 124 +- ivy/functional/ivy/manipulation.py | 824 +-- ivy/functional/ivy/meta.py | 399 +- ivy/functional/ivy/nest.py | 1716 ++--- ivy/functional/ivy/norms.py | 24 +- ivy/functional/ivy/random.py | 466 +- ivy/functional/ivy/searching.py | 150 +- ivy/functional/ivy/set.py | 210 +- ivy/functional/ivy/sorting.py | 214 +- ivy/functional/ivy/statistical.py | 1554 ++--- ivy/functional/ivy/utility.py | 26 +- 182 files changed, 36023 insertions(+), 36140 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f4c4c298876fb..34c2ff59c9c14 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,6 +26,6 @@ repos: # Exclude everything in frontends except __init__.py, and func_wrapper.py exclude: 'ivy/functional/(frontends|backends)/(?!.*/func_wrapper\.py$).*(?!__init__\.py$)' - repo: https://github.com/unifyai/lint-hook - rev: 27646397c5390f644a645f439535b1061b9c0105 + rev: 5abf78187bb7a5f839174ce4febd4e63f5d758f6 hooks: - id: ivy-lint diff --git a/ivy/functional/backends/jax/activations.py b/ivy/functional/backends/jax/activations.py index 6d980aaf6022e..7732873e41b6f 100644 --- a/ivy/functional/backends/jax/activations.py +++ b/ivy/functional/backends/jax/activations.py @@ -23,6 +23,10 @@ def gelu( return jax.nn.gelu(x, approximate) +def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.hard_swish(x) + + def leaky_relu( x: JaxArray, /, @@ -34,6 +38,18 @@ def leaky_relu( return jnp.asarray(jnp.where(x > 0, x, jnp.multiply(x, alpha)), x.dtype) +def log_softmax( + x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None +): + if axis is None: + axis = -1 + return jax.nn.log_softmax(x, axis) + + +def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): + return x * jnp.tanh(jax.nn.softplus(x)) + + def relu( x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None ) -> JaxArray: @@ -78,19 +94,3 @@ def softplus( if threshold is not None: return jnp.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -def log_softmax( - x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None -): - if axis is None: - axis = -1 - return jax.nn.log_softmax(x, axis) - - -def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): - return x * jnp.tanh(jax.nn.softplus(x)) - - -def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.hard_swish(x) diff --git a/ivy/functional/backends/jax/creation.py b/ivy/functional/backends/jax/creation.py index 4fbc39f493756..45b8ad89b441f 100644 --- a/ivy/functional/backends/jax/creation.py +++ b/ivy/functional/backends/jax/creation.py @@ -79,6 +79,19 @@ def asarray( return jnp.asarray(obj, dtype=dtype) +def copy_array( + x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + x = ( + jax.core.ShapedArray(x.shape, x.dtype) + if isinstance(x, jax.core.ShapedArray) + else jnp.array(x) + ) + if to_ivy_array: + return ivy.to_ivy(x) + return x + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -127,6 +140,15 @@ def from_dlpack(x, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jax.dlpack.from_dlpack(capsule) +def frombuffer( + buffer: bytes, + dtype: Optional[jnp.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> JaxArray: + return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -234,6 +256,42 @@ def meshgrid( return jnp.meshgrid(*arrays, sparse=sparse, indexing=indexing) +def one_hot( + indices: JaxArray, + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[jnp.dtype] = None, + device: jaxlib.xla_extension.Device, + out: Optional[JaxArray] = None, +) -> JaxArray: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = jnp.float32 + else: + if not on_none: + dtype = jnp.array(on_value).dtype + elif not off_none: + dtype = jnp.array(off_value).dtype + + res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] + res = res.reshape(list(indices.shape) + [depth]) + + if not on_none and not off_none: + res = jnp.where(res == 1, on_value, off_value) + + if axis is not None: + res = jnp.moveaxis(res, -1, axis) + + return res + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -263,6 +321,17 @@ def triu(x: JaxArray, /, *, k: int = 0, out: Optional[JaxArray] = None) -> JaxAr return jnp.triu(x, k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: jaxlib.xla_extension.Device, +) -> Tuple[JaxArray]: + return jnp.triu_indices(n=n_rows, k=k, m=n_cols) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -289,72 +358,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - x = ( - jax.core.ShapedArray(x.shape, x.dtype) - if isinstance(x, jax.core.ShapedArray) - else jnp.array(x) - ) - if to_ivy_array: - return ivy.to_ivy(x) - return x - - -def one_hot( - indices: JaxArray, - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - device: jaxlib.xla_extension.Device, - out: Optional[JaxArray] = None, -) -> JaxArray: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = jnp.float32 - else: - if not on_none: - dtype = jnp.array(on_value).dtype - elif not off_none: - dtype = jnp.array(off_value).dtype - - res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] - res = res.reshape(list(indices.shape) + [depth]) - - if not on_none and not off_none: - res = jnp.where(res == 1, on_value, off_value) - - if axis is not None: - res = jnp.moveaxis(res, -1, axis) - - return res - - -def frombuffer( - buffer: bytes, - dtype: Optional[jnp.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> JaxArray: - return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: jaxlib.xla_extension.Device, -) -> Tuple[JaxArray]: - return jnp.triu_indices(n=n_rows, k=k, m=n_cols) diff --git a/ivy/functional/backends/jax/data_type.py b/ivy/functional/backends/jax/data_type.py index e25e54e745bbe..9a53044c244f3 100644 --- a/ivy/functional/backends/jax/data_type.py +++ b/ivy/functional/backends/jax/data_type.py @@ -9,6 +9,26 @@ from ivy.functional.backends.jax import JaxArray from ivy.functional.ivy.data_type import _handle_nestable_dtype_info +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} ivy_dtype_dict = { jnp.dtype("int8"): "int8", jnp.dtype("int16"): "int16", @@ -41,7 +61,6 @@ jnp.complex128: "complex128", jnp.bool_: "bool", } - native_dtype_dict = { "int8": jnp.dtype("int8"), "int16": jnp.dtype("int16"), @@ -60,27 +79,6 @@ "bool": jnp.dtype("bool"), } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} - class Finfo: def __init__(self, jnp_finfo: jnp.finfo): @@ -110,69 +108,6 @@ def smallest_normal(self): return float(self._jnp_finfo.tiny) -# Array API Standard # -# -------------------# - - -def astype( - x: JaxArray, - dtype: jnp.dtype, - /, - *, - copy: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - ivy.utils.assertions._check_jax_x64_flag(dtype) - if x.dtype == dtype: - return jnp.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: - try: - return jnp.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: JaxArray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return jnp.broadcast_to(x.reshape(-1), shape) - return jnp.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return Finfo(jnp.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return jnp.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return jnp.result_type(arrays_and_dtypes) - - result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = jnp.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -245,6 +180,45 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: JaxArray, + dtype: jnp.dtype, + /, + *, + copy: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + ivy.utils.assertions._check_jax_x64_flag(dtype) + if x.dtype == dtype: + return jnp.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: + try: + return jnp.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: JaxArray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return jnp.broadcast_to(x.reshape(-1), shape) + return jnp.broadcast_to(x, shape) + + def dtype(x: Union[JaxArray, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -264,6 +238,20 @@ def dtype_bits(dtype_in: Union[jnp.dtype, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return Finfo(jnp.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return jnp.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -271,3 +259,13 @@ def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return jnp.result_type(arrays_and_dtypes) + + result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = jnp.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/jax/device.py b/ivy/functional/backends/jax/device.py index 5158c9a0809e5..9f41e96c4be8c 100644 --- a/ivy/functional/backends/jax/device.py +++ b/ivy/functional/backends/jax/device.py @@ -15,54 +15,44 @@ ) -# Helpers # -# --------# +# noinspection PyMethodMayBeStatic +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._save_dir = os.path.join(self._save_dir, "profile") + def start(self): + jax.profiler.start_trace(self._save_dir) -def _to_array(x): - if isinstance(x, jax.interpreters.ad.JVPTracer): - return _to_array(x.primal) - elif isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): - return _to_array(x.aval) - elif isinstance(x, jax.interpreters.batching.BatchTracer): - return _to_array(x.val) - return x + def stop(self): + jax.profiler.stop_trace() + def __enter__(self): + self.start() -# API # -# ----# + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() -def dev( - x: JaxArray, - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, jaxlib.xla_extension.Device]: - if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): - return "" +# --- Helpers --- # +# --------------- # + + +def _dev_is_available(base_dev): try: - dv = _to_array(x).device_buffer.device - dv = dv() - except Exception: - dv = jax.devices()[0] - if as_native: - return dv - return as_ivy_dev(dv) + jax.devices(base_dev) + return True + except RuntimeError: + return False -def to_device( - x: JaxArray, - device: jaxlib.xla_extension.Device, - /, - *, - stream: Optional[int] = None, - out: Optional[JaxArray] = None, -): - if device is not None: - cur_dev = as_native_dev(dev(x)) - if cur_dev != device: - x = jax.device_put(x, as_native_dev(device)) +def _to_array(x): + if isinstance(x, jax.interpreters.ad.JVPTracer): + return _to_array(x.primal) + elif isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): + return _to_array(x.aval) + elif isinstance(x, jax.interpreters.batching.BatchTracer): + return _to_array(x.val) return x @@ -76,6 +66,10 @@ def _to_device(x, device=None): return x +# --- Main --- # +# ------------ # + + def as_ivy_dev(device, /): if isinstance(device, str): return ivy.Device(device) @@ -99,30 +93,44 @@ def as_native_dev(device, /): return jax.devices(device)[idx] -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( - *args, device_shifting_dev=device_shifting_dev, **kwargs - ) - with jax.default_device(device_shifting_dev): - return fn(*args, **kwargs) - - def clear_cached_mem_on_dev(device: str, /): return None -def _dev_is_available(base_dev): +# API # +# ----# + + +def dev( + x: JaxArray, + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, jaxlib.xla_extension.Device]: + if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): + return "" try: - jax.devices(base_dev) - return True - except RuntimeError: - return False + dv = _to_array(x).device_buffer.device + dv = dv() + except Exception: + dv = jax.devices()[0] + if as_native: + return dv + return as_ivy_dev(dv) def gpu_is_available() -> bool: return _dev_is_available("gpu") +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( + *args, device_shifting_dev=device_shifting_dev, **kwargs + ) + with jax.default_device(device_shifting_dev): + return fn(*args, **kwargs) + + def num_gpus() -> int: try: return len(jax.devices("gpu")) @@ -130,24 +138,20 @@ def num_gpus() -> int: return 0 -def tpu_is_available() -> bool: - return _dev_is_available("tpu") - - -# noinspection PyMethodMayBeStatic -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._save_dir = os.path.join(self._save_dir, "profile") - - def start(self): - jax.profiler.start_trace(self._save_dir) - - def stop(self): - jax.profiler.stop_trace() +def to_device( + x: JaxArray, + device: jaxlib.xla_extension.Device, + /, + *, + stream: Optional[int] = None, + out: Optional[JaxArray] = None, +): + if device is not None: + cur_dev = as_native_dev(dev(x)) + if cur_dev != device: + x = jax.device_put(x, as_native_dev(device)) + return x - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + return _dev_is_available("tpu") diff --git a/ivy/functional/backends/jax/elementwise.py b/ivy/functional/backends/jax/elementwise.py index 6de81f65a16eb..77f304a382779 100644 --- a/ivy/functional/backends/jax/elementwise.py +++ b/ivy/functional/backends/jax/elementwise.py @@ -16,6 +16,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _abs_variant_sign(x): + return jnp.where(x != 0, x / jnp.abs(x), 0) + + +# --- Main --- # +# ------------ # + + def abs( x: Union[float, JaxArray], /, @@ -51,6 +63,16 @@ def add( return jnp.add(x1, x2) +def angle( + z: JaxArray, + /, + *, + deg: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.angle(z, deg=deg) + + def asin(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.arcsin(x) @@ -156,6 +178,10 @@ def cosh(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.cosh(x) +def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.deg2rad(x) + + def divide( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -183,10 +209,28 @@ def equal( return jnp.equal(x1, x2) +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.scipy.special.erf(x) + + def exp(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.exp(x) +def exp2( + x: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.power(2, x) + + def expm1(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.expm1(x) @@ -221,6 +265,29 @@ def fmin( return jnp.fmin(x1, x2) +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def fmod( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmod(x1, x2) + + +def gcd( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.gcd(x1, x2) + + def greater( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -243,6 +310,15 @@ def greater_equal( return jnp.greater_equal(x1, x2) +def imag( + val: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.imag(val) + + def isfinite(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isfinite(x) @@ -268,6 +344,10 @@ def isnan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isnan(x) +def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.isreal(x) + + def lcm(x1: JaxArray, x2: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: x1, x2 = promote_types_of_inputs(x1, x2) return jnp.lcm(x1, x2) @@ -353,6 +433,34 @@ def logical_xor( return jnp.logical_xor(x1, x2) +def maximum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 >= x2, x1, x2) + return jnp.maximum(x1, x2) + + +def minimum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 <= x2, x1, x2) + return jnp.minimum(x1, x2) + + def multiply( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -411,6 +519,20 @@ def pow( return jnp.power(x1, x2) +def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.rad2deg(x) + + +def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.real(x) + + +def reciprocal( + x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.reciprocal(x) + + @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def remainder( x1: Union[float, JaxArray], @@ -442,10 +564,6 @@ def round( return ret -def _abs_variant_sign(x): - return jnp.where(x != 0, x / jnp.abs(x), 0) - - def sign( x: JaxArray, /, *, np_variant: Optional[bool] = True, out: Optional[JaxArray] = None ) -> JaxArray: @@ -486,6 +604,16 @@ def subtract( return jnp.subtract(x1, x2) +def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.tan(x) + + +def tanh( + x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.tanh(x) + + def trapz( y: JaxArray, /, @@ -498,129 +626,9 @@ def trapz( return jnp.trapz(y, x=x, dx=dx, axis=axis) -def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.tan(x) - - -def tanh( - x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.tanh(x) - - @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def trunc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: if "int" in str(x.dtype): return x else: return jnp.trunc(x) - - -def exp2( - x: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.power(2, x) - - -def imag( - val: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.imag(val) - - -def angle( - z: JaxArray, - /, - *, - deg: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.angle(z, deg=deg) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.scipy.special.erf(x) - - -def maximum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 >= x2, x1, x2) - return jnp.maximum(x1, x2) - - -def minimum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 <= x2, x1, x2) - return jnp.minimum(x1, x2) - - -def reciprocal( - x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.reciprocal(x) - - -def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.deg2rad(x) - - -def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.rad2deg(x) - - -def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.isreal(x) - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def fmod( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmod(x1, x2) - - -def gcd( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.gcd(x1, x2) - - -def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.real(x) diff --git a/ivy/functional/backends/jax/experimental/activations.py b/ivy/functional/backends/jax/experimental/activations.py index 788126f4b44cf..8d270c6c8458a 100644 --- a/ivy/functional/backends/jax/experimental/activations.py +++ b/ivy/functional/backends/jax/experimental/activations.py @@ -8,6 +8,15 @@ import ivy +def elu( + x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None +) -> JaxArray: + ret = jax.nn.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret + + def logit( x: JaxArray, /, @@ -22,6 +31,10 @@ def logit( return jnp.log(x / (1 - x)) +def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.log_sigmoid(input) + + def relu6(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: relu6_func = jax.nn.relu6 @@ -38,20 +51,6 @@ def custom_grad_func(x_and_grad, one): return new_func(x).astype(x.dtype) -def thresholded_relu( - x: JaxArray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.where(x > threshold, x, 0).astype(x.dtype) - - -def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.log_sigmoid(input) - - def selu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: ret = jax.nn.selu(x).astype(x.dtype) if ivy.exists(out): @@ -66,10 +65,11 @@ def silu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return ret -def elu( - x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None +def thresholded_relu( + x: JaxArray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[JaxArray] = None, ) -> JaxArray: - ret = jax.nn.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret + return jnp.where(x > threshold, x, 0).astype(x.dtype) diff --git a/ivy/functional/backends/jax/experimental/creation.py b/ivy/functional/backends/jax/experimental/creation.py index 2bbc0b1b895e6..343a8dfaf9b49 100644 --- a/ivy/functional/backends/jax/experimental/creation.py +++ b/ivy/functional/backends/jax/experimental/creation.py @@ -9,27 +9,23 @@ from ivy.functional.backends.jax import JaxArray import ivy -# Array API Standard # -# ------------------ # - -def vorbis_window( - window_length: JaxArray, +def blackman_window( + size: int, + /, *, - dtype: jnp.dtype = jnp.float32, + periodic: bool = True, + dtype: Optional[jnp.dtype] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.array( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, + if size < 2: + return jnp.ones([size], dtype=dtype) + if periodic: + count = jnp.arange(size) / size + else: + count = jnp.linspace(start=0, stop=size, num=size) + return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( + 0.08 * jnp.cos(2 * jnp.pi * 2 * count) ) @@ -77,6 +73,14 @@ def tril_indices( return jnp.tril_indices(n=n_rows, k=k, m=n_cols) +def trilu( + x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + if upper: + return jnp.triu(x, k) + return jnp.tril(x, k) + + def unsorted_segment_min( data: JaxArray, segment_ids: JaxArray, @@ -104,28 +108,25 @@ def unsorted_segment_sum( return jax.ops.segment_sum(data, segment_ids, num_segments) -def blackman_window( - size: int, - /, +# Array API Standard # +# ------------------ # + + +def vorbis_window( + window_length: JaxArray, *, - periodic: bool = True, - dtype: Optional[jnp.dtype] = None, + dtype: jnp.dtype = jnp.float32, out: Optional[JaxArray] = None, ) -> JaxArray: - if size < 2: - return jnp.ones([size], dtype=dtype) - if periodic: - count = jnp.arange(size) / size - else: - count = jnp.linspace(start=0, stop=size, num=size) - return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( - 0.08 * jnp.cos(2 * jnp.pi * 2 * count) + return jnp.array( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, ) - - -def trilu( - x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - if upper: - return jnp.triu(x, k) - return jnp.tril(x, k) diff --git a/ivy/functional/backends/jax/experimental/elementwise.py b/ivy/functional/backends/jax/experimental/elementwise.py index f6517ca26c95d..c05fecdf8c937 100644 --- a/ivy/functional/backends/jax/experimental/elementwise.py +++ b/ivy/functional/backends/jax/experimental/elementwise.py @@ -19,41 +19,65 @@ jax_ArrayLike = Union[JaxArray, Number] -def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.sinc(x) +# --- Helpers --- # +# --------------- # -@with_supported_dtypes( - {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version -) -def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jlax.lgamma(x) +# def gradient( +# x: JaxArray, +# /, +# *, +# spacing: Optional[Union[int, list, tuple]] = 1, +# axis: Optional[Union[int, list, tuple]] = None, +# edge_order: Optional[int] = 1, +# ) -> Union[JaxArray, List[JaxArray]]: +# if type(spacing) == int: +# return jnp.gradient(x, spacing, axis=axis) +# return jnp.gradient(x, *spacing, axis=axis) -def fmax( +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim + + +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis + + +# --- Main --- # +# ------------ # + + +def allclose( x1: JaxArray, x2: JaxArray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmax(x1, x2) +) -> bool: + return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) -def float_power( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], +def conj( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): - out_dtype = jnp.complex128 - else: - out_dtype = jnp.float64 - return jnp.float_power(x1, x2).astype(out_dtype) + return jnp.conj(x) def copysign( @@ -86,65 +110,6 @@ def count_nonzero( return jnp.array(jnp.count_nonzero(a, axis=axis, keepdims=keepdims), dtype=dtype) -def nansum( - x: JaxArray, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[jnp.dtype] = None, - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - if isinstance(axis, list): - axis = tuple(axis) - return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -def isclose( - a: JaxArray, - b: JaxArray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -def signbit( - x: Union[JaxArray, float, int, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.signbit(x) - - -def hypot( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.hypot(x1, x2) - - -def allclose( - x1: JaxArray, - x2: JaxArray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[JaxArray] = None, -) -> bool: - return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - - def diff( x: JaxArray, /, @@ -163,74 +128,54 @@ def diff( return jnp.diff(x, n=n, axis=axis, prepend=prepend, append=append) -def fix( +def digamma( x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.fix(x, out=out) + return js.special.digamma(x) -def nextafter( - x1: JaxArray, - x2: JaxArray, +def fix( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.nextafter(x1, x2) + return jnp.fix(x, out=out) -def zeta( - x: JaxArray, - q: JaxArray, +def float_power( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) - inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) - temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) - ret = js.special.zeta(x, q) - ret = ret.at[nan_indices].set(jnp.nan) - ret = ret.at[inf_indices].set(jnp.inf) - return ret - - -# def gradient( -# x: JaxArray, -# /, -# *, -# spacing: Optional[Union[int, list, tuple]] = 1, -# axis: Optional[Union[int, list, tuple]] = None, -# edge_order: Optional[int] = 1, -# ) -> Union[JaxArray, List[JaxArray]]: -# if type(spacing) == int: -# return jnp.gradient(x, spacing, axis=axis) -# return jnp.gradient(x, *spacing, axis=axis) + x1, x2 = promote_types_of_inputs(x1, x2) + if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): + out_dtype = jnp.complex128 + else: + out_dtype = jnp.float64 + return jnp.float_power(x1, x2).astype(out_dtype) -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim +def fmax( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmax(x1, x2) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +def frexp( + x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None +) -> Tuple[JaxArray, JaxArray]: + return jnp.frexp(x) def gradient( @@ -422,18 +367,27 @@ def gradient( return outvals -def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - x, y = promote_types_of_inputs(x, y) - return js.special.xlogy(x, y) +def hypot( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.hypot(x1, x2) -def conj( - x: JaxArray, +def isclose( + a: JaxArray, + b: JaxArray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.conj(x) + return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) def ldexp( @@ -442,10 +396,11 @@ def ldexp( return jnp.ldexp(x1, x2) -def frexp( - x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None -) -> Tuple[JaxArray, JaxArray]: - return jnp.frexp(x) +@with_supported_dtypes( + {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version +) +def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jlax.lgamma(x) def modf( @@ -457,10 +412,63 @@ def modf( return jnp.modf(x) -def digamma( +def nansum( x: JaxArray, /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[jnp.dtype] = None, + keepdims: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return js.special.digamma(x) + if isinstance(axis, list): + axis = tuple(axis) + return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +def nextafter( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.nextafter(x1, x2) + + +def signbit( + x: Union[JaxArray, float, int, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.signbit(x) + + +def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.sinc(x) + + +def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + x, y = promote_types_of_inputs(x, y) + return js.special.xlogy(x, y) + + +def zeta( + x: JaxArray, + q: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) + inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) + temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) + ret = js.special.zeta(x, q) + ret = ret.at[nan_indices].set(jnp.nan) + ret = ret.at[inf_indices].set(jnp.inf) + return ret diff --git a/ivy/functional/backends/jax/experimental/layers.py b/ivy/functional/backends/jax/experimental/layers.py index e65256730b950..6ea8306098c5d 100644 --- a/ivy/functional/backends/jax/experimental/layers.py +++ b/ivy/functional/backends/jax/experimental/layers.py @@ -23,6 +23,10 @@ from ivy.functional.backends.jax.experimental.manipulation import _to_nested_tuple +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # determine depth pooling _, _, depth_pooling = _depth_max_pooling_helper( @@ -48,239 +52,8 @@ def _pad_str_to_list(inputs, dims, padding, strides, new_window_shape): return pad_list -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - # This function assumes that param validation is already done - window_shape = tuple(window_shape) - strides = (1,) + strides + (1,) if len(strides) == dim else strides - dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape - if isinstance(dilation, int): - dilation = (1,) + (dilation,) * dim + (1,) - else: - dilation = (1,) + tuple(dilation) + (1,) - - is_single_input = False - if inputs.ndim == len(dims) - 1: - # add singleton batch dimension because lax.reduce_window always - # needs a batch dimension. - inputs = inputs[None] - is_single_input = True - - assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" - - # shape of window after dilation - new_window_shape = tuple( - [ - window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) - for i in range(1, len(dims) - 1) - ] - ) - inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( - inputs, window_shape, strides, dim, data_format="channel_last" - ) - if not depth_pooling: - # manually creating padding list - if isinstance(padding, str): - pad_list = _pad_str_to_list( - inputs, dims, padding, strides, new_window_shape - ) - else: - if isinstance(padding, int): - padding = [(padding,) * 2] * dim - pad_list = [(0, 0)] + list(padding) + [(0, 0)] - - if ceil_mode: - c = [] - for i in range(len(dims) - 2): - pad_list[i + 1], ceil = _padding_ceil_mode( - inputs.shape[i + 1], - new_window_shape[i], - pad_list[i + 1], - strides[i + 1], - True, - ) - c.append(ceil) - - if count_include_pad: - # manually pad inputs with 0 if ceil_mode is True - # because they're not counted in average calculation - if ceil_mode: - ceil = [(0, c[i]) for i in range(len(dims) - 2)] - for i in range(len(dims) - 2): - pad_list[i + 1] = ( - pad_list[i + 1][0], - pad_list[i + 1][1] - ceil[i][1], - ) - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - inputs = jnp.pad( - inputs, - [(0, 0)] + ceil + [(0, 0)], - mode="constant", - constant_values=0.0, - ) - else: - # manually pad inputs with 1s - # because they are counted in average calculation - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - pad_list = [(0, 0)] * len(pad_list) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - pad_list = [(0, 0)] * (dim + 2) - - if not ivy.is_array(inputs): - # if dtype is not set here, jax casts it to float64 - inputs = jnp.array(inputs, dtype=jnp.float32) - if not ivy.is_array(init): - init = jnp.array(init, dtype=jnp.float32) - promoted_type = jnp.promote_types(inputs.dtype, init.dtype) - inputs = inputs.astype(promoted_type) - init = init.astype(promoted_type) - y = jlax.reduce_window( - inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation - ) - if is_single_input: - y = jnp.squeeze(y, axis=0) - return y - - -def max_pool1d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCW": - x = jnp.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCW": - res = jnp.transpose(res, (0, 2, 1)) - return res - - -def max_pool2d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCHW": - x = jnp.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCHW": - return jnp.transpose(res, (0, 3, 1, 2)) - - return res - - -def max_pool3d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - if data_format == "NCDHW": - x = jnp.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCDHW": - res = jnp.transpose(res, (0, 4, 1, 2, 3)) - - return res +# --- Main --- # +# ------------ # def avg_pool1d( @@ -512,60 +285,13 @@ def dct( return dct_out -def idct( +def dropout1d( x: JaxArray, + prob: float, /, *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - -def fft( - x: JaxArray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return jnp.fft.fft(x, n, dim, norm) - - -def dropout1d( - x: JaxArray, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", + training: bool = True, + data_format: str = "NWC", out: Optional[JaxArray] = None, ) -> JaxArray: if training: @@ -639,6 +365,225 @@ def dropout3d( return res +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version +) +def embedding( + weights: JaxArray, + indices: JaxArray, + /, + *, + max_norm: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + + embeddings = jnp.take(weights, indices, axis=0) + if max_norm is not None: + norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = jnp.where( + norms > max_norm, embeddings * max_norm / norms, embeddings + ) + embeddings = jnp.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings + ) + return embeddings + + +def fft( + x: JaxArray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return jnp.fft.fft(x, n, dim, norm) + + +def fft2( + x: JaxArray, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_elem_in_list( + norm, + ["backward", "ortho", "forward"], + message=f"Unrecognized normalization mode {norm}", + ) + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " + ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) + + +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + # This function assumes that param validation is already done + window_shape = tuple(window_shape) + strides = (1,) + strides + (1,) if len(strides) == dim else strides + dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape + if isinstance(dilation, int): + dilation = (1,) + (dilation,) * dim + (1,) + else: + dilation = (1,) + tuple(dilation) + (1,) + + is_single_input = False + if inputs.ndim == len(dims) - 1: + # add singleton batch dimension because lax.reduce_window always + # needs a batch dimension. + inputs = inputs[None] + is_single_input = True + + assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" + + # shape of window after dilation + new_window_shape = tuple( + [ + window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) + for i in range(1, len(dims) - 1) + ] + ) + inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( + inputs, window_shape, strides, dim, data_format="channel_last" + ) + if not depth_pooling: + # manually creating padding list + if isinstance(padding, str): + pad_list = _pad_str_to_list( + inputs, dims, padding, strides, new_window_shape + ) + else: + if isinstance(padding, int): + padding = [(padding,) * 2] * dim + pad_list = [(0, 0)] + list(padding) + [(0, 0)] + + if ceil_mode: + c = [] + for i in range(len(dims) - 2): + pad_list[i + 1], ceil = _padding_ceil_mode( + inputs.shape[i + 1], + new_window_shape[i], + pad_list[i + 1], + strides[i + 1], + True, + ) + c.append(ceil) + + if count_include_pad: + # manually pad inputs with 0 if ceil_mode is True + # because they're not counted in average calculation + if ceil_mode: + ceil = [(0, c[i]) for i in range(len(dims) - 2)] + for i in range(len(dims) - 2): + pad_list[i + 1] = ( + pad_list[i + 1][0], + pad_list[i + 1][1] - ceil[i][1], + ) + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + inputs = jnp.pad( + inputs, + [(0, 0)] + ceil + [(0, 0)], + mode="constant", + constant_values=0.0, + ) + else: + # manually pad inputs with 1s + # because they are counted in average calculation + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + pad_list = [(0, 0)] * len(pad_list) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + pad_list = [(0, 0)] * (dim + 2) + + if not ivy.is_array(inputs): + # if dtype is not set here, jax casts it to float64 + inputs = jnp.array(inputs, dtype=jnp.float32) + if not ivy.is_array(init): + init = jnp.array(init, dtype=jnp.float32) + promoted_type = jnp.promote_types(inputs.dtype, init.dtype) + inputs = inputs.astype(promoted_type) + init = init.astype(promoted_type) + y = jlax.reduce_window( + inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation + ) + if is_single_input: + y = jnp.squeeze(y, axis=0) + return y + + +def idct( + x: JaxArray, + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + def ifft( x: JaxArray, dim: int, @@ -671,6 +616,17 @@ def ifft( return jnp.fft.ifft(x, n, dim, norm) +def ifftn( + x: JaxArray, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.fft.ifftn(x, s, axes, norm) + + def interpolate( x: JaxArray, size: Union[Sequence[int], int], @@ -713,19 +669,129 @@ def interpolate( ) -interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (align_corners is None or not align_corners) - and mode - not in [ - "area", - "nearest", - "nd", - "tf_area", - "mitchellcubic", - "gaussian", - "bicubic", - ] -) +def max_pool1d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCW": + x = jnp.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCW": + res = jnp.transpose(res, (0, 2, 1)) + return res + + +def max_pool2d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCHW": + x = jnp.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCHW": + return jnp.transpose(res, (0, 3, 1, 2)) + + return res + + +def max_pool3d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NCDHW": + x = jnp.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCDHW": + res = jnp.transpose(res, (0, 4, 1, 2, 3)) + + return res def reduce_window( @@ -763,79 +829,6 @@ def reduce_window( ) -def fft2( - x: JaxArray, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_elem_in_list( - norm, - ["backward", "ortho", "forward"], - message=f"Unrecognized normalization mode {norm}", - ) - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" - ) - return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) - - -def ifftn( - x: JaxArray, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.fft.ifftn(x, s, axes, norm) - - -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version -) -def embedding( - weights: JaxArray, - indices: JaxArray, - /, - *, - max_norm: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - - embeddings = jnp.take(weights, indices, axis=0) - if max_norm is not None: - norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = jnp.where( - norms > max_norm, embeddings * max_norm / norms, embeddings - ) - embeddings = jnp.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings - ) - return embeddings - - @with_unsupported_dtypes({"0.4.14 and below": ("float16", "complex")}, backend_version) def rfftn( x: JaxArray, @@ -867,3 +860,18 @@ def rfftn( if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") return jnp.fft.rfftn(x, s, axes, norm).astype(jnp.complex128) + + +interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (align_corners is None or not align_corners) + and mode + not in [ + "area", + "nearest", + "nd", + "tf_area", + "mitchellcubic", + "gaussian", + "bicubic", + ] +) diff --git a/ivy/functional/backends/jax/experimental/linear_algebra.py b/ivy/functional/backends/jax/experimental/linear_algebra.py index 897254757e7ba..214010733ee24 100644 --- a/ivy/functional/backends/jax/experimental/linear_algebra.py +++ b/ivy/functional/backends/jax/experimental/linear_algebra.py @@ -10,6 +10,31 @@ from ivy.utils.exceptions import IvyNotImplementedException +dot.support_native_out = True + + +def adjoint( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return jnp.conjugate(jnp.transpose(x, axes=axes)) + + +def cond( + x: JaxArray, + /, + *, + p: Optional[Union[int, str, None]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.cond(x, p=p) + + def diagflat( x: JaxArray, /, @@ -86,23 +111,14 @@ def diagflat( return ret -def kron( +def dot( a: JaxArray, b: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.kron(a, b) - - -def matrix_exp( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jla.expm(x) + return jnp.dot(a, b) def eig( @@ -120,35 +136,14 @@ def eigvals(x: JaxArray, /) -> JaxArray: return jnp.linalg.eigvals(x) -def adjoint( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return jnp.conjugate(jnp.transpose(x, axes=axes)) - - -def multi_dot( - x: Sequence[JaxArray], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.multi_dot(x) - - -def cond( - x: JaxArray, +def kron( + a: JaxArray, + b: JaxArray, /, *, - p: Optional[Union[int, str, None]] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.linalg.cond(x, p=p) + return jnp.kron(a, b) def lu_factor( @@ -161,14 +156,19 @@ def lu_factor( raise IvyNotImplementedException() -def dot( - a: JaxArray, - b: JaxArray, +def matrix_exp( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dot(a, b) + return jla.expm(x) -dot.support_native_out = True +def multi_dot( + x: Sequence[JaxArray], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.multi_dot(x) diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index 7258a7ae7ebe7..e38a8343ba23d 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -20,102 +20,180 @@ from ivy.functional.backends.jax import JaxArray -def moveaxis( - a: JaxArray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], +# --- Helpers --- # +# --------------- # + + +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x + + +def _to_nested_tuple(nested_list): + ret = () + if hasattr(nested_list, "__iter__"): + for inner_list in nested_list: + if hasattr(inner_list, "__iter__"): + ret += (tuple(inner_list),) + else: + ret += (inner_list,) + return ret + if ret == (): + return nested_list + + +# --- Main --- # +# ------------ # + + +def atleast_1d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_1d(*arys) + + +def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: + return jnp.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_3d(*arys) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return jnp.broadcast_shapes(*shapes) + + +def concat_from_sequence( + input_sequence: Union[Tuple[JaxArray], List[JaxArray]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.moveaxis(a, source, destination) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = jnp.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = jnp.stack(input_sequence, axis=axis) + return ret -def heaviside( - x1: JaxArray, - x2: JaxArray, +def dsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], + /, + *, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def dstack( + arrays: Sequence[JaxArray], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.heaviside(x1, x2) + return jnp.dstack(arrays) -def flipud( - m: JaxArray, +def expand( + x: JaxArray, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.flipud(m) + shape = list(shape) + if len(shape) > len(x.shape): + x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape[i] + return jnp.broadcast_to(x, tuple(shape)) -def vstack( - arrays: Sequence[JaxArray], +def fill_diagonal( + a: JaxArray, + v: Union[int, float], /, *, - out: Optional[JaxArray] = None, + wrap: bool = False, ) -> JaxArray: - return jnp.vstack(arrays) + shape = jnp.array(a.shape) + end = None + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (jnp.cumprod(shape[:-1])).sum() + a = jnp.reshape(a, (-1,)) + a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) + a = jnp.reshape(a, shape) + return a -def hstack( - arrays: Sequence[JaxArray], +def fliplr( + m: JaxArray, /, *, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.hstack(arrays) + return jnp.fliplr(m) -def rot90( +def flipud( m: JaxArray, /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[JaxArray] = None, ) -> JaxArray: - if isinstance(axes, list): - axes = tuple(axes) - return jnp.rot90(m, k, axes) + return jnp.flipud(m) -def top_k( - x: JaxArray, - k: int, +def heaviside( + x1: JaxArray, + x2: JaxArray, /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[JaxArray, JaxArray]] = None, -) -> Tuple[JaxArray, JaxArray]: - k = min(k, x.shape[axis]) - if not largest: - indices = jnp.argsort(x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - else: - indices = jnp.argsort(-x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - if not sorted: - indices = jnp.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) - val = jnp.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.heaviside(x1, x2) -def fliplr( - m: JaxArray, +def hsplit( + ary: JaxArray, + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +def hstack( + arrays: Sequence[JaxArray], + /, + *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.fliplr(m) + return jnp.hstack(arrays) def i0( @@ -127,21 +205,16 @@ def i0( return jnp.i0(x) -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x - - -def _to_nested_tuple(nested_list): - ret = () - if hasattr(nested_list, "__iter__"): - for inner_list in nested_list: - if hasattr(inner_list, "__iter__"): - ret += (tuple(inner_list),) - else: - ret += (inner_list,) - return ret - if ret == (): - return nested_list +def moveaxis( + a: JaxArray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.moveaxis(a, source, destination) def pad( @@ -228,57 +301,18 @@ def pad( return ret -def vsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def dsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], +def rot90( + m: JaxArray, /, *, copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_1d(*arys) - - -def dstack( - arrays: Sequence[JaxArray], - /, - *, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dstack(arrays) - - -def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: - return jnp.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_3d(*arys) + if isinstance(axes, list): + axes = tuple(axes) + return jnp.rot90(m, k, axes) def take_along_axis( @@ -298,56 +332,28 @@ def take_along_axis( return jnp.take_along_axis(arr, indices, axis, mode=mode) -def hsplit( - ary: JaxArray, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return jnp.broadcast_shapes(*shapes) - - -def expand( +def top_k( x: JaxArray, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - shape = list(shape) - if len(shape) > len(x.shape): - x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape[i] - return jnp.broadcast_to(x, tuple(shape)) - - -def concat_from_sequence( - input_sequence: Union[Tuple[JaxArray], List[JaxArray]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = jnp.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = jnp.stack(input_sequence, axis=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[JaxArray, JaxArray]] = None, +) -> Tuple[JaxArray, JaxArray]: + k = min(k, x.shape[axis]) + if not largest: + indices = jnp.argsort(x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + else: + indices = jnp.argsort(-x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + if not sorted: + indices = jnp.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) + val = jnp.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) def unique_consecutive( @@ -393,22 +399,24 @@ def unique_consecutive( ) -def fill_diagonal( - a: JaxArray, - v: Union[int, float], +def vsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], /, *, - wrap: bool = False, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Sequence[JaxArray], + /, + *, + out: Optional[JaxArray] = None, ) -> JaxArray: - shape = jnp.array(a.shape) - end = None - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (jnp.cumprod(shape[:-1])).sum() - a = jnp.reshape(a, (-1,)) - a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) - a = jnp.reshape(a, shape) - return a + return jnp.vstack(arrays) diff --git a/ivy/functional/backends/jax/experimental/random.py b/ivy/functional/backends/jax/experimental/random.py index cfec5e0f496ce..0b75f787a241c 100644 --- a/ivy/functional/backends/jax/experimental/random.py +++ b/ivy/functional/backends/jax/experimental/random.py @@ -15,26 +15,27 @@ from ivy.func_wrapper import with_unsupported_dtypes from .. import backend_version -# Extra # -# ----- # - -# dirichlet -def dirichlet( - alpha: Union[JaxArray, float, Sequence[float]], - /, +def bernoulli( + probs: Union[float, JaxArray], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + logits: Optional[Union[float, JaxArray]] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: Optional[jaxlib.xla_extension.Device] = None, dtype: Optional[jnp.dtype] = None, seed: Optional[int] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - if seed is not None: + if seed: rng_input = jax.random.PRNGKey(seed) else: RNG_, rng_input = jax.random.split(_getRNG()) _setRNG(RNG_) - return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) + if logits is not None: + probs = jax.nn.softmax(logits, axis=-1) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return jax.random.bernoulli(rng_input, probs, shape=shape) def beta( @@ -56,6 +57,28 @@ def beta( return jax.random.beta(rng_input, a, b, shape, dtype) +# Extra # +# ----- # + + +# dirichlet +def dirichlet( + alpha: Union[JaxArray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[jnp.dtype] = None, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if seed is not None: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) + + @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, JaxArray], @@ -105,25 +128,3 @@ def poisson( else: ret = jax.random.poisson(rng_input, lam, shape=list_shape).astype(dtype) return ret - - -def bernoulli( - probs: Union[float, JaxArray], - *, - logits: Optional[Union[float, JaxArray]] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: Optional[jaxlib.xla_extension.Device] = None, - dtype: Optional[jnp.dtype] = None, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - if logits is not None: - probs = jax.nn.softmax(logits, axis=-1) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return jax.random.bernoulli(rng_input, probs, shape=shape) diff --git a/ivy/functional/backends/jax/experimental/statistical.py b/ivy/functional/backends/jax/experimental/statistical.py index 64f1bfc031973..b391273e1b628 100644 --- a/ivy/functional/backends/jax/experimental/statistical.py +++ b/ivy/functional/backends/jax/experimental/statistical.py @@ -9,6 +9,198 @@ from ..statistical import _infer_dtype +# --- Helpers --- # +# --------------- # + + +def __find_cummax_indices( + x: JaxArray, + axis: int = 0, +) -> JaxArray: + n, indice, indices = 0, [], [] + + if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + else: + z_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + else: + indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) + else: + n, indices = 0, [] + for idx, y in enumerate(x): + if idx == 0 or x[n] <= y: + n = idx + indices.append(n) + + return jnp.asarray(indices, dtype="int64") + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + +# --- Main --- # +# ------------ # + + +def bincount( + x: JaxArray, + /, + *, + weights: Optional[JaxArray] = None, + minlength: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + if weights is not None: + ret = jnp.bincount(x, weights=weights, minlength=minlength) + ret = ret.astype(weights.dtype) + else: + ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) + return ret + + +def corrcoef( + x: JaxArray, + /, + *, + y: Optional[JaxArray] = None, + rowvar: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.corrcoef(x, y=y, rowvar=rowvar) + + +def cov( + x1: JaxArray, + x2: JaxArray = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[JaxArray] = None, + aweights: Optional[JaxArray] = None, + dtype: Optional[jnp.dtype] = None, +) -> JaxArray: + if not dtype: + x1 = jnp.asarray(x1, dtype=jnp.float64) + + if jnp.ndim(x1) > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if jnp.ndim(x2) > 2: + raise ValueError("x2 has more than 2 dimensions") + + if fweights is not None: + fweights = jnp.asarray(fweights, dtype=jnp.int64) + + return jnp.cov( + m=x1, + y=x2, + rowvar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, + ) + + +def cummax( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> Tuple[JaxArray, JaxArray]: + if x.dtype in (jnp.bool_, jnp.float16): + x = x.astype(jnp.float64) + elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): + x = x.astype(jnp.int64) + elif x.dtype in (jnp.complex128, jnp.complex64): + x = jnp.real(x).astype(jnp.float64) + + if exclusive or (reverse and exclusive): + if exclusive and reverse: + indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) + x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + x, indices = jnp.concatenate( + (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 + ), jnp.concatenate( + (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 + ) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + indices = __find_cummax_indices(x, axis=axis) + res = jlax.cummax(x, axis=axis) + return res, indices + + if reverse: + y = jnp.flip(x, axis=axis) + indices = __find_cummax_indices(y, axis=axis) + indices = jnp.flip(indices, axis=axis) + else: + indices = __find_cummax_indices(x, axis=axis) + return jlax.cummax(x, axis, reverse=reverse), indices + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cummin( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if axis < 0: + axis = axis + len(x.shape) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + return jlax.cummin(x, axis, reverse=reverse).astype(dtype) + + @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16",)}, backend_version, @@ -120,6 +312,16 @@ def histogram( return ret +def igamma( + a: JaxArray, + /, + *, + x: JaxArray, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jlax.igamma(a=a, x=x) + + @with_unsupported_dtypes( {"0.4.14 and below": ("complex64", "complex128")}, backend_version ) @@ -162,34 +364,6 @@ def nanmean( return jnp.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) -def quantile( - a: JaxArray, - q: Union[float, JaxArray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - interpolation = "nearest" if interpolation == "nearest_jax" else interpolation - return jnp.quantile( - a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out - ) - - -def corrcoef( - x: JaxArray, - /, - *, - y: Optional[JaxArray] = None, - rowvar: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.corrcoef(x, y=y, rowvar=rowvar) - - def nanmedian( input: JaxArray, /, @@ -219,184 +393,18 @@ def nanmedian( ) -def bincount( - x: JaxArray, - /, - *, - weights: Optional[JaxArray] = None, - minlength: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - if weights is not None: - ret = jnp.bincount(x, weights=weights, minlength=minlength) - ret = ret.astype(weights.dtype) - else: - ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) - return ret - - -def cov( - x1: JaxArray, - x2: JaxArray = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[JaxArray] = None, - aweights: Optional[JaxArray] = None, - dtype: Optional[jnp.dtype] = None, -) -> JaxArray: - if not dtype: - x1 = jnp.asarray(x1, dtype=jnp.float64) - - if jnp.ndim(x1) > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if jnp.ndim(x2) > 2: - raise ValueError("x2 has more than 2 dimensions") - - if fweights is not None: - fweights = jnp.asarray(fweights, dtype=jnp.int64) - - return jnp.cov( - m=x1, - y=x2, - rowvar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - ) - - -def cummax( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> Tuple[JaxArray, JaxArray]: - if x.dtype in (jnp.bool_, jnp.float16): - x = x.astype(jnp.float64) - elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): - x = x.astype(jnp.int64) - elif x.dtype in (jnp.complex128, jnp.complex64): - x = jnp.real(x).astype(jnp.float64) - - if exclusive or (reverse and exclusive): - if exclusive and reverse: - indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) - x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - x, indices = jnp.concatenate( - (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 - ), jnp.concatenate( - (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 - ) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - indices = __find_cummax_indices(x, axis=axis) - res = jlax.cummax(x, axis=axis) - return res, indices - - if reverse: - y = jnp.flip(x, axis=axis) - indices = __find_cummax_indices(y, axis=axis) - indices = jnp.flip(indices, axis=axis) - else: - indices = __find_cummax_indices(x, axis=axis) - return jlax.cummax(x, axis, reverse=reverse), indices - - -def __find_cummax_indices( - x: JaxArray, - axis: int = 0, -) -> JaxArray: - n, indice, indices = 0, [], [] - - if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - else: - z_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - else: - indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) - else: - n, indices = 0, [] - for idx, y in enumerate(x): - if idx == 0 or x[n] <= y: - n = idx - indices.append(n) - - return jnp.asarray(indices, dtype="int64") - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cummin( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if axis < 0: - axis = axis + len(x.shape) - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - return jlax.cummin(x, axis, reverse=reverse).astype(dtype) - - -def igamma( +def quantile( a: JaxArray, + q: Union[float, JaxArray], /, *, - x: JaxArray, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return jlax.igamma(a=a, x=x) + axis = tuple(axis) if isinstance(axis, list) else axis + interpolation = "nearest" if interpolation == "nearest_jax" else interpolation + return jnp.quantile( + a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ) diff --git a/ivy/functional/backends/jax/general.py b/ivy/functional/backends/jax/general.py index 9221319cfeaed..1315cef38980a 100644 --- a/ivy/functional/backends/jax/general.py +++ b/ivy/functional/backends/jax/general.py @@ -22,33 +22,8 @@ from . import backend_version -def container_types(): - flat_mapping_spec = importlib.util.find_spec( - "FlatMapping", "haiku._src.data_structures" - ) - if not flat_mapping_spec: - from haiku._src.data_structures import FlatMapping - else: - FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) - return [FlatMapping] - - -def current_backend_str() -> str: - return "jax" - - -def is_native_array(x, /, *, exclusive=False): - if exclusive: - return isinstance(x, NativeArray) - return isinstance( - x, - ( - NativeArray, - jax.interpreters.ad.JVPTracer, - jax.core.ShapedArray, - jax.interpreters.partial_eval.DynamicJaxprTracer, - ), - ) +# --- Helpers --- # +# --------------- # def _mask_to_index(query, x): @@ -61,63 +36,35 @@ def _mask_to_index(query, x): return jnp.where(query), expected_shape -def get_item( - x: JaxArray, - /, - query: Union[JaxArray, Tuple], - *, - copy: bool = None, -) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return jnp.array([], dtype=x.dtype) - else: - return jnp.expand_dims(x, 0) - query, _ = _mask_to_index(query, x) - elif isinstance(query, list): - query = (query,) - return x.__getitem__(query) +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view -def set_item( - x: JaxArray, - query: Union[JaxArray, Tuple], - val: JaxArray, - /, - *, - copy: Optional[bool] = False, -) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - query, expected_shape = _mask_to_index(query, x) - val = _broadcast_to(val, expected_shape)._data - ret = x.at[query].set(val) - if copy: - return ret - return ivy.inplace_update(x, _to_device(ret)) +# --- Main --- # +# ------------ # def array_equal(x0: JaxArray, x1: JaxArray, /) -> bool: return bool(jnp.array_equal(x0, x1)) -@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) -def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return np.array(_to_array(x)) - else: - return np.asarray(_to_array(x)) - - -def to_scalar(x: JaxArray, /) -> Number: - if isinstance(x, Number): - return x +def container_types(): + flat_mapping_spec = importlib.util.find_spec( + "FlatMapping", "haiku._src.data_structures" + ) + if not flat_mapping_spec: + from haiku._src.data_structures import FlatMapping else: - return _to_array(x).item() + FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) + return [FlatMapping] -def to_list(x: JaxArray, /) -> list: - return _to_array(x).tolist() +def current_backend_str() -> str: + return "jax" def gather( @@ -152,6 +99,36 @@ def gather( return result +def gather_nd( + params: JaxArray, + indices: JaxArray, + /, + *, + batch_dims: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = jnp.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -183,34 +160,23 @@ def gather_nd_helper(params, indices): return ret -def gather_nd( - params: JaxArray, - indices: JaxArray, +def get_item( + x: JaxArray, /, + query: Union[JaxArray, Tuple], *, - batch_dims: int = 0, - out: Optional[JaxArray] = None, + copy: bool = None, ) -> JaxArray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return jnp.array([], dtype=x.dtype) else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = jnp.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result + return jnp.expand_dims(x, 0) + query, _ = _mask_to_index(query, x) + elif isinstance(query, list): + query = (query,) + return x.__getitem__(query) def get_num_dims(x: JaxArray, /, *, as_array: bool = False) -> Union[JaxArray, int]: @@ -288,18 +254,40 @@ def inplace_update( return val -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view - - def inplace_variables_supported(): return False +def is_native_array(x, /, *, exclusive=False): + if exclusive: + return isinstance(x, NativeArray) + return isinstance( + x, + ( + NativeArray, + jax.interpreters.ad.JVPTracer, + jax.core.ShapedArray, + jax.interpreters.partial_eval.DynamicJaxprTracer, + ), + ) + + +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def isin( + elements: JaxArray, + test_elements: JaxArray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> JaxArray: + return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) + + +def itemsize(x: JaxArray) -> int: + return x.itemsize + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -342,9 +330,6 @@ def scatter_flat( return target -scatter_flat.support_native_out = True - - def scatter_nd( indices: JaxArray, updates: JaxArray, @@ -392,7 +377,21 @@ def scatter_nd( return target -scatter_nd.support_native_out = True +def set_item( + x: JaxArray, + query: Union[JaxArray, Tuple], + val: JaxArray, + /, + *, + copy: Optional[bool] = False, +) -> JaxArray: + if ivy.is_array(query) and ivy.is_bool_dtype(query): + query, expected_shape = _mask_to_index(query, x) + val = _broadcast_to(val, expected_shape)._data + ret = x.at[query].set(val) + if copy: + return ret + return ivy.inplace_update(x, _to_device(ret)) def shape( @@ -407,6 +406,25 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: JaxArray, /) -> list: + return _to_array(x).tolist() + + +@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) +def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return np.array(_to_array(x)) + else: + return np.asarray(_to_array(x)) + + +def to_scalar(x: JaxArray, /) -> Number: + if isinstance(x, Number): + return x + else: + return _to_array(x).item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -418,17 +436,5 @@ def vmap( ) -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def isin( - elements: JaxArray, - test_elements: JaxArray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> JaxArray: - return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) - - -def itemsize(x: JaxArray) -> int: - return x.itemsize +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True diff --git a/ivy/functional/backends/jax/gradients.py b/ivy/functional/backends/jax/gradients.py index a5e6d789a070b..dcd85bc2b395e 100644 --- a/ivy/functional/backends/jax/gradients.py +++ b/ivy/functional/backends/jax/gradients.py @@ -18,19 +18,8 @@ ) -# ToDo: modify these functions to track whether variable() has been called -def variable(x, /): - return x - - -def is_variable(x, /, *, exclusive=False): - if exclusive: - return False - return isinstance(x, NativeArray) - - -def variable_data(x: JaxArray, /) -> JaxArray: - return x +# --- Helpers --- # +# --------------- # def _forward_fn( @@ -68,6 +57,10 @@ def _forward_fn( return ret_values +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: JaxArray, @@ -127,21 +120,18 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - value, grad = jax.value_and_grad(grad_fn)(xs) - return ivy.to_ivy(value), ivy.to_ivy(grad) - +def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): + grad_fn = lambda x_in: ivy.to_native(func(x_in)) + callback_fn = lambda x_in: ivy.to_ivy( + jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) + ) return callback_fn -def stop_gradient( - x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - return jlax.stop_gradient(x) +def is_variable(x, /, *, exclusive=False): + if exclusive: + return False + return isinstance(x, NativeArray) def jac(func: Callable): @@ -158,9 +148,27 @@ def jac(func: Callable): return callback_fn -def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): - grad_fn = lambda x_in: ivy.to_native(func(x_in)) - callback_fn = lambda x_in: ivy.to_ivy( - jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) - ) +def stop_gradient( + x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + return jlax.stop_gradient(x) + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + value, grad = jax.value_and_grad(grad_fn)(xs) + return ivy.to_ivy(value), ivy.to_ivy(grad) + return callback_fn + + +# ToDo: modify these functions to track whether variable() has been called +def variable(x, /): + return x + + +def variable_data(x: JaxArray, /) -> JaxArray: + return x diff --git a/ivy/functional/backends/jax/layers.py b/ivy/functional/backends/jax/layers.py index ff4c264439da7..c27e92a2680eb 100644 --- a/ivy/functional/backends/jax/layers.py +++ b/ivy/functional/backends/jax/layers.py @@ -16,41 +16,18 @@ ) -def _transpose_padding_helper(k, s, padding, dilation, diff=0): - k = (k - 1) * dilation + 1 - if padding == "SAME": - pad_len = k + s - 2 - pad_len -= diff - if s > k - 1: - pad_a = k - 1 - else: - pad_a = int(jnp.ceil(pad_len / 2)) - else: - pad_len = k + s - 2 + max(k - s, 0) - pad_a = k - 1 - pad_b = pad_len - pad_a - return pad_a, pad_b +# --- Helpers --- # +# --------------- # -def _get_tranpose_padding( - x_shape, filter_shape, strides, padding, dims, dilations, output_shape -): - new_shape = [ - _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) - for i in range(dims) - ] - if output_shape is None: - output_shape = [x_shape[0], *new_shape, filter_shape[-1]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] - shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] - pad_list = [ - _transpose_padding_helper( - filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] - ) - for i in range(dims) - ] - return pad_list +def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): + first = True if filter_format == "channel_first" else False + if dims == 1: + return "WIO" if not first else "OIW" + if dims == 2: + return "HWIO" if not first else "OIHW" + elif dims == 3: + return "DHWIO" if not first else "OIDHW" def _get_new_padding_before_conv( @@ -94,6 +71,47 @@ def _get_new_padding_before_conv( return padding +def _get_tranpose_padding( + x_shape, filter_shape, strides, padding, dims, dilations, output_shape +): + new_shape = [ + _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) + for i in range(dims) + ] + if output_shape is None: + output_shape = [x_shape[0], *new_shape, filter_shape[-1]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] + shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] + pad_list = [ + _transpose_padding_helper( + filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] + ) + for i in range(dims) + ] + return pad_list + + +def _transpose_padding_helper(k, s, padding, dilation, diff=0): + k = (k - 1) * dilation + 1 + if padding == "SAME": + pad_len = k + s - 2 + pad_len -= diff + if s > k - 1: + pad_a = k - 1 + else: + pad_a = int(jnp.ceil(pad_len / 2)) + else: + pad_len = k + s - 2 + max(k - s, 0) + pad_a = k - 1 + pad_b = pad_len - pad_a + return pad_a, pad_b + + +# --- Main --- # +# ------------ # + + def conv1d( x: JaxArray, filters: JaxArray, @@ -231,37 +249,6 @@ def conv2d_transpose( return res -def depthwise_conv2d( - x: JaxArray, - filters: JaxArray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[JaxArray] = None, -) -> JaxArray: - strides = [strides] * 2 if isinstance(strides, int) else strides - strides = [strides[1], strides[2]] if len(strides) == 4 else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters - cn = filters.shape[-1] - filters = jnp.expand_dims(filters, -2) - return jlax.conv_general_dilated( - x, - filters, - strides, - padding, - None, - dilations, - (data_format, "HWIO", data_format), - feature_group_count=cn, - ) - - def conv3d( x: JaxArray, filters: JaxArray, @@ -330,16 +317,6 @@ def conv3d_transpose( return res -def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): - first = True if filter_format == "channel_first" else False - if dims == 1: - return "WIO" if not first else "OIW" - if dims == 2: - return "HWIO" if not first else "OIHW" - elif dims == 3: - return "DHWIO" if not first else "OIDHW" - - def conv_general_dilated( x: JaxArray, filters: JaxArray, @@ -455,3 +432,34 @@ def conv_general_transpose( if data_format == "channel_first": return jnp.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +def depthwise_conv2d( + x: JaxArray, + filters: JaxArray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[JaxArray] = None, +) -> JaxArray: + strides = [strides] * 2 if isinstance(strides, int) else strides + strides = [strides[1], strides[2]] if len(strides) == 4 else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters + cn = filters.shape[-1] + filters = jnp.expand_dims(filters, -2) + return jlax.conv_general_dilated( + x, + filters, + strides, + padding, + None, + dilations, + (data_format, "HWIO", data_format), + feature_group_count=cn, + ) diff --git a/ivy/functional/backends/jax/linear_algebra.py b/ivy/functional/backends/jax/linear_algebra.py index 383a9d032abc6..64c80f2339090 100644 --- a/ivy/functional/backends/jax/linear_algebra.py +++ b/ivy/functional/backends/jax/linear_algebra.py @@ -58,13 +58,19 @@ def det(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.linalg.det(x) -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] - ) - eigenvalues, eigenvectors = jnp.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def diag( + x: JaxArray, + /, + *, + k: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.diag(x, k=k) @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) @@ -92,15 +98,13 @@ def diagonal( return ret -def tensorsolve( - x1: JaxArray, - x2: JaxArray, - /, - *, - axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.tensorsolve(x1, x2, axes) +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] + ) + eigenvalues, eigenvectors = jnp.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) @with_unsupported_dtypes( @@ -388,6 +392,17 @@ def tensordot( return jnp.tensordot(x1, x2, axes) +def tensorsolve( + x1: JaxArray, + x2: JaxArray, + /, + *, + axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.tensorsolve(x1, x2, axes) + + @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -404,6 +419,21 @@ def trace( return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, + backend_version, +) +def vander( + x: JaxArray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def vecdot( x1: JaxArray, x2: JaxArray, /, *, axis: int = -1, out: Optional[JaxArray] = None @@ -438,36 +468,6 @@ def vector_norm( return jnp.sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def diag( - x: JaxArray, - /, - *, - k: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.diag(x, k=k) - - -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, - backend_version, -) -def vander( - x: JaxArray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes( { "0.4.14 and below": ( diff --git a/ivy/functional/backends/jax/manipulation.py b/ivy/functional/backends/jax/manipulation.py index 61e638b9509cb..4ebc42a93b9c7 100644 --- a/ivy/functional/backends/jax/manipulation.py +++ b/ivy/functional/backends/jax/manipulation.py @@ -12,10 +12,61 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x +# --- Main --- # +# ------------ # + + +def clip( + x: JaxArray, + x_min: Union[Number, JaxArray], + x_max: Union[Number, JaxArray], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + if ( + hasattr(x_min, "dtype") + and hasattr(x_max, "dtype") + and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) + ): + if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( + jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float32) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + elif ( + jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) + ) and ( + jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float64) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + else: + promoted_type = jnp.promote_types(x.dtype, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x.astype(promoted_type) + # jnp.clip isn't used because of inconsistent gradients + x = jnp.where(x > x_max, x_max, x) + return jnp.where(x < x_min, x_min, x) + + # Array API Standard # # -------------------# @@ -42,6 +93,18 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) +@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) +def constant_pad( + x: JaxArray, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) + + def expand_dims( x: JaxArray, /, @@ -79,6 +142,17 @@ def permute_dims( return jnp.transpose(x, axes) +def repeat( + x: JaxArray, + /, + repeats: Union[int, Iterable[int]], + *, + axis: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.repeat(x, repeats, axis) + + def reshape( x: JaxArray, /, @@ -114,38 +188,6 @@ def roll( return jnp.roll(x, shift, axis) -def squeeze( - x: JaxArray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - else: - ret = jnp.squeeze(x, axis=axis) - return ret - - -def stack( - arrays: Union[Tuple[JaxArray], List[JaxArray]], - /, - *, - axis: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - try: - return jnp.stack(arrays, axis=axis) - except ValueError as error: - raise ivy.utils.exceptions.IvyIndexError(error) - - # Extra # # ------# @@ -184,76 +226,54 @@ def split( return jnp.split(x, num_or_size_splits, axis) -def repeat( +def squeeze( x: JaxArray, /, - repeats: Union[int, Iterable[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.repeat(x, repeats, axis) - - -def tile( - x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.tile(x, repeats) + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + else: + ret = jnp.squeeze(x, axis=axis) + return ret -def clip( - x: JaxArray, - x_min: Union[Number, JaxArray], - x_max: Union[Number, JaxArray], +def stack( + arrays: Union[Tuple[JaxArray], List[JaxArray]], /, *, + axis: int = 0, out: Optional[JaxArray] = None, ) -> JaxArray: - if ( - hasattr(x_min, "dtype") - and hasattr(x_max, "dtype") - and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) - ): - if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( - jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float32) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - elif ( - jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) - ) and ( - jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float64) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - else: - promoted_type = jnp.promote_types(x.dtype, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x.astype(promoted_type) - # jnp.clip isn't used because of inconsistent gradients - x = jnp.where(x > x_max, x_max, x) - return jnp.where(x < x_min, x_min, x) + try: + return jnp.stack(arrays, axis=axis) + except ValueError as error: + raise ivy.utils.exceptions.IvyIndexError(error) -@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) -def constant_pad( +def swapaxes( x: JaxArray, + axis0: int, + axis1: int, /, - pad_width: List[List[int]], *, - value: Number = 0.0, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) + return jnp.swapaxes(x, axis0, axis1) + + +def tile( + x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.tile(x, repeats) def unstack( @@ -278,15 +298,3 @@ def zero_pad( x: JaxArray, /, pad_width: List[List[int]], *, out: Optional[JaxArray] = None ): return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=0) - - -def swapaxes( - x: JaxArray, - axis0: int, - axis1: int, - /, - *, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.swapaxes(x, axis0, axis1) diff --git a/ivy/functional/backends/jax/random.py b/ivy/functional/backends/jax/random.py index 53aedb4e811e4..38b599f8ef405 100644 --- a/ivy/functional/backends/jax/random.py +++ b/ivy/functional/backends/jax/random.py @@ -26,12 +26,8 @@ def __init__(self): self.key = jax.random.PRNGKey(0) -RNG = RNGWrapper() - - -def _setRNG(key): - global RNG - RNG.key = key +# --- Helpers --- # +# --------------- # def _getRNG(): @@ -39,47 +35,13 @@ def _getRNG(): return RNG.key -def random_uniform( - *, - low: Union[float, JaxArray] = 0.0, - high: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - shape = _check_bounds_and_get_shape(low, high, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.uniform( - rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 - ).astype(dtype) +def _setRNG(key): + global RNG + RNG.key = key -def random_normal( - *, - mean: Union[float, JaxArray] = 0.0, - std: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) @@ -152,6 +114,49 @@ def randint( return jax.random.randint(rng_input, shape, low, high, dtype) +def random_normal( + *, + mean: Union[float, JaxArray] = 0.0, + std: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean + + +def random_uniform( + *, + low: Union[float, JaxArray] = 0.0, + high: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + shape = _check_bounds_and_get_shape(low, high, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.uniform( + rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 + ).astype(dtype) + + def seed(*, seed_value: int = 0) -> None: _setRNG(jax.random.PRNGKey(seed_value)) return @@ -176,3 +181,6 @@ def shuffle( # jax.random.shuffle is deprecated; identical behaviour reproduced with # jax.random.permutation return jax.random.permutation(key=rng_input, x=x, axis=axis, independent=True) + + +RNG = RNGWrapper() diff --git a/ivy/functional/backends/jax/searching.py b/ivy/functional/backends/jax/searching.py index 794436773c74a..e5f6a320ef428 100644 --- a/ivy/functional/backends/jax/searching.py +++ b/ivy/functional/backends/jax/searching.py @@ -64,6 +64,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.argwhere(x) + + def nonzero( x: JaxArray, /, @@ -90,16 +103,3 @@ def where( ) -> JaxArray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(jnp.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.argwhere(x) diff --git a/ivy/functional/backends/jax/sorting.py b/ivy/functional/backends/jax/sorting.py index 5111eca178530..5ccd39bef98f6 100644 --- a/ivy/functional/backends/jax/sorting.py +++ b/ivy/functional/backends/jax/sorting.py @@ -26,20 +26,15 @@ def argsort( ) -def sort( - x: JaxArray, +# msort +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def msort( + a: Union[JaxArray, list, tuple], /, *, - axis: int = -1, - descending: bool = False, - stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - kind = "stable" if stable else "quicksort" - ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) - if descending: - ret = jnp.asarray(jnp.flip(ret, axis=axis)) - return ret + return jnp.sort(a, axis=0, kind="mergesort") def searchsorted( @@ -79,12 +74,17 @@ def searchsorted( return ret.astype(ret_dtype) -# msort -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def msort( - a: Union[JaxArray, list, tuple], +def sort( + x: JaxArray, /, *, + axis: int = -1, + descending: bool = False, + stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.sort(a, axis=0, kind="mergesort") + kind = "stable" if stable else "quicksort" + ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) + if descending: + ret = jnp.asarray(jnp.flip(ret, axis=axis)) + return ret diff --git a/ivy/functional/backends/jax/statistical.py b/ivy/functional/backends/jax/statistical.py index b03be2a367141..b6af324bafde4 100644 --- a/ivy/functional/backends/jax/statistical.py +++ b/ivy/functional/backends/jax/statistical.py @@ -9,20 +9,102 @@ from ivy.functional.backends.jax import JaxArray from . import backend_version -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -def min( + +def _infer_dtype(dtype: jnp.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cumprod( x: JaxArray, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return jnp.cumprod(x, axis, dtype=dtype) + elif exclusive and reverse: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + return jnp.flip(x, axis=(axis,)) + + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumprod(x, -1, dtype=dtype) + return jnp.swapaxes(x, axis, -1) + else: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + return jnp.flip(x, axis=axis) + + +def cumsum( + x: JaxArray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + res = jnp.flip(x, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumsum(x, -1, dtype=dtype) + res = jnp.swapaxes(x, axis, -1) + elif reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + res = jnp.flip(x, axis=axis) + return res + return jnp.cumsum(x, axis, dtype=dtype) + + +def einsum( + equation: str, *operands: JaxArray, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.einsum(equation, *operands) def max( @@ -49,11 +131,20 @@ def mean( return jnp.mean(x, axis=axis, keepdims=keepdims) -def _infer_dtype(dtype: jnp.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +# Array API Standard # +# -------------------# + + +def min( + x: JaxArray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + axis = tuple(axis) if isinstance(axis, list) else axis + return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) def prod( @@ -133,85 +224,3 @@ def var( x.dtype, copy=False, ) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cumprod( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return jnp.cumprod(x, axis, dtype=dtype) - elif exclusive and reverse: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - return jnp.flip(x, axis=(axis,)) - - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumprod(x, -1, dtype=dtype) - return jnp.swapaxes(x, axis, -1) - else: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - return jnp.flip(x, axis=axis) - - -def cumsum( - x: JaxArray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - res = jnp.flip(x, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumsum(x, -1, dtype=dtype) - res = jnp.swapaxes(x, axis, -1) - elif reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - res = jnp.flip(x, axis=axis) - return res - return jnp.cumsum(x, axis, dtype=dtype) - - -def einsum( - equation: str, *operands: JaxArray, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.einsum(equation, *operands) diff --git a/ivy/functional/backends/mxnet/activations.py b/ivy/functional/backends/mxnet/activations.py index 8efd486936334..a6f3fc0e9b713 100644 --- a/ivy/functional/backends/mxnet/activations.py +++ b/ivy/functional/backends/mxnet/activations.py @@ -32,6 +32,14 @@ def leaky_relu( return mx.nd.LeakyReLU(x, slope=alpha) +def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): + raise IvyNotImplementedException() + + +def mish(x: None, /, *, out: Optional[None] = None) -> None: + raise IvyNotImplementedException() + + def relu(x: None, /, *, complex_mode="jax", out: Optional[None] = None) -> None: return mx.nd.relu(x) @@ -71,11 +79,3 @@ def softplus( if threshold is not None: return mx.nd.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): - raise IvyNotImplementedException() - - -def mish(x: None, /, *, out: Optional[None] = None) -> None: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/creation.py b/ivy/functional/backends/mxnet/creation.py index 632b1daca8f9e..417bea7a45a17 100644 --- a/ivy/functional/backends/mxnet/creation.py +++ b/ivy/functional/backends/mxnet/creation.py @@ -60,7 +60,15 @@ def asarray( return ret -array = asarray +def copy_array( + x: Union[(None, mx.ndarray.NDArray)], + *, + to_ivy_array: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() def empty( @@ -107,6 +115,15 @@ def from_dlpack( raise IvyNotImplementedException() +def frombuffer( + buffer: bytes, + dtype: Optional[None] = float, + count: Optional[int] = (-1), + offset: Optional[int] = 0, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def full( shape: Union[(ivy.NativeShape, Sequence[int])], fill_value: Union[(int, float, bool)], @@ -154,6 +171,21 @@ def meshgrid( raise IvyNotImplementedException() +def one_hot( + indices: Union[(None, mx.ndarray.NDArray)], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[None] = None, + device: str, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def ones( shape: Optional[ivy.NativeShape] = None, *, @@ -195,6 +227,12 @@ def triu( raise IvyNotImplementedException() +def triu_indices( + n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str +) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + def zeros( *size: Union[(int, Sequence[int])], shape: Optional[ivy.NativeShape] = None, @@ -220,42 +258,4 @@ def zeros_like( return ivy.to_device(ret, device) -def copy_array( - x: Union[(None, mx.ndarray.NDArray)], - *, - to_ivy_array: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() - - -def one_hot( - indices: Union[(None, mx.ndarray.NDArray)], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[None] = None, - device: str, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def frombuffer( - buffer: bytes, - dtype: Optional[None] = float, - count: Optional[int] = (-1), - offset: Optional[int] = 0, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def triu_indices( - n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str -) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() +array = asarray diff --git a/ivy/functional/backends/mxnet/data_type.py b/ivy/functional/backends/mxnet/data_type.py index 1ec0e1174ceec..6d15845c7c19f 100644 --- a/ivy/functional/backends/mxnet/data_type.py +++ b/ivy/functional/backends/mxnet/data_type.py @@ -5,6 +5,18 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from ivy.utils.exceptions import IvyNotImplementedException +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "u1": "uint8", +} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int32"): "int32", @@ -34,19 +46,6 @@ "bool": np.bool_, } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "u1": "uint8", -} - class Finfo: def __init__(self, mx_finfo: mx.np.finfo): @@ -84,57 +83,6 @@ def __repr__(self): return repr(self._mx_finfo) -def astype( - x: Union[(None, mx.ndarray.NDArray)], - dtype: Union[(None, str)], - /, - *, - copy: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return mx.nd.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays( - *arrays: Union[(None, mx.ndarray.NDArray)] -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_to( - x: Union[(None, mx.ndarray.NDArray)], - /, - shape: Union[(ivy.NativeShape, Sequence[int])], - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -@_handle_nestable_dtype_info -def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: - if isinstance(type, mx.ndarray.NDArray): - type = type.dtype - return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: - # using np.iinfo as mx use np dtypes and mxnet iinfo not provided - if isinstance(type, mx.ndarray.NDArray): - type = type.asnumpy().dtype - return np.iinfo(ivy.as_native_dtype(type)) - - -def result_type( - *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] -) -> ivy.Dtype: - raise IvyNotImplementedException() - - def as_ivy_dtype( dtype_in: Union[(str, int, float, complex, bool, np.dtype)], / ) -> ivy.Dtype: @@ -189,6 +137,36 @@ def as_native_dtype(dtype_in: Union[(None, str, bool, int, float, np.dtype)]) -> ) +def astype( + x: Union[(None, mx.ndarray.NDArray)], + dtype: Union[(None, str)], + /, + *, + copy: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return mx.nd.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays( + *arrays: Union[(None, mx.ndarray.NDArray)] +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def broadcast_to( + x: Union[(None, mx.ndarray.NDArray)], + /, + shape: Union[(ivy.NativeShape, Sequence[int])], + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def dtype( x: Union[(None, mx.ndarray.NDArray, np.ndarray)], *, as_native: bool = False ) -> ivy.Dtype: @@ -201,5 +179,26 @@ def dtype_bits(dtype_in: Union[(None, str, np.dtype)], /) -> int: raise IvyNotImplementedException() +@_handle_nestable_dtype_info +def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: + if isinstance(type, mx.ndarray.NDArray): + type = type.dtype + return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: + # using np.iinfo as mx use np dtypes and mxnet iinfo not provided + if isinstance(type, mx.ndarray.NDArray): + type = type.asnumpy().dtype + return np.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[(None, str)], /) -> bool: raise IvyNotImplementedException() + + +def result_type( + *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] +) -> ivy.Dtype: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/device.py b/ivy/functional/backends/mxnet/device.py index 309621811a71e..2a47dc7848916 100644 --- a/ivy/functional/backends/mxnet/device.py +++ b/ivy/functional/backends/mxnet/device.py @@ -11,23 +11,21 @@ from ivy.utils.exceptions import IvyNotImplementedException -def dev( - x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False -) -> Union[(ivy.Device, str)]: - if as_native: - return x.context - return as_ivy_dev(x.context) +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + raise IvyNotImplementedException() + def start(self): + raise IvyNotImplementedException() -def to_device( - x: Union[(None, mx.ndarray.NDArray)], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return x.as_in_context(as_native_dev(device)) + def stop(self): + raise IvyNotImplementedException() + + def __enter__(self): + raise IvyNotImplementedException() + + def __exit__(self, exc_type, exc_val, exc_tb): + raise IvyNotImplementedException() def as_ivy_dev(device): @@ -62,8 +60,12 @@ def clear_cached_mem_on_dev(device: str, /): raise IvyNotImplementedException() -def num_gpus() -> int: - return mx.context.num_gpus() +def dev( + x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False +) -> Union[(ivy.Device, str)]: + if as_native: + return x.context + return as_ivy_dev(x.context) def gpu_is_available() -> bool: @@ -72,22 +74,20 @@ def gpu_is_available() -> bool: return False -def tpu_is_available() -> bool: - return False - - -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - raise IvyNotImplementedException() +def num_gpus() -> int: + return mx.context.num_gpus() - def start(self): - raise IvyNotImplementedException() - def stop(self): - raise IvyNotImplementedException() +def to_device( + x: Union[(None, mx.ndarray.NDArray)], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return x.as_in_context(as_native_dev(device)) - def __enter__(self): - raise IvyNotImplementedException() - def __exit__(self, exc_type, exc_val, exc_tb): - raise IvyNotImplementedException() +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/mxnet/elementwise.py b/ivy/functional/backends/mxnet/elementwise.py index 42a8a2788b605..f3a98ade3f1cf 100644 --- a/ivy/functional/backends/mxnet/elementwise.py +++ b/ivy/functional/backends/mxnet/elementwise.py @@ -46,6 +46,16 @@ def add( return mx.nd.add(mx.nd.multiply(x1, alpha), x2) +def angle( + input: Union[(None, mx.ndarray.NDArray)], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def asin( x: Union[(None, mx.ndarray.NDArray)], /, @@ -178,6 +188,15 @@ def cosh( return mx.nd.cosh(x) +def deg2rad( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def divide( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -203,6 +222,15 @@ def equal( raise IvyNotImplementedException() +def erf( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def exp( x: Union[(None, mx.ndarray.NDArray)], /, @@ -212,6 +240,15 @@ def exp( return mx.nd.exp(x) +def exp2( + x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def expm1( x: Union[(None, mx.ndarray.NDArray)], /, @@ -250,6 +287,16 @@ def fmin( raise IvyNotImplementedException() +def fmod( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def greater( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -270,6 +317,15 @@ def greater_equal( return mx.nd.greater_equal(x1, x2) +def imag( + val: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def isfinite( x: Union[(None, mx.ndarray.NDArray)], /, @@ -299,6 +355,15 @@ def isnan( raise IvyNotImplementedException() +def isreal( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def lcm( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -424,6 +489,28 @@ def logical_xor( return mx.nd.logical_xor(x1, x2) +def maximum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def minimum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def multiply( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -472,6 +559,33 @@ def pow( return mx.nd.power(x1, x2) +def rad2deg( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def real( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def reciprocal( + x: Union[(float, None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.reciprocal(x) + + def remainder( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -573,18 +687,6 @@ def subtract( return mx.nd.subtract(x1, x2) -def trapz( - y: Union[(None, mx.ndarray.NDArray)], - /, - *, - x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - dx: float = 1.0, - axis: int = (-1), - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def tan( x: Union[(None, mx.ndarray.NDArray)], /, @@ -604,121 +706,19 @@ def tanh( return mx.nd.tanh(x) -def trunc( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def imag( - val: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def angle( - input: Union[(None, mx.ndarray.NDArray)], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def exp2( - x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def erf( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def maximum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def minimum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def reciprocal( - x: Union[(float, None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.reciprocal(x) - - -def deg2rad( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def rad2deg( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def isreal( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def fmod( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def trapz( + y: Union[(None, mx.ndarray.NDArray)], /, *, + x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + dx: float = 1.0, + axis: int = (-1), out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def real( +def trunc( x: Union[(None, mx.ndarray.NDArray)], /, *, diff --git a/ivy/functional/backends/mxnet/experimental/activations.py b/ivy/functional/backends/mxnet/experimental/activations.py index 2ab8f7443e0af..74eac3b3ada9c 100644 --- a/ivy/functional/backends/mxnet/experimental/activations.py +++ b/ivy/functional/backends/mxnet/experimental/activations.py @@ -14,9 +14,7 @@ def logit( raise IvyNotImplementedException() -def thresholded_relu( - x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None -) -> None: +def logsigmoid(input: None) -> None: raise IvyNotImplementedException() @@ -24,13 +22,15 @@ def relu6(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def logsigmoid(input: None) -> None: +def selu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def selu(x: None, /, *, out: Optional[None] = None) -> None: +def silu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def silu(x: None, /, *, out: Optional[None] = None) -> None: +def thresholded_relu( + x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None +) -> None: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/creation.py b/ivy/functional/backends/mxnet/experimental/creation.py index 633a0125504d0..3942f38dae26c 100644 --- a/ivy/functional/backends/mxnet/experimental/creation.py +++ b/ivy/functional/backends/mxnet/experimental/creation.py @@ -5,42 +5,44 @@ from ivy.utils.exceptions import IvyNotImplementedException -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def blackman_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def kaiser_bessel_derived_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def vorbis_window( - window_length: Union[(None, mx.ndarray.NDArray)], +def kaiser_bessel_derived_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, *, - dtype: None = np.float32, + dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def hann_window( - size: int, - /, - *, +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: @@ -53,12 +55,10 @@ def tril_indices( raise IvyNotImplementedException() -def blackman_window( - size: int, - /, +def vorbis_window( + window_length: Union[(None, mx.ndarray.NDArray)], *, - periodic: bool = True, - dtype: Optional[None] = None, + dtype: None = np.float32, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/elementwise.py b/ivy/functional/backends/mxnet/experimental/elementwise.py index ce69fe9e91ed3..c6b93dab5e8a7 100644 --- a/ivy/functional/backends/mxnet/experimental/elementwise.py +++ b/ivy/functional/backends/mxnet/experimental/elementwise.py @@ -7,41 +7,21 @@ from .. import backend_version -@with_supported_dtypes( - {"1.9.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.np.log(mx.npx.gamma(x)) - - -def sinc( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def fmax( +def allclose( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> bool: raise IvyNotImplementedException() -def float_power( - x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], - x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def conj( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -71,21 +51,23 @@ def count_nonzero( raise IvyNotImplementedException() -def nansum( - x: Union[(None, mx.ndarray.NDArray)], +def diff( + x: Union[(None, mx.ndarray.NDArray, list, tuple)], /, *, - axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, - dtype: Optional[None] = None, - keepdims: bool = False, + n: int = 1, + axis: int = (-1), + prepend: Optional[ + Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] + ] = None, + append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def gcd( - x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], - x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def fix( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -93,21 +75,19 @@ def gcd( raise IvyNotImplementedException() -def isclose( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def float_power( + x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def signbit( - x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], +def fmax( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -115,9 +95,20 @@ def signbit( raise IvyNotImplementedException() -def hypot( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def frexp( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[ + Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] + ] = None, +) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: + raise IvyNotImplementedException() + + +def gcd( + x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], + x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -125,31 +116,43 @@ def hypot( raise IvyNotImplementedException() -def allclose( +def gradient( + x: None, + /, + *, + spacing: Union[(int, list, tuple)] = 1, + axis: Optional[Union[(int, list, tuple)]] = None, + edge_order: int = 1, +) -> Union[(None, List[None])]: + raise IvyNotImplementedException() + + +def hypot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> bool: +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def fix( - x: Union[(None, mx.ndarray.NDArray)], +def isclose( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def nextafter( +def ldexp( x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray, int)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -157,24 +160,21 @@ def nextafter( raise IvyNotImplementedException() -def diff( - x: Union[(None, mx.ndarray.NDArray, list, tuple)], +@with_supported_dtypes( + {"1.9.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( + x: Union[(None, mx.ndarray.NDArray)], /, *, - n: int = 1, - axis: int = (-1), - prepend: Optional[ - Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] - ] = None, - append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + return mx.np.log(mx.npx.gamma(x)) -def zeta( +def modf( x: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -182,20 +182,21 @@ def zeta( raise IvyNotImplementedException() -def gradient( - x: None, +def nansum( + x: Union[(None, mx.ndarray.NDArray)], /, *, - spacing: Union[(int, list, tuple)] = 1, - axis: Optional[Union[(int, list, tuple)]] = None, - edge_order: int = 1, -) -> Union[(None, List[None])]: + axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, + dtype: Optional[None] = None, + keepdims: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def xlogy( - x: Union[(None, mx.ndarray.NDArray)], - y: Union[(None, mx.ndarray.NDArray)], +def nextafter( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -203,8 +204,8 @@ def xlogy( raise IvyNotImplementedException() -def conj( - x: Union[(None, mx.ndarray.NDArray)], +def signbit( + x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -212,9 +213,8 @@ def conj( raise IvyNotImplementedException() -def ldexp( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray, int)], +def sinc( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -222,19 +222,19 @@ def ldexp( raise IvyNotImplementedException() -def frexp( +def xlogy( x: Union[(None, mx.ndarray.NDArray)], + y: Union[(None, mx.ndarray.NDArray)], /, *, - out: Optional[ - Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] - ] = None, -) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def modf( +def zeta( x: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, diff --git a/ivy/functional/backends/mxnet/experimental/layers.py b/ivy/functional/backends/mxnet/experimental/layers.py index 84fab7a4ba311..4d9b6a781898e 100644 --- a/ivy/functional/backends/mxnet/experimental/layers.py +++ b/ivy/functional/backends/mxnet/experimental/layers.py @@ -7,70 +7,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - raise IvyNotImplementedException() - - -def max_pool1d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: Union[str, int, Tuple[int]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool2d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: Union[str, int, Tuple[int], Tuple[int, int]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool3d( - x: mx.nd.NDArray, - kernel: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - strides: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - padding: Union[str, int, Tuple[int], Tuple[int, int, int]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - def avg_pool1d( x: mx.nd.NDArray, kernel: Union[int, Tuple[int]], @@ -131,18 +67,6 @@ def dct( raise IvyNotImplementedException() -def fft( - x: mx.nd.NDArray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - def dropout1d( x: mx.nd.NDArray, prob: float, @@ -179,6 +103,33 @@ def dropout3d( raise IvyNotImplementedException() +def fft( + x: mx.nd.NDArray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + raise IvyNotImplementedException() + + def ifft( x: mx.nd.NDArray, dim: int, @@ -231,3 +182,52 @@ def interpolate( out: Optional[mx.nd.NDArray] = None, ): raise IvyNotImplementedException() + + +def max_pool1d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: Union[str, int, Tuple[int]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool2d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: Union[str, int, Tuple[int], Tuple[int, int]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool3d( + x: mx.nd.NDArray, + kernel: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + strides: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + padding: Union[str, int, Tuple[int], Tuple[int, int, int]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/linear_algebra.py b/ivy/functional/backends/mxnet/experimental/linear_algebra.py index 8212bc3c54f76..a5da1d277cf8d 100644 --- a/ivy/functional/backends/mxnet/experimental/linear_algebra.py +++ b/ivy/functional/backends/mxnet/experimental/linear_algebra.py @@ -4,24 +4,25 @@ from ivy.utils.exceptions import IvyNotImplementedException -def eigh_tridiagonal( - alpha: Union[(None, mx.ndarray.NDArray)], - beta: Union[(None, mx.ndarray.NDArray)], +dot.support_native_out = True + + +def adjoint( + x: Union[(None, mx.ndarray.NDArray)], /, *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] - ] = None, - tol: Optional[float] = None, -) -> Union[ - ( - None, - mx.ndarray.NDArray, - Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], - ) -]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def cond( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + p: Optional[Union[(None, int, str)]] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -39,31 +40,43 @@ def diagflat( raise IvyNotImplementedException() -def kron( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def dot( + a: mx.ndarray.NDArray, + b: mx.ndarray.NDArray, /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + out: Optional[mx.ndarray.NDArray] = None, +) -> mx.ndarray.NDArray: + return mx.symbol.dot(a, b, out=out) -def matrix_exp( +def eig( x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> Tuple[None]: raise IvyNotImplementedException() -def eig( - x: Union[(None, mx.ndarray.NDArray)], +def eigh_tridiagonal( + alpha: Union[(None, mx.ndarray.NDArray)], + beta: Union[(None, mx.ndarray.NDArray)], /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Tuple[None]: + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] + ] = None, + tol: Optional[float] = None, +) -> Union[ + ( + None, + mx.ndarray.NDArray, + Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], + ) +]: raise IvyNotImplementedException() @@ -73,8 +86,9 @@ def eigvals( raise IvyNotImplementedException() -def adjoint( - x: Union[(None, mx.ndarray.NDArray)], +def kron( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -82,33 +96,19 @@ def adjoint( raise IvyNotImplementedException() -def multi_dot( - x: Sequence[Union[(None, mx.ndarray.NDArray)]], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: - raise IvyNotImplementedException() - - -def cond( +def matrix_exp( x: Union[(None, mx.ndarray.NDArray)], /, *, - p: Optional[Union[(None, int, str)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dot( - a: mx.ndarray.NDArray, - b: mx.ndarray.NDArray, +def multi_dot( + x: Sequence[Union[(None, mx.ndarray.NDArray)]], /, *, - out: Optional[mx.ndarray.NDArray] = None, -) -> mx.ndarray.NDArray: - return mx.symbol.dot(a, b, out=out) - - -dot.support_native_out = True + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> None: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/manipulation.py b/ivy/functional/backends/mxnet/experimental/manipulation.py index 66b46e813ae47..aa8ac9077a4a0 100644 --- a/ivy/functional/backends/mxnet/experimental/manipulation.py +++ b/ivy/functional/backends/mxnet/experimental/manipulation.py @@ -5,39 +5,50 @@ from ivy.utils.exceptions import IvyNotImplementedException -def moveaxis( - a: Union[(None, mx.ndarray.NDArray)], - source: Union[(int, Sequence[int])], - destination: Union[(int, Sequence[int])], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +def atleast_1d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def heaviside( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def atleast_2d( + *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def atleast_3d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: + raise IvyNotImplementedException() + + +def concat_from_sequence( + input_sequence: Union[(Tuple[None], List[None])], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def flipud( - m: Union[(None, mx.ndarray.NDArray)], +def dsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def vstack( +def dstack( arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, @@ -46,41 +57,28 @@ def vstack( raise IvyNotImplementedException() -def hstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def expand( + x: Union[(None, mx.ndarray.NDArray)], + shape: Union[(List[int], List[Tuple])], /, *, + copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def rot90( +def fliplr( m: Union[(None, mx.ndarray.NDArray)], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[(int, int)] = (0, 1), - out: Union[(None, mx.ndarray.NDArray)] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def top_k( - x: None, - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[(None, None)]] = None, -) -> Tuple[(None, None)]: - raise IvyNotImplementedException() - - -def fliplr( +def flipud( m: Union[(None, mx.ndarray.NDArray)], /, *, @@ -90,8 +88,9 @@ def fliplr( raise IvyNotImplementedException() -def i0( - x: Union[(None, mx.ndarray.NDArray)], +def heaviside( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -99,7 +98,7 @@ def i0( raise IvyNotImplementedException() -def vsplit( +def hsplit( ary: Union[(None, mx.ndarray.NDArray)], indices_or_sections: Union[(int, Tuple[(int, ...)])], /, @@ -109,24 +108,17 @@ def vsplit( raise IvyNotImplementedException() -def dsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def hstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def atleast_1d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def i0( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -134,15 +126,27 @@ def dstack( raise IvyNotImplementedException() -def atleast_2d( - *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: +def moveaxis( + a: Union[(None, mx.ndarray.NDArray)], + source: Union[(int, Sequence[int])], + destination: Union[(int, Sequence[int])], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def atleast_3d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: +def rot90( + m: Union[(None, mx.ndarray.NDArray)], + /, + *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[(int, int)] = (0, 1), + out: Union[(None, mx.ndarray.NDArray)] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -158,37 +162,33 @@ def take_along_axis( raise IvyNotImplementedException() -def hsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def top_k( + x: None, + k: int, /, *, - copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[(None, None)]] = None, +) -> Tuple[(None, None)]: raise IvyNotImplementedException() -def expand( - x: Union[(None, mx.ndarray.NDArray)], - shape: Union[(List[int], List[Tuple])], +def vsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def concat_from_sequence( - input_sequence: Union[(Tuple[None], List[None])], +def vstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - new_axis: int = 0, - axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/norms.py b/ivy/functional/backends/mxnet/experimental/norms.py index 19fea95289a8d..ee8cb935f70f0 100644 --- a/ivy/functional/backends/mxnet/experimental/norms.py +++ b/ivy/functional/backends/mxnet/experimental/norms.py @@ -4,16 +4,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def l2_normalize( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: Optional[int] = None, - out: Optional[None] = None, -) -> None: - raise IvyNotImplementedException() - - def batch_norm( x: Union[(None, mx.ndarray.NDArray)], mean: Union[(None, mx.ndarray.NDArray)], @@ -58,6 +48,16 @@ def instance_norm( raise IvyNotImplementedException() +def l2_normalize( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[int] = None, + out: Optional[None] = None, +) -> None: + raise IvyNotImplementedException() + + def lp_normalize( x: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/experimental/random.py b/ivy/functional/backends/mxnet/experimental/random.py index 8c3c4dfc0554d..36c313891b167 100644 --- a/ivy/functional/backends/mxnet/experimental/random.py +++ b/ivy/functional/backends/mxnet/experimental/random.py @@ -5,14 +5,15 @@ from ivy.utils.exceptions import IvyNotImplementedException -def dirichlet( - alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], - /, +def bernoulli( + probs: Union[(float, None, mx.ndarray.NDArray)], *, - size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + logits: Union[(float, None, mx.ndarray.NDArray)] = None, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + device: str, + dtype: None, seed: Optional[int] = None, - dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -31,6 +32,18 @@ def beta( raise IvyNotImplementedException() +def dirichlet( + alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], + /, + *, + size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + seed: Optional[int] = None, + dtype: Optional[None] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def gamma( alpha: Union[(float, None, mx.ndarray.NDArray)], beta: Union[(float, None, mx.ndarray.NDArray)], @@ -55,16 +68,3 @@ def poisson( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def bernoulli( - probs: Union[(float, None, mx.ndarray.NDArray)], - *, - logits: Union[(float, None, mx.ndarray.NDArray)] = None, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - device: str, - dtype: None, - seed: Optional[int] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/statistical.py b/ivy/functional/backends/mxnet/experimental/statistical.py index 252e27dd08e8c..985ab019ce0e1 100644 --- a/ivy/functional/backends/mxnet/experimental/statistical.py +++ b/ivy/functional/backends/mxnet/experimental/statistical.py @@ -4,6 +4,43 @@ from ivy.utils.exceptions import IvyNotImplementedException +def bincount( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + minlength: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def corrcoef( + x: None, + /, + *, + y: None, + rowvar: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> None: + raise IvyNotImplementedException() + + +def cov( + x1: None, + x2: None = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[None] = None, + aweights: Optional[None] = None, + dtype: Optional[type] = None, +) -> None: + raise IvyNotImplementedException() + + def histogram( a: None, /, @@ -44,30 +81,6 @@ def nanmean( raise IvyNotImplementedException() -def quantile( - a: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, float)], - /, - *, - axis: Optional[Union[(int, Sequence[int])]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def corrcoef( - x: None, - /, - *, - y: None, - rowvar: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: - raise IvyNotImplementedException() - - def nanmedian( input: Union[(None, mx.ndarray.NDArray)], /, @@ -79,27 +92,14 @@ def nanmedian( raise IvyNotImplementedException() -def bincount( - x: Union[(None, mx.ndarray.NDArray)], +def quantile( + a: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, float)], /, *, - weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - minlength: int = 0, + axis: Optional[Union[(int, Sequence[int])]] = None, + interpolation: str = "linear", + keepdims: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def cov( - x1: None, - x2: None = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[None] = None, - aweights: Optional[None] = None, - dtype: Optional[type] = None, -) -> None: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/general.py b/ivy/functional/backends/mxnet/general.py index c469b517c47c1..34eebe8331057 100644 --- a/ivy/functional/backends/mxnet/general.py +++ b/ivy/functional/backends/mxnet/general.py @@ -5,10 +5,26 @@ from ivy.utils.exceptions import IvyNotImplementedException +def container_types(): + return [] + + def current_backend_str() -> str: return "mxnet" +def gather( + x: mx.ndarray.NDArray, + indices: mx.ndarray.NDArray, + /, + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[mx.ndarray.NDArray] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def is_native_array( x: Union[(None, mx.ndarray.NDArray)], /, @@ -21,6 +37,10 @@ def is_native_array( return isinstance(x, mx.ndarray.NDArray) or isinstance(x, np.ndarray) +def itemsize(x: mx.ndarray.NDArray, /) -> int: + return x.asnumpy().itemsize + + def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: if copy: if x.shape == (): @@ -29,23 +49,3 @@ def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: return x.copy().asnumpy() else: return x.asnumpy() - - -def itemsize(x: mx.ndarray.NDArray, /) -> int: - return x.asnumpy().itemsize - - -def container_types(): - return [] - - -def gather( - x: mx.ndarray.NDArray, - indices: mx.ndarray.NDArray, - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[mx.ndarray.NDArray] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/gradients.py b/ivy/functional/backends/mxnet/gradients.py index dd3e9041601be..28189a4412a33 100644 --- a/ivy/functional/backends/mxnet/gradients.py +++ b/ivy/functional/backends/mxnet/gradients.py @@ -8,18 +8,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def variable(x, /): - return x - - -def is_variable(x, /, *, exclusive=False): - return isinstance(x, mx.ndarray.NDArray) - - -def variable_data(x, /): - raise IvyNotImplementedException() - - def execute_with_gradients( func, xs, @@ -32,17 +20,29 @@ def execute_with_gradients( raise IvyNotImplementedException() -def value_and_grad(func): +def grad(func, argnums=0): raise IvyNotImplementedException() +def is_variable(x, /, *, exclusive=False): + return isinstance(x, mx.ndarray.NDArray) + + def jac(func): raise IvyNotImplementedException() -def grad(func, argnums=0): +def stop_gradient(x, /, *, preserve_type=True, out=None): raise IvyNotImplementedException() -def stop_gradient(x, /, *, preserve_type=True, out=None): +def value_and_grad(func): + raise IvyNotImplementedException() + + +def variable(x, /): + return x + + +def variable_data(x, /): raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/layers.py b/ivy/functional/backends/mxnet/layers.py index 1ae0560d0c356..3937442087024 100644 --- a/ivy/functional/backends/mxnet/layers.py +++ b/ivy/functional/backends/mxnet/layers.py @@ -66,20 +66,6 @@ def conv2d_transpose( raise IvyNotImplementedException() -def depthwise_conv2d( - x: Union[(None, mx.ndarray.NDArray)], - filters: Union[(None, mx.ndarray.NDArray)], - strides: Union[(int, Tuple[(int, int)])], - padding: Union[(str, int, Sequence[Tuple[(int, int)]])], - /, - *, - data_format: str = "NHWC", - dilations: Union[(int, Tuple[(int, int)])] = 1, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def conv3d( x: Union[(None, mx.ndarray.NDArray)], filters: Union[(None, mx.ndarray.NDArray)], @@ -145,3 +131,17 @@ def conv_general_transpose( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def depthwise_conv2d( + x: Union[(None, mx.ndarray.NDArray)], + filters: Union[(None, mx.ndarray.NDArray)], + strides: Union[(int, Tuple[(int, int)])], + padding: Union[(str, int, Sequence[Tuple[(int, int)]])], + /, + *, + data_format: str = "NHWC", + dilations: Union[(int, Tuple[(int, int)])] = 1, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/linear_algebra.py b/ivy/functional/backends/mxnet/linear_algebra.py index 5077fd3a24bc2..52682499037c2 100644 --- a/ivy/functional/backends/mxnet/linear_algebra.py +++ b/ivy/functional/backends/mxnet/linear_algebra.py @@ -38,6 +38,16 @@ def det( return mx.nd.linalg.det(x) +def diag( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + k: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def diagonal( x: Union[(None, mx.ndarray.NDArray)], /, @@ -249,6 +259,17 @@ def trace( raise IvyNotImplementedException() +def vander( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def vecdot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -273,27 +294,6 @@ def vector_norm( raise IvyNotImplementedException() -def diag( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - k: int = 0, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def vander( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def vector_to_skew_symmetric_matrix( vector: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/manipulation.py b/ivy/functional/backends/mxnet/manipulation.py index feb262db2a152..49dd4aa56f818 100644 --- a/ivy/functional/backends/mxnet/manipulation.py +++ b/ivy/functional/backends/mxnet/manipulation.py @@ -6,6 +6,17 @@ from ivy.utils.exceptions import IvyNotImplementedException +def clip( + x: Union[(None, mx.ndarray.NDArray)], + x_min: Union[(Number, None, mx.ndarray.NDArray)], + x_max: Union[(Number, None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.clip(x, x_min, x_max) + + def concat( xs: Union[(Tuple[(None, ...)], List[None])], /, @@ -16,6 +27,12 @@ def concat( raise IvyNotImplementedException() +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): + raise IvyNotImplementedException() + + def expand_dims( x: Union[(None, mx.ndarray.NDArray)], /, @@ -49,46 +66,36 @@ def permute_dims( raise IvyNotImplementedException() -def reshape( +def repeat( x: Union[(None, mx.ndarray.NDArray)], /, - shape: Union[(ivy.NativeShape, Sequence[int])], + repeats: Union[(int, List[int])], *, - copy: Optional[bool] = None, - order: str = "C", - allowzero: bool = True, + axis: int = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def roll( +def reshape( x: Union[(None, mx.ndarray.NDArray)], /, - shift: Union[(int, Sequence[int])], + shape: Union[(ivy.NativeShape, Sequence[int])], *, - axis: Optional[Union[(int, Sequence[int])]] = None, + copy: Optional[bool] = None, + order: str = "C", + allowzero: bool = True, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def squeeze( +def roll( x: Union[(None, mx.ndarray.NDArray)], /, + shift: Union[(int, Sequence[int])], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.squeeze(x, axis=axis) - - -def stack( - arrays: Union[(Tuple[None], List[None])], - /, - *, - axis: int = 0, + axis: Optional[Union[(int, Sequence[int])]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -106,36 +113,24 @@ def split( raise IvyNotImplementedException() -def repeat( +def squeeze( x: Union[(None, mx.ndarray.NDArray)], /, - repeats: Union[(int, List[int])], *, - axis: int = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + return mx.nd.squeeze(x, axis=axis) -def tile( - x: Union[(None, mx.ndarray.NDArray)], +def stack( + arrays: Union[(Tuple[None], List[None])], /, - repeats: Sequence[int], *, + axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.tile(x, repeats) - - -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): - raise IvyNotImplementedException() - - -def zero_pad( - x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): raise IvyNotImplementedException() @@ -151,15 +146,14 @@ def swapaxes( raise IvyNotImplementedException() -def clip( +def tile( x: Union[(None, mx.ndarray.NDArray)], - x_min: Union[(Number, None, mx.ndarray.NDArray)], - x_max: Union[(Number, None, mx.ndarray.NDArray)], /, + repeats: Sequence[int], *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.clip(x, x_min, x_max) + return mx.nd.tile(x, repeats) def unstack( @@ -171,3 +165,9 @@ def unstack( keepdims: bool = False, ) -> List[None]: raise IvyNotImplementedException() + + +def zero_pad( + x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/random.py b/ivy/functional/backends/mxnet/random.py index 4f1e25f4763d5..97a4c3959de99 100644 --- a/ivy/functional/backends/mxnet/random.py +++ b/ivy/functional/backends/mxnet/random.py @@ -11,12 +11,14 @@ from ivy.utils.exceptions import IvyNotImplementedException -def random_uniform( +def multinomial( + population_size: int, + num_samples: int, + /, *, - low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, - shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, - dtype: None, + batch_size: int = 1, + probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + replace: bool = True, device: str, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -24,42 +26,40 @@ def random_uniform( raise IvyNotImplementedException() -def random_normal( +def randint( + low: Union[(float, None, mx.ndarray.NDArray)], + high: Union[(float, None, mx.ndarray.NDArray)], + /, *, - mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - dtype: None, - seed: Optional[int] = None, device: str, + dtype: Optional[Union[(None, ivy.Dtype)]] = None, + seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def multinomial( - population_size: int, - num_samples: int, - /, +def random_normal( *, - batch_size: int = 1, - probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - replace: bool = True, - device: str, + mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + dtype: None, seed: Optional[int] = None, + device: str, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def randint( - low: Union[(float, None, mx.ndarray.NDArray)], - high: Union[(float, None, mx.ndarray.NDArray)], - /, +def random_uniform( *, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, + shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, + dtype: None, device: str, - dtype: Optional[Union[(None, ivy.Dtype)]] = None, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: diff --git a/ivy/functional/backends/mxnet/searching.py b/ivy/functional/backends/mxnet/searching.py index ed20dee3fccc9..8e758dae93f55 100644 --- a/ivy/functional/backends/mxnet/searching.py +++ b/ivy/functional/backends/mxnet/searching.py @@ -33,6 +33,15 @@ def argmin( raise IvyNotImplementedException() +def argwhere( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def nonzero( x: Union[(None, mx.ndarray.NDArray)], /, @@ -53,12 +62,3 @@ def where( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def argwhere( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/sorting.py b/ivy/functional/backends/mxnet/sorting.py index 356c4b9414403..85954c0a9e91b 100644 --- a/ivy/functional/backends/mxnet/sorting.py +++ b/ivy/functional/backends/mxnet/sorting.py @@ -17,18 +17,6 @@ def argsort( raise IvyNotImplementedException() -def sort( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: int = (-1), - descending: bool = False, - stable: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def msort( a: Union[(None, mx.ndarray.NDArray, list, tuple)], /, @@ -49,3 +37,15 @@ def searchsorted( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def sort( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: int = (-1), + descending: bool = False, + stable: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/statistical.py b/ivy/functional/backends/mxnet/statistical.py index 67a24d9922d61..c688052b9fc2e 100644 --- a/ivy/functional/backends/mxnet/statistical.py +++ b/ivy/functional/backends/mxnet/statistical.py @@ -7,12 +7,34 @@ import ivy -def min( +def cumprod( x: Union[(None, mx.ndarray.NDArray)], /, *, - axis: Optional[Union[(int, Sequence[int])]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def cumsum( + x: Union[(None, mx.ndarray.NDArray)], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def einsum( + equation: str, + *operands: Union[(None, mx.ndarray.NDArray)], out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -40,6 +62,17 @@ def mean( raise IvyNotImplementedException() +def min( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[Union[(int, Sequence[int])]] = None, + keepdims: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def prod( x: Union[(None, mx.ndarray.NDArray)], /, @@ -110,36 +143,3 @@ def var( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def cumprod( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def cumsum( - x: Union[(None, mx.ndarray.NDArray)], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def einsum( - equation: str, - *operands: Union[(None, mx.ndarray.NDArray)], - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/numpy/activations.py b/ivy/functional/backends/numpy/activations.py index 2e230498b9574..f7ef8d66048b9 100644 --- a/ivy/functional/backends/numpy/activations.py +++ b/ivy/functional/backends/numpy/activations.py @@ -9,14 +9,34 @@ from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array +relu.support_native_out = True +softmax.support_native_out = True +softplus.support_native_out = True +log_softmax.support_native_out = True +mish.support_native_out = True +hardswish.support_native_out = True + + @_scalar_output_to_0d_array -def relu( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +def gelu( + x: np.ndarray, + /, + *, + approximate: bool = False, + complex_mode="jax", + out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.maximum(x, 0, out=out, dtype=x.dtype) + if approximate: + ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) + else: + ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) + return ivy.astype(ret, x.dtype, copy=False) -relu.support_native_out = True +@_scalar_output_to_0d_array +def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) + return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) def leaky_relu( @@ -31,19 +51,36 @@ def leaky_relu( @_scalar_output_to_0d_array -def gelu( - x: np.ndarray, - /, - *, - approximate: bool = False, - complex_mode="jax", - out: Optional[np.ndarray] = None, +def log_softmax( + x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None ) -> np.ndarray: - if approximate: - ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) - else: - ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) - return ivy.astype(ret, x.dtype, copy=False) + if axis is None: + axis = -1 + x_max = np.max(x, axis=axis, keepdims=True) + if x_max.ndim > 0: + x_max[~np.isfinite(x_max)] = 0 + elif not np.isfinite(x_max): + x_max = 0 + exp_tmp = np.exp(x - x_max) + + with np.errstate(divide="ignore"): + s = np.sum(exp_tmp, axis=axis, keepdims=True) + ret = np.log(s) + + ret = x - x_max - ret + return ret + + +@_scalar_output_to_0d_array +def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return x * np.tanh(np.log1p(np.exp(x))) + + +@_scalar_output_to_0d_array +def relu( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.maximum(x, 0, out=out, dtype=x.dtype) def sigmoid(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -60,9 +97,6 @@ def softmax( return np.divide(exp_x, np.sum(exp_x, axis=axis, keepdims=True), out=out) -softmax.support_native_out = True - - @_scalar_output_to_0d_array def softplus( x: np.ndarray, @@ -92,47 +126,3 @@ def softplus( if threshold is not None: return np.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -softplus.support_native_out = True - - -@_scalar_output_to_0d_array -def log_softmax( - x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None -) -> np.ndarray: - if axis is None: - axis = -1 - x_max = np.max(x, axis=axis, keepdims=True) - if x_max.ndim > 0: - x_max[~np.isfinite(x_max)] = 0 - elif not np.isfinite(x_max): - x_max = 0 - exp_tmp = np.exp(x - x_max) - - with np.errstate(divide="ignore"): - s = np.sum(exp_tmp, axis=axis, keepdims=True) - ret = np.log(s) - - ret = x - x_max - ret - return ret - - -log_softmax.support_native_out = True - - -@_scalar_output_to_0d_array -def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return x * np.tanh(np.log1p(np.exp(x))) - - -mish.support_native_out = True - - -@_scalar_output_to_0d_array -def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) - return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) - - -hardswish.support_native_out = True diff --git a/ivy/functional/backends/numpy/creation.py b/ivy/functional/backends/numpy/creation.py index 7600f3cf9dba0..0eb26a54646d9 100644 --- a/ivy/functional/backends/numpy/creation.py +++ b/ivy/functional/backends/numpy/creation.py @@ -64,6 +64,17 @@ def asarray( return np.copy(ret) if copy else ret +def copy_array( + x: np.ndarray, + *, + to_ivy_array: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -107,6 +118,17 @@ def from_dlpack(x, /, *, out: Optional[np.ndarray] = None): return np.from_dlpack(x) +def frombuffer( + buffer: bytes, + dtype: Optional[np.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> np.ndarray: + if isinstance(dtype, list): + dtype = np.dtype(dtype[0]) + return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -170,68 +192,6 @@ def meshgrid( return np.meshgrid(*arrays, sparse=sparse, indexing=indexing) -def ones( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.ones(shape, dtype), device=device) - - -def ones_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.ones_like(x, dtype=dtype), device=device) - - -def tril( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tril(x, k) - - -def triu( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.triu(x, k) - - -def zeros( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.zeros(shape, dtype), device=device) - - -def zeros_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.zeros_like(x, dtype=dtype), device=device) - - -# Extra # -# ------# - - -array = asarray - - -def copy_array( - x: np.ndarray, - *, - to_ivy_array: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() - - def one_hot( indices: np.ndarray, depth: int, @@ -268,15 +228,32 @@ def one_hot( return res -def frombuffer( - buffer: bytes, - dtype: Optional[np.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, +def ones( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(dtype, list): - dtype = np.dtype(dtype[0]) - return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + return _to_device(np.ones(shape, dtype), device=device) + + +def ones_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.ones_like(x, dtype=dtype), device=device) + + +def tril( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tril(x, k) + + +def triu( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.triu(x, k) def triu_indices( @@ -290,3 +267,26 @@ def triu_indices( return tuple( _to_device(np.asarray(np.triu_indices(n=n_rows, k=k, m=n_cols)), device=device) ) + + +def zeros( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return _to_device(np.zeros(shape, dtype), device=device) + + +def zeros_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.zeros_like(x, dtype=dtype), device=device) + + +# Extra # +# ------# + + +array = asarray diff --git a/ivy/functional/backends/numpy/data_type.py b/ivy/functional/backends/numpy/data_type.py index 317cc71c8608d..172989ff82968 100644 --- a/ivy/functional/backends/numpy/data_type.py +++ b/ivy/functional/backends/numpy/data_type.py @@ -9,6 +9,26 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from . import backend_version +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int16"): "int16", @@ -40,7 +60,6 @@ np.complex128: "complex128", np.bool_: "bool", } - native_dtype_dict = { "int8": np.dtype("int8"), "int16": np.dtype("int16"), @@ -58,27 +77,6 @@ "bool": np.dtype("bool"), } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} - class Finfo: def __init__(self, np_finfo: np.finfo): @@ -108,68 +106,6 @@ def smallest_normal(self): return float(self._np_finfo.tiny) -# Array API Standard # -# -------------------# - - -def astype( - x: np.ndarray, - dtype: np.dtype, - /, - *, - copy: bool = True, - out: Optional[ivy.Array] = None, -) -> np.ndarray: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return np.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: - try: - return np.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def broadcast_to( - x: np.ndarray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return np.broadcast_to(x.reshape([-1]), shape) - return np.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype - return Finfo(np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype - return np.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return np.result_type(arrays_and_dtypes) - result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = np.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -238,6 +174,45 @@ def as_native_dtype(dtype_in: Union[np.dtype, str, bool, int, float], /) -> np.d ) +# Array API Standard # +# -------------------# + + +def astype( + x: np.ndarray, + dtype: np.dtype, + /, + *, + copy: bool = True, + out: Optional[ivy.Array] = None, +) -> np.ndarray: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return np.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: + try: + return np.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def broadcast_to( + x: np.ndarray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return np.broadcast_to(x.reshape([-1]), shape) + return np.broadcast_to(x, shape) + + def dtype(x: np.ndarray, *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.to_native(x).dtype @@ -257,6 +232,20 @@ def dtype_bits(dtype_in: Union[np.dtype, str], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype + return Finfo(np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype + return np.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -264,3 +253,12 @@ def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return np.result_type(arrays_and_dtypes) + result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = np.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/numpy/device.py b/ivy/functional/backends/numpy/device.py index b10c8119f8cc1..4576b42831bc5 100644 --- a/ivy/functional/backends/numpy/device.py +++ b/ivy/functional/backends/numpy/device.py @@ -11,34 +11,30 @@ from ivy.functional.ivy.device import Profiler as BaseProfiler -def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: - if as_native: - return "cpu" - return as_ivy_dev("cpu") - - -def as_ivy_dev(device: str, /): - return ivy.Device("cpu") - - -def as_native_dev(device: str, /): - return "cpu" - - -def clear_cached_mem_on_dev(device: str, /): - return None +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper numpy profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + def start(self): + self._start_time = time.perf_counter() -def tpu_is_available() -> bool: - return False + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + def __enter__(self): + self.start() -def num_gpus() -> int: - return 0 + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() -def gpu_is_available() -> bool: - return False +# --- Helpers --- # +# --------------- # # private version of to_device to be used in backend implementations @@ -60,6 +56,40 @@ def _to_device(x: np.ndarray, device=None) -> np.ndarray: return x +# --- Main --- # +# ------------ # + + +def as_ivy_dev(device: str, /): + return ivy.Device("cpu") + + +def as_native_dev(device: str, /): + return "cpu" + + +def clear_cached_mem_on_dev(device: str, /): + return None + + +def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: + if as_native: + return "cpu" + return as_ivy_dev("cpu") + + +def gpu_is_available() -> bool: + return False + + +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return fn(*args, **kwargs) + + +def num_gpus() -> int: + return 0 + + def to_device( x: np.ndarray, device: str, @@ -85,27 +115,5 @@ def to_device( return x -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return fn(*args, **kwargs) - - -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper numpy profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None - - def start(self): - self._start_time = time.perf_counter() - - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) - - def __enter__(self): - self.start() - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/numpy/elementwise.py b/ivy/functional/backends/numpy/elementwise.py index 6c4459bbaaf4f..5d643eb77e7a3 100644 --- a/ivy/functional/backends/numpy/elementwise.py +++ b/ivy/functional/backends/numpy/elementwise.py @@ -10,6 +10,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _abs_variant_sign(x): + return np.divide(x, np.abs(x), where=x != 0) + + +# --- Main --- # +# ------------ # + + @_scalar_output_to_0d_array def abs( x: Union[float, np.ndarray], @@ -20,25 +32,16 @@ def abs( return np.absolute(x, out=out) -abs.support_native_out = True - - @_scalar_output_to_0d_array def acos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccos(x, out=out) -acos.support_native_out = True - - @_scalar_output_to_0d_array def acosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccosh(x, out=out) -acosh.support_native_out = True - - @_scalar_output_to_0d_array def add( x1: Union[float, np.ndarray], @@ -55,7 +58,14 @@ def add( return np.add(x1, x2, out=out) -add.support_native_out = True +def angle( + z: np.ndarray, + /, + *, + deg: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.angle(z, deg=deg) @_scalar_output_to_0d_array @@ -63,25 +73,16 @@ def asin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsin(x, out=out) -asin.support_native_out = True - - @_scalar_output_to_0d_array def asinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsinh(x, out=out) -asinh.support_native_out = True - - @_scalar_output_to_0d_array def atan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctan(x, out=out) -atan.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def atan2( @@ -91,17 +92,11 @@ def atan2( return np.arctan2(x1, x2, out=out) -atan2.support_native_out = True - - @_scalar_output_to_0d_array def atanh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctanh(x, out=out) -atanh.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_and( @@ -115,9 +110,6 @@ def bitwise_and( return np.bitwise_and(x1, x2, out=out) -bitwise_and.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_invert( @@ -126,9 +118,6 @@ def bitwise_invert( return np.invert(x, out=out) -bitwise_invert.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_left_shift( @@ -142,9 +131,6 @@ def bitwise_left_shift( return np.left_shift(x1, x2, out=out) -bitwise_left_shift.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_or( @@ -158,9 +144,6 @@ def bitwise_or( return np.bitwise_or(x1, x2, out=out) -bitwise_or.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_right_shift( @@ -174,9 +157,6 @@ def bitwise_right_shift( return np.right_shift(x1, x2, out=out) -bitwise_right_shift.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_xor( @@ -190,9 +170,6 @@ def bitwise_xor( return np.bitwise_xor(x1, x2, out=out) -bitwise_xor.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) @_scalar_output_to_0d_array def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -205,24 +182,21 @@ def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -ceil.support_native_out = True - - @_scalar_output_to_0d_array def cos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cos(x, out=out) -cos.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) @_scalar_output_to_0d_array def cosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cosh(x, out=out) -cosh.support_native_out = True +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.deg2rad(x, out=out) @_scalar_output_to_0d_array @@ -242,9 +216,6 @@ def divide( return ret -divide.support_native_out = True - - @_scalar_output_to_0d_array def equal( x1: Union[float, np.ndarray], @@ -257,15 +228,36 @@ def equal( return np.equal(x1, x2, out=out) -equal.support_native_out = True +# Extra # +# ------# @_scalar_output_to_0d_array -def exp(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.exp(x, out=out) +def erf(x, /, *, out: Optional[np.ndarray] = None): + a1 = 0.254829592 + a2 = -0.284496736 + a3 = 1.421413741 + a4 = -1.453152027 + a5 = 1.061405429 + p = 0.3275911 + sign = np.sign(x) + x = np.abs(x) -exp.support_native_out = True + # A&S formula 7.1.26 + t = 1.0 / (1.0 + p * x) + y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) + ret = sign * y + if hasattr(x, "dtype"): + ret = np.asarray(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + +@_scalar_output_to_0d_array +def exp(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.exp(x, out=out) def exp2( @@ -277,17 +269,11 @@ def exp2( return np.exp2(x, out=out) -exp2.support_native_out = True - - @_scalar_output_to_0d_array def expm1(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.expm1(x, out=out) -expm1.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -300,9 +286,6 @@ def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -floor.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor_divide( @@ -337,7 +320,32 @@ def fmin( ) -fmin.support_native_out = True +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def fmod( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmod( + x1, + x2, + out=None, + ) + + +def gcd( + x1: Union[np.ndarray, int, list, tuple], + x2: Union[np.ndarray, float, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.gcd(x1, x2, out=out) @_scalar_output_to_0d_array @@ -352,9 +360,6 @@ def greater( return np.greater(x1, x2, out=out) -greater.support_native_out = True - - @_scalar_output_to_0d_array def greater_equal( x1: Union[float, np.ndarray], @@ -367,7 +372,13 @@ def greater_equal( return np.greater_equal(x1, x2, out=out) -greater_equal.support_native_out = True +def imag( + val: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.imag(val) @_scalar_output_to_0d_array @@ -375,9 +386,6 @@ def isfinite(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarra return np.isfinite(x, out=out) -isfinite.support_native_out = True - - @_scalar_output_to_0d_array def isinf( x: np.ndarray, @@ -401,7 +409,9 @@ def isnan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.isnan(x, out=out) -isnan.support_native_out = True +@_scalar_output_to_0d_array +def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.isreal(x) @_scalar_output_to_0d_array @@ -420,9 +430,6 @@ def lcm( ) -lcm.support_native_out = True - - @_scalar_output_to_0d_array def less( x1: Union[float, np.ndarray], @@ -435,9 +442,6 @@ def less( return np.less(x1, x2, out=out) -less.support_native_out = True - - @_scalar_output_to_0d_array def less_equal( x1: Union[float, np.ndarray], @@ -450,41 +454,26 @@ def less_equal( return np.less_equal(x1, x2, out=out) -less_equal.support_native_out = True - - @_scalar_output_to_0d_array def log(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log(x, out=out) -log.support_native_out = True - - @_scalar_output_to_0d_array def log10(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log10(x, out=out) -log10.support_native_out = True - - @_scalar_output_to_0d_array def log1p(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log1p(x, out=out) -log1p.support_native_out = True - - @_scalar_output_to_0d_array def log2(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log2(x, out=out) -log2.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def logaddexp( @@ -494,9 +483,6 @@ def logaddexp( return np.logaddexp(x1, x2, out=out) -logaddexp.support_native_out = True - - def logaddexp2( x1: Union[np.ndarray, int, list, tuple], x2: Union[np.ndarray, int, list, tuple], @@ -511,9 +497,6 @@ def logaddexp2( return np.logaddexp2(x1, x2, out=out) -logaddexp2.support_native_out = True - - @_scalar_output_to_0d_array def logical_and( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -521,17 +504,11 @@ def logical_and( return np.logical_and(x1, x2, out=out) -logical_and.support_native_out = True - - @_scalar_output_to_0d_array def logical_not(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.logical_not(x, out=out) -logical_not.support_native_out = True - - @_scalar_output_to_0d_array def logical_or( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -539,9 +516,6 @@ def logical_or( return np.logical_or(x1, x2, out=out) -logical_or.support_native_out = True - - @_scalar_output_to_0d_array def logical_xor( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -549,7 +523,40 @@ def logical_xor( return np.logical_xor(x1, x2, out=out) -logical_xor.support_native_out = True +@_scalar_output_to_0d_array +def maximum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +): + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 >= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.maximum(x1, x2, out=out) + + +@_scalar_output_to_0d_array +def minimum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 <= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.minimum(x1, x2, out=out) @_scalar_output_to_0d_array @@ -564,7 +571,17 @@ def multiply( return np.multiply(x1, x2, out=out) -multiply.support_native_out = True +def nan_to_num( + x: np.ndarray, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) @_scalar_output_to_0d_array @@ -574,9 +591,6 @@ def negative( return np.negative(x, out=out) -negative.support_native_out = True - - @_scalar_output_to_0d_array def not_equal( x1: Union[float, np.ndarray], @@ -589,9 +603,6 @@ def not_equal( return np.not_equal(x1, x2, out=out) -not_equal.support_native_out = True - - @_scalar_output_to_0d_array def positive( x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None @@ -599,9 +610,6 @@ def positive( return np.positive(x, out=out) -positive.support_native_out = True - - @_scalar_output_to_0d_array def pow( x1: Union[float, np.ndarray], @@ -614,16 +622,31 @@ def pow( return np.power(x1, x2, out=out) -pow.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def remainder( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, +def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.rad2deg(x, out=out) + + +def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.real(x) + + +@_scalar_output_to_0d_array +def reciprocal( + x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + numerator = np.ones_like(x) + return np.true_divide(numerator, x, out=out) + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def remainder( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, modulus: bool = True, out: Optional[np.ndarray] = None, ) -> np.ndarray: @@ -637,9 +660,6 @@ def remainder( return np.remainder(x1, x2, out=out) -remainder.support_native_out = True - - @_scalar_output_to_0d_array def round( x: np.ndarray, /, *, decimals: int = 0, out: Optional[np.ndarray] = None @@ -653,13 +673,6 @@ def round( return ret -round.support_native_out = True - - -def _abs_variant_sign(x): - return np.divide(x, np.abs(x), where=x != 0) - - @_scalar_output_to_0d_array def sign( x: np.ndarray, @@ -673,41 +686,26 @@ def sign( return np.sign(x, out=out) -sign.support_native_out = True - - @_scalar_output_to_0d_array def sin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sin(x, out=out) -sin.support_native_out = True - - @_scalar_output_to_0d_array def sinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sinh(x, out=out) -sinh.support_native_out = True - - @_scalar_output_to_0d_array def sqrt(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sqrt(x, out=out) -sqrt.support_native_out = True - - @_scalar_output_to_0d_array def square(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.square(x, out=out) -square.support_native_out = True - - @_scalar_output_to_0d_array def subtract( x1: Union[float, np.ndarray], @@ -725,7 +723,16 @@ def subtract( return np.subtract(x1, x2, out=out) -subtract.support_native_out = True +@_scalar_output_to_0d_array +def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.tan(x, out=out) + + +@_scalar_output_to_0d_array +def tanh( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tanh(x, out=out) @_scalar_output_to_0d_array @@ -741,27 +748,6 @@ def trapz( return np.trapz(y, x=x, dx=dx, axis=axis) -trapz.support_native_out = False - - -@_scalar_output_to_0d_array -def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.tan(x, out=out) - - -tan.support_native_out = True - - -@_scalar_output_to_0d_array -def tanh( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tanh(x, out=out) - - -tanh.support_native_out = True - - @_scalar_output_to_0d_array def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: if "int" in str(x.dtype): @@ -773,192 +759,74 @@ def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret +abs.support_native_out = True +acos.support_native_out = True +acosh.support_native_out = True +add.support_native_out = True +asin.support_native_out = True +asinh.support_native_out = True +atan.support_native_out = True +atan2.support_native_out = True +atanh.support_native_out = True +bitwise_and.support_native_out = True +bitwise_invert.support_native_out = True +bitwise_left_shift.support_native_out = True +bitwise_or.support_native_out = True +bitwise_right_shift.support_native_out = True +bitwise_xor.support_native_out = True +ceil.support_native_out = True +cos.support_native_out = True +cosh.support_native_out = True +divide.support_native_out = True +equal.support_native_out = True +exp.support_native_out = True +exp2.support_native_out = True +expm1.support_native_out = True +floor.support_native_out = True +fmin.support_native_out = True +greater.support_native_out = True +greater_equal.support_native_out = True +isfinite.support_native_out = True +isnan.support_native_out = True +lcm.support_native_out = True +less.support_native_out = True +less_equal.support_native_out = True +log.support_native_out = True +log10.support_native_out = True +log1p.support_native_out = True +log2.support_native_out = True +logaddexp.support_native_out = True +logaddexp2.support_native_out = True +logical_and.support_native_out = True +logical_not.support_native_out = True +logical_or.support_native_out = True +logical_xor.support_native_out = True +multiply.support_native_out = True +negative.support_native_out = True +not_equal.support_native_out = True +positive.support_native_out = True +pow.support_native_out = True +remainder.support_native_out = True +round.support_native_out = True +sign.support_native_out = True +sin.support_native_out = True +sinh.support_native_out = True +sqrt.support_native_out = True +square.support_native_out = True +subtract.support_native_out = True +trapz.support_native_out = False +tan.support_native_out = True +tanh.support_native_out = True trunc.support_native_out = True - - -# Extra # -# ------# - - -@_scalar_output_to_0d_array -def erf(x, /, *, out: Optional[np.ndarray] = None): - a1 = 0.254829592 - a2 = -0.284496736 - a3 = 1.421413741 - a4 = -1.453152027 - a5 = 1.061405429 - p = 0.3275911 - - sign = np.sign(x) - x = np.abs(x) - - # A&S formula 7.1.26 - t = 1.0 / (1.0 + p * x) - y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) - ret = sign * y - if hasattr(x, "dtype"): - ret = np.asarray(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - erf.support_native_out = True - - -@_scalar_output_to_0d_array -def maximum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -): - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 >= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.maximum(x1, x2, out=out) - - maximum.support_native_out = True - - -@_scalar_output_to_0d_array -def minimum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 <= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.minimum(x1, x2, out=out) - - minimum.support_native_out = True - - -@_scalar_output_to_0d_array -def reciprocal( - x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - numerator = np.ones_like(x) - return np.true_divide(numerator, x, out=out) - - reciprocal.support_native_out = True - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.deg2rad(x, out=out) - - deg2rad.support_native_out = True - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.rad2deg(x, out=out) - - rad2deg.support_native_out = True - - -@_scalar_output_to_0d_array -def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.isreal(x) - - isreal.support_native_out = False - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def fmod( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmod( - x1, - x2, - out=None, - ) - - fmod.support_native_out = True - - -def angle( - z: np.ndarray, - /, - *, - deg: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.angle(z, deg=deg) - - angle.support_native_out = False - - -def gcd( - x1: Union[np.ndarray, int, list, tuple], - x2: Union[np.ndarray, float, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.gcd(x1, x2, out=out) - - gcd.support_native_out = True - - -def imag( - val: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.imag(val) - - imag.support_native_out = False - - -def nan_to_num( - x: np.ndarray, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) - - nan_to_num.support_native_out = False - - -def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.real(x) diff --git a/ivy/functional/backends/numpy/experimental/activations.py b/ivy/functional/backends/numpy/experimental/activations.py index d3f8282a22de2..0ae99da3621d4 100644 --- a/ivy/functional/backends/numpy/experimental/activations.py +++ b/ivy/functional/backends/numpy/experimental/activations.py @@ -10,6 +10,24 @@ from . import backend_version +thresholded_relu.support_native_out = True +relu6.support_native_out = True +selu.support_native_out = True +silu.support_native_out = True +elu.support_native_out = True + + +@_scalar_output_to_0d_array +def elu( + x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None +) -> np.ndarray: + # exp = np.expm1(x) + ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret + + def logit( x: np.ndarray, /, @@ -28,18 +46,10 @@ def logit( return ret +@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) @_scalar_output_to_0d_array -def thresholded_relu( - x: np.ndarray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.where(x > threshold, x, 0).astype(x.dtype) - - -thresholded_relu.support_native_out = True +def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return -(np.log1p(np.exp(-(input)))) @_scalar_output_to_0d_array @@ -47,15 +57,6 @@ def relu6(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.minimum(np.maximum(x, 0, dtype=x.dtype), 6, out=out, dtype=x.dtype) -relu6.support_native_out = True - - -@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) -@_scalar_output_to_0d_array -def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return -(np.log1p(np.exp(-(input)))) - - @_scalar_output_to_0d_array def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: alpha = 1.6732632423543772848170429916717 @@ -66,9 +67,6 @@ def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -selu.support_native_out = True - - @_scalar_output_to_0d_array def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: ret = np.asarray(x * (1 / (1 + np.exp(-x)))) @@ -80,18 +78,12 @@ def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.asarray(x * (1 / (1 + np.exp(-x)))).astype(x.dtype) -silu.support_native_out = True - - @_scalar_output_to_0d_array -def elu( - x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None +def thresholded_relu( + x: np.ndarray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - # exp = np.expm1(x) - ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret - - -elu.support_native_out = True + return np.where(x > threshold, x, 0).astype(x.dtype) diff --git a/ivy/functional/backends/numpy/experimental/creation.py b/ivy/functional/backends/numpy/experimental/creation.py index c2e01f77c55e8..150d49fce0060 100644 --- a/ivy/functional/backends/numpy/experimental/creation.py +++ b/ivy/functional/backends/numpy/experimental/creation.py @@ -7,37 +7,32 @@ from ivy.functional.backends.numpy.device import _to_device import ivy -# Array API Standard # -# -------------------# - - -def vorbis_window( - window_length: np.ndarray, - *, - dtype: np.dtype = np.float32, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - result = [] - for i in range(1, window_length * 2, 2): - temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) - result.append(round(temp, 8)) - return np.array(result, dtype=dtype) - vorbis_window.support_native_out = False +hann_window.support_native_out = False +kaiser_window.support_native_out = False +blackman_window.support_native_out = False -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, +def blackman_window( + size: int, /, *, - device: str, -) -> Tuple[np.ndarray, ...]: - return tuple( - _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) - ) + periodic: bool = True, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if size < 2: + return np.ones([size], dtype=dtype) + if periodic: + count = np.arange(size) / size + else: + count = np.linspace(start=0, stop=size, num=size) + + return ( + (0.42 - 0.5 * np.cos(2 * np.pi * count)) + + (0.08 * np.cos(2 * np.pi * 2 * count)) + ).astype(dtype) def hann_window( @@ -57,7 +52,12 @@ def hann_window( return (0.5 - 0.5 * np.cos(2 * np.pi * count)).astype(dtype) -hann_window.support_native_out = False +def indices( + dimensions: Sequence, + dtype: np.dtype = np.int64, + sparse: bool = False, +) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: + return np.indices(dimensions, dtype=dtype, sparse=sparse) def kaiser_window( @@ -76,15 +76,30 @@ def kaiser_window( return np.kaiser(M=window_length + 1, beta=beta)[:-1].astype(dtype) -kaiser_window.support_native_out = False +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[np.ndarray, ...]: + return tuple( + _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) + ) -def indices( - dimensions: Sequence, - dtype: np.dtype = np.int64, - sparse: bool = False, -) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: - return np.indices(dimensions, dtype=dtype, sparse=sparse) +def trilu( + x: np.ndarray, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if upper: + return np.triu(x, k) + return np.tril(x, k) def unsorted_segment_min( @@ -113,30 +128,6 @@ def unsorted_segment_min( return res -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if size < 2: - return np.ones([size], dtype=dtype) - if periodic: - count = np.arange(size) / size - else: - count = np.linspace(start=0, stop=size, num=size) - - return ( - (0.42 - 0.5 * np.cos(2 * np.pi * count)) - + (0.08 * np.cos(2 * np.pi * 2 * count)) - ).astype(dtype) - - -blackman_window.support_native_out = False - - def unsorted_segment_sum( data: np.ndarray, segment_ids: np.ndarray, @@ -160,14 +151,18 @@ def unsorted_segment_sum( return res -def trilu( - x: np.ndarray, - /, +# Array API Standard # +# -------------------# + + +def vorbis_window( + window_length: np.ndarray, *, - k: int = 0, - upper: bool = True, + dtype: np.dtype = np.float32, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if upper: - return np.triu(x, k) - return np.tril(x, k) + result = [] + for i in range(1, window_length * 2, 2): + temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) + result.append(round(temp, 8)) + return np.array(result, dtype=dtype) diff --git a/ivy/functional/backends/numpy/experimental/elementwise.py b/ivy/functional/backends/numpy/experimental/elementwise.py index cf91193091296..04479cf5508c6 100644 --- a/ivy/functional/backends/numpy/experimental/elementwise.py +++ b/ivy/functional/backends/numpy/experimental/elementwise.py @@ -9,49 +9,96 @@ from . import backend_version -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.sinc(x).astype(x.dtype) +fmax.support_native_out = True +float_power.support_native_out = True +copysign.support_native_out = True +count_nonzero.support_native_out = False +nansum.support_native_out = True +isclose.support_native_out = False +signbit.support_native_out = True +diff.support_native_out = False +allclose.support_native_out = False +fix.support_native_out = True +nextafter.support_natvie_out = True +zeta.support_native_out = False +# --- LGAMMA --- # +LANCZOS_N = 13 +kBaseLanczosCoeff = 0.99999999999980993227684700473478 +kLanczosCoefficients = np.array( + [ + 676.520368121885098567009190444019, + -1259.13921672240287047156078755283, + 771.3234287776530788486528258894, + -176.61502916214059906584551354, + 12.507343278686904814458936853, + -0.13857109526572011689554707, + 9.984369578019570859563e-6, + 1.50563273514931155834e-7, + ] +) +# ---digamma---# +kLanczosGamma = 7 # aka g +lanczos_den_coeffs = np.array( + [ + 0.0, + 39916800.0, + 120543840.0, + 150917976.0, + 105258076.0, + 45995730.0, + 13339535.0, + 2637558.0, + 357423.0, + 32670.0, + 1925.0, + 66.0, + 1.0, + ] +) +lanczos_g = 6.024680040776729583740234375 +lanczos_num_coeffs = np.array( + [ + 23531376880.410759688572007674451636754734846804940, + 42919803642.649098768957899047001988850926355848959, + 35711959237.355668049440185451547166705960488635843, + 17921034426.037209699919755754458931112671403265390, + 6039542586.3520280050642916443072979210699388420708, + 1439720407.3117216736632230727949123939715485786772, + 248874557.86205415651146038641322942321632125127801, + 31426415.585400194380614231628318205362874684987640, + 2876370.6289353724412254090516208496135991145378768, + 186056.26539522349504029498971604569928220784236328, + 8071.6720023658162106380029022722506138218516325024, + 210.82427775157934587250973392071336271166969580291, + 2.5066282746310002701649081771338373386264310793408, + ] +) @_scalar_output_to_0d_array -def fmax( +def allclose( x1: np.ndarray, x2: np.ndarray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmax( - x1, - x2, - out=None, - where=True, - casting="same_kind", - order="K", - dtype=None, - subok=True, - ) - - -fmax.support_native_out = True +) -> bool: + return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) -@_scalar_output_to_0d_array -def float_power( - x1: Union[np.ndarray, float, list, tuple], - x2: Union[np.ndarray, float, list, tuple], +def conj( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.float_power(x1, x2, out=out) - - -float_power.support_native_out = True + ret = np.conj(x, out=out) + if x.dtype == bool: + return ret.astype("bool") + return ret @_scalar_output_to_0d_array @@ -69,9 +116,6 @@ def copysign( return np.copysign(x1, x2, out=out) -copysign.support_native_out = True - - @_scalar_output_to_0d_array def count_nonzero( a: np.ndarray, @@ -90,67 +134,6 @@ def count_nonzero( return ret.astype(dtype) -count_nonzero.support_native_out = False - - -def nansum( - x: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[np.dtype] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -nansum.support_native_out = True - - -def isclose( - a: np.ndarray, - b: np.ndarray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - if np.isscalar(ret): - return np.array(ret, dtype="bool") - return ret - - -isclose.support_native_out = False - - -def signbit( - x: Union[np.ndarray, float, int, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.signbit(x, out=out) - - -signbit.support_native_out = True - - -def hypot( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.hypot(x1, x2) - - def diff( x: Union[np.ndarray, list, tuple], /, @@ -159,165 +142,11 @@ def diff( axis: int = -1, prepend: Optional[Union[np.ndarray, int, float, list, tuple]] = None, append: Optional[Union[np.ndarray, int, float, list, tuple]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - prepend = prepend if prepend is not None else np._NoValue - append = append if append is not None else np._NoValue - return np.diff(x, n=n, axis=axis, prepend=prepend, append=append) - - -diff.support_native_out = False - - -@_scalar_output_to_0d_array -def allclose( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[np.ndarray] = None, -) -> bool: - return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -allclose.support_native_out = False - - -def fix( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.fix(x, out=out) - - -fix.support_native_out = True - - -def nextafter( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nextafter(x1, x2) - - -nextafter.support_natvie_out = True - - -def zeta( - x: np.ndarray, - q: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) - inf_indices = np.logical_or(temp, np.equal(x, 1)) - temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - nan_indices = np.logical_or(temp, np.less(x, 1)) - n, res = 1, 1 / q**x - while n < 10000: - term = 1 / (q + n) ** x - n, res = n + 1, res + term - ret = np.round(res, decimals=4) - ret[nan_indices] = np.nan - ret[inf_indices] = np.inf - return ret - - -zeta.support_native_out = False - - -def gradient( - x: np.ndarray, - /, - *, - spacing: Union[int, list, tuple] = 1, - axis: Optional[Union[int, list, tuple]] = None, - edge_order: int = 1, -) -> Union[np.ndarray, List[np.ndarray]]: - if type(spacing) in (int, float): - return np.gradient(x, spacing, axis=axis, edge_order=edge_order) - return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) - - -def xlogy( - x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - x, y = promote_types_of_inputs(x, y) - if (x == 0).all(): - return 0.0 - else: - return x * np.log(y) - - -def conj( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ret = np.conj(x, out=out) - if x.dtype == bool: - return ret.astype("bool") - return ret - - -def ldexp( - x1: np.ndarray, - x2: Union[np.ndarray, int, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.ldexp(x1, x2, out=out) - - -def frexp( - x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None -) -> Tuple[np.ndarray, np.ndarray]: - if out is None: - return np.frexp(x, out=(None, None)) - else: - return np.frexp(x, out=out) - - -def modf( - x: np.ndarray, - /, - *, - out: Optional[Tuple[np.ndarray, np.ndarray]] = None, -) -> np.ndarray: - if out: - return np.modf(x, out=out) - return np.modf(x) - - -# ---digamma---# -kLanczosGamma = 7 # aka g -kBaseLanczosCoeff = 0.99999999999980993227684700473478 -kLanczosCoefficients = np.array( - [ - 676.520368121885098567009190444019, - -1259.13921672240287047156078755283, - 771.3234287776530788486528258894, - -176.61502916214059906584551354, - 12.507343278686904814458936853, - -0.13857109526572011689554707, - 9.984369578019570859563e-6, - 1.50563273514931155834e-7, - ] -) + out: Optional[np.ndarray] = None, +) -> np.ndarray: + prepend = prepend if prepend is not None else np._NoValue + append = append if append is not None else np._NoValue + return np.diff(x, n=n, axis=axis, prepend=prepend, append=append) def digamma( @@ -364,64 +193,94 @@ def digamma( ) -# --- LGAMMA --- # -LANCZOS_N = 13 -lanczos_g = 6.024680040776729583740234375 -lanczos_num_coeffs = np.array( - [ - 23531376880.410759688572007674451636754734846804940, - 42919803642.649098768957899047001988850926355848959, - 35711959237.355668049440185451547166705960488635843, - 17921034426.037209699919755754458931112671403265390, - 6039542586.3520280050642916443072979210699388420708, - 1439720407.3117216736632230727949123939715485786772, - 248874557.86205415651146038641322942321632125127801, - 31426415.585400194380614231628318205362874684987640, - 2876370.6289353724412254090516208496135991145378768, - 186056.26539522349504029498971604569928220784236328, - 8071.6720023658162106380029022722506138218516325024, - 210.82427775157934587250973392071336271166969580291, - 2.5066282746310002701649081771338373386264310793408, - ] -) -lanczos_den_coeffs = np.array( - [ - 0.0, - 39916800.0, - 120543840.0, - 150917976.0, - 105258076.0, - 45995730.0, - 13339535.0, - 2637558.0, - 357423.0, - 32670.0, - 1925.0, - 66.0, - 1.0, - ] -) +def fix( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.fix(x, out=out) -def sinpi(x): - y = np.abs(x) % 2.0 - n = np.round(2.0 * y) - assert 0 <= n and n <= 4 +@_scalar_output_to_0d_array +def float_power( + x1: Union[np.ndarray, float, list, tuple], + x2: Union[np.ndarray, float, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.float_power(x1, x2, out=out) - if n == 0: - r = np.sin(np.pi * y) - elif n == 1: - r = np.cos(np.pi * (y - 0.5)) - elif n == 2: - r = np.sin(np.pi * (1.0 - y)) - elif n == 3: - r = -np.cos(np.pi * (y - 1.5)) - elif n == 4: - r = np.sin(np.pi * (y - 2.0)) + +@_scalar_output_to_0d_array +def fmax( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmax( + x1, + x2, + out=None, + where=True, + casting="same_kind", + order="K", + dtype=None, + subok=True, + ) + + +def frexp( + x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None +) -> Tuple[np.ndarray, np.ndarray]: + if out is None: + return np.frexp(x, out=(None, None)) else: - raise Exception("Unreachable code") + return np.frexp(x, out=out) - return np.copysign(1.0, x) * r + +def gradient( + x: np.ndarray, + /, + *, + spacing: Union[int, list, tuple] = 1, + axis: Optional[Union[int, list, tuple]] = None, + edge_order: int = 1, +) -> Union[np.ndarray, List[np.ndarray]]: + if type(spacing) in (int, float): + return np.gradient(x, spacing, axis=axis, edge_order=edge_order) + return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) + + +def hypot( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.hypot(x1, x2) + + +def isclose( + a: np.ndarray, + b: np.ndarray, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + if np.isscalar(ret): + return np.array(ret, dtype="bool") + return ret def lanczos_sum(x): @@ -440,6 +299,16 @@ def lanczos_sum(x): return num / den +def ldexp( + x1: np.ndarray, + x2: Union[np.ndarray, int, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.ldexp(x1, x2, out=out) + + # TODO: Replace with native lgamma implementation when available def lgamma( x: np.ndarray, @@ -480,3 +349,108 @@ def func(x): # Vectorize 'func' for element-wise operations on 'x', output matching 'x' dtype. vfunc = np.vectorize(func, otypes=[x.dtype]) return vfunc(x) + + +def modf( + x: np.ndarray, + /, + *, + out: Optional[Tuple[np.ndarray, np.ndarray]] = None, +) -> np.ndarray: + if out: + return np.modf(x, out=out) + return np.modf(x) + + +def nansum( + x: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[np.dtype] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +def nextafter( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nextafter(x1, x2) + + +def signbit( + x: Union[np.ndarray, float, int, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.signbit(x, out=out) + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.sinc(x).astype(x.dtype) + + +def sinpi(x): + y = np.abs(x) % 2.0 + n = np.round(2.0 * y) + assert 0 <= n and n <= 4 + + if n == 0: + r = np.sin(np.pi * y) + elif n == 1: + r = np.cos(np.pi * (y - 0.5)) + elif n == 2: + r = np.sin(np.pi * (1.0 - y)) + elif n == 3: + r = -np.cos(np.pi * (y - 1.5)) + elif n == 4: + r = np.sin(np.pi * (y - 2.0)) + else: + raise Exception("Unreachable code") + + return np.copysign(1.0, x) * r + + +def xlogy( + x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + x, y = promote_types_of_inputs(x, y) + if (x == 0).all(): + return 0.0 + else: + return x * np.log(y) + + +def zeta( + x: np.ndarray, + q: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) + inf_indices = np.logical_or(temp, np.equal(x, 1)) + temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + nan_indices = np.logical_or(temp, np.less(x, 1)) + n, res = 1, 1 / q**x + while n < 10000: + term = 1 / (q + n) ** x + n, res = n + 1, res + term + ret = np.round(res, decimals=4) + ret[nan_indices] = np.nan + ret[inf_indices] = np.inf + return ret diff --git a/ivy/functional/backends/numpy/experimental/layers.py b/ivy/functional/backends/numpy/experimental/layers.py index 757e434c7b834..4360cf4099908 100644 --- a/ivy/functional/backends/numpy/experimental/layers.py +++ b/ivy/functional/backends/numpy/experimental/layers.py @@ -19,6 +19,10 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): kernel, strides, depth_pooling = _depth_max_pooling_helper( x.shape, kernel, strides, dims=dims, data_format=data_format @@ -28,82 +32,77 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def max_pool1d( +def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): + if isinstance(padding, str): + pad_specific = [ + _handle_padding(x_shape[i], strides[i], kernel[i], padding) + for i in range(dim) + ] + padding = [ + (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) + for i in range(dim) + ] + else: + pad_specific = [sum(padding[i]) for i in range(dim)] + + c = [] + if ceil_mode: + for i in range(dim): + padding[i], c_i = _padding_ceil_mode( + x_shape[i], kernel[i], padding[i], strides[i], True + ) + c.append(c_i) + pad_specific[i] = sum(padding[i]) + return padding, pad_specific, c + + +# --- Main --- # +# ------------ # + + +def avg_pool1d( x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, /, *, data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, + count_include_pad: bool = False, ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + if isinstance(kernel, int): + kernel = [kernel] + elif len(kernel) == 1: + kernel = [kernel[0]] - if data_format == "NCW": + if isinstance(strides, int): + strides = [strides] + elif len(strides) == 1: + strides = [strides[0]] + + if data_format in ("NCW", "NCL"): x = np.swapaxes(x, 1, 2) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" + x_shape = x.shape[1:-1] + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 1 ) - x_shape = x.shape[1:2] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - if dilation[0] > 1: - filters = _add_dilations(filters, dilation[0], axis=0, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_list = [ - (pad_w // 2, pad_w - pad_w // 2), - ] - if ceil_mode: - pad_list[0] = _padding_ceil_mode( - x_shape[0], kernel[0], pad_list[0], strides[0] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) x_shape = x.shape new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] + new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] new_strides = ( x.strides[0], x.strides[1] * strides[0], @@ -111,101 +110,79 @@ def max_pool1d( x.strides[2], ) - # B x OW x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - # B x OW x KW x I - sub_matrices = np.where( - filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf - ) + res = np.mean(sub_matrices, axis=2) - res = sub_matrices.max(axis=(2)) + if (not count_include_pad or ceil_mode) and any(pad_specific): + if not count_include_pad: + num_padded_values = np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[0], + "n": x.shape[1] - pad_specific[0], + "k": kernel[0], + "s": strides[0], + }, + unique={ + "i": np.arange(res.shape[1]), + }, + ), + dtype=res.dtype, + ) + else: + num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) + num_padded_values[-1] = c[0] + res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) + + if data_format in ("NCW", "NCL"): + return res.swapaxes(1, 2) - if depth_pooling: - res = np.swapaxes(res, 1, 2) - if data_format == "NCW": - res = np.swapaxes(res, 1, 2) return res -def max_pool2d( +def avg_pool2d( x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, /, *, data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, + count_include_pad: bool = False, ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + if isinstance(kernel, int): + kernel = [kernel] * 2 + elif len(kernel) == 1: + kernel = [kernel[0]] * 2 + + if isinstance(strides, int): + strides = [strides] * 2 + elif len(strides) == 1: + strides = [strides[0]] * 2 if data_format == "NCHW": x = np.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) x_shape = list(x.shape[1:3]) - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_list = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 2 + ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) x_shape = x.shape new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 @@ -219,195 +196,90 @@ def max_pool2d( x.strides[2], x.strides[3], ) - # B x OH x OW x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - # B x OH x OW x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf - ) - # B x OH x OW x O - res = sub_matrices.max(axis=(3, 4)) - - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 1)) - if data_format == "NCHW": - return np.transpose(res, (0, 3, 1, 2)) - return res - - -def max_pool3d( - x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCDHW": - x = np.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - x_shape = x.shape[1:4] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) - pad_list = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override + else: + res = np.mean(sub_matrices, axis=(3, 4)) + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): + if not count_include_pad: + num_padded_values = [ + np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[i], + "n": x.shape[i + 1] - pad_specific[i], + "k": kernel[i], + "s": strides[i], + }, + unique={ + "i": np.arange(res.shape[i + 1]), + }, + ), + dtype=res.dtype, ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, + for i in range(2) + ] + else: + num_padded_values = [] + for i in range(2): + num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) + num_pad[-1] = c[i] + num_padded_values.append(num_pad) + num_padded_values1 = num_padded_values[0][:, None] + num_padded_values2 = num_padded_values[1][None, :] + num_padded_values = ( + num_padded_values1 * kernel[1] + + num_padded_values2 * kernel[0] + - num_padded_values1 * num_padded_values2 ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - x_shape = x.shape - new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 - new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[3] * strides[2], - x.strides[1], - x.strides[2], - x.strides[3], - x.strides[4], - ) - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf - ) - - # B x OD x OH x OW x O - res = sub_matrices.max(axis=(4, 5, 6)) + kernel_mul = np.prod(kernel) + res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 4, 1)) - if data_format == "NCDHW": - return np.transpose(res, (0, 4, 1, 2, 3)) + if data_format == "NCHW": + return np.transpose(res, (0, 3, 1, 2)) return res -def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): - if isinstance(padding, str): - pad_specific = [ - _handle_padding(x_shape[i], strides[i], kernel[i], padding) - for i in range(dim) - ] - padding = [ - (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) - for i in range(dim) - ] - else: - pad_specific = [sum(padding[i]) for i in range(dim)] - - c = [] - if ceil_mode: - for i in range(dim): - padding[i], c_i = _padding_ceil_mode( - x_shape[i], kernel[i], padding[i], strides[i], True - ) - c.append(c_i) - pad_specific[i] = sum(padding[i]) - return padding, pad_specific, c - - -def avg_pool1d( +def avg_pool3d( x: np.ndarray, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], padding: str, /, *, - data_format: str = "NWC", + data_format: str = "NDHWC", count_include_pad: bool = False, ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: if isinstance(kernel, int): - kernel = [kernel] + kernel = [kernel] * 3 elif len(kernel) == 1: - kernel = [kernel[0]] + kernel = [kernel[0]] * 3 if isinstance(strides, int): - strides = [strides] + strides = [strides] * 3 elif len(strides) == 1: - strides = [strides[0]] + strides = [strides[0]] * 3 - if data_format in ("NCW", "NCL"): - x = np.swapaxes(x, 1, 2) + if data_format == "NCDHW": + x = np.transpose(x, (0, 2, 3, 4, 1)) - x_shape = x.shape[1:-1] + x_shape = list(x.shape[1:4]) padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 1 + x_shape, kernel, strides, padding, ceil_mode, 3 ) x = np.pad( @@ -421,228 +293,36 @@ def avg_pool1d( ) x_shape = x.shape - new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] + new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 + new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] new_strides = ( x.strides[0], x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[3] * strides[2], x.strides[1], x.strides[2], + x.strides[3], + x.strides[4], ) - + # B x OH x OW x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - res = np.mean(sub_matrices, axis=2) + # B x OH x OW x O + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override + else: + res = np.mean(sub_matrices, axis=(4, 5, 6)) - if (not count_include_pad or ceil_mode) and any(pad_specific): - if not count_include_pad: - num_padded_values = np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[0], - "n": x.shape[1] - pad_specific[0], - "k": kernel[0], - "s": strides[0], - }, - unique={ - "i": np.arange(res.shape[1]), - }, - ), - dtype=res.dtype, - ) - else: - num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) - num_padded_values[-1] = c[0] - res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) - - if data_format in ("NCW", "NCL"): - return res.swapaxes(1, 2) - - return res - - -def avg_pool2d( - x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 2 - elif len(kernel) == 1: - kernel = [kernel[0]] * 2 - - if isinstance(strides, int): - strides = [strides] * 2 - elif len(strides) == 1: - strides = [strides[0]] * 2 - - if data_format == "NCHW": - x = np.transpose(x, (0, 2, 3, 1)) - - x_shape = list(x.shape[1:3]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 2 - ) - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) - - x_shape = x.shape - new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[1], - x.strides[2], - x.strides[3], - ) - # B x OH x OW x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(3, 4)) - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): - if not count_include_pad: - num_padded_values = [ - np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[i], - "n": x.shape[i + 1] - pad_specific[i], - "k": kernel[i], - "s": strides[i], - }, - unique={ - "i": np.arange(res.shape[i + 1]), - }, - ), - dtype=res.dtype, - ) - for i in range(2) - ] - else: - num_padded_values = [] - for i in range(2): - num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) - num_pad[-1] = c[i] - num_padded_values.append(num_pad) - num_padded_values1 = num_padded_values[0][:, None] - num_padded_values2 = num_padded_values[1][None, :] - num_padded_values = ( - num_padded_values1 * kernel[1] - + num_padded_values2 * kernel[0] - - num_padded_values1 * num_padded_values2 - ) - kernel_mul = np.prod(kernel) - res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) - - if data_format == "NCHW": - return np.transpose(res, (0, 3, 1, 2)) - return res - - -def avg_pool3d( - x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 3 - elif len(kernel) == 1: - kernel = [kernel[0]] * 3 - - if isinstance(strides, int): - strides = [strides] * 3 - elif len(strides) == 1: - strides = [strides[0]] * 3 - - if data_format == "NCDHW": - x = np.transpose(x, (0, 2, 3, 4, 1)) - - x_shape = list(x.shape[1:4]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 3 - ) - - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) - - x_shape = x.shape - new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 - new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[3] * strides[2], - x.strides[1], - x.strides[2], - x.strides[3], - x.strides[4], - ) - # B x OH x OW x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(4, 5, 6)) - - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): if not count_include_pad: num_padded_values = [ np.array( @@ -687,43 +367,6 @@ def avg_pool3d( return res -def fft( - x: np.ndarray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.dtype in [np.uint64, np.int64, np.float64, np.complex128]: - out_dtype = np.complex128 - else: - out_dtype = np.complex64 - return np.fft.fft(x, n, dim, norm).astype(out_dtype) - - @with_supported_dtypes({"1.25.2 and below": ("float32", "float64")}, backend_version) def dct( x: np.ndarray, @@ -820,20 +463,6 @@ def dct( return dct_out.astype(np.float32) if cast_final else dct_out -def idct( - x: np.ndarray, - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - def dropout1d( x: np.ndarray, prob: float, @@ -911,22 +540,48 @@ def dropout3d( return res -def ifft( - x: np.ndarray, - dim: int, +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def embedding( + weights: np.ndarray, + indices: np.ndarray, + /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + max_norm: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + + embeddings = np.take(weights, indices, axis=0) + if max_norm is not None: + norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = np.where( + norms > max_norm, embeddings * max_norm / norms, embeddings ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( + embeddings = np.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings + ) + return embeddings + + +def fft( + x: np.ndarray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( f"Invalid dim {dim}, expecting ranging" " from {-len(x.shape)} to {len(x.shape)-1} " ) @@ -940,7 +595,11 @@ def ifft( ) if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return np.asarray(np.fft.ifft(x, n, dim, norm), dtype=x.dtype) + if x.dtype in [np.uint64, np.int64, np.float64, np.complex128]: + out_dtype = np.complex128 + else: + out_dtype = np.complex64 + return np.fft.fft(x, n, dim, norm).astype(out_dtype) def fft2( @@ -978,6 +637,52 @@ def fft2( return np.fft.fft2(x, s, dim, norm).astype(np.complex128) +def idct( + x: np.ndarray, + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def ifft( + x: np.ndarray, + dim: int, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return np.asarray(np.fft.ifft(x, n, dim, norm), dtype=x.dtype) + + def ifftn( x: np.ndarray, s: Optional[Union[int, Tuple[int]]] = None, @@ -989,29 +694,332 @@ def ifftn( return np.fft.ifftn(x, s, axes, norm).astype(x.dtype) -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def embedding( - weights: np.ndarray, - indices: np.ndarray, +def max_pool1d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - max_norm: Optional[int] = None, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims ) - embeddings = np.take(weights, indices, axis=0) - if max_norm is not None: - norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = np.where( - norms > max_norm, embeddings * max_norm / norms, embeddings + if data_format == "NCW": + x = np.swapaxes(x, 1, 2) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides ) - embeddings = np.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - return embeddings + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = x.shape[1:2] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + if dilation[0] > 1: + filters = _add_dilations(filters, dilation[0], axis=0, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_list = [ + (pad_w // 2, pad_w - pad_w // 2), + ] + if ceil_mode: + pad_list[0] = _padding_ceil_mode( + x_shape[0], kernel[0], pad_list[0], strides[0] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[1], + x.strides[2], + ) + + # B x OW x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OW x KW x I + sub_matrices = np.where( + filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + res = sub_matrices.max(axis=(2)) + + if depth_pooling: + res = np.swapaxes(res, 1, 2) + if data_format == "NCW": + res = np.swapaxes(res, 1, 2) + return res + + +def max_pool2d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCHW": + x = np.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = list(x.shape[1:3]) + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_list = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[1], + x.strides[2], + x.strides[3], + ) + + # B x OH x OW x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OH x OW x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + # B x OH x OW x O + res = sub_matrices.max(axis=(3, 4)) + + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 1)) + if data_format == "NCHW": + return np.transpose(res, (0, 3, 1, 2)) + return res + + +def max_pool3d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCDHW": + x = np.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = x.shape[1:4] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) + pad_list = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 + new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[3] * strides[2], + x.strides[1], + x.strides[2], + x.strides[3], + x.strides[4], + ) + # B x OD x OH x OW x KD x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OD x OH x OW x KD x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + # B x OD x OH x OW x O + res = sub_matrices.max(axis=(4, 5, 6)) + + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 4, 1)) + if data_format == "NCDHW": + return np.transpose(res, (0, 4, 1, 2, 3)) + return res def rfftn( diff --git a/ivy/functional/backends/numpy/experimental/linear_algebra.py b/ivy/functional/backends/numpy/experimental/linear_algebra.py index 915b0b4ad9111..54029e3dd0321 100644 --- a/ivy/functional/backends/numpy/experimental/linear_algebra.py +++ b/ivy/functional/backends/numpy/experimental/linear_algebra.py @@ -10,6 +10,36 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size +kron.support_native_out = False +eig.support_native_out = False +eigvals.support_native_out = False +multi_dot.support_native_out = True +cond.support_native_out = False +dot.support_native_out = True + + +def adjoint( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return np.conjugate(np.transpose(x, axes=axes)) + + +def cond( + x: np.ndarray, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[np.ndarray] = None, +) -> Any: + return np.linalg.cond(x, p=p) + + def diagflat( x: np.ndarray, /, @@ -84,34 +114,14 @@ def diagflat( return ret -def kron( +def dot( a: np.ndarray, b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.kron(a, b) - - -kron.support_native_out = False - - -@with_supported_dtypes( - {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, - backend_version, -) -def matrix_exp( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - eig_vals, eig_vecs = np.linalg.eig(x) - exp_diag = np.exp(eig_vals) - exp_diag_mat = np.diag(exp_diag) - exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) - return exp_mat.astype(x.dtype) + return np.dot(a, b, out=out) def eig( @@ -126,9 +136,6 @@ def eig( return e.astype(complex), v.astype(complex) -eig.support_native_out = False - - def eigvals(x: np.ndarray, /) -> np.ndarray: if ivy.dtype(x) == ivy.float16: x = x.astype(np.float32) @@ -136,44 +143,14 @@ def eigvals(x: np.ndarray, /) -> np.ndarray: return e.astype(complex) -eigvals.support_native_out = False - - -def adjoint( - x: np.ndarray, +def kron( + a: np.ndarray, + b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return np.conjugate(np.transpose(x, axes=axes)) - - -def multi_dot( - x: Sequence[np.ndarray], - /, - *, - out: Optional[np.array] = None, -) -> np.ndarray: - return np.linalg.multi_dot(x, out=out) - - -multi_dot.support_native_out = True - - -def cond( - x: np.ndarray, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[np.ndarray] = None, -) -> Any: - return np.linalg.cond(x, p=p) - - -cond.support_native_out = False + return np.kron(a, b) def lu_factor( @@ -186,14 +163,27 @@ def lu_factor( raise IvyNotImplementedException() -def dot( - a: np.ndarray, - b: np.ndarray, +@with_supported_dtypes( + {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, + backend_version, +) +def matrix_exp( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.dot(a, b, out=out) + eig_vals, eig_vecs = np.linalg.eig(x) + exp_diag = np.exp(eig_vals) + exp_diag_mat = np.diag(exp_diag) + exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) + return exp_mat.astype(x.dtype) -dot.support_native_out = True +def multi_dot( + x: Sequence[np.ndarray], + /, + *, + out: Optional[np.array] = None, +) -> np.ndarray: + return np.linalg.multi_dot(x, out=out) diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 53436987e2187..0c03925ae8aa2 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -19,103 +19,152 @@ from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array -def moveaxis( - a: np.ndarray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.moveaxis(a, source, destination) +# --- Helpers --- # +# --------------- # -moveaxis.support_native_out = False +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x -def heaviside( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.heaviside( - x1, - x2, - out=out, - ) +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = np.full(new_shape, padding_value, dtype=operand.dtype) + src_indices = np.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) -heaviside.support_native_out = True + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = np.pad(padded, pad_width, constant_values=padding_value) + return padded -def flipud( - m: np.ndarray, - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.flipud(m) +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] -flipud.support_native_out = False +# --- Main --- # +# ------------ # -def vstack( - arrays: Sequence[np.ndarray], + +def atleast_1d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_1d(*arys) + + +def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: + return np.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_3d(*arys) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return np.broadcast_shapes(*shapes) + + +def concat_from_sequence( + input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.vstack(arrays) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = np.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = np.stack(input_sequence, axis=axis) + return ret -def hstack( +def dsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def dstack( arrays: Sequence[np.ndarray], /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.hstack(arrays) + return np.dstack(arrays) -def rot90( - m: np.ndarray, +def expand( + x: np.ndarray, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.rot90(m, k, axes) + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) + return np.broadcast_to(x, tuple(shape)) -def top_k( - x: np.ndarray, - k: int, +def fill_diagonal( + a: np.ndarray, + v: Union[int, float], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[np.ndarray, np.ndarray]] = None, -) -> Tuple[np.ndarray, np.ndarray]: - k = min(k, x.shape[axis]) - if not largest: - indices = np.argsort(x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - else: - indices = np.argsort(-x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - if not sorted: - indices = np.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) - val = np.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + wrap: bool = False, +) -> np.ndarray: + np.fill_diagonal(a, v, wrap=wrap) + return a def fliplr( @@ -128,71 +177,70 @@ def fliplr( return np.fliplr(m) -fliplr.support_native_out = False - - -def i0( - x: np.ndarray, +def flipud( + m: np.ndarray, /, *, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.i0(x) - + return np.flipud(m) -i0.support_native_out = False +def heaviside( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.heaviside( + x1, + x2, + out=out, + ) -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x +def hsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] +def hstack( + arrays: Sequence[np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.hstack(arrays) -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = np.full(new_shape, padding_value, dtype=operand.dtype) - src_indices = np.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array +def i0( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.i0(x) - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = np.pad(padded, pad_width, constant_values=padding_value) - return padded +def moveaxis( + a: np.ndarray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.moveaxis(a, source, destination) def pad( @@ -268,57 +316,16 @@ def pad( ) -def vsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Sequence[int], np.ndarray], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def dsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], +def rot90( + m: np.ndarray, /, *, copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_1d(*arys) - - -def dstack( - arrays: Sequence[np.ndarray], - /, - *, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.dstack(arrays) - - -def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: - return np.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_3d(*arys) + return np.rot90(m, k, axes) @_scalar_output_to_0d_array @@ -366,63 +373,28 @@ def take_along_axis( return np.take_along_axis(arr, indices, axis) -def hsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -take_along_axis.support_native_out = False - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return np.broadcast_shapes(*shapes) - - -broadcast_shapes.support_native_out = False - - -def expand( +def top_k( x: np.ndarray, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) - return np.broadcast_to(x, tuple(shape)) - - -expand.support_native_out = False - - -def concat_from_sequence( - input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = np.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = np.stack(input_sequence, axis=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[np.ndarray, np.ndarray]] = None, +) -> Tuple[np.ndarray, np.ndarray]: + k = min(k, x.shape[axis]) + if not largest: + indices = np.argsort(x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + else: + indices = np.argsort(-x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + if not sorted: + indices = np.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) + val = np.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) def unique_consecutive( @@ -468,12 +440,34 @@ def unique_consecutive( ) -def fill_diagonal( - a: np.ndarray, - v: Union[int, float], +def vsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Sequence[int], np.ndarray], /, *, - wrap: bool = False, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Sequence[np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - np.fill_diagonal(a, v, wrap=wrap) - return a + return np.vstack(arrays) + + +moveaxis.support_native_out = False +heaviside.support_native_out = True +flipud.support_native_out = False +fliplr.support_native_out = False +i0.support_native_out = False +take_along_axis.support_native_out = False +broadcast_shapes.support_native_out = False +expand.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/random.py b/ivy/functional/backends/numpy/experimental/random.py index 16293d4eac407..e2af3988496df 100644 --- a/ivy/functional/backends/numpy/experimental/random.py +++ b/ivy/functional/backends/numpy/experimental/random.py @@ -10,24 +10,26 @@ ) -# dirichlet -def dirichlet( - alpha: Union[np.ndarray, float, Sequence[float]], - /, +dirichlet.support_native_out = False + + +def bernoulli( + probs: Union[float, np.ndarray], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + logits: Optional[Union[float, np.ndarray]] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, dtype: Optional[np.dtype] = None, seed: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - size = size if size is not None else len(alpha) - dtype = dtype if dtype is not None else np.float64 if seed is not None: np.random.seed(seed) - return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) - - -dirichlet.support_native_out = False + if logits is not None: + probs = np.asarray(ivy.softmax(logits), dtype=dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) def beta( @@ -47,6 +49,23 @@ def beta( return np.asarray(np.random.beta(alpha, beta, shape), dtype=dtype) +# dirichlet +def dirichlet( + alpha: Union[np.ndarray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[np.dtype] = None, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + size = size if size is not None else len(alpha) + dtype = dtype if dtype is not None else np.float64 + if seed is not None: + np.random.seed(seed) + return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) + + def gamma( alpha: Union[float, np.ndarray], beta: Union[float, np.ndarray], @@ -87,22 +106,3 @@ def poisson( else: ret = np.random.poisson(lam, shape) return np.asarray(ret, dtype=dtype) - - -def bernoulli( - probs: Union[float, np.ndarray], - *, - logits: Optional[Union[float, np.ndarray]] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, - dtype: Optional[np.dtype] = None, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if seed is not None: - np.random.seed(seed) - if logits is not None: - probs = np.asarray(ivy.softmax(logits), dtype=dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) diff --git a/ivy/functional/backends/numpy/experimental/searching.py b/ivy/functional/backends/numpy/experimental/searching.py index 23811e75b4077..297c9cdba2db2 100644 --- a/ivy/functional/backends/numpy/experimental/searching.py +++ b/ivy/functional/backends/numpy/experimental/searching.py @@ -7,6 +7,9 @@ from . import backend_version +unravel_index.support_native_out = False + + @with_supported_dtypes({"1.25.2 and below": ("int32", "int64")}, backend_version) def unravel_index( indices: np.ndarray, @@ -17,6 +20,3 @@ def unravel_index( ) -> Tuple[np.ndarray]: ret = np.asarray(np.unravel_index(indices, shape), dtype=np.int32) return tuple(ret) - - -unravel_index.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/sorting.py b/ivy/functional/backends/numpy/experimental/sorting.py index cd734f8e69835..fd04f1f400214 100644 --- a/ivy/functional/backends/numpy/experimental/sorting.py +++ b/ivy/functional/backends/numpy/experimental/sorting.py @@ -3,6 +3,9 @@ from typing import Optional, Union +lexsort.support_native_out = False + + # invert_permutation def invert_permutation( x: Union[np.ndarray, list, tuple], @@ -20,6 +23,3 @@ def lexsort( keys: np.ndarray, /, *, axis: int = -1, out: Optional[np.ndarray] = None ) -> np.ndarray: return np.asarray(np.lexsort(keys, axis=axis)) - - -lexsort.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/statistical.py b/ivy/functional/backends/numpy/experimental/statistical.py index 12c3565148ec7..6efebf5a93ba6 100644 --- a/ivy/functional/backends/numpy/experimental/statistical.py +++ b/ivy/functional/backends/numpy/experimental/statistical.py @@ -8,194 +8,87 @@ from copy import deepcopy -@with_unsupported_dtypes( - {"1.25.2 and below": ("bfloat16",)}, - backend_version, -) -def histogram( - a: np.ndarray, - /, - *, - bins: Optional[Union[int, np.ndarray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[np.dtype] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[np.ndarray] = None, - density: Optional[bool] = False, - out: Optional[np.ndarray] = None, -) -> Tuple[np.ndarray]: - min_a = np.min(a) - max_a = np.max(a) - if isinstance(bins, np.ndarray) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - elif isinstance(bins, int): - range = (min_a, max_a) - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - if bins.size < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - bins_out = bins.copy() - if extend_lower_interval and min_a < bins[0]: - bins[0] = min_a - if extend_upper_interval and max_a > bins[-1]: - bins[-1] = max_a - if a.ndim > 0 and axis is not None: - inverted_shape_dims = list(np.flip(np.arange(a.ndim))) - if isinstance(axis, int): - axis = [axis] - shape_axes = 1 - for dimension in axis: - inverted_shape_dims.remove(dimension) - inverted_shape_dims.append(dimension) - shape_axes *= a.shape[dimension] - a_along_axis_1d = ( - a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) - ) - if weights is None: - ret = [] - for a_1d in a_along_axis_1d: - ret_1d = np.histogram( - a_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - else: - weights_along_axis_1d = ( - weights.transpose(inverted_shape_dims) - .flatten() - .reshape((-1, shape_axes)) - ) - ret = [] - for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): - ret_1d = np.histogram( - a_1d, - weights=weights_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - out_shape = list(a.shape) - for dimension in sorted(axis, reverse=True): - del out_shape[dimension] - out_shape.insert(0, len(bins) - 1) - ret = np.array(ret) - ret = ret.flatten() - index = np.zeros(len(out_shape), dtype=int) - ret_shaped = np.zeros(out_shape) - dim = 0 - i = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - while list(index) != list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - dim_full_flag = False - while index[dim] == out_shape[dim] - 1: - index[dim] = 0 - dim += 1 - dim_full_flag = True - index[dim] += 1 - i += 1 - if dim_full_flag: - dim = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - ret = ret_shaped - else: - ret = np.histogram( - a=a, bins=bins, range=range, weights=weights, density=density - )[0] - if dtype: - ret = ret.astype(dtype) - bins_out = np.array(bins_out).astype(dtype) - # TODO: weird error when returning bins: return ret, bins_out - return ret +# --- Helpers --- # +# --------------- # -def median( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if out is not None: - out = np.reshape(out, input.shape) - ret = np.median( - input, - axis=axis, - keepdims=keepdims, - out=out, - ) - if input.dtype in [np.uint64, np.int64, np.float64]: - return ret.astype(np.float64) - elif input.dtype in [np.float16]: - return ret.astype(input.dtype) - else: - return ret.astype(np.float32) - - -median.support_native_out = True - - -def nanmean( - a: np.ndarray, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, +def __find_cummax_indices( + x: np.ndarray, + axis: int = 0, ) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) + indices = [] + if x[0] is np.ndarray: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + else: + indice_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices.fill(0) + indice_list = sorted(indice_list, key=lambda i: i[1]) + for y, y_index in indice_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + else: + indices[y_index] = n1[tuple(multi_index[1:])] + else: + n = 0 + for index1, ret1 in enumerate(x): + if x[n] <= ret1 or index1 == 0: + n = index1 + indices.append(n) + return np.array(indices, dtype=np.int64) -nanmean.support_native_out = True +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] -def _validate_quantile(q): - if isinstance(q, float): - q = np.asarray(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) else: - if not (np.all(0 <= q) and np.all(q <= 1)): - return False - return True - - -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] + indices.append((lst, tuple(prefix))) + return indices - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + axis = tuple(axis) if isinstance(axis, list) else axis - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis + return np.quantile( + x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ).astype(x.dtype) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) def _handle_axis(a, q, fn, keepdims=False, axis=None): @@ -261,89 +154,39 @@ def _quantile(a, q, axis=None): return out.astype(ret_dtype) -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - axis = tuple(axis) if isinstance(axis, list) else axis - - return np.quantile( - x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out - ).astype(x.dtype) - else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - - -def quantile( - a: np.ndarray, - q: Union[float, np.ndarray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - interpolation: str = "linear", - out: Optional[np.ndarray] = None, -) -> np.ndarray: - # quantile method in numpy backend, always return an array with dtype=float64. - # in other backends, the output is the same dtype as the input. - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - out=out, - ) +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + if len(axis) == 0: + raise ValueError("Axis can't be empty!") -def corrcoef( - x: np.ndarray, - /, - *, - y: Optional[np.ndarray] = None, - rowvar: bool = True, - dtype: np.dtype = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dtype = dtype if dtype is not None else np.float64 + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") - return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis -@with_unsupported_dtypes( - {"1.25.0 and below": ("bfloat16",)}, - backend_version, -) -def nanmedian( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out - ) +def _validate_quantile(q): + if isinstance(q, float): + q = np.asarray(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (np.all(0 <= q) and np.all(q <= 1)): + return False + return True -nanmedian.support_native_out = True +# --- Main --- # +# ------------ # def bincount( @@ -363,7 +206,18 @@ def bincount( return ret -bincount.support_native_out = False +def corrcoef( + x: np.ndarray, + /, + *, + y: Optional[np.ndarray] = None, + rowvar: bool = True, + dtype: np.dtype = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dtype = dtype if dtype is not None else np.float64 + + return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) def cov( @@ -390,9 +244,6 @@ def cov( ) -cov.support_native_out = False - - def cummax( x: np.ndarray, /, @@ -438,58 +289,6 @@ def cummax( return np.maximum.accumulate(x, axis=axis, dtype=x.dtype), indices -def __find_cummax_indices( - x: np.ndarray, - axis: int = 0, -) -> np.ndarray: - indices = [] - if x[0] is np.ndarray: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - - else: - indice_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices.fill(0) - indice_list = sorted(indice_list, key=lambda i: i[1]) - for y, y_index in indice_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - else: - indices[y_index] = n1[tuple(multi_index[1:])] - else: - n = 0 - for index1, ret1 in enumerate(x): - if x[n] <= ret1 or index1 == 0: - n = index1 - indices.append(n) - return np.array(indices, dtype=np.int64) - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) def cummin( x: np.ndarray, @@ -513,6 +312,121 @@ def cummin( return np.flip(x, axis=axis) +@with_unsupported_dtypes( + {"1.25.2 and below": ("bfloat16",)}, + backend_version, +) +def histogram( + a: np.ndarray, + /, + *, + bins: Optional[Union[int, np.ndarray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[np.dtype] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[np.ndarray] = None, + density: Optional[bool] = False, + out: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray]: + min_a = np.min(a) + max_a = np.max(a) + if isinstance(bins, np.ndarray) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + elif isinstance(bins, int): + range = (min_a, max_a) + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + if bins.size < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + bins_out = bins.copy() + if extend_lower_interval and min_a < bins[0]: + bins[0] = min_a + if extend_upper_interval and max_a > bins[-1]: + bins[-1] = max_a + if a.ndim > 0 and axis is not None: + inverted_shape_dims = list(np.flip(np.arange(a.ndim))) + if isinstance(axis, int): + axis = [axis] + shape_axes = 1 + for dimension in axis: + inverted_shape_dims.remove(dimension) + inverted_shape_dims.append(dimension) + shape_axes *= a.shape[dimension] + a_along_axis_1d = ( + a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) + ) + if weights is None: + ret = [] + for a_1d in a_along_axis_1d: + ret_1d = np.histogram( + a_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + else: + weights_along_axis_1d = ( + weights.transpose(inverted_shape_dims) + .flatten() + .reshape((-1, shape_axes)) + ) + ret = [] + for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): + ret_1d = np.histogram( + a_1d, + weights=weights_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + out_shape = list(a.shape) + for dimension in sorted(axis, reverse=True): + del out_shape[dimension] + out_shape.insert(0, len(bins) - 1) + ret = np.array(ret) + ret = ret.flatten() + index = np.zeros(len(out_shape), dtype=int) + ret_shaped = np.zeros(out_shape) + dim = 0 + i = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + while list(index) != list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + dim_full_flag = False + while index[dim] == out_shape[dim] - 1: + index[dim] = 0 + dim += 1 + dim_full_flag = True + index[dim] += 1 + i += 1 + if dim_full_flag: + dim = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + ret = ret_shaped + else: + ret = np.histogram( + a=a, bins=bins, range=range, weights=weights, density=density + )[0] + if dtype: + ret = ret.astype(dtype) + bins_out = np.array(bins_out).astype(dtype) + # TODO: weird error when returning bins: return ret, bins_out + return ret + + def igamma( a: np.ndarray, /, @@ -528,3 +442,89 @@ def igamma_cal(a, x): igamma_vec = np.vectorize(igamma_cal) return igamma_vec(a, x).astype(a.dtype) + + +def median( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if out is not None: + out = np.reshape(out, input.shape) + ret = np.median( + input, + axis=axis, + keepdims=keepdims, + out=out, + ) + if input.dtype in [np.uint64, np.int64, np.float64]: + return ret.astype(np.float64) + elif input.dtype in [np.float16]: + return ret.astype(input.dtype) + else: + return ret.astype(np.float32) + + +def nanmean( + a: np.ndarray, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) + + +@with_unsupported_dtypes( + {"1.25.0 and below": ("bfloat16",)}, + backend_version, +) +def nanmedian( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + ) + + +def quantile( + a: np.ndarray, + q: Union[float, np.ndarray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + interpolation: str = "linear", + out: Optional[np.ndarray] = None, +) -> np.ndarray: + # quantile method in numpy backend, always return an array with dtype=float64. + # in other backends, the output is the same dtype as the input. + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + out=out, + ) + + +median.support_native_out = True +nanmean.support_native_out = True +nanmedian.support_native_out = True +bincount.support_native_out = False +cov.support_native_out = False diff --git a/ivy/functional/backends/numpy/general.py b/ivy/functional/backends/numpy/general.py index e4c5a8d75ccac..f7d07be482b5d 100644 --- a/ivy/functional/backends/numpy/general.py +++ b/ivy/functional/backends/numpy/general.py @@ -17,6 +17,11 @@ from ...ivy.general import _broadcast_to +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True +isin.support_native_out = True + + def array_equal(x0: np.ndarray, x1: np.ndarray, /) -> bool: return np.array_equal(x0, x1) @@ -29,49 +34,6 @@ def current_backend_str() -> str: return "numpy" -@_scalar_output_to_0d_array -def get_item( - x: np.ndarray, - /, - query: Union[np.ndarray, Tuple], - *, - copy: bool = None, -) -> np.ndarray: - return x.__getitem__(query) - - -@_scalar_output_to_0d_array -def set_item( - x: np.ndarray, - query: Union[np.ndarray, Tuple], - val: np.ndarray, - /, - *, - copy: Optional[bool] = False, -) -> np.ndarray: - if copy: - x = np.copy(x) - x.__setitem__(query, val) - return x - - -def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return x.copy() - else: - return x - - -def to_scalar(x: np.ndarray, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - -def to_list(x: np.ndarray, /) -> list: - return x.tolist() - - def gather( params: np.ndarray, indices: np.ndarray, @@ -104,6 +66,36 @@ def gather( return _to_device(result) +def gather_nd( + params: np.ndarray, + indices: np.ndarray, + /, + *, + batch_dims: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, np.asarray(i, indices.dtype)) + result.append(r) + result = np.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return _to_device(result) + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -135,34 +127,15 @@ def gather_nd_helper(params, indices): return res -def gather_nd( - params: np.ndarray, - indices: np.ndarray, +@_scalar_output_to_0d_array +def get_item( + x: np.ndarray, /, + query: Union[np.ndarray, Tuple], *, - batch_dims: int = 0, - out: Optional[np.ndarray] = None, + copy: bool = None, ) -> np.ndarray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, np.asarray(i, indices.dtype)) - result.append(r) - result = np.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return _to_device(result) + return x.__getitem__(query) def get_num_dims(x, /, *, as_array=False): @@ -244,6 +217,27 @@ def is_native_array(x, /, *, exclusive=False): return False +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def isin( + elements: np.ndarray, + test_elements: np.ndarray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> np.ndarray: + return np.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + +def itemsize(x: np.ndarray) -> int: + return x.itemsize + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -288,9 +282,6 @@ def scatter_flat( return target -scatter_flat.support_native_out = True - - def scatter_nd( indices: np.ndarray, updates: np.ndarray, @@ -332,7 +323,19 @@ def scatter_nd( return _to_device(target) -scatter_nd.support_native_out = True +@_scalar_output_to_0d_array +def set_item( + x: np.ndarray, + query: Union[np.ndarray, Tuple], + val: np.ndarray, + /, + *, + copy: Optional[bool] = False, +) -> np.ndarray: + if copy: + x = np.copy(x) + x.__setitem__(query, val) + return x def shape( @@ -347,6 +350,23 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: np.ndarray, /) -> list: + return x.tolist() + + +def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return x.copy() + else: + return x + + +def to_scalar(x: np.ndarray, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -431,27 +451,3 @@ def _vmap(*args): return res return _vmap - - -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def isin( - elements: np.ndarray, - test_elements: np.ndarray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> np.ndarray: - return np.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - -isin.support_native_out = True - - -def itemsize(x: np.ndarray) -> int: - return x.itemsize diff --git a/ivy/functional/backends/numpy/gradients.py b/ivy/functional/backends/numpy/gradients.py index 87be747e053b8..73621665e3d54 100644 --- a/ivy/functional/backends/numpy/gradients.py +++ b/ivy/functional/backends/numpy/gradients.py @@ -6,24 +6,6 @@ import ivy -def variable(x, /): - logging.warning( - "NumPy does not support autograd, declaring a 'variable' " - "is identical to declaring an 'array' when using numpy backend." - ) - return x - - -def is_variable(x, /, *, exclusive=False): - # NumPy does not support autograd, checking if x is a variable does have any meaning - # for NumPy. Return False. - return False - - -def variable_data(x, /): - return x - - def execute_with_gradients( func, xs, @@ -42,23 +24,29 @@ def execute_with_gradients( return func_ret, None -def value_and_grad(func): +def grad(func, argnums=0): logging.warning( - "NumPy does not support autograd, 'value_and_grad' " + "NumPy does not support autograd, 'grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grads = ivy.nested_map( + grad = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return y, grads + return grad return grad_fn +def is_variable(x, /, *, exclusive=False): + # NumPy does not support autograd, checking if x is a variable does have any meaning + # for NumPy. Return False. + return False + + def jac(func): logging.warning( "NumPy does not support autograd, 'jac' " @@ -74,26 +62,38 @@ def grad_fn(xs): return grad_fn -def grad(func, argnums=0): +def stop_gradient(x, /, *, preserve_type=True, out=None): logging.warning( - "NumPy does not support autograd, 'grad' " + "NumPy does not support autograd, 'stop_gradient' " + "has no effect on the array, as gradients are not supported in the first place." + ) + return x + + +def value_and_grad(func): + logging.warning( + "NumPy does not support autograd, 'value_and_grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grad = ivy.nested_map( + grads = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return grad + return y, grads return grad_fn -def stop_gradient(x, /, *, preserve_type=True, out=None): +def variable(x, /): logging.warning( - "NumPy does not support autograd, 'stop_gradient' " - "has no effect on the array, as gradients are not supported in the first place." + "NumPy does not support autograd, declaring a 'variable' " + "is identical to declaring an 'array' when using numpy backend." ) return x + + +def variable_data(x, /): + return x diff --git a/ivy/functional/backends/numpy/helpers.py b/ivy/functional/backends/numpy/helpers.py index b5ae02d3a09f0..110cbebe2a53c 100644 --- a/ivy/functional/backends/numpy/helpers.py +++ b/ivy/functional/backends/numpy/helpers.py @@ -1,23 +1,27 @@ -import functools -from typing import Callable -import numpy as np - - -def _scalar_output_to_0d_array(function: Callable) -> Callable: - """ - Convert scalar outputs to 0d arrays. - - Sometimes NumPy functions return scalars e.g. `np.add` does when the - inputs are both 0 dimensional. - - We use this wrapper to handle such cases, and convert scalar outputs - to 0d arrays, since the array API standard dictates outputs must be - arrays. - """ - - @functools.wraps(function) - def new_function(*args, **kwargs): - ret = function(*args, **kwargs) - return np.asarray(ret) if np.isscalar(ret) else ret - - return new_function +import functools +from typing import Callable +import numpy as np + + +# --- Helpers --- # +# --------------- # + + +def _scalar_output_to_0d_array(function: Callable) -> Callable: + """ + Convert scalar outputs to 0d arrays. + + Sometimes NumPy functions return scalars e.g. `np.add` does when the + inputs are both 0 dimensional. + + We use this wrapper to handle such cases, and convert scalar outputs + to 0d arrays, since the array API standard dictates outputs must be + arrays. + """ + + @functools.wraps(function) + def new_function(*args, **kwargs): + ret = function(*args, **kwargs) + return np.asarray(ret) if np.isscalar(ret) else ret + + return new_function diff --git a/ivy/functional/backends/numpy/layers.py b/ivy/functional/backends/numpy/layers.py index 52ad223304735..ca278076265d0 100644 --- a/ivy/functional/backends/numpy/layers.py +++ b/ivy/functional/backends/numpy/layers.py @@ -14,6 +14,10 @@ ) +# --- Helpers --- # +# --------------- # + + def _add_dilations(x, dilations, axis, values=0): return np.insert( x, @@ -113,6 +117,10 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters +# --- Main --- # +# ------------ # + + def conv1d( x: np.ndarray, filters: np.ndarray, @@ -280,62 +288,6 @@ def conv2d_transpose( return res -def depthwise_conv2d( - x: np.ndarray, - filters: np.ndarray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[np.ndarray] = None, -): - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - if data_format == "NHWC": - x = np.transpose(x, (3, 0, 1, 2)) - else: - x = np.transpose(x, (1, 0, 2, 3)) - filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters - filters = np.transpose(filters, (2, 0, 1)) - filters = np.expand_dims(filters, (-1, -2)) - filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) - filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) - if isinstance(padding, str): - if padding == "VALID": - out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) - out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) - else: - out_height = np.ceil(float(x.shape[2]) / float(strides[0])) - out_width = np.ceil(float(x.shape[3]) / float(strides[1])) - else: - out_height = np.ceil( - float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) - / float(strides[0]) - ) - out_width = np.ceil( - float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) - / float(strides[1]) - ) - if data_format == "NHWC": - outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) - else: - outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) - x = np.expand_dims(x, -1) - for i in range(x.shape[0]): - output = conv2d( - x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations - ) - if data_format == "NHWC": - outputs = np.append(outputs, output, axis=-1) - else: - outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) - return outputs - - def conv3d( x: np.ndarray, filters: np.ndarray, @@ -550,3 +502,59 @@ def conv_general_transpose( if data_format == "channel_first": return np.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +def depthwise_conv2d( + x: np.ndarray, + filters: np.ndarray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[np.ndarray] = None, +): + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + if data_format == "NHWC": + x = np.transpose(x, (3, 0, 1, 2)) + else: + x = np.transpose(x, (1, 0, 2, 3)) + filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters + filters = np.transpose(filters, (2, 0, 1)) + filters = np.expand_dims(filters, (-1, -2)) + filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) + filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) + if isinstance(padding, str): + if padding == "VALID": + out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) + out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) + else: + out_height = np.ceil(float(x.shape[2]) / float(strides[0])) + out_width = np.ceil(float(x.shape[3]) / float(strides[1])) + else: + out_height = np.ceil( + float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) + / float(strides[0]) + ) + out_width = np.ceil( + float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) + / float(strides[1]) + ) + if data_format == "NHWC": + outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) + else: + outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) + x = np.expand_dims(x, -1) + for i in range(x.shape[0]): + output = conv2d( + x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations + ) + if data_format == "NHWC": + outputs = np.append(outputs, output, axis=-1) + else: + outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) + return outputs diff --git a/ivy/functional/backends/numpy/linear_algebra.py b/ivy/functional/backends/numpy/linear_algebra.py index 9f9fcdbc17f1b..8f87923cbdbfc 100644 --- a/ivy/functional/backends/numpy/linear_algebra.py +++ b/ivy/functional/backends/numpy/linear_algebra.py @@ -15,6 +15,12 @@ from . import backend_version +matmul.support_native_out = True +outer.support_native_out = True +trace.support_native_out = True +vector_to_skew_symmetric_matrix.support_native_out = True + + # Array API Standard # # -------------------# @@ -52,6 +58,20 @@ def det(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.linalg.det(x) +# Extra # +# ----- # + + +def diag( + x: np.ndarray, + /, + *, + k: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.diag(x, k=k) + + def diagonal( x: np.ndarray, /, @@ -64,6 +84,15 @@ def diagonal( return np.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) +@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) +def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] + ) + eigenvalues, eigenvectors = np.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def eigh( x: np.ndarray, /, *, UPLO: str = "L", out: Optional[np.ndarray] = None @@ -137,9 +166,6 @@ def matmul( return ret -matmul.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("float16", "bfloat16")}, backend_version) def matrix_norm( @@ -214,9 +240,6 @@ def outer( return np.outer(x1, x2, out=out) -outer.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def pinv( x: np.ndarray, @@ -303,27 +326,27 @@ def svdvals(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray return np.linalg.svd(x, compute_uv=False) -def tensorsolve( +def tensordot( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + axes: Union[int, Tuple[List[int], List[int]]] = 2, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.linalg.tensorsolve(x1, x2, axes=axes) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return np.tensordot(x1, x2, axes=axes) -def tensordot( +def tensorsolve( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return np.tensordot(x1, x2, axes=axes) + return np.linalg.tensorsolve(x1, x2, axes=axes) @_scalar_output_to_0d_array @@ -340,7 +363,16 @@ def trace( return np.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) -trace.support_native_out = True +@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) +def vander( + x: np.ndarray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.vander(x, N=N, increasing=increasing).astype(x.dtype) def vecdot( @@ -355,15 +387,6 @@ def vecdot( return np.tensordot(x1, x2, axes=(axis, axis)) -@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) -def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] - ) - eigenvalues, eigenvectors = np.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - def vector_norm( x: np.ndarray, /, @@ -404,32 +427,6 @@ def vector_norm( return res -# Extra # -# ----- # - - -def diag( - x: np.ndarray, - /, - *, - k: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.diag(x, k=k) - - -@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) -def vander( - x: np.ndarray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.vander(x, N=N, increasing=increasing).astype(x.dtype) - - @with_unsupported_dtypes( { "1.25.2 and below": ( @@ -457,6 +454,3 @@ def vector_to_skew_symmetric_matrix( row3 = np.concatenate((-a2s, a1s, zs), -1) # BS x 3 x 3 return np.concatenate((row1, row2, row3), -2, out=out) - - -vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/numpy/manipulation.py b/ivy/functional/backends/numpy/manipulation.py index f258de2b5c59d..8e614f3d7335a 100644 --- a/ivy/functional/backends/numpy/manipulation.py +++ b/ivy/functional/backends/numpy/manipulation.py @@ -10,10 +10,42 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x +# --- Main --- # +# ------------ # + + +def as_strided( + x: np.ndarray, + shape: Union[ivy.NativeShape, Sequence[int]], + strides: Sequence[int], + /, +) -> np.ndarray: + return np.lib.stride_tricks.as_strided( + x, + shape=shape, + strides=strides, + ) + + +def clip( + x: np.ndarray, + x_min: Union[Number, np.ndarray], + x_max: Union[Number, np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) + + # Array API Standard # # -------------------# @@ -41,7 +73,15 @@ def concat( return ivy.astype(ret, highest_dtype, copy=False) -concat.support_native_out = True +def constant_pad( + x: np.ndarray, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) def expand_dims( @@ -88,6 +128,18 @@ def permute_dims( return np.transpose(x, axes) +@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) +def repeat( + x: np.ndarray, + /, + repeats: Union[int, List[int]], + *, + axis: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.repeat(x, repeats, axis) + + def reshape( x: np.ndarray, /, @@ -118,38 +170,6 @@ def roll( return np.roll(x, shift, axis) -def squeeze( - x: np.ndarray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - return np.squeeze(x, axis=axis) - - -def stack( - arrays: Union[Tuple[np.ndarray], List[np.ndarray]], - /, - *, - axis: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.stack(arrays, axis, out=out) - - -stack.support_native_out = True - - # Extra # # ------# @@ -191,39 +211,33 @@ def split( return np.split(x, num_or_size_splits, axis) -@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) -def repeat( +def squeeze( x: np.ndarray, /, - repeats: Union[int, List[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.repeat(x, repeats, axis) - - -def tile( - x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tile(x, repeats) + if isinstance(axis, list): + axis = tuple(axis) + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + return np.squeeze(x, axis=axis) -def constant_pad( - x: np.ndarray, +def stack( + arrays: Union[Tuple[np.ndarray], List[np.ndarray]], /, - pad_width: List[List[int]], *, - value: Number = 0.0, + axis: int = 0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) - - -def zero_pad( - x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None -): - return np.pad(_flat_array_to_1_dim_array(x), pad_width) + return np.stack(arrays, axis, out=out) def swapaxes( @@ -238,6 +252,12 @@ def swapaxes( return np.swapaxes(x, axis0, axis1) +def tile( + x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tile(x, repeats) + + def unstack( x: np.ndarray, /, @@ -259,28 +279,12 @@ def unstack( return [np.squeeze(item, axis) for item in x_split] -def clip( - x: np.ndarray, - x_min: Union[Number, np.ndarray], - x_max: Union[Number, np.ndarray], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) +def zero_pad( + x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None +): + return np.pad(_flat_array_to_1_dim_array(x), pad_width) +concat.support_native_out = True +stack.support_native_out = True clip.support_native_out = True - - -def as_strided( - x: np.ndarray, - shape: Union[ivy.NativeShape, Sequence[int]], - strides: Sequence[int], - /, -) -> np.ndarray: - return np.lib.stride_tricks.as_strided( - x, - shape=shape, - strides=strides, - ) diff --git a/ivy/functional/backends/numpy/random.py b/ivy/functional/backends/numpy/random.py index bba1591aa039a..2f8db47cedd99 100644 --- a/ivy/functional/backends/numpy/random.py +++ b/ivy/functional/backends/numpy/random.py @@ -14,42 +14,6 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, np.ndarray] = 0.0, - high: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, - seed: Optional[int] = None, -) -> np.ndarray: - if seed: - np.random.seed(seed) - shape = _check_bounds_and_get_shape(low, high, shape).shape - return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) - - -def random_normal( - *, - mean: Union[float, np.ndarray] = 0.0, - std: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: np.dtype, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - if seed: - np.random.seed(seed) - return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) - @with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) def multinomial( @@ -110,6 +74,43 @@ def randint( return np.random.randint(low, high, shape, dtype=dtype) +def random_normal( + *, + mean: Union[float, np.ndarray] = 0.0, + std: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: np.dtype, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + if seed: + np.random.seed(seed) + return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, np.ndarray] = 0.0, + high: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, + seed: Optional[int] = None, +) -> np.ndarray: + if seed: + np.random.seed(seed) + shape = _check_bounds_and_get_shape(low, high, shape).shape + return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) + + def seed(*, seed_value: int = 0) -> None: np.random.seed(seed_value) return diff --git a/ivy/functional/backends/numpy/searching.py b/ivy/functional/backends/numpy/searching.py index c5b2b7194025c..604c06c861b77 100644 --- a/ivy/functional/backends/numpy/searching.py +++ b/ivy/functional/backends/numpy/searching.py @@ -60,6 +60,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.argwhere(x) + + def nonzero( x: np.ndarray, /, @@ -95,16 +108,3 @@ def where( ) -> np.ndarray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(np.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.argwhere(x) diff --git a/ivy/functional/backends/numpy/sorting.py b/ivy/functional/backends/numpy/sorting.py index aa3d84358220b..b3d949062b87a 100644 --- a/ivy/functional/backends/numpy/sorting.py +++ b/ivy/functional/backends/numpy/sorting.py @@ -8,6 +8,9 @@ from . import backend_version +msort.support_native_out = False + + def argsort( x: np.ndarray, /, @@ -25,22 +28,6 @@ def argsort( ) -def sort( - x: np.ndarray, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - kind = "stable" if stable else "quicksort" - ret = np.asarray(np.sort(x, axis=axis, kind=kind)) - if descending: - ret = np.asarray((np.flip(ret, axis))) - return ret - - # msort @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def msort( @@ -49,9 +36,6 @@ def msort( return np.msort(a) -msort.support_native_out = False - - def searchsorted( x: np.ndarray, v: np.ndarray, @@ -89,3 +73,19 @@ def searchsorted( else: ret = np.searchsorted(x, v, side=side, sorter=sorter) return ret.astype(ret_dtype) + + +def sort( + x: np.ndarray, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + kind = "stable" if stable else "quicksort" + ret = np.asarray(np.sort(x, axis=axis, kind=kind)) + if descending: + ret = np.asarray((np.flip(ret, axis))) + return ret diff --git a/ivy/functional/backends/numpy/statistical.py b/ivy/functional/backends/numpy/statistical.py index 419ef9355d336..248318c8104ac 100644 --- a/ivy/functional/backends/numpy/statistical.py +++ b/ivy/functional/backends/numpy/statistical.py @@ -10,23 +10,100 @@ from ivy.utils.einsum_parser import legalise_einsum_expr -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -def min( +def _infer_dtype(dtype: np.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) +def cumprod( x: np.ndarray, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[np.dtype] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - axis = tuple(axis) if isinstance(axis, list) else axis - return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return np.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + return np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumprod(x, -1, dtype=dtype) + return np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + return np.flip(x, axis=axis) -min.support_native_out = True +def cumsum( + x: np.ndarray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + + if exclusive or reverse: + if exclusive and reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + res = np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumsum(x, -1, dtype=dtype) + res = np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + res = np.flip(x, axis=axis) + return res + return np.cumsum(x, axis, dtype=dtype, out=out) + + +@_scalar_output_to_0d_array +def einsum( + equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None +) -> np.ndarray: + equation = legalise_einsum_expr(*[equation, *operands]) + return np.einsum(equation, *operands, out=out) def max( @@ -41,9 +118,6 @@ def max( return np.asarray(np.amax(a=x, axis=axis, keepdims=keepdims, out=out)) -max.support_native_out = True - - @_scalar_output_to_0d_array def mean( x: np.ndarray, @@ -59,14 +133,20 @@ def mean( ) -mean.support_native_out = True +# Array API Standard # +# -------------------# -def _infer_dtype(dtype: np.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +def min( + x: np.ndarray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + axis = tuple(axis) if isinstance(axis, list) else axis + return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) def prod( @@ -85,9 +165,6 @@ def prod( return np.asarray(np.prod(a=x, axis=axis, dtype=dtype, keepdims=keepdims, out=out)) -prod.support_native_out = True - - def std( x: np.ndarray, /, @@ -101,9 +178,6 @@ def std( return np.asarray(np.std(x, axis=axis, ddof=correction, keepdims=keepdims, out=out)) -std.support_native_out = True - - def sum( x: np.ndarray, /, @@ -127,9 +201,6 @@ def sum( ) -sum.support_native_out = True - - @_scalar_output_to_0d_array def var( x: np.ndarray, @@ -164,94 +235,13 @@ def var( ) +min.support_native_out = True +max.support_native_out = True +mean.support_native_out = True +prod.support_native_out = True +std.support_native_out = True +sum.support_native_out = True var.support_native_out = True - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) -def cumprod( - x: np.ndarray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return np.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - return np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumprod(x, -1, dtype=dtype) - return np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - return np.flip(x, axis=axis) - - cumprod.support_native_out = True - - -def cumsum( - x: np.ndarray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - - if exclusive or reverse: - if exclusive and reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - res = np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumsum(x, -1, dtype=dtype) - res = np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - res = np.flip(x, axis=axis) - return res - return np.cumsum(x, axis, dtype=dtype, out=out) - - cumsum.support_native_out = True - - -@_scalar_output_to_0d_array -def einsum( - equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None -) -> np.ndarray: - equation = legalise_einsum_expr(*[equation, *operands]) - return np.einsum(equation, *operands, out=out) - - einsum.support_native_out = True diff --git a/ivy/functional/backends/numpy/utility.py b/ivy/functional/backends/numpy/utility.py index de5597f952940..6904438177374 100644 --- a/ivy/functional/backends/numpy/utility.py +++ b/ivy/functional/backends/numpy/utility.py @@ -6,6 +6,10 @@ import ivy +all.support_native_out = True +any.support_native_out = True + + def all( x: np.ndarray, /, @@ -20,9 +24,6 @@ def all( raise ivy.utils.exceptions.IvyIndexError(error) -all.support_native_out = True - - def any( x: np.ndarray, /, @@ -32,6 +33,3 @@ def any( out: Optional[np.ndarray] = None, ) -> np.ndarray: return np.asarray(np.any(x, axis=axis, keepdims=keepdims, out=out)) - - -any.support_native_out = True diff --git a/ivy/functional/backends/paddle/activations.py b/ivy/functional/backends/paddle/activations.py index d090f69a66a4b..6979b59f13f8b 100644 --- a/ivy/functional/backends/paddle/activations.py +++ b/ivy/functional/backends/paddle/activations.py @@ -30,14 +30,38 @@ ] -def relu( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version +) +def gelu( + x: paddle.Tensor, + /, + *, + approximate: bool = False, + complex_mode="jax", + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: + if paddle.is_complex(x): + sqrt_2_over_pi = 0.7978845608 + # the other magic number comes directly from the formula in + # https://doi.org/10.48550/arXiv.1606.08415 + return ( + 0.5 + * x + * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) + ) if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return paddle.complex(F.relu(x.real()), F.relu(x.imag())) - return F.relu(x.cast("float32")).cast(x.dtype) - return F.relu(x) + return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) + return F.gelu(x, approximate=approximate) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def hardswish( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return F.hardswish(x) @with_unsupported_device_and_dtypes( @@ -62,28 +86,47 @@ def leaky_relu( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) -def gelu( +def log_softmax( x: paddle.Tensor, /, *, - approximate: bool = False, - complex_mode="jax", + axis: Optional[int] = None, out: Optional[paddle.Tensor] = None, +): + if axis is None: + axis = -1 + x_max = paddle_backend.max(x, axis=axis, keepdims=True) + x_max = paddle_backend.where( + paddle_backend.isfinite(x_max), + x_max, + paddle.zeros(shape=x_max.shape).astype(x_max.dtype), + ) + exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) + + s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) + ret = paddle_backend.log(s) + ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) + return ret + + +def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in unsupported_dtypes: + if paddle.is_complex(x): + return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) + return F.mish(x.cast("float32")).cast(x.dtype) + return F.mish(x) + + +def relu( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if paddle.is_complex(x): - sqrt_2_over_pi = 0.7978845608 - # the other magic number comes directly from the formula in - # https://doi.org/10.48550/arXiv.1606.08415 - return ( - 0.5 - * x - * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) - ) if x.dtype in unsupported_dtypes: - return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) - return F.gelu(x, approximate=approximate) + if paddle.is_complex(x): + return paddle.complex(F.relu(x.real()), F.relu(x.imag())) + return F.relu(x.cast("float32")).cast(x.dtype) + return F.relu(x) def sigmoid( @@ -142,46 +185,3 @@ def softplus( if threshold is not None: return ivy.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def log_softmax( - x: paddle.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -): - if axis is None: - axis = -1 - x_max = paddle_backend.max(x, axis=axis, keepdims=True) - x_max = paddle_backend.where( - paddle_backend.isfinite(x_max), - x_max, - paddle.zeros(shape=x_max.shape).astype(x_max.dtype), - ) - exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) - - s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) - ret = paddle_backend.log(s) - ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) - return ret - - -def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) - return F.mish(x.cast("float32")).cast(x.dtype) - return F.mish(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def hardswish( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return F.hardswish(x) diff --git a/ivy/functional/backends/paddle/creation.py b/ivy/functional/backends/paddle/creation.py index 7b01cae902d53..c216ae2522833 100644 --- a/ivy/functional/backends/paddle/creation.py +++ b/ivy/functional/backends/paddle/creation.py @@ -25,6 +25,114 @@ from . import backend_version from paddle.device import core + +# --- Helpers --- # +# --------------- # + + +def _differentiable_linspace(start, stop, num, *, dtype=None): + start = ivy.to_native(start) + num = paddle.to_tensor(num, stop_gradient=False) + if num == 1: + return paddle_backend.expand_dims(start, axis=0) + n_m_1 = paddle_backend.subtract(num, 1) + increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) + increment_tiled = paddle_backend.repeat(increment, n_m_1) + increments = paddle_backend.multiply( + increment_tiled, + paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), + ) + if isinstance(start, int) or start.ndim == 0: + start = paddle_backend.expand_dims(start, axis=0) + res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) + return res.cast(dtype) + + +def _linspace_helper(start, stop, num, axis=None, *, dtype=None): + num = num.detach().item() if isinstance(num, paddle.Tensor) else num + start_is_array = isinstance(start, paddle.Tensor) + stop_is_array = isinstance(stop, paddle.Tensor) + linspace_method = paddle.linspace + sos_shape = [] + if start_is_array: + start_shape = start.shape + sos_shape = start_shape + if num == 1: + if axis is not None: + return paddle_backend.expand_dims(start, axis=axis) + else: + return paddle_backend.expand_dims(start, axis=-1) + start = start.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not start.stop_gradient else paddle.linspace + ) + if stop_is_array: + stop_shape = stop.shape + sos_shape = stop_shape + if num == 1: + return ( + paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start + ) + stop = stop.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not stop.stop_gradient else paddle.linspace + ) + if start_is_array and stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=-1) + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = paddle_backend.subtract(stop, start) + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [ + linspace_method(strt, stp, num) + for strt, stp in zip( + paddle_backend.unstack(start, keepdims=True), + paddle_backend.unstack(stop, keepdims=True), + ) + ] + elif start_is_array and not stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=axis) + diff = stop - start + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(paddle.ones(start.shape).astype(start.dtype) * stop) + else: + res = [linspace_method(strt, stop, num) for strt in start] + elif not start_is_array and stop_is_array: + if num < stop.shape[0]: + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = stop - start + inc = diff / (num - 1) + res = [paddle.ones(stop.shape).astype(stop.dtype) * start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [linspace_method(start, stp, num) for stp in stop] + else: + return linspace_method(start, stop, num, dtype=dtype) + res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) + if axis is not None: + ndim = res.ndim + perm = list(range(ndim - 1)) + perm.insert(axis % (ndim + 1), ndim - 1) + res = paddle_backend.permute_dims(res, perm) + return res + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +# --- Main --- # +# ------------ # + + # Array API Standard # # -------------------# @@ -108,6 +216,21 @@ def asarray( return paddle.to_tensor(obj, dtype=dtype, place=device) +def copy_array( + x: paddle.Tensor, + *, + to_ivy_array: Optional[bool] = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if 0 in x.shape: + new_arr = paddle.empty(x.shape, dtype=x.dtype) + else: + new_arr = x.clone() + if to_ivy_array: + return ivy.to_ivy(new_arr) + return new_arr + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -198,6 +321,46 @@ def from_dlpack(x, /, *, out: Optional[paddle.Tensor] = None): return paddle.utils.dlpack.from_dlpack(x_d) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def frombuffer( + buffer: bytes, + dtype: Optional[paddle.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> paddle.Tensor: + dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) + if str(dtype) == "bool": + dtype_bytes = 1 + dtype_str = str(dtype) + struct_format = { + "bool": "?", + "int8": "b", + "int16": "h", + "int32": "i", + "int64": "q", + "uint8": "B", + "float16": "e", + "float32": "f", + "float64": "d", + } + ret = [] + for i in range(0, len(buffer), dtype_bytes): + x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) + ret = ret + list(x) + if offset > 0: + offset = int(offset / dtype_bytes) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + ret = paddle.to_tensor(ret, dtype=dtype) + + return ret + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -237,105 +400,6 @@ def full_like( ) -def _linspace_helper(start, stop, num, axis=None, *, dtype=None): - num = num.detach().item() if isinstance(num, paddle.Tensor) else num - start_is_array = isinstance(start, paddle.Tensor) - stop_is_array = isinstance(stop, paddle.Tensor) - linspace_method = paddle.linspace - sos_shape = [] - if start_is_array: - start_shape = start.shape - sos_shape = start_shape - if num == 1: - if axis is not None: - return paddle_backend.expand_dims(start, axis=axis) - else: - return paddle_backend.expand_dims(start, axis=-1) - start = start.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not start.stop_gradient else paddle.linspace - ) - if stop_is_array: - stop_shape = stop.shape - sos_shape = stop_shape - if num == 1: - return ( - paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start - ) - stop = stop.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not stop.stop_gradient else paddle.linspace - ) - if start_is_array and stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=-1) - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = paddle_backend.subtract(stop, start) - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [ - linspace_method(strt, stp, num) - for strt, stp in zip( - paddle_backend.unstack(start, keepdims=True), - paddle_backend.unstack(stop, keepdims=True), - ) - ] - elif start_is_array and not stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=axis) - diff = stop - start - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(paddle.ones(start.shape).astype(start.dtype) * stop) - else: - res = [linspace_method(strt, stop, num) for strt in start] - elif not start_is_array and stop_is_array: - if num < stop.shape[0]: - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = stop - start - inc = diff / (num - 1) - res = [paddle.ones(stop.shape).astype(stop.dtype) * start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [linspace_method(start, stp, num) for stp in stop] - else: - return linspace_method(start, stop, num, dtype=dtype) - res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) - if axis is not None: - ndim = res.ndim - perm = list(range(ndim - 1)) - perm.insert(axis % (ndim + 1), ndim - 1) - res = paddle_backend.permute_dims(res, perm) - return res - - -def _differentiable_linspace(start, stop, num, *, dtype=None): - start = ivy.to_native(start) - num = paddle.to_tensor(num, stop_gradient=False) - if num == 1: - return paddle_backend.expand_dims(start, axis=0) - n_m_1 = paddle_backend.subtract(num, 1) - increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) - increment_tiled = paddle_backend.repeat(increment, n_m_1) - increments = paddle_backend.multiply( - increment_tiled, - paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), - ) - if isinstance(start, int) or start.ndim == 0: - start = paddle_backend.expand_dims(start, axis=0) - res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) - return res.cast(dtype) - - -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint16", "bfloat16", "float16")}}, backend_version ) @@ -441,6 +505,58 @@ def meshgrid( return res +def one_hot( + indices: paddle.Tensor, + depth: int, + /, + *, + on_value: Optional[paddle.Tensor] = None, + off_value: Optional[paddle.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[paddle.dtype] = None, + device: core.Place, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + on_none = on_value is None + off_none = off_value is None + expand_ret = False + if indices.ndim == 0: + expand_ret = True + indices = indices.cast("int64").unsqueeze(0) + if dtype is None: + if on_none and off_none: + dtype = paddle.float32 + else: + if not on_none: + dtype = paddle.to_tensor(on_value).dtype + elif not off_none: + dtype = paddle.to_tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = ( + paddle.to_tensor(1.0, dtype="float32") + if on_none + else paddle.to_tensor(on_value, dtype="float32") + ) + off_value = ( + paddle.to_tensor(0.0, dtype="float32") + if off_none + else paddle.to_tensor(off_value, dtype="float32") + ) + + res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) + + if not on_none or not off_none: + res = paddle.where(res == 1, on_value, off_value) + + if axis is not None: + res = paddle.moveaxis(res, -1, axis) + if expand_ret: + res = res.squeeze() + return res.cast(dtype) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -500,6 +616,22 @@ def triu( return paddle.triu(x=x, diagonal=k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: Optional[int] = 0, + /, + *, + device: core.Place, +) -> Tuple[paddle.Tensor]: + # special case due to inconsistent behavior when n_cols=1 and n_rows=0 + if n_cols == 1 and n_rows == 0: + return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( + [], place=device, dtype="int64" + ) + return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -526,126 +658,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: paddle.Tensor, - *, - to_ivy_array: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if 0 in x.shape: - new_arr = paddle.empty(x.shape, dtype=x.dtype) - else: - new_arr = x.clone() - if to_ivy_array: - return ivy.to_ivy(new_arr) - return new_arr - - -def one_hot( - indices: paddle.Tensor, - depth: int, - /, - *, - on_value: Optional[paddle.Tensor] = None, - off_value: Optional[paddle.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[paddle.dtype] = None, - device: core.Place, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - on_none = on_value is None - off_none = off_value is None - expand_ret = False - if indices.ndim == 0: - expand_ret = True - indices = indices.cast("int64").unsqueeze(0) - if dtype is None: - if on_none and off_none: - dtype = paddle.float32 - else: - if not on_none: - dtype = paddle.to_tensor(on_value).dtype - elif not off_none: - dtype = paddle.to_tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = ( - paddle.to_tensor(1.0, dtype="float32") - if on_none - else paddle.to_tensor(on_value, dtype="float32") - ) - off_value = ( - paddle.to_tensor(0.0, dtype="float32") - if off_none - else paddle.to_tensor(off_value, dtype="float32") - ) - - res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) - - if not on_none or not off_none: - res = paddle.where(res == 1, on_value, off_value) - - if axis is not None: - res = paddle.moveaxis(res, -1, axis) - if expand_ret: - res = res.squeeze() - return res.cast(dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def frombuffer( - buffer: bytes, - dtype: Optional[paddle.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> paddle.Tensor: - dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) - if str(dtype) == "bool": - dtype_bytes = 1 - dtype_str = str(dtype) - struct_format = { - "bool": "?", - "int8": "b", - "int16": "h", - "int32": "i", - "int64": "q", - "uint8": "B", - "float16": "e", - "float32": "f", - "float64": "d", - } - ret = [] - for i in range(0, len(buffer), dtype_bytes): - x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) - ret = ret + list(x) - if offset > 0: - offset = int(offset / dtype_bytes) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - ret = paddle.to_tensor(ret, dtype=dtype) - - return ret - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: Optional[int] = 0, - /, - *, - device: core.Place, -) -> Tuple[paddle.Tensor]: - # special case due to inconsistent behavior when n_cols=1 and n_rows=0 - if n_cols == 1 and n_rows == 0: - return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( - [], place=device, dtype="int64" - ) - return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) diff --git a/ivy/functional/backends/paddle/data_type.py b/ivy/functional/backends/paddle/data_type.py index 1a67457c36079..bb4550a92dcc8 100644 --- a/ivy/functional/backends/paddle/data_type.py +++ b/ivy/functional/backends/paddle/data_type.py @@ -22,7 +22,6 @@ paddle.complex128: "complex128", paddle.bool: "bool", } - native_dtype_dict = { "int8": paddle.int8, "int16": paddle.int16, @@ -102,6 +101,51 @@ def __repr__(self): ) +# Extra # +# ------# + + +def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: + if dtype_in is int: + return ivy.default_int_dtype() + if dtype_in is float: + return ivy.default_float_dtype() + if dtype_in is complex: + return ivy.default_complex_dtype() + if dtype_in is bool: + return ivy.Dtype("bool") + if isinstance(dtype_in, str): + if dtype_in in native_dtype_dict: + return ivy.Dtype(dtype_in) + else: + raise ivy.utils.exceptions.IvyException( + "Cannot convert to ivy dtype." + f" {dtype_in} is not supported by Paddle backend." + ) + return ivy.Dtype(ivy_dtype_dict[dtype_in]) + + +def as_native_dtype( + dtype_in: Union[paddle.dtype, str, bool, int, float] +) -> paddle.dtype: + if dtype_in is int: + return ivy.default_int_dtype(as_native=True) + if dtype_in is float: + return ivy.default_float_dtype(as_native=True) + if dtype_in is complex: + return ivy.default_complex_dtype(as_native=True) + if dtype_in is bool: + return paddle.bool + if not isinstance(dtype_in, str): + return dtype_in + if dtype_in in native_dtype_dict.keys(): + return native_dtype_dict[ivy.Dtype(dtype_in)] + else: + raise ivy.utils.exceptions.IvyException( + f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." + ) + + # Array API Standard # # -------------------# @@ -172,6 +216,26 @@ def broadcast_to( return paddle.broadcast_to(x, shape) +def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: + if as_native: + return ivy.to_native(x).dtype + return as_ivy_dtype(x.dtype) + + +def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: + dtype_str = as_ivy_dtype(dtype_in) + if "bool" in dtype_str: + return 1 + return int( + dtype_str.replace("paddle.", "") + .replace("uint", "") + .replace("int", "") + .replace("bfloat", "") + .replace("float", "") + .replace("complex", "") + ) + + @_handle_nestable_dtype_info def finfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Finfo: if isinstance(type, paddle.Tensor): @@ -195,75 +259,6 @@ def iinfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Iinfo: return Iinfo(np.iinfo(type)) -def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: - return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) - - -# Extra # -# ------# - - -def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: - if dtype_in is int: - return ivy.default_int_dtype() - if dtype_in is float: - return ivy.default_float_dtype() - if dtype_in is complex: - return ivy.default_complex_dtype() - if dtype_in is bool: - return ivy.Dtype("bool") - if isinstance(dtype_in, str): - if dtype_in in native_dtype_dict: - return ivy.Dtype(dtype_in) - else: - raise ivy.utils.exceptions.IvyException( - "Cannot convert to ivy dtype." - f" {dtype_in} is not supported by Paddle backend." - ) - return ivy.Dtype(ivy_dtype_dict[dtype_in]) - - -def as_native_dtype( - dtype_in: Union[paddle.dtype, str, bool, int, float] -) -> paddle.dtype: - if dtype_in is int: - return ivy.default_int_dtype(as_native=True) - if dtype_in is float: - return ivy.default_float_dtype(as_native=True) - if dtype_in is complex: - return ivy.default_complex_dtype(as_native=True) - if dtype_in is bool: - return paddle.bool - if not isinstance(dtype_in, str): - return dtype_in - if dtype_in in native_dtype_dict.keys(): - return native_dtype_dict[ivy.Dtype(dtype_in)] - else: - raise ivy.utils.exceptions.IvyException( - f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." - ) - - -def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: - if as_native: - return ivy.to_native(x).dtype - return as_ivy_dtype(x.dtype) - - -def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: - dtype_str = as_ivy_dtype(dtype_in) - if "bool" in dtype_str: - return 1 - return int( - dtype_str.replace("paddle.", "") - .replace("uint", "") - .replace("int", "") - .replace("bfloat", "") - .replace("float", "") - .replace("complex", "") - ) - - def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -271,3 +266,7 @@ def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: + return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) diff --git a/ivy/functional/backends/paddle/device.py b/ivy/functional/backends/paddle/device.py index c8e9d9ba385d9..f8295eaaf3d1d 100644 --- a/ivy/functional/backends/paddle/device.py +++ b/ivy/functional/backends/paddle/device.py @@ -13,29 +13,26 @@ from paddle.device import core -# API # -# ----# +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper Paddle profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + def start(self): + self._start_time = time.perf_counter() -def dev( - x: paddle.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, core.Place]: - return x.place if as_native else as_ivy_dev(x.place) + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + def __enter__(self): + self.start() -def to_device( - x: paddle.Tensor, - device: core.Place, - /, - *, - stream: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - device = as_native_dev(device) - if device.is_cpu_place(): - return x.cpu() - elif device.is_gpu_place(): - return x.cuda(device.gpu_device_id()) + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() def as_ivy_dev(device: core.Place, /): @@ -73,29 +70,30 @@ def as_native_dev( return native_dev -def clear_mem_on_dev(device: core.Place, /): +def clear_cached_mem_on_dev(device: str, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -def clear_cached_mem_on_dev(device: str, /): +def clear_mem_on_dev(device: core.Place, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -def num_gpus() -> int: - return paddle.device.cuda.device_count() +# API # +# ----# -def gpu_is_available() -> bool: - return bool(paddle.device.cuda.device_count()) +def dev( + x: paddle.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, core.Place]: + return x.place if as_native else as_ivy_dev(x.place) -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - return False +def gpu_is_available() -> bool: + return bool(paddle.device.cuda.device_count()) def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): @@ -112,23 +110,25 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return ret -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper Paddle profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None +def num_gpus() -> int: + return paddle.device.cuda.device_count() - def start(self): - self._start_time = time.perf_counter() - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) +def to_device( + x: paddle.Tensor, + device: core.Place, + /, + *, + stream: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + device = as_native_dev(device) + if device.is_cpu_place(): + return x.cpu() + elif device.is_gpu_place(): + return x.cuda(device.gpu_device_id()) - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/paddle/elementwise.py b/ivy/functional/backends/paddle/elementwise.py index f4664e17d527a..cbba2fecde68e 100644 --- a/ivy/functional/backends/paddle/elementwise.py +++ b/ivy/functional/backends/paddle/elementwise.py @@ -12,187 +12,155 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. +def _determine_sqrt_dtype_cast( + dtype: Type[paddle.Tensor], +) -> Tuple[Optional[str], Optional[str]]: + """ + Determine the appropriate casting dtype for sqrt operations. + + Returns: + (intermediate_dtype, output_dtype) + """ + + cast_and_return_float32_dtype = { + paddle.int8, + paddle.int16, + paddle.int32, + paddle.uint8, + paddle.bool, + } + + if dtype in cast_and_return_float32_dtype: + return "float32", "float32" + elif dtype == paddle.int64: + return "float64", "float64" + elif dtype == paddle.float16: + return "float32", "float16" + elif dtype == paddle.bfloat16: + return "float32", "bfloat16" + else: + return None, None + + def _elementwise_helper(x1, x2): x1, x2 = ivy.promote_types_of_inputs(x1, x2) x1, x2 = paddle_backend.broadcast_arrays(x1, x2) return x1, x2, x1.dtype -def add( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], +# --- Main --- # +# ------------ # + + +def abs( + x: Union[float, paddle.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() + if x.dtype in [ paddle.int8, + paddle.int16, paddle.uint8, paddle.float16, - paddle.bool, paddle.bfloat16, + paddle.bool, ]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if alpha not in (1, None): - x2 = paddle_backend.multiply(x2, alpha) - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.add(x1, x2).astype(ret_dtype) - - -def bitwise_xor( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_xor(x1, x2) - - -def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: - return paddle.expm1(x) - return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) - - -def bitwise_invert( - x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.bitwise_not(x) + return paddle.abs(x.astype("float32")).astype(x.dtype) + return paddle.abs(x) @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex64", - "complex128", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def isfinite( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.isfinite(x) - - -def isinf( - x: paddle.Tensor, - /, - *, - detect_positive: bool = True, - detect_negative: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if detect_negative and detect_positive: - return paddle.isinf(x) - - if detect_negative: - return paddle_backend.equal(x, float("-inf")) - - if detect_positive: - return paddle_backend.equal(x, float("inf")) - - return paddle.zeros(shape=x.shape, dtype=bool) +def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.acos(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa + s1 = paddle_backend.sqrt(1 - x) + s2 = paddle_backend.sqrt(1 + x) + return paddle.complex( + 2.0 * paddle.atan2(s1.real(), s2.real()), + paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), + ) + return paddle.acos(x) -def equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - diff = paddle_backend.subtract(x1, x2) - ret = paddle_backend.logical_and( - paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) - ) - # ret result is sufficient for all cases except where the value is +/-INF of NaN - return paddle_backend.where( - paddle_backend.isnan(diff), - ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), - ret, - ) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.acosh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa + s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) + s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) + return paddle.complex( + paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), + 2.0 * paddle.atan2(s1.imag(), s2.real()), + ) + return paddle.acosh(x) -def less_equal( +def add( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - if paddle.is_complex(x1): - real = paddle.less_equal(x1.real(), x2.real()) - imag = paddle.less_equal(x1.imag(), x2.imag()) - return paddle_backend.logical_and(real, imag) - return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_equal(x1, x2) - - -def bitwise_and( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_and(x1, x2) - - -def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ + if x1.dtype in [ paddle.int8, - paddle.int16, - paddle.int32, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, paddle.bool, + paddle.bfloat16, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) - return paddle.ceil(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.ceil(x.astype("float64")).astype(x_dtype) - return paddle.ceil(x) + x1, x2 = x1.astype("float32"), x2.astype("float32") + if alpha not in (1, None): + x2 = paddle_backend.multiply(x2, alpha) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.add(x1, x2).astype(ret_dtype) -def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) - return paddle.floor(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.floor(x.astype("float64")).astype(x_dtype) - return paddle.floor(x) +def angle( + input: paddle.Tensor, + /, + *, + deg: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + result = paddle.angle(input) + if deg: + result = paddle.rad2deg(result) + return result @with_unsupported_device_and_dtypes( @@ -243,16 +211,10 @@ def asinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def sign( - x: paddle.Tensor, - /, - *, - np_variant: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: +def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -260,69 +222,33 @@ def sign( paddle.int64, paddle.uint8, paddle.float16, - paddle.bfloat16, - paddle.bool, ]: - return paddle.sgn(x.astype("float32")).astype(x.dtype) - return paddle.sgn(x) - - -# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. -def _determine_sqrt_dtype_cast( - dtype: Type[paddle.Tensor], -) -> Tuple[Optional[str], Optional[str]]: - """ - Determine the appropriate casting dtype for sqrt operations. - - Returns: - (intermediate_dtype, output_dtype) - """ - - cast_and_return_float32_dtype = { - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.bool, - } - - if dtype in cast_and_return_float32_dtype: - return "float32", "float32" - elif dtype == paddle.int64: - return "float64", "float64" - elif dtype == paddle.float16: - return "float32", "float16" - elif dtype == paddle.bfloat16: - return "float32", "bfloat16" - else: - return None, None - - -def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - """Calculate the square root with type handling.""" - - if paddle.is_complex(x): - angle = paddle.angle(x) - return paddle.complex( - paddle.cos(angle / 2), paddle.sin(angle / 2) - ) * paddle.sqrt(paddle.abs(x)) - - if x.dtype in {paddle.float32, paddle.float64}: - return paddle.sqrt(x) + ret_dtype = x.dtype + return paddle.atan(x.astype("float32")).astype(ret_dtype) + if x.dtype in [paddle.complex64, paddle.complex128]: + atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) + return paddle.atan(x) - intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) - if intermediate_dtype: - result = paddle.sqrt(x.astype(intermediate_dtype)) - return result.astype(output_dtype) - raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + backend_version, +) +def atan2( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.atan2(x1, x2).astype(ret_dtype) @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -332,22 +258,83 @@ def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.float16, ]: ret_dtype = x.dtype - return paddle.cosh(x.astype("float32")).astype(ret_dtype) + return paddle.atanh(x.astype("float32")).astype(ret_dtype) if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) - ) - return paddle.cosh(x) + return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) + return paddle.atanh(x) -def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def bitwise_and( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_and(x1, x2) + + +def bitwise_invert( + x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.bitwise_not(x) + + +def bitwise_left_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( + ret_dtype + ) + + +def bitwise_or( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_or(x1, x2) + + +def bitwise_right_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( + ret_dtype + ) + + +def bitwise_xor( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_xor(x1, x2) + + +def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ paddle.int8, paddle.int16, paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, @@ -355,15 +342,18 @@ def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.bool, ]: if paddle.is_complex(x): - base = paddle.to_tensor(10.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log10(x.astype("float32")).astype(x.dtype) - return paddle.log10(x) + return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) + return paddle.ceil(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.ceil(x.astype("float64")).astype(x_dtype) + return paddle.ceil(x) -def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def cos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -371,20 +361,24 @@ def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - base = paddle.to_tensor(2.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log2(x.astype("float32")).astype(x.dtype) - return paddle.log2(x) + ret_dtype = x.dtype + return paddle.cos(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.cos(re) * paddle.cosh(im), + -paddle.sin(re) * paddle.sinh(im), + ) + return paddle.cos(x) -def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -392,32 +386,27 @@ def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - return paddle_backend.log(x + 1) - return paddle.log1p(x.astype("float32")).astype(x.dtype) - return paddle.log1p(x) + ret_dtype = x.dtype + return paddle.cosh(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) + ) + return paddle.cosh(x) -def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) - return paddle.isnan(x.astype("float32")) - return paddle.isnan(x) +def deg2rad( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: + return paddle.deg2rad(x.astype("float32")).astype(x.dtype) + return paddle.deg2rad(x) -def less( +def divide( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -425,17 +414,14 @@ def less( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - real = paddle.less_than(x1.real(), x2.real()) - imag = paddle.less_than(x1.imag(), x2.imag()) - return logical_and(real, imag) - return paddle.less_than(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_than(x1, x2) + if x1.dtype in [paddle.float16, paddle.bfloat16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if not (ivy.is_float_dtype(ret_dtype) or ivy.is_complex_dtype(ret_dtype)): + ret_dtype = ivy.default_float_dtype(as_native=True) + return (x1 / x2).astype(ret_dtype) -def multiply( +def equal( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -443,49 +429,81 @@ def multiply( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.multiply(x1, x2).astype(ret_dtype) - + diff = paddle_backend.subtract(x1, x2) + ret = paddle_backend.logical_and( + paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) + ) + # ret result is sufficient for all cases except where the value is +/-INF of NaN + return paddle_backend.where( + paddle_backend.isnan(diff), + ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), + ret, + ) + + +# Extra # +# ------# + @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, backend_version, ) -def cos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + # TODO: add support for complex x, supported in scipy only atm + if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: + return paddle.erf(x.astype("float32")).astype(x.dtype) + return paddle.erf(x) + + +def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + return paddle.exp(x) + if paddle.is_complex(x): + return paddle.multiply( + paddle.exp(x.real()), + paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), + ) + return paddle_backend.pow(math.e, x).astype(x.dtype) + + +def exp2( + x: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.pow(2, x) + + +def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: + return paddle.expm1(x) + return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) + + +def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ paddle.int8, paddle.int16, paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - ret_dtype = x.dtype - return paddle.cos(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.cos(re) * paddle.cosh(im), - -paddle.sin(re) * paddle.sinh(im), - ) - return paddle.cos(x) - - -def logical_not( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: if paddle.is_complex(x): - return paddle.logical_and( - paddle.logical_not(x.real()), paddle.logical_not(x.imag()) - ) - return paddle.logical_not(x.astype("float32")) - return paddle.logical_not(x) + return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) + return paddle.floor(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.floor(x.astype("float64")).astype(x_dtype) + return paddle.floor(x) -def divide( +def floor_divide( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -493,11 +511,9 @@ def divide( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.float16, paddle.bfloat16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if not (ivy.is_float_dtype(ret_dtype) or ivy.is_complex_dtype(ret_dtype)): - ret_dtype = ivy.default_float_dtype(as_native=True) - return (x1 / x2).astype(ret_dtype) + if x1.dtype in [paddle.int32, paddle.int64]: + return paddle.floor_divide(x1, x2) + return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) @with_supported_dtypes( @@ -516,6 +532,32 @@ def fmin( return paddle.fmin(x1, x2) +def fmod( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) + return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version +) +def gcd( + x1: Union[paddle.Tensor, int, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.gcd(x1, x2) + + def greater( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], @@ -553,179 +595,121 @@ def greater_equal( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, backend_version, ) -def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.acos(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa - s1 = paddle_backend.sqrt(1 - x) - s2 = paddle_backend.sqrt(1 + x) - return paddle.complex( - 2.0 * paddle.atan2(s1.real(), s2.real()), - paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), - ) - return paddle.acos(x) +def imag( + val: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.imag(val) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex64", + "complex128", + "bool", + ) + } + }, backend_version, ) -def logical_xor( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def isfinite( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # with the XOR logic - # if paddle.is_complex(x1): - # return paddle.logical_xor( - # paddle.logical_xor(x1.real(), x2.real()), - # paddle.logical_xor(x1.imag(), x2.imag()), - # ) - return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_xor(x1, x2) + return paddle.isfinite(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def logical_and( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def isinf( + x: paddle.Tensor, + /, + *, + detect_positive: bool = True, + detect_negative: bool = True, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # if paddle.is_complex(x1): - # return paddle.logical_and( - # paddle.logical_and(x1.real(), x2.real()), - # paddle.logical_and(x1.imag(), x2.imag()), - # ) - return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_and(x1, x2) + if detect_negative and detect_positive: + return paddle.isinf(x) + if detect_negative: + return paddle_backend.equal(x, float("-inf")) -def logical_or( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - return paddle.logical_or( - paddle.logical_or(x1.real(), x2.real()), - paddle.logical_or(x1.imag(), x2.imag()), - ) - return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_or(x1, x2) + if detect_positive: + return paddle_backend.equal(x, float("inf")) + return paddle.zeros(shape=x.shape, dtype=bool) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + +def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, - paddle.int32, - paddle.int64, paddle.uint8, - paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - return paddle.acosh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa - s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) - s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) - return paddle.complex( - paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), - 2.0 * paddle.atan2(s1.imag(), s2.real()), - ) - return paddle.acosh(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.sin(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) - ) - return paddle.sin(x) - - -def negative( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - if x.dtype == paddle.bool: - return paddle.logical_not(x) - return paddle.neg(x) + if paddle.is_complex(x): + return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) + return paddle.isnan(x.astype("float32")) + return paddle.isnan(x) -def not_equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def isreal( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - return paddle.logical_not(paddle_backend.equal(x1, x2)) + if paddle.is_complex(x): + return paddle.logical_not(x.imag().astype(bool)) + else: + return paddle.ones_like(x, dtype="bool") @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, backend_version, ) -def tanh( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +def lcm( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.tanh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - tanh_a = paddle.tanh(paddle.real(x)) - tan_b = paddle.tan(paddle.imag(x)) - return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) - return paddle.tanh(x) + x1_dtype = x1.dtype + x2_dtype = x2.dtype + if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): + return paddle.cast( + paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), + paddle.int16, + ) + elif x1_dtype != x2_dtype: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.lcm(x1, x2) -def floor_divide( +def less( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -733,27 +717,36 @@ def floor_divide( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int32, paddle.int64]: - return paddle.floor_divide(x1, x2) - return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + real = paddle.less_than(x1.real(), x2.real()) + imag = paddle.less_than(x1.imag(), x2.imag()) + return logical_and(real, imag) + return paddle.less_than(x1.astype("float32"), x2.astype("float32")) + return paddle.less_than(x1, x2) -def bitwise_or( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], + +def less_equal( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_or(x1, x2) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + if paddle.is_complex(x1): + real = paddle.less_equal(x1.real(), x2.real()) + imag = paddle.less_equal(x1.imag(), x2.imag()) + return paddle_backend.logical_and(real, imag) + return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) + return paddle.less_equal(x1, x2) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + +def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -761,112 +754,43 @@ def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - ret_dtype = x.dtype - return paddle.sinh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) - ) - return paddle.sinh(x) - - -def positive( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - return x.clone() - - -def square( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.square(x) - if paddle.is_complex(x): - return paddle.complex( - paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), - 2.0 * paddle.real(x) * paddle.imag(x), - ) - return paddle_backend.pow(x, 2).astype(x.dtype) + if paddle.is_complex(x): + return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) + return paddle.log(x.astype("float32")).astype(x.dtype) + return paddle.log(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version -) -def pow( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ +def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, paddle.bool, ]: - return paddle.pow(x1.astype("float32"), x2.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x1): - # https://math.stackexchange.com/questions/476968/complex-power-of-a-complex-number - r = paddle.abs(x1) - theta = paddle.angle(x1) - res_mag = paddle.pow(r, x2.real()) / paddle.exp(x2.imag() * theta) - res_ang = paddle.log(r) * x2.imag() + theta * x2.real() - result = res_mag * paddle.complex(paddle.cos(res_ang), paddle.sin(res_ang)) - return result.astype(ret_dtype) - return paddle.pow(x1, x2) - - -def round( - x: paddle.Tensor, /, *, decimals: int = 0, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - def _np_round(x, decimals): - # this is a logic to mimic np.round behaviour - # which rounds odd numbers up and even numbers down at limits like 0.5 - eps = 1e-6 * paddle.sign(x) - - # check if the integer is even or odd - candidate_ints = paddle_backend.remainder(paddle_backend.trunc(x), 2.0).astype( - bool - ) - # check if the fraction is exactly half - candidate_fractions = paddle_backend.equal( - paddle_backend.abs(paddle_backend.subtract(x, paddle_backend.trunc(x))), - 0.5, - ) - x = paddle_backend.where( - paddle.logical_and(~candidate_ints, candidate_fractions), - x - eps, - x, - ) - factor = paddle_backend.pow(10.0, decimals).astype(x.dtype) - factor_denom = ivy.where(ivy.isinf(x), 1.0, factor) - return paddle_backend.divide( - paddle.round(paddle_backend.multiply(x, factor)), factor_denom - ) - - if x.dtype not in [paddle.float32, paddle.float64]: if paddle.is_complex(x): - return paddle.complex( - _np_round(x.real(), decimals), _np_round(x.imag(), decimals) - ) - return _np_round(x.astype("float32"), decimals).astype(x.dtype) - return _np_round(x, decimals).astype(x.dtype) + base = paddle.to_tensor(10.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log10(x.astype("float32")).astype(x.dtype) + return paddle.log10(x) -def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, @@ -874,79 +798,30 @@ def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.bool, ]: if paddle.is_complex(x): - return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) - return paddle.trunc(x.astype("float32")).astype(x.dtype) - return paddle.trunc(x) + return paddle_backend.log(x + 1) + return paddle.log1p(x.astype("float32")).astype(x.dtype) + return paddle.log1p(x) -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32")}, - backend_version, -) -def trapz( - y: paddle.Tensor, - /, - *, - x: Optional[paddle.Tensor] = None, - dx: Optional[float] = 1.0, - axis: Optional[int] = -1, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x is None: - d = dx - else: - if x.ndim == 1: - d = paddle.diff(x) - # reshape to correct shape - shape = [1] * y.ndim - shape[axis] = d.shape[0] - d = d.reshape(shape) - else: - d = paddle.diff(x, axis=axis) - - slice1 = [slice(None)] * y.ndim - slice2 = [slice(None)] * y.ndim - - slice1[axis] = slice(1, None) - slice2[axis] = slice(None, -1) - - with ivy.ArrayMode(False): - if y.shape[axis] < 2: - return ivy.zeros_like(ivy.squeeze(y, axis=axis)) - ret = ivy.sum( - ivy.divide( - ivy.multiply( - d, - ivy.add( - ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) - ), - ), - 2.0, - ), - axis=axis, - ) - - return ret - - -def abs( - x: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() +def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, - paddle.bfloat16, + paddle.complex64, + paddle.complex128, paddle.bool, ]: - return paddle.abs(x.astype("float32")).astype(x.dtype) - return paddle.abs(x) + if paddle.is_complex(x): + base = paddle.to_tensor(2.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log2(x.astype("float32")).astype(x.dtype) + return paddle.log2(x) @with_unsupported_device_and_dtypes( @@ -977,264 +852,224 @@ def logaddexp2( @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.real(x) +def logical_and( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # if paddle.is_complex(x1): + # return paddle.logical_and( + # paddle.logical_and(x1.real(), x2.real()), + # paddle.logical_and(x1.imag(), x2.imag()), + # ) + return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_and(x1, x2) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.tan(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) - return paddle.tan(x) +def logical_not( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + return paddle.logical_and( + paddle.logical_not(x.real()), paddle.logical_not(x.imag()) + ) + return paddle.logical_not(x.astype("float32")) + return paddle.logical_not(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.atan(x.astype("float32")).astype(ret_dtype) - if x.dtype in [paddle.complex64, paddle.complex128]: - atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) - return paddle.atan(x) +def logical_or( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + return paddle.logical_or( + paddle.logical_or(x1.real(), x2.real()), + paddle.logical_or(x1.imag(), x2.imag()), + ) + return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_or(x1, x2) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def atan2( +def logical_xor( x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.atan2(x1, x2).astype(ret_dtype) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # with the XOR logic + # if paddle.is_complex(x1): + # return paddle.logical_xor( + # paddle.logical_xor(x1.real(), x2.real()), + # paddle.logical_xor(x1.imag(), x2.imag()), + # ) + return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_xor(x1, x2) -def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def maximum( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + use_where: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [ paddle.int8, paddle.int16, - paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) - return paddle.log(x.astype("float32")).astype(x.dtype) - return paddle.log(x) - - -def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.exp(x) - if paddle.is_complex(x): - return paddle.multiply( - paddle.exp(x.real()), - paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), - ) - return paddle_backend.pow(math.e, x).astype(x.dtype) - - -def exp2( - x: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.pow(2, x) + if paddle.is_complex(x1): + use_where = True + else: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if use_where: + return paddle_backend.where( + paddle_backend.greater_equal(x1, x2), x1, x2 + ).astype(ret_dtype) + return paddle.maximum(x1, x2).astype(ret_dtype) -def subtract( +def minimum( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, + use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.float16, paddle.bool]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if alpha not in (1, None): - x2 = paddle_backend.multiply(x2, alpha) - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.subtract(x1, x2).astype(ret_dtype) + if x1.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x1): + use_where = True + else: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if use_where: + return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( + ret_dtype + ) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def remainder( + return paddle.minimum(x1, x2).astype(ret_dtype) + + +def multiply( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - modulus: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if not modulus: - res = paddle_backend.divide(x1, x2) - res_floored = paddle_backend.where( - paddle_backend.greater_equal(res, 0.0), - paddle_backend.floor(res), - paddle_backend.ceil(res), - ) - diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) - return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.remainder(x1, x2).astype(ret_dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.atanh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) - return paddle.atanh(x) + return paddle.multiply(x1, x2).astype(ret_dtype) -def bitwise_right_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], +def nan_to_num( + x: paddle.Tensor, /, *, + copy: Optional[bool] = True, + nan: Optional[Union[float, int]] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( - ret_dtype - ) + with ivy.ArrayMode(False): + if ivy.is_int_dtype(x): + if posinf is None: + posinf = ivy.iinfo(x).max + if neginf is None: + neginf = ivy.iinfo(x).min + elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): + if posinf is None: + posinf = ivy.finfo(x).max + if neginf is None: + neginf = ivy.finfo(x).min + ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret > 0), + paddle.to_tensor(posinf, dtype=x.dtype), + ret, + ) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret < 0), + paddle.to_tensor(neginf, dtype=x.dtype), + ret, + ) + if copy: + return ret.clone() + else: + x = ret + return x -def bitwise_left_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def negative( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( - ret_dtype - ) - - -# Extra # -# ------# - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - # TODO: add support for complex x, supported in scipy only atm - if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: - return paddle.erf(x.astype("float32")).astype(x.dtype) - return paddle.erf(x) + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + if x.dtype == paddle.bool: + return paddle.logical_not(x) + return paddle.neg(x) -def minimum( +def not_equal( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.logical_not(paddle_backend.equal(x1, x2)) - if use_where: - return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( - ret_dtype - ) - return paddle.minimum(x1, x2).astype(ret_dtype) +def positive( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + return x.clone() -def maximum( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version +) +def pow( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) @@ -1243,35 +1078,18 @@ def maximum( paddle.int16, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if use_where: - return paddle_backend.where( - paddle_backend.greater_equal(x1, x2), x1, x2 - ).astype(ret_dtype) - return paddle.maximum(x1, x2).astype(ret_dtype) - - -def reciprocal( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return paddle.reciprocal(x) - return paddle_backend.divide(1, x) - - -def deg2rad( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: - return paddle.deg2rad(x.astype("float32")).astype(x.dtype) - return paddle.deg2rad(x) + return paddle.pow(x1.astype("float32"), x2.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x1): + # https://math.stackexchange.com/questions/476968/complex-power-of-a-complex-number + r = paddle.abs(x1) + theta = paddle.angle(x1) + res_mag = paddle.pow(r, x2.real()) / paddle.exp(x2.imag() * theta) + res_ang = paddle.log(r) * x2.imag() + theta * x2.real() + result = res_mag * paddle.complex(paddle.cos(res_ang), paddle.sin(res_ang)) + return result.astype(ret_dtype) + return paddle.pow(x1, x2) def rad2deg( @@ -1282,148 +1100,338 @@ def rad2deg( return paddle.rad2deg(x) -def trunc_divide( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.trunc(paddle_backend.divide(x1, x2)) - - -def isreal( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if paddle.is_complex(x): - return paddle.logical_not(x.imag().astype(bool)) - else: - return paddle.ones_like(x, dtype="bool") +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, + backend_version, +) +def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.real(x) -def fmod( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, +def reciprocal( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) - return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) + if x.dtype in [paddle.float32, paddle.float64]: + return paddle.reciprocal(x) + return paddle_backend.divide(1, x) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, backend_version, ) -def lcm( - x1: paddle.Tensor, - x2: paddle.Tensor, +def remainder( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, + modulus: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1_dtype = x1.dtype - x2_dtype = x2.dtype - if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): - return paddle.cast( - paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), - paddle.int16, + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if not modulus: + res = paddle_backend.divide(x1, x2) + res_floored = paddle_backend.where( + paddle_backend.greater_equal(res, 0.0), + paddle_backend.floor(res), + paddle_backend.ceil(res), ) - elif x1_dtype != x2_dtype: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.lcm(x1, x2) + diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) + return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) + + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.remainder(x1, x2).astype(ret_dtype) -def angle( - input: paddle.Tensor, +def round( + x: paddle.Tensor, /, *, decimals: int = 0, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + def _np_round(x, decimals): + # this is a logic to mimic np.round behaviour + # which rounds odd numbers up and even numbers down at limits like 0.5 + eps = 1e-6 * paddle.sign(x) + + # check if the integer is even or odd + candidate_ints = paddle_backend.remainder(paddle_backend.trunc(x), 2.0).astype( + bool + ) + # check if the fraction is exactly half + candidate_fractions = paddle_backend.equal( + paddle_backend.abs(paddle_backend.subtract(x, paddle_backend.trunc(x))), + 0.5, + ) + x = paddle_backend.where( + paddle.logical_and(~candidate_ints, candidate_fractions), + x - eps, + x, + ) + factor = paddle_backend.pow(10.0, decimals).astype(x.dtype) + factor_denom = ivy.where(ivy.isinf(x), 1.0, factor) + return paddle_backend.divide( + paddle.round(paddle_backend.multiply(x, factor)), factor_denom + ) + + if x.dtype not in [paddle.float32, paddle.float64]: + if paddle.is_complex(x): + return paddle.complex( + _np_round(x.real(), decimals), _np_round(x.imag(), decimals) + ) + return _np_round(x.astype("float32"), decimals).astype(x.dtype) + return _np_round(x, decimals).astype(x.dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def sign( + x: paddle.Tensor, /, *, - deg: Optional[bool] = None, + np_variant: Optional[bool] = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - result = paddle.angle(input) - if deg: - result = paddle.rad2deg(result) - return result + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.bfloat16, + paddle.bool, + ]: + return paddle.sgn(x.astype("float32")).astype(x.dtype) + return paddle.sgn(x) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, ) -def gcd( - x1: Union[paddle.Tensor, int, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], +def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.sin(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) + ) + return paddle.sin(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + ret_dtype = x.dtype + return paddle.sinh(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) + ) + return paddle.sinh(x) + + +def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + """Calculate the square root with type handling.""" + + if paddle.is_complex(x): + angle = paddle.angle(x) + return paddle.complex( + paddle.cos(angle / 2), paddle.sin(angle / 2) + ) * paddle.sqrt(paddle.abs(x)) + + if x.dtype in {paddle.float32, paddle.float64}: + return paddle.sqrt(x) + + intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) + if intermediate_dtype: + result = paddle.sqrt(x.astype(intermediate_dtype)) + return result.astype(output_dtype) + + raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") + + +def square( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + return paddle.square(x) + if paddle.is_complex(x): + return paddle.complex( + paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), + 2.0 * paddle.real(x) * paddle.imag(x), + ) + return paddle_backend.pow(x, 2).astype(x.dtype) + + +def subtract( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.gcd(x1, x2) + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.float16, paddle.bool]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if alpha not in (1, None): + x2 = paddle_backend.multiply(x2, alpha) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.subtract(x1, x2).astype(ret_dtype) @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def imag( - val: paddle.Tensor, +def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + ret_dtype = x.dtype + return paddle.tan(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) + return paddle.tan(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def tanh( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.tanh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + tanh_a = paddle.tanh(paddle.real(x)) + tan_b = paddle.tan(paddle.imag(x)) + return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) + return paddle.tanh(x) + + +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32")}, + backend_version, +) +def trapz( + y: paddle.Tensor, /, *, + x: Optional[paddle.Tensor] = None, + dx: Optional[float] = 1.0, + axis: Optional[int] = -1, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.imag(val) + if x is None: + d = dx + else: + if x.ndim == 1: + d = paddle.diff(x) + # reshape to correct shape + shape = [1] * y.ndim + shape[axis] = d.shape[0] + d = d.reshape(shape) + else: + d = paddle.diff(x, axis=axis) + + slice1 = [slice(None)] * y.ndim + slice2 = [slice(None)] * y.ndim + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) -def nan_to_num( - x: paddle.Tensor, + with ivy.ArrayMode(False): + if y.shape[axis] < 2: + return ivy.zeros_like(ivy.squeeze(y, axis=axis)) + ret = ivy.sum( + ivy.divide( + ivy.multiply( + d, + ivy.add( + ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) + ), + ), + 2.0, + ), + axis=axis, + ) + + return ret + + +def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) + return paddle.trunc(x.astype("float32")).astype(x.dtype) + return paddle.trunc(x) + + +def trunc_divide( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, - copy: Optional[bool] = True, - nan: Optional[Union[float, int]] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - if ivy.is_int_dtype(x): - if posinf is None: - posinf = ivy.iinfo(x).max - if neginf is None: - neginf = ivy.iinfo(x).min - elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): - if posinf is None: - posinf = ivy.finfo(x).max - if neginf is None: - neginf = ivy.finfo(x).min - ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret > 0), - paddle.to_tensor(posinf, dtype=x.dtype), - ret, - ) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret < 0), - paddle.to_tensor(neginf, dtype=x.dtype), - ret, - ) - if copy: - return ret.clone() - else: - x = ret - return x + return paddle_backend.trunc(paddle_backend.divide(x1, x2)) diff --git a/ivy/functional/backends/paddle/experimental/activations.py b/ivy/functional/backends/paddle/experimental/activations.py index 81bc6bdf25cb7..069bf40aeac5b 100644 --- a/ivy/functional/backends/paddle/experimental/activations.py +++ b/ivy/functional/backends/paddle/experimental/activations.py @@ -9,6 +9,24 @@ from . import backend_version +def elu( + x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.elu(x, alpha=alpha) + + if paddle.is_complex(x): + ret = ( + paddle_backend.where( + paddle_backend.greater(x, 0), + x, + paddle_backend.multiply(alpha, paddle_backend.expm1(x)), + ), + ) + return ret + return F.elu(x.cast("float32"), alpha).cast(x.dtype) + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -33,28 +51,6 @@ def logit(x: paddle.Tensor, /, *, eps: Optional[float] = None, out=None): ).cast(x.dtype) -def thresholded_relu( - x: paddle.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.thresholded_relu(x, threshold=threshold) - return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( - x.dtype - ) - - -def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.relu6(x) - if paddle.is_complex(x): - return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) - return F.relu6(x.cast("float32")).cast(x.dtype) - - def logsigmoid( input: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: @@ -69,6 +65,14 @@ def logsigmoid( return F.log_sigmoid(input.cast("float32")).cast(input.dtype) +def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.relu6(x) + if paddle.is_complex(x): + return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) + return F.relu6(x.cast("float32")).cast(x.dtype) + + def selu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: return F.selu(x) @@ -95,19 +99,15 @@ def silu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. return F.silu(x.cast("float32")).cast(x.dtype) -def elu( - x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None +def thresholded_relu( + x: paddle.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = 0, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: - return F.elu(x, alpha=alpha) - - if paddle.is_complex(x): - ret = ( - paddle_backend.where( - paddle_backend.greater(x, 0), - x, - paddle_backend.multiply(alpha, paddle_backend.expm1(x)), - ), - ) - return ret - return F.elu(x.cast("float32"), alpha).cast(x.dtype) + return F.thresholded_relu(x, threshold=threshold) + return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( + x.dtype + ) diff --git a/ivy/functional/backends/paddle/experimental/creation.py b/ivy/functional/backends/paddle/experimental/creation.py index 04d15e026a29d..09ab86a05324d 100644 --- a/ivy/functional/backends/paddle/experimental/creation.py +++ b/ivy/functional/backends/paddle/experimental/creation.py @@ -15,6 +15,11 @@ import ivy from .. import backend_version + +# --- Helpers --- # +# --------------- # + + # noinspection PyProtectedMember # Helpers for calculating Window Functions # ---------------------------------------- @@ -29,39 +34,28 @@ def _kaiser_window(window_length, beta): ) / paddle_backend.i0(beta) -# Array API Standard # -# -------------------# +# --- Main --- # +# ------------ # -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def blackman_window( + size: int, + /, *, + periodic: Optional[bool] = True, dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if window_length < 2: - return paddle.ones([window_length], dtype=dtype) - if periodic is False: - return _kaiser_window(window_length, beta).cast(dtype) + if size < 2: + return paddle.ones([size], dtype=dtype) + if periodic: + count = paddle.arange(size) / size else: - return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) - - -def vorbis_window( - window_length: paddle.Tensor, - *, - dtype: Optional[paddle.dtype] = paddle.float32, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if window_length == 0: - return paddle.to_tensor([], dtype=dtype) - i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) - pi = paddle.full(shape=i.shape, fill_value=math.pi) - return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( - dtype - ) + count = paddle.linspace(start=0, stop=size, num=size) + return ( + (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) + + (0.08 * paddle.cos(2 * math.pi * 2 * count)) + ).cast(dtype) def hann_window( @@ -81,6 +75,26 @@ def hann_window( return (0.5 - 0.5 * paddle.cos(2 * math.pi * count)).cast(dtype) +# Array API Standard # +# -------------------# + + +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if window_length < 2: + return paddle.ones([window_length], dtype=dtype) + if periodic is False: + return _kaiser_window(window_length, beta).cast(dtype) + else: + return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) + + def tril_indices( n_rows: int, n_cols: Optional[int] = None, @@ -101,6 +115,32 @@ def tril_indices( ) +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex", + ) + } + }, + backend_version, +) +def trilu( + x: paddle.Tensor, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if upper: + return paddle.triu(x=x, diagonal=k) + return paddle.tril(x=x, diagonal=k) + + @with_supported_dtypes( {"2.4.2 and below": ("float64", "float32", "int32", "int64")}, backend_version, @@ -134,26 +174,6 @@ def unsorted_segment_min( return res -def blackman_window( - size: int, - /, - *, - periodic: Optional[bool] = True, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if size < 2: - return paddle.ones([size], dtype=dtype) - if periodic: - count = paddle.arange(size) / size - else: - count = paddle.linspace(start=0, stop=size, num=size) - return ( - (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) - + (0.08 * paddle.cos(2 * math.pi * 2 * count)) - ).cast(dtype) - - def unsorted_segment_sum( data: paddle.Tensor, segment_ids: paddle.Tensor, @@ -188,27 +208,16 @@ def unsorted_segment_sum( return res -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex", - ) - } - }, - backend_version, -) -def trilu( - x: paddle.Tensor, - /, +def vorbis_window( + window_length: paddle.Tensor, *, - k: int = 0, - upper: bool = True, + dtype: Optional[paddle.dtype] = paddle.float32, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if upper: - return paddle.triu(x=x, diagonal=k) - return paddle.tril(x=x, diagonal=k) + if window_length == 0: + return paddle.to_tensor([], dtype=dtype) + i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) + pi = paddle.full(shape=i.shape, fill_value=math.pi) + return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( + dtype + ) diff --git a/ivy/functional/backends/paddle/experimental/elementwise.py b/ivy/functional/backends/paddle/experimental/elementwise.py index 8fe58b9e1519a..95b5be21b6aad 100644 --- a/ivy/functional/backends/paddle/experimental/elementwise.py +++ b/ivy/functional/backends/paddle/experimental/elementwise.py @@ -17,79 +17,102 @@ from .. import backend_version -@with_supported_dtypes( - {"2.5.1 and below": ("float32", "float64")}, - backend_version, -) -def lgamma( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.lgamma(x) +_BERNOULLI_COEFS = [ + 12, + -720, + 30240, + -1209600, + 47900160, + -1307674368000 / 691, + 74724249600, + -10670622842880000 / 3617, + 5109094217170944000 / 43867, + -802857662698291200000 / 174611, + 14101100039391805440000 / 77683, + -1693824136731743669452800000 / 236364091, + 186134520519971831808000000 / 657931, + -37893265687455865519472640000000 / 3392780147, + 759790291646040068357842010112000000 / 1723168255201, + -134196726836183700385281186201600000000 / 7709321041217, +] -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, - backend_version, -) -def fmax( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x1.dtype != x2.dtype: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.fmax(x1, x2) +# --- Helpers --- # +# --------------- # -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) - return paddle.divide(paddle.sin(y), y) +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim -def float_power( - x1: Union[paddle.Tensor, float, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1 = paddle.cast(x1, dtype="float64") - x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power - return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis -def frexp( - x: Union[paddle.Tensor, Number], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException +def _np_ndim(x): + return ivy.array(x).ndim -def ldexp( - x1: Union[paddle.Tensor, Number], - x2: Union[paddle.Tensor, Number], +# --- Main --- # +# ------------ # + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def allclose( + x1: paddle.Tensor, + x2: paddle.Tensor, /, *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - out_dtype = x1.dtype - x1, x2 = promote_types_of_inputs(x1, x2) - with ivy.ArrayMode(False): - if ivy.any(ivy.less(x2, 0)): - pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 - neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 - ret = ivy.multiply(ivy.pow(2, pos_exp), x1) - ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) - else: - ret = ivy.multiply(ivy.pow(2, x2), x1) - return ivy.astype(ret, out_dtype, copy=False) +) -> bool: + return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) + + +@with_supported_dtypes( + { + "2.5.1 and below": ( + "complex64", + "complex128", + "float32", + "float64", + "int32", + "int64", + ) + }, + backend_version, +) +def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.conj(x) def copysign( @@ -107,38 +130,17 @@ def copysign( return result -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version -) -def nansum( - x: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[paddle.dtype] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) - if result.shape == [1]: - result = paddle.fluid.layers.squeeze(result, [0]) - return result - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def isclose( +def count_nonzero( a: paddle.Tensor, - b: paddle.Tensor, /, *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, + axis: Optional[Union[int, list, tuple]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + with ivy.ArrayMode(False): + return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) def diff( @@ -167,185 +169,69 @@ def _tensor(val): ) -def signbit( - x: Union[paddle.Tensor, float, int, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.less( - paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 - ) - - -def hypot( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -@with_unsupported_device_and_dtypes( +@with_supported_dtypes( { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def allclose( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> bool: - return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) - - -def fix( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.trunc(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def nextafter( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - with ivy.ArrayMode(False): - eps = ivy.finfo(x1.dtype).eps - return ivy.where( - ivy.equal(x1, x2), - x2, - ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), + "2.5.0 and below": ( + "float32", + "float64", ) - - -_BERNOULLI_COEFS = [ - 12, - -720, - 30240, - -1209600, - 47900160, - -1307674368000 / 691, - 74724249600, - -10670622842880000 / 3617, - 5109094217170944000 / 43867, - -802857662698291200000 / 174611, - 14101100039391805440000 / 77683, - -1693824136731743669452800000 / 236364091, - 186134520519971831808000000 / 657931, - -37893265687455865519472640000000 / 3392780147, - 759790291646040068357842010112000000 / 1723168255201, - -134196726836183700385281186201600000000 / 7709321041217, -] - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "uint16", - "float16", - "bool", - ) - } }, - backend_version, -) -def zeta( - x: paddle.Tensor, - q: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - s, a = ivy.promote_types_of_inputs(x, q) - s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) - N = M = ( - paddle.to_tensor(8.0, dtype="float32") - if q.dtype == paddle.float32 - else paddle.to_tensor(8.0, dtype="float64") - ) - assert M <= len(_BERNOULLI_COEFS) - k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) - S = paddle.sum((a_ + k) ** -s_, -1) - Q = ivy.divide((q + N) ** (1 - x), x - 1) - T0 = (q + N) ** -x - m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) - s_over_a = (s_ + m) / (a_ + N) - s_over_a = ivy.where( - s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a - ) - T1 = paddle.cumprod(s_over_a, -1)[..., ::2] - # t=np.array(T1) - T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) - coefs = paddle.unsqueeze( - paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), - tuple(range(a.ndim)), - ) - T1 = T1 / coefs - T = T0 * (0.5 + paddle.sum(T1, -1)) - ans = S + Q + T - mask = x < 1 - ans[mask] = ivy.nan - return ans + backend_version, +) +def digamma( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.digamma(x) -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim +def fix( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.trunc(x) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +def float_power( + x1: Union[paddle.Tensor, float, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1 = paddle.cast(x1, dtype="float64") + x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power + return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) -def _np_ndim(x): - return ivy.array(x).ndim +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, + backend_version, +) +def fmax( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x1.dtype != x2.dtype: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.fmax(x1, x2) + + +def frexp( + x: Union[paddle.Tensor, Number], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException @with_supported_dtypes( @@ -578,47 +464,60 @@ def gradient( return outvals -def xlogy( - x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def hypot( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x, y, ret_dtype = _elementwise_helper(x, y) - with ivy.ArrayMode(False): - x_ok = ivy.not_equal(x, 0.0) - safe_x = ivy.where(x_ok, x, 1.0) - safe_y = ivy.where(x_ok, y, 1.0) - return ivy.where( - x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) - ).cast(ret_dtype) + raise IvyNotImplementedException() -def count_nonzero( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def isclose( a: paddle.Tensor, + b: paddle.Tensor, + /, + *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + +def ldexp( + x1: Union[paddle.Tensor, Number], + x2: Union[paddle.Tensor, Number], /, *, - axis: Optional[Union[int, list, tuple]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: + out_dtype = x1.dtype + x1, x2 = promote_types_of_inputs(x1, x2) with ivy.ArrayMode(False): - return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) + if ivy.any(ivy.less(x2, 0)): + pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 + neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 + ret = ivy.multiply(ivy.pow(2, pos_exp), x1) + ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) + else: + ret = ivy.multiply(ivy.pow(2, x2), x1) + return ivy.astype(ret, out_dtype, copy=False) @with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - "float32", - "float64", - "int32", - "int64", - ) - }, + {"2.5.1 and below": ("float32", "float64")}, backend_version, ) -def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.conj(x) +def lgamma( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.lgamma(x) def modf( @@ -628,19 +527,128 @@ def modf( return paddle.modf(x, out=out) -@with_supported_dtypes( - { - "2.5.0 and below": ( - "float32", - "float64", +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version +) +def nansum( + x: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[paddle.dtype] = None, + keepdims: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) + if result.shape == [1]: + result = paddle.fluid.layers.squeeze(result, [0]) + return result + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def nextafter( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + with ivy.ArrayMode(False): + eps = ivy.finfo(x1.dtype).eps + return ivy.where( + ivy.equal(x1, x2), + x2, + ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), ) + + +def signbit( + x: Union[paddle.Tensor, float, int, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.less( + paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 + ) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) + return paddle.divide(paddle.sin(y), y) + + +def xlogy( + x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x, y, ret_dtype = _elementwise_helper(x, y) + with ivy.ArrayMode(False): + x_ok = ivy.not_equal(x, 0.0) + safe_x = ivy.where(x_ok, x, 1.0) + safe_y = ivy.where(x_ok, y, 1.0) + return ivy.where( + x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) + ).cast(ret_dtype) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "float16", + "bool", + ) + } }, backend_version, ) -def digamma( +def zeta( x: paddle.Tensor, + q: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.digamma(x) + with ivy.ArrayMode(False): + s, a = ivy.promote_types_of_inputs(x, q) + s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) + N = M = ( + paddle.to_tensor(8.0, dtype="float32") + if q.dtype == paddle.float32 + else paddle.to_tensor(8.0, dtype="float64") + ) + assert M <= len(_BERNOULLI_COEFS) + k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) + S = paddle.sum((a_ + k) ** -s_, -1) + Q = ivy.divide((q + N) ** (1 - x), x - 1) + T0 = (q + N) ** -x + m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) + s_over_a = (s_ + m) / (a_ + N) + s_over_a = ivy.where( + s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a + ) + T1 = paddle.cumprod(s_over_a, -1)[..., ::2] + # t=np.array(T1) + T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) + coefs = paddle.unsqueeze( + paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), + tuple(range(a.ndim)), + ) + T1 = T1 / coefs + T = T0 * (0.5 + paddle.sum(T1, -1)) + ans = S + Q + T + mask = x < 1 + ans[mask] = ivy.nan + return ans diff --git a/ivy/functional/backends/paddle/experimental/layers.py b/ivy/functional/backends/paddle/experimental/layers.py index fafd86275beb3..4fdba46985ad9 100644 --- a/ivy/functional/backends/paddle/experimental/layers.py +++ b/ivy/functional/backends/paddle/experimental/layers.py @@ -13,6 +13,11 @@ ) from .. import backend_version + +# --- Helpers --- # +# --------------- # + + # local @@ -26,6 +31,255 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling +# --- Main --- # +# ------------ # + + +def adaptive_max_pool2d( + input: paddle.Tensor, output_size: Union[Sequence[int], int] +) -> paddle.Tensor: + squeeze = input.ndim == 3 + x = paddle.unsqueeze(input, axis=0) if squeeze else input + ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) + return paddle.squeeze(ret, axis=0) if squeeze else ret + + +def avg_pool1d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool2d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, + /, + *, + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool3d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, + /, + *, + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def dct( + x: paddle.Tensor, + /, + *, + type: Optional[Literal[1, 2, 3, 4]] = 2, + n: Optional[int] = None, + axis: Optional[int] = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout1d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 3 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout2d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 4 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout3d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 5 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +def embedding( + weights: paddle.Tensor, + indices: paddle.Tensor, + /, + *, + max_norm: Optional[int] = None, + out=None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def fft( + x: paddle.Tensor, + dim: int, + /, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if not isinstance(dim, int): + raise IvyValueError(f"Expecting instead of {type(dim)}") + + if n is None: + n = x.shape[dim] + + if dim < -x.ndim or dim >= x.ndim: + raise IvyValueError( + f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" + ) + + if not isinstance(n, int): + raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") + + if n <= 1: + raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") + + valid_norm_modes = ["backward", "ortho", "forward"] + if norm not in valid_norm_modes: + raise IvyValueError( + f"Unrecognized normalization mode {norm}, expecting one of" + f" {valid_norm_modes}" + ) + + if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: + x = x.cast(paddle.complex128) + else: + x = x.cast(paddle.complex64) + + return paddle.fft.fft(x, n, dim, norm=norm) + + +@with_supported_dtypes( + { + "2.5.1 and below": ( + "complex64", + "complex128", + ) + }, + backend_version, +) +def fft2( + x: paddle.Tensor, + *, + dim: Optional[Union[int, Tuple[int]]] = None, + norm: Optional[str] = "backward", + s: Optional[Union[int, Tuple[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + res = paddle.fft.fft2(x, s, dim, norm) + return res.astype("complex128") + + +def ifft( + x: paddle.Tensor, + dim: int, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def ifftn( + x: paddle.Tensor, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: Optional[str] = "backward", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.fft.ifftn(x, s, axes, norm) + + +def interpolate( + x: paddle.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +): + raise IvyNotImplementedException() + + @with_supported_device_and_dtypes( { "2.5.1 and below": { @@ -237,230 +491,6 @@ def max_pool3d( return res -def avg_pool1d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool2d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool3d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def dct( - x: paddle.Tensor, - /, - *, - type: Optional[Literal[1, 2, 3, 4]] = 2, - n: Optional[int] = None, - axis: Optional[int] = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def fft( - x: paddle.Tensor, - dim: int, - /, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if not isinstance(dim, int): - raise IvyValueError(f"Expecting instead of {type(dim)}") - - if n is None: - n = x.shape[dim] - - if dim < -x.ndim or dim >= x.ndim: - raise IvyValueError( - f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" - ) - - if not isinstance(n, int): - raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") - - if n <= 1: - raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") - - valid_norm_modes = ["backward", "ortho", "forward"] - if norm not in valid_norm_modes: - raise IvyValueError( - f"Unrecognized normalization mode {norm}, expecting one of" - f" {valid_norm_modes}" - ) - - if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: - x = x.cast(paddle.complex128) - else: - x = x.cast(paddle.complex64) - - return paddle.fft.fft(x, n, dim, norm=norm) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout1d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 3 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout2d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 4 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout3d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 5 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -def ifft( - x: paddle.Tensor, - dim: int, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def embedding( - weights: paddle.Tensor, - indices: paddle.Tensor, - /, - *, - max_norm: Optional[int] = None, - out=None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def interpolate( - x: paddle.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -): - raise IvyNotImplementedException() - - -def adaptive_max_pool2d( - input: paddle.Tensor, output_size: Union[Sequence[int], int] -) -> paddle.Tensor: - squeeze = input.ndim == 3 - x = paddle.unsqueeze(input, axis=0) if squeeze else input - ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) - return paddle.squeeze(ret, axis=0) if squeeze else ret - - -def ifftn( - x: paddle.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.fft.ifftn(x, s, axes, norm) - - @with_unsupported_dtypes( {"2.5.1 and below": ("bfloat16", "float16", "complex64", "complex128", "bool")}, backend_version, @@ -475,24 +505,3 @@ def rfftn( ) -> paddle.Tensor: result = paddle.fft.rfftn(x, s, axes, norm) return result.astype("complex128") - - -@with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - ) - }, - backend_version, -) -def fft2( - x: paddle.Tensor, - *, - dim: Optional[Union[int, Tuple[int]]] = None, - norm: Optional[str] = "backward", - s: Optional[Union[int, Tuple[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - res = paddle.fft.fft2(x, s, dim, norm) - return res.astype("complex128") diff --git a/ivy/functional/backends/paddle/experimental/linear_algebra.py b/ivy/functional/backends/paddle/experimental/linear_algebra.py index 00f9e583ad2e7..e6f3e99f2bc4a 100644 --- a/ivy/functional/backends/paddle/experimental/linear_algebra.py +++ b/ivy/functional/backends/paddle/experimental/linear_algebra.py @@ -12,6 +12,29 @@ from .. import backend_version +dot.support_native_out = True + + +def adjoint( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + _check_valid_dimension_size(x) + return paddle.moveaxis(x, -2, -1).conj() + + +def cond( + x: paddle.Tensor, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[paddle.Tensor] = None, +) -> Any: + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) @@ -45,28 +68,14 @@ def diagflat( )(diag) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version -) -def kron( +def dot( a: paddle.Tensor, b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.kron(a, b) - - -def matrix_exp( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # TODO: this is elementwise exp, should be changed to matrix exp ASAP - # return paddle.exp(x) - raise IvyNotImplementedException() + return paddle.dot(a, b, out=out) def eig( @@ -79,24 +88,17 @@ def eigvals(x: paddle.Tensor, /) -> paddle.Tensor: return paddle.linalg.eig(x)[0] -def adjoint( - x: paddle.Tensor, +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version +) +def kron( + a: paddle.Tensor, + b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - _check_valid_dimension_size(x) - return paddle.moveaxis(x, -2, -1).conj() - - -def cond( - x: paddle.Tensor, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[paddle.Tensor] = None, -) -> Any: - raise IvyNotImplementedException() + return paddle.kron(a, b) def lu_factor( @@ -109,17 +111,15 @@ def lu_factor( raise IvyNotImplementedException() -def dot( - a: paddle.Tensor, - b: paddle.Tensor, +def matrix_exp( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.dot(a, b, out=out) - - -dot.support_native_out = True + # TODO: this is elementwise exp, should be changed to matrix exp ASAP + # return paddle.exp(x) + raise IvyNotImplementedException() @with_supported_device_and_dtypes( diff --git a/ivy/functional/backends/paddle/experimental/losses.py b/ivy/functional/backends/paddle/experimental/losses.py index 70441eb4655f3..54f155697c070 100644 --- a/ivy/functional/backends/paddle/experimental/losses.py +++ b/ivy/functional/backends/paddle/experimental/losses.py @@ -26,20 +26,21 @@ }, backend_version, ) -def l1_loss( +def huber_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, ) -> paddle.Tensor: - return F.l1_loss(input, target, reduction=reduction) + return paddle.fluid.layers.huber_loss(input, target, delta=delta) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( + "float16", "int8", "int16", "int32", @@ -53,24 +54,20 @@ def l1_loss( }, backend_version, ) -def smooth_l1_loss( +def l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - beta: Optional[float] = 1.0, reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return paddle.nn.functional.smooth_l1_loss( - input, target, reduction=reduction, beta=beta - ) + return F.l1_loss(input, target, reduction=reduction) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( - "float16", "int8", "int16", "int32", @@ -84,11 +81,14 @@ def smooth_l1_loss( }, backend_version, ) -def huber_loss( +def smooth_l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - delta: Optional[float] = 1.0, + beta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return paddle.fluid.layers.huber_loss(input, target, delta=delta) + return paddle.nn.functional.smooth_l1_loss( + input, target, reduction=reduction, beta=beta + ) diff --git a/ivy/functional/backends/paddle/experimental/manipulation.py b/ivy/functional/backends/paddle/experimental/manipulation.py index 59e781c937591..c195b7894d40f 100644 --- a/ivy/functional/backends/paddle/experimental/manipulation.py +++ b/ivy/functional/backends/paddle/experimental/manipulation.py @@ -44,7 +44,6 @@ -3.04682672343198398683e-1, 6.76795274409476084995e-1, ] - _i0B = [ -7.23318048787475395456e-18, -4.83050448594418207126e-18, @@ -74,197 +73,174 @@ ] -def moveaxis( - a: paddle.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if isinstance(source, tuple): - source = list(source) - if isinstance(destination, tuple): - source = list(destination) - if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) - return paddle.moveaxis(a, source, destination) +def atleast_1d( + *arys: paddle.Tensor, copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 1: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=0)) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res + + +def atleast_2d( + *arys: paddle.Tensor, copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 2: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, ) -def heaviside( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.heaviside(x1, x2) +def atleast_3d( + *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim == 0: + result = ary.reshape((1, 1, 1)) + elif ary.ndim == 1: + result = ary[None, :, None] + elif ary.ndim == 2: + result = ary[:, :, None] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return res -def flipud( - m: paddle.Tensor, +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + def _broadcast_shape(s1, s2): + len_1 = len(s1) + len_2 = len(s2) + if len_1 == 0: + return () if len_2 == 0 else s2 + elif len_1 != 0 and len_2 == 0: + return s1 + else: + return paddle.broadcast_shape(s1, s2) + + if len(shapes) == 0: + raise ValueError("shapes=[] must be non-empty") + elif len(shapes) == 1: + return shapes[0] + result = _broadcast_shape(shapes[0], shapes[1]) + for i in range(2, len(shapes)): + result = _broadcast_shape(result, shapes[i]) + # paddle outputs -1 if the output dimension is 0 + result = [0 if dim == -1 else dim for dim in result] + return tuple(result) + + +def concat_from_sequence( + input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) - return paddle.flip(m, axis=0) + with ivy.ArrayMode(False): + if new_axis == 0: + return ivy.concat(input_sequence, axis=axis) + elif new_axis == 1: + return ivy.stack(input_sequence, axis=axis) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int16", "float16")}}, - backend_version, -) -def vstack( - arrays: Sequence[paddle.Tensor], +def dsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], /, *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=0) - else: - return ivy.stack(arrays, axis=0) + copy: Optional[bool] = None, +) -> List[paddle.Tensor]: + if ary.ndim < 3: + raise ivy.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def dstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=1) + arrays = ivy.atleast_2d(*arrays) + if not isinstance(arrays, list): + arrays = [arrays] + if arrays[0].ndim < 3: + return ivy.stack(arrays, axis=-1) else: - return ivy.concat(arrays, axis=0) + return ivy.concat(arrays, axis=2) -@with_supported_device_and_dtypes( - { - "2.5.1 and above": { - "cpu": ( - "bool", - "int32", - "int64", - "float32", - "float64", - ), - "gpu": ("float16",), - }, - }, - backend_version, -) -def rot90( - m: paddle.Tensor, +def expand( + x: paddle.Tensor, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: Optional[int] = 1, - axes: Optional[Tuple[int, int]] = (0, 1), out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) - return paddle.rot90(m, k=k, axes=axes) + return paddle_backend.broadcast_to(x, shape) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) -def top_k( - x: paddle.Tensor, - k: int, - /, - *, - axis: int = -1, - largest: Optional[bool] = True, - sorted: bool = True, - out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] - ) - with ivy.ArrayMode(False): - indices = ivy.argsort(x, axis=axis, descending=largest) - indices = paddle.index_select(indices, paddle.arange(end=k), axis) - if not sorted: - indices = paddle.sort(indices, axis=axis) - val = ivy.take_along_axis(x, indices, axis) - return topk_res(val, indices) - - -def fliplr( - m: paddle.Tensor, - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) - return paddle.flip(m, axis=1) - - -def i0( - x: paddle.Tensor, +def fill_diagonal( + a: paddle.Tensor, + v: Union[int, float], /, *, - out: Optional[paddle.Tensor] = None, + wrap: bool = False, ) -> paddle.Tensor: - def _i0_1(x): - return paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), - ) - - def _i0_2(x): - return paddle_backend.divide( - paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl( - paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B - ), - ), - paddle_backend.sqrt(x), - ) - - def _chbevl(x, vals): - b0 = vals[0] - b1 = 0.0 - - for i in range(1, len(vals)): - b2 = b1 - b1 = b0 - b0 = paddle_backend.add( - paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] - ) - return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) + shape = a.shape + max_end = paddle.prod(paddle.to_tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() + end = max_end if end > max_end else end + a = paddle.reshape(a, (-1,)) + w = paddle.zeros(a.shape, dtype=bool) + ins = paddle.arange(0, max_end) + steps = paddle.arange(0, end, step) - x = paddle_backend.abs(x) - return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) + for i in steps: + i = ins == i + w = paddle.logical_or(w, i) + v = paddle.to_tensor(v, dtype=a.dtype) + a = paddle.where(w, v, a) + a = paddle.reshape(a, shape) + return a def flatten( @@ -306,105 +282,165 @@ def _flatten(x, start_dim, end_dim): return _flatten(x, start_dim, end_dim) -def vsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], +def fliplr( + m: paddle.Tensor, /, *, copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) + return paddle.flip(m, axis=1) -def dsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], +def flipud( + m: paddle.Tensor, /, *, copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 3: - raise ivy.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) + return paddle.flip(m, axis=0) -def atleast_1d( - *arys: paddle.Tensor, copy: Optional[bool] = None +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def heaviside( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.heaviside(x1, x2) + + +def hsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, ) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 1: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=0)) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def dstack( +def hstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - arrays = ivy.atleast_2d(*arrays) - if not isinstance(arrays, list): - arrays = [arrays] - if arrays[0].ndim < 3: - return ivy.stack(arrays, axis=-1) + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=1) else: - return ivy.concat(arrays, axis=2) + return ivy.concat(arrays, axis=0) -def atleast_2d( - *arys: paddle.Tensor, copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 2: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res +def i0( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + def _i0_1(x): + return paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), + ) + + def _i0_2(x): + return paddle_backend.divide( + paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl( + paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B + ), + ), + paddle_backend.sqrt(x), + ) + def _chbevl(x, vals): + b0 = vals[0] + b1 = 0.0 -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, + for i in range(1, len(vals)): + b2 = b1 + b1 = b0 + b0 = paddle_backend.add( + paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] + ) + return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) + + x = paddle_backend.abs(x) + return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) + + +def moveaxis( + a: paddle.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if isinstance(source, tuple): + source = list(source) + if isinstance(destination, tuple): + source = list(destination) + if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) + return paddle.moveaxis(a, source, destination) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and above": { + "cpu": ( + "bool", + "int32", + "int64", + "float32", + "float64", + ), + "gpu": ("float16",), + }, + }, backend_version, ) -def atleast_3d( - *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim == 0: - result = ary.reshape((1, 1, 1)) - elif ary.ndim == 1: - result = ary[None, :, None] - elif ary.ndim == 2: - result = ary[:, :, None] - else: - result = ary - res.append(result) - if len(res) == 1: - return res[0] - else: - return res +def rot90( + m: paddle.Tensor, + /, + *, + copy: Optional[bool] = None, + k: Optional[int] = 1, + axes: Optional[Tuple[int, int]] = (0, 1), + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) + return paddle.rot90(m, k=k, axes=axes) @with_unsupported_device_and_dtypes( @@ -484,65 +520,31 @@ def take_along_axis( return paddle.take_along_axis(arr, indices, axis) -def hsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - def _broadcast_shape(s1, s2): - len_1 = len(s1) - len_2 = len(s2) - if len_1 == 0: - return () if len_2 == 0 else s2 - elif len_1 != 0 and len_2 == 0: - return s1 - else: - return paddle.broadcast_shape(s1, s2) - - if len(shapes) == 0: - raise ValueError("shapes=[] must be non-empty") - elif len(shapes) == 1: - return shapes[0] - result = _broadcast_shape(shapes[0], shapes[1]) - for i in range(2, len(shapes)): - result = _broadcast_shape(result, shapes[i]) - # paddle outputs -1 if the output dimension is 0 - result = [0 if dim == -1 else dim for dim in result] - return tuple(result) - - -def expand( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def top_k( x: paddle.Tensor, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.broadcast_to(x, shape) - - -def concat_from_sequence( - input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: + axis: int = -1, + largest: Optional[bool] = True, + sorted: bool = True, + out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] + ) with ivy.ArrayMode(False): - if new_axis == 0: - return ivy.concat(input_sequence, axis=axis) - elif new_axis == 1: - return ivy.stack(input_sequence, axis=axis) + indices = ivy.argsort(x, axis=axis, descending=largest) + indices = paddle.index_select(indices, paddle.arange(end=k), axis) + if not sorted: + indices = paddle.sort(indices, axis=axis) + val = ivy.take_along_axis(x, indices, axis) + return topk_res(val, indices) @with_unsupported_device_and_dtypes( @@ -608,35 +610,32 @@ def unique_consecutive( ) +def vsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], + /, + *, + copy: Optional[bool] = None, +) -> List[paddle.Tensor]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) + + @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version + {"2.5.1 and below": {"cpu": ("int16", "float16")}}, + backend_version, ) -def fill_diagonal( - a: paddle.Tensor, - v: Union[int, float], +def vstack( + arrays: Sequence[paddle.Tensor], /, *, - wrap: bool = False, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - shape = a.shape - max_end = paddle.prod(paddle.to_tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() - end = max_end if end > max_end else end - a = paddle.reshape(a, (-1,)) - w = paddle.zeros(a.shape, dtype=bool) - ins = paddle.arange(0, max_end) - steps = paddle.arange(0, end, step) - - for i in steps: - i = ins == i - w = paddle.logical_or(w, i) - v = paddle.to_tensor(v, dtype=a.dtype) - a = paddle.where(w, v, a) - a = paddle.reshape(a, shape) - return a + with ivy.ArrayMode(False): + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=0) + else: + return ivy.stack(arrays, axis=0) diff --git a/ivy/functional/backends/paddle/experimental/norms.py b/ivy/functional/backends/paddle/experimental/norms.py index c1fe25a638596..6694cd81aca28 100644 --- a/ivy/functional/backends/paddle/experimental/norms.py +++ b/ivy/functional/backends/paddle/experimental/norms.py @@ -7,6 +7,13 @@ from . import backend_version +batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( + (x.ndim > 1 and x.ndim < 6) + and (scale is not None and scale.ndim == 1) + and (offset is not None and offset.ndim == 1) +) + + # TODO: add support for the rest of the dtypes # use numpy implementation with ivy functions @with_unsupported_device_and_dtypes( @@ -97,11 +104,27 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( - (x.ndim > 1 and x.ndim < 6) - and (scale is not None and scale.ndim == 1) - and (offset is not None and offset.ndim == 1) -) +def instance_norm( + x: paddle.Tensor, + mean: paddle.Tensor, + variance: paddle.Tensor, + /, + *, + scale: Optional[paddle.Tensor] = None, + offset: Optional[paddle.Tensor] = None, + training: Optional[bool] = False, + eps: Optional[float] = 1e-5, + momentum: Optional[float] = 1e-1, + data_format: Optional[str] = "NSC", + out: Optional[ + Tuple[ + paddle.Tensor, + paddle.Tensor, + paddle.Tensor, + ] + ] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: + raise IvyNotImplementedException() def l1_normalize( @@ -132,29 +155,6 @@ def l2_normalize( raise IvyNotImplementedException() -def instance_norm( - x: paddle.Tensor, - mean: paddle.Tensor, - variance: paddle.Tensor, - /, - *, - scale: Optional[paddle.Tensor] = None, - offset: Optional[paddle.Tensor] = None, - training: Optional[bool] = False, - eps: Optional[float] = 1e-5, - momentum: Optional[float] = 1e-1, - data_format: Optional[str] = "NSC", - out: Optional[ - Tuple[ - paddle.Tensor, - paddle.Tensor, - paddle.Tensor, - ] - ] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: - raise IvyNotImplementedException() - - def lp_normalize( x: paddle.Tensor, /, *, p: float = 2, axis: int = None, out: paddle.Tensor = None ) -> paddle.Tensor: diff --git a/ivy/functional/backends/paddle/experimental/random.py b/ivy/functional/backends/paddle/experimental/random.py index 542f665e4bb10..1c95f96b39179 100644 --- a/ivy/functional/backends/paddle/experimental/random.py +++ b/ivy/functional/backends/paddle/experimental/random.py @@ -12,6 +12,70 @@ from paddle.device import core from ivy import with_supported_device_and_dtypes + +# bernoulli +@with_supported_device_and_dtypes( + { + "2.5.0 and above": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + }, + "2.4.2 and below": { + "cpu": ( + "float32", + "float64", + ), + "gpu": ("float16", "float32", "float64"), + }, + }, + backend_version, +) +def bernoulli( + probs: Union[float, paddle.Tensor], + *, + logits: Union[float, paddle.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: core.Place, + dtype: paddle.dtype, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + if probs is not None: + probs = probs + elif logits is not None: + probs = ivy.softmax(logits) + probs = paddle.cast(probs, dtype) + probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs + probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) + sample = paddle.bernoulli(probs) + return to_device(sample, device) + + +# beta +def beta( + alpha: Union[float, paddle.Tensor], + beta: Union[float, paddle.Tensor], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + device: core.Place = None, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + shape = _check_bounds_and_get_shape(alpha, beta, shape) + dtype = paddle.float32 if dtype is None else dtype + beta = paddle.cast(beta, alpha.dtype) + dist = paddle.distribution.Beta(alpha, beta) + sample = dist.sample(shape) + sample = paddle.cast(sample, dtype) + return to_device(sample, device) if device is not None else sample + + # dirichlet @@ -51,29 +115,6 @@ def dirichlet( return res -# beta -def beta( - alpha: Union[float, paddle.Tensor], - beta: Union[float, paddle.Tensor], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - device: core.Place = None, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - shape = _check_bounds_and_get_shape(alpha, beta, shape) - dtype = paddle.float32 if dtype is None else dtype - beta = paddle.cast(beta, alpha.dtype) - dist = paddle.distribution.Beta(alpha, beta) - sample = dist.sample(shape) - sample = paddle.cast(sample, dtype) - return to_device(sample, device) if device is not None else sample - - def gamma( alpha: Union[float, paddle.Tensor], beta: Union[float, paddle.Tensor], @@ -99,43 +140,3 @@ def poisson( out: Optional[paddle.Tensor] = None, ): raise IvyNotImplementedException() - - -# bernoulli -@with_supported_device_and_dtypes( - { - "2.5.0 and above": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - }, - "2.4.2 and below": { - "cpu": ( - "float32", - "float64", - ), - "gpu": ("float16", "float32", "float64"), - }, - }, - backend_version, -) -def bernoulli( - probs: Union[float, paddle.Tensor], - *, - logits: Union[float, paddle.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: core.Place, - dtype: paddle.dtype, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - if probs is not None: - probs = probs - elif logits is not None: - probs = ivy.softmax(logits) - probs = paddle.cast(probs, dtype) - probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs - probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) - sample = paddle.bernoulli(probs) - return to_device(sample, device) diff --git a/ivy/functional/backends/paddle/experimental/statistical.py b/ivy/functional/backends/paddle/experimental/statistical.py index 0cc4e40739733..8047578332080 100644 --- a/ivy/functional/backends/paddle/experimental/statistical.py +++ b/ivy/functional/backends/paddle/experimental/statistical.py @@ -11,120 +11,120 @@ from . import backend_version -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def median( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # keepdims is set to True because in versions up to 2.5.1 - # there was a problem when the axis was defined and it was the - # only axis in the tensor so it needs to be handled manually - - ret_dtype = input.dtype - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(input): - ret = paddle.complex( - paddle.median(input.real(), axis=axis, keepdim=True), - paddle.median(input.imag(), axis=axis, keepdim=True), - ) - else: - ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) - else: - ret = paddle.median(input, axis=axis, keepdim=True) - if not keepdims: - ret = paddle_backend.squeeze(ret, axis=axis) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == input.ndim: - axis = None - if (input.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) +# --- Helpers --- # +# --------------- # -def nanmean( - a: paddle.Tensor, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - ret_dtype = dtype if dtype is not None else a.dtype - a = a.cast( - ret_dtype - ) # this is necessary to match other FWs behaviour which cast before calculation - if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(a): - ret = paddle.complex( - paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), - paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), - ) +def __find_cummax( + x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None +) -> Tuple[paddle.Tensor, paddle.Tensor]: + indices = [] + values = [] + x_dtype = x.dtype if dtype is None else dtype + if ( + isinstance(x.tolist()[0], list) + and len(x[0].shape) >= 1 + and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) + ): + if axis >= 1: + if not isinstance(x, list): + x = x.tolist() + for ret1 in x: + value, indice = __find_cummax( + paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype + ) + indices.append(indice) + values.append(value) else: - ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] else: - ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) - - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == a.ndim: - axis = None - if (a.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) - + if not isinstance(x, list): + x = x.tolist() + n = 0 + for idx, y in enumerate(x): + if x[n] > y: + values.append(x[n]) + elif x[n] <= y or idx == 0: + n = idx + values.append(y) + indices.append(n) -def _validate_quantile(q): - if isinstance(q, float): - q = paddle.to_tensor(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(x, paddle.Tensor): + return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( + indices, dtype="int64" + ) else: - if not (paddle.all(0 <= q) and paddle.all(q <= 1)): - return False - return True + return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] - if len(axis) == 0: - raise ValueError("Axis can't be empty!") + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis +def _compute_quantile_wrapper( + x, + q, + axis=None, + keepdims=False, + interpolation="linear", +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation not in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + return _handle_axis( + x, + q, + _quantile, + keepdims=keepdims, + axis=axis, + interpolation=interpolation, + ) def _handle_axis(a, q, fn, keepdims=False, axis=None, interpolation="nearest"): @@ -216,154 +216,39 @@ def _quantile(a, q, axis=None, interpolation="nearest"): return out.astype(ret_dtype) -def _compute_quantile_wrapper( - x, - q, - axis=None, - keepdims=False, - interpolation="linear", -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation not in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - return _handle_axis( - x, - q, - _quantile, - keepdims=keepdims, - axis=axis, - interpolation=interpolation, - ) +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + if len(axis) == 0: + raise ValueError("Axis can't be empty!") -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bfloat16", - "complex64", - "complex128", - ) - } - }, - backend_version, -) -def quantile( - a: paddle.Tensor, - q: Union[paddle.Tensor, float], - /, - *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: Optional[bool] = False, - interpolation: Optional[str] = "linear", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - x=a, - q=q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) - + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") -def corrcoef( - x: paddle.Tensor, - /, - *, - y: Optional[paddle.Tensor] = None, - rowvar: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis -def histogram( - a: paddle.Tensor, - /, - *, - bins: Optional[Union[int, paddle.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[paddle.Tensor] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[paddle.Tensor] = None, - density: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> Tuple[paddle.Tensor]: - if range is None: - min_range = 0 - max_range = 0 +def _validate_quantile(q): + if isinstance(q, float): + q = paddle.to_tensor(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False else: - min_range = range[0] - max_range = range[1] - return paddle.histogram(a, bins=bins, min=min_range, max=max_range) - - -def nanmedian( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - overwrite_input: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if dtype is None: - dtype = input.dtype - input = input.cast("float32") - paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - + if not (paddle.all(0 <= q) and paddle.all(q <= 1)): + return False + return True -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bool", - ) - } - }, - backend_version, -) -def unravel_index( - indices: paddle.Tensor, - shape: Tuple[int], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if indices.ndim == 0: - indices = indices.unsqueeze(0) - coord = [] - indices = indices - for dim in reversed(shape): - coord.append((indices % dim).astype("int32")) - indices = paddle.floor(indices / dim) - return tuple(reversed(coord)) +# --- Main --- # +# ------------ # @with_unsupported_device_and_dtypes( @@ -397,34 +282,15 @@ def bincount( ) -def igamma( - a: paddle.Tensor, +def corrcoef( + x: paddle.Tensor, /, *, - x: paddle.Tensor, + y: Optional[paddle.Tensor] = None, + rowvar: Optional[bool] = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - results = [] - ret_dtype = a.dtype if out is None else out.dtype - if paddle.float16 in [a.dtype, x.dtype]: - a = a.astype("float32") - x = x.astype("float32") - - for ai, xi in zip(a.flatten(), x.flatten()): - ai = ai.astype("float64") - xi = xi.astype("float64") - - def integrand(t): - return paddle.exp(-t) * paddle.pow(t, ai - 1) - - intervals = paddle.linspace(0, xi, 10001).astype("float64") - interval_width = xi / 10000 - values = integrand(intervals) - integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) - result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) - results.append(result) - - return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) + raise IvyNotImplementedException() def cov( @@ -536,88 +402,6 @@ def cummax( return ivy.flip(x, axis=axis), ivy.flip(indices, axis=axis) -def __find_cummax( - x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None -) -> Tuple[paddle.Tensor, paddle.Tensor]: - indices = [] - values = [] - x_dtype = x.dtype if dtype is None else dtype - if ( - isinstance(x.tolist()[0], list) - and len(x[0].shape) >= 1 - and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) - ): - if axis >= 1: - if not isinstance(x, list): - x = x.tolist() - for ret1 in x: - value, indice = __find_cummax( - paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype - ) - indices.append(indice) - values.append(value) - else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - else: - if not isinstance(x, list): - x = x.tolist() - n = 0 - for idx, y in enumerate(x): - if x[n] > y: - values.append(x[n]) - elif x[n] <= y or idx == 0: - n = idx - values.append(y) - indices.append(n) - - if isinstance(x, paddle.Tensor): - return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( - indices, dtype="int64" - ) - else: - return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16")}}, backend_version, @@ -645,3 +429,227 @@ def cummin( if reverse: cummin_x = paddle.flip(cummin_x, axis=[axis]) return cummin_x.cast(dtype) + + +def histogram( + a: paddle.Tensor, + /, + *, + bins: Optional[Union[int, paddle.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[paddle.Tensor] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[paddle.Tensor] = None, + density: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> Tuple[paddle.Tensor]: + if range is None: + min_range = 0 + max_range = 0 + else: + min_range = range[0] + max_range = range[1] + return paddle.histogram(a, bins=bins, min=min_range, max=max_range) + + +def igamma( + a: paddle.Tensor, + /, + *, + x: paddle.Tensor, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + results = [] + ret_dtype = a.dtype if out is None else out.dtype + if paddle.float16 in [a.dtype, x.dtype]: + a = a.astype("float32") + x = x.astype("float32") + + for ai, xi in zip(a.flatten(), x.flatten()): + ai = ai.astype("float64") + xi = xi.astype("float64") + + def integrand(t): + return paddle.exp(-t) * paddle.pow(t, ai - 1) + + intervals = paddle.linspace(0, xi, 10001).astype("float64") + interval_width = xi / 10000 + values = integrand(intervals) + integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) + result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) + results.append(result) + + return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def median( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # keepdims is set to True because in versions up to 2.5.1 + # there was a problem when the axis was defined and it was the + # only axis in the tensor so it needs to be handled manually + + ret_dtype = input.dtype + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(input): + ret = paddle.complex( + paddle.median(input.real(), axis=axis, keepdim=True), + paddle.median(input.imag(), axis=axis, keepdim=True), + ) + else: + ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) + else: + ret = paddle.median(input, axis=axis, keepdim=True) + if not keepdims: + ret = paddle_backend.squeeze(ret, axis=axis) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == input.ndim: + axis = None + if (input.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + +def nanmean( + a: paddle.Tensor, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + ret_dtype = dtype if dtype is not None else a.dtype + a = a.cast( + ret_dtype + ) # this is necessary to match other FWs behaviour which cast before calculation + if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(a): + ret = paddle.complex( + paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), + paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), + ) + else: + ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) + else: + ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) + + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == a.ndim: + axis = None + if (a.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + +def nanmedian( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + overwrite_input: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if dtype is None: + dtype = input.dtype + input = input.cast("float32") + paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bfloat16", + "complex64", + "complex128", + ) + } + }, + backend_version, +) +def quantile( + a: paddle.Tensor, + q: Union[paddle.Tensor, float], + /, + *, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: Optional[bool] = False, + interpolation: Optional[str] = "linear", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + x=a, + q=q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bool", + ) + } + }, + backend_version, +) +def unravel_index( + indices: paddle.Tensor, + shape: Tuple[int], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if indices.ndim == 0: + indices = indices.unsqueeze(0) + coord = [] + indices = indices + for dim in reversed(shape): + coord.append((indices % dim).astype("int32")) + indices = paddle.floor(indices / dim) + + return tuple(reversed(coord)) diff --git a/ivy/functional/backends/paddle/general.py b/ivy/functional/backends/paddle/general.py index f4258c4f452b9..31ad5024f28f7 100644 --- a/ivy/functional/backends/paddle/general.py +++ b/ivy/functional/backends/paddle/general.py @@ -13,24 +13,8 @@ from ivy.utils.exceptions import _check_inplace_update_support -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, paddle.Tensor): - if exclusive and not x.stop_gradient: - return False - return True - return False - - -def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: - return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) - - -def container_types(): - return [] - - -def current_backend_str() -> str: - return "paddle" +# --- Helpers --- # +# --------------- # def _check_query(query): @@ -45,59 +29,20 @@ def _check_query(query): ) -def get_item( - x: paddle.Tensor, - /, - query: Union[paddle.Tensor, Tuple], - *, - copy: bool = None, -) -> paddle.Tensor: - dtype = x.dtype - if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: - ret = x.cast("float32").__getitem__(query).cast(dtype) - elif dtype in [paddle.complex64, paddle.complex128]: - ret = paddle.complex( - x.real().__getitem__(query), - x.imag().__getitem__(query), - ) - else: - ret = x.__getitem__(query) - return ret - - -get_item.partial_mixed_handler = ( - lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape -) +# --- Main --- # +# ------------ # -def to_numpy( - x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif paddle.is_tensor(x): - if copy: - return np.array(x) - else: - return np.asarray(x) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") +def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: + return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) -def to_scalar(x: paddle.Tensor, /) -> Number: - if isinstance(x, (Number, complex)): - return x - return x.item() +def container_types(): + return [] -def to_list(x: paddle.Tensor, /) -> list: - return x.tolist() +def current_backend_str() -> str: + return "paddle" def gather( @@ -287,6 +232,26 @@ def gather_nd( return out +def get_item( + x: paddle.Tensor, + /, + query: Union[paddle.Tensor, Tuple], + *, + copy: bool = None, +) -> paddle.Tensor: + dtype = x.dtype + if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: + ret = x.cast("float32").__getitem__(query).cast(dtype) + elif dtype in [paddle.complex64, paddle.complex128]: + ret = paddle.complex( + x.real().__getitem__(query), + x.imag().__getitem__(query), + ) + else: + ret = x.__getitem__(query) + return ret + + def get_num_dims( x: paddle.Tensor, /, *, as_array: bool = False ) -> Union[paddle.Tensor, int]: @@ -356,6 +321,48 @@ def inplace_variables_supported(): return False +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, paddle.Tensor): + if exclusive and not x.stop_gradient: + return False + return True + return False + + +def isin( + elements: paddle.Tensor, + test_elements: paddle.Tensor, + /, + *, + assume_unique: Optional[bool] = False, + invert: Optional[bool] = False, +) -> paddle.Tensor: + input_shape = elements.shape + if elements.ndim == 0: + elements = paddle_backend.expand_dims(elements, axis=0) + if test_elements.ndim == 0: + test_elements = paddle_backend.expand_dims(test_elements, axis=0) + if not assume_unique: + test_elements = paddle_backend.unique_values(test_elements) + + elements = elements.reshape([-1]) + test_elements = test_elements.reshape([-1]) + + output = paddle_backend.any( + paddle_backend.equal( + paddle_backend.expand_dims(elements, axis=-1), test_elements + ), + axis=-1, + ) + return paddle_backend.logical_xor( + paddle_backend.reshape(output, input_shape), invert + ) + + +def itemsize(x: paddle.Tensor) -> int: + return x.element_size() + + def multiprocessing(context=None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -512,6 +519,36 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: paddle.Tensor, /) -> list: + return x.tolist() + + +def to_numpy( + x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif paddle.is_tensor(x): + if copy: + return np.array(x) + else: + return np.asarray(x) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") + + +def to_scalar(x: paddle.Tensor, /) -> Number: + if isinstance(x, (Number, complex)): + return x + return x.item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -595,35 +632,6 @@ def _vmap(*args, **kwargs): return _vmap -def isin( - elements: paddle.Tensor, - test_elements: paddle.Tensor, - /, - *, - assume_unique: Optional[bool] = False, - invert: Optional[bool] = False, -) -> paddle.Tensor: - input_shape = elements.shape - if elements.ndim == 0: - elements = paddle_backend.expand_dims(elements, axis=0) - if test_elements.ndim == 0: - test_elements = paddle_backend.expand_dims(test_elements, axis=0) - if not assume_unique: - test_elements = paddle_backend.unique_values(test_elements) - - elements = elements.reshape([-1]) - test_elements = test_elements.reshape([-1]) - - output = paddle_backend.any( - paddle_backend.equal( - paddle_backend.expand_dims(elements, axis=-1), test_elements - ), - axis=-1, - ) - return paddle_backend.logical_xor( - paddle_backend.reshape(output, input_shape), invert - ) - - -def itemsize(x: paddle.Tensor) -> int: - return x.element_size() +get_item.partial_mixed_handler = ( + lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape +) diff --git a/ivy/functional/backends/paddle/gradients.py b/ivy/functional/backends/paddle/gradients.py index 9faf1f8d2dca0..f7a7fee303be4 100644 --- a/ivy/functional/backends/paddle/gradients.py +++ b/ivy/functional/backends/paddle/gradients.py @@ -20,24 +20,45 @@ ) -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = x.astype(ivy.default_float_dtype()) - if not x.is_leaf: - ret = x.detach() - ret.stop_gradient = False +# --- Helpers --- # +# --------------- # + + +def _get_jac_one_arg_fn(grad_fn, xs, out_idx): + nested_indices = iter(ivy.all_nested_indices(xs)) + + def one_arg_fn(x): + idx = next(nested_indices) + new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x + ret = grad_fn(new_xs) + for i in out_idx: + ret = ret[i] return ret - ret = paddle_backend.copy_array(x) - ret.stop_gradient = False - return ret + return one_arg_fn -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, paddle.Tensor) and not x.stop_gradient +def _get_one_out_fn(grad_fn, xs, fn_ret): + out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) -def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: - return x.value() + def one_out_fn(o): + out_idx = next(out_nested_indices) + out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape + one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) + jacobian = ivy.nested_map( + xs, + lambda x: jacobian_to_ivy( + paddle.incubate.autograd.Jacobian( + one_arg_fn, ivy.to_native(x.expand_dims()) + ), + x.shape, + out_shape, + ), + shallow=False, + ) + return jacobian + + return one_out_fn def _grad_func(y, xs, retain_grads): @@ -103,6 +124,10 @@ def grad_(x): return grads +# --- Main --- # +# ------------ # + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -159,100 +184,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = paddle.grad(y, x, allow_unused=True)[0] - grad = grad if grad is not None else paddle.zeros_like(x) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def stop_gradient( - x: Optional[paddle.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[paddle.Tensor] = None, -): - is_var = is_variable(x) - x.stop_gradient = True - if is_var and preserve_type: - return variable(x) - return x - - -def _get_jac_one_arg_fn(grad_fn, xs, out_idx): - nested_indices = iter(ivy.all_nested_indices(xs)) - - def one_arg_fn(x): - idx = next(nested_indices) - new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x - ret = grad_fn(new_xs) - for i in out_idx: - ret = ret[i] - return ret - - return one_arg_fn - - -def _get_one_out_fn(grad_fn, xs, fn_ret): - out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) - - def one_out_fn(o): - out_idx = next(out_nested_indices) - out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape - one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) - jacobian = ivy.nested_map( - xs, - lambda x: jacobian_to_ivy( - paddle.incubate.autograd.Jacobian( - one_arg_fn, ivy.to_native(x.expand_dims()) - ), - x.shape, - out_shape, - ), - shallow=False, - ) - return jacobian - - return one_out_fn - - -def jacobian_to_ivy(jacobian, in_shape, out_shape): - jac_ivy = ivy.to_ivy(jacobian[:]) - jac_shape = out_shape + in_shape - jac_reshaped = jac_ivy.reshape(jac_shape) - return jac_reshaped - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(xs): - fn_ret = grad_fn(xs) - one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) - jacobian = ivy.nested_map(fn_ret, one_out_fn) - return jacobian - - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -290,5 +221,82 @@ def _inner(x): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, paddle.Tensor) and not x.stop_gradient + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(xs): + fn_ret = grad_fn(xs) + one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) + jacobian = ivy.nested_map(fn_ret, one_out_fn) + return jacobian + + return callback_fn + + +def jacobian_to_ivy(jacobian, in_shape, out_shape): + jac_ivy = ivy.to_ivy(jacobian[:]) + jac_shape = out_shape + in_shape + jac_reshaped = jac_ivy.reshape(jac_shape) + return jac_reshaped + + +def stop_gradient( + x: Optional[paddle.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[paddle.Tensor] = None, +): + is_var = is_variable(x) + x.stop_gradient = True + if is_var and preserve_type: + return variable(x) + return x + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = paddle.grad(y, x, allow_unused=True)[0] + grad = grad if grad is not None else paddle.zeros_like(x) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = x.astype(ivy.default_float_dtype()) + if not x.is_leaf: + ret = x.detach() + ret.stop_gradient = False + return ret + ret = paddle_backend.copy_array(x) + ret.stop_gradient = False + return ret + + +def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: + return x.value() + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/paddle/layers.py b/ivy/functional/backends/paddle/layers.py index f70c374bc70de..462d2b96f23e8 100644 --- a/ivy/functional/backends/paddle/layers.py +++ b/ivy/functional/backends/paddle/layers.py @@ -19,8 +19,8 @@ from . import backend_version -def _is_list_or_tuple(inp): - return isinstance(inp, (list, tuple)) +# --- Helpers --- # +# --------------- # def _convert_to_list(value, n, name="padding", _type=int): @@ -37,6 +37,27 @@ def _convert_to_list(value, n, name="padding", _type=int): return value_list +def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): + if filter_format == "channel_first": + filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) + + # adding dilation in input + x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations + for i in range(dims): + if x_dilations[i] > 1: + h = x.shape[1 + i] + new_height = h + (h - 1) * (x_dilations[i] - 1) + h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] + x = paddle_backend.swapaxes(x, 1 + i, -1) + x = paddle.matmul(x, h) + x = paddle_backend.swapaxes(x, -1, 1 + i) + return x, filters + + +def _is_list_or_tuple(inp): + return isinstance(inp, (list, tuple)) + + def _pad_before_conv(x, filters, strides, padding, dims, dilations, data_format): dilations = _convert_to_list(dilations, dims, "dilations") strides = _convert_to_list(strides, dims, "strides") @@ -122,21 +143,8 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding -def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): - if filter_format == "channel_first": - filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) - - # adding dilation in input - x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations - for i in range(dims): - if x_dilations[i] > 1: - h = x.shape[1 + i] - new_height = h + (h - 1) * (x_dilations[i] - 1) - h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] - x = paddle_backend.swapaxes(x, 1 + i, -1) - x = paddle.matmul(x, h) - x = paddle_backend.swapaxes(x, -1, 1 + i) - return x, filters +# --- Main --- # +# ------------ # def conv1d( @@ -259,21 +267,6 @@ def conv2d_transpose( return res -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: paddle.Tensor, - filters: paddle.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: Optional[str] = "NHWC", - dilations: Optional[Union[int, Tuple[int, int]]] = 1, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, @@ -494,3 +487,18 @@ def conv_general_transpose( if data_format == "channel_last": res = res.transpose(0, *range(2, dims + 2), 1) return res + + +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: paddle.Tensor, + filters: paddle.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: Optional[str] = "NHWC", + dilations: Optional[Union[int, Tuple[int, int]]] = 1, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/paddle/linear_algebra.py b/ivy/functional/backends/paddle/linear_algebra.py index 8a99ab77110fc..61e2f5744c1dc 100644 --- a/ivy/functional/backends/paddle/linear_algebra.py +++ b/ivy/functional/backends/paddle/linear_algebra.py @@ -108,6 +108,33 @@ def det(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T return ret +# Extra # +# ----- # + + +def diag( + x: paddle.Tensor, + /, + *, + k: int = 0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) + ) + return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) + return paddle.diag(x, offset=k) + + def diagonal( x: paddle.Tensor, /, @@ -136,6 +163,16 @@ def diagonal( return paddle.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) +def eig( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> Tuple[paddle.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] + ) + eigenvalues, eigenvectors = paddle.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + def eigh( x: paddle.Tensor, /, @@ -318,16 +355,6 @@ def matrix_norm( return ret -def eig( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> Tuple[paddle.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] - ) - eigenvalues, eigenvectors = paddle.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -423,18 +450,6 @@ def pinv( return paddle.linalg.pinv(x, rcond=rtol) -def tensorsolve( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # Implemented as a composite function in ivy.functional.ivy.linear_algebra - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -560,6 +575,18 @@ def tensordot( return ret.squeeze().cast(ret_dtype) if x1.ndim == axes else ret.cast(ret_dtype) +def tensorsolve( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # Implemented as a composite function in ivy.functional.ivy.linear_algebra + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( { "2.5.1 and below": { @@ -588,6 +615,28 @@ def trace( return ret.squeeze() if x.ndim <= 2 else ret +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, + backend_version, +) +def vander( + x: paddle.Tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + return paddle.pow( + paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), + paddle.arange(start, stop, step, dtype=x.dtype), + ) + + def vecdot( x1: paddle.Tensor, x2: paddle.Tensor, @@ -633,55 +682,6 @@ def vector_norm( ) -# Extra # -# ----- # - - -def diag( - x: paddle.Tensor, - /, - *, - k: int = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) - ) - return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) - return paddle.diag(x, offset=k) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, - backend_version, -) -def vander( - x: paddle.Tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - return paddle.pow( - paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), - paddle.arange(start, stop, step, dtype=x.dtype), - ) - - @with_unsupported_dtypes( {"2.5.1 and below": ("unsigned", "int8", "int16", "float16")}, backend_version, diff --git a/ivy/functional/backends/paddle/manipulation.py b/ivy/functional/backends/paddle/manipulation.py index a52336b3c46a0..0902034d2362a 100644 --- a/ivy/functional/backends/paddle/manipulation.py +++ b/ivy/functional/backends/paddle/manipulation.py @@ -14,6 +14,33 @@ from ...ivy.manipulation import _calculate_out_shape +# --- Helpers --- # +# --------------- # + + +def _reshape_fortran_paddle(x, shape): + if len(x.shape) > 0: + x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) + return paddle_backend.permute_dims( + paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] + ) + + +# --- Main --- # +# ------------ # + + +def clip( + x: paddle.Tensor, + x_min: Union[Number, paddle.Tensor], + x_max: Union[Number, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) + + # Array API Standard # # -------------------# @@ -54,6 +81,35 @@ def concat( return ret +def constant_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + paddings = [] + pad_width = list(pad_width) + for item in pad_width: + if len(item) != 2: + raise ivy.utils.exceptions.IvyException("Length of each item should be 2") + else: + paddings.append(item[0]) + paddings.append(item[1]) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.nn.functional.pad( + x.cast("float32"), pad=paddings, value=value + ).cast(x.dtype) + return paddle.nn.functional.pad(x=x, pad=paddings, value=value) + + def expand_dims( x: paddle.Tensor, /, @@ -97,12 +153,50 @@ def permute_dims( return paddle.transpose(x, axes) -def _reshape_fortran_paddle(x, shape): - if len(x.shape) > 0: - x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) - return paddle_backend.permute_dims( - paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] - ) +def repeat( + x: paddle.Tensor, + /, + repeats: Union[int, Iterable[int]], + *, + axis: int = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # handle the case when repeats contains 0 as paddle doesn't support it + if (isinstance(repeats, Number) and repeats == 0) or ( + isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 + ): + if axis is None: + return paddle.to_tensor([], dtype=x.dtype) + else: + shape = x.shape + shape[axis] = 0 + return paddle.zeros(shape=shape).cast(x.dtype) + + if isinstance(repeats, paddle.Tensor) and repeats.size == 1: + repeats = repeats.item() + + if axis is not None: + axis = axis % x.ndim + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), + paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), + ) + + return paddle.repeat_interleave( + x.cast("float32"), repeats=repeats, axis=axis + ).cast(x.dtype) + + return paddle.repeat_interleave(x, repeats=repeats, axis=axis) def reshape( @@ -170,6 +264,62 @@ def roll( return paddle.roll(x, shift, axis) +# Extra # +# ------# + + +def split( + x: paddle.Tensor, + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, + axis: Optional[int] = 0, + with_remainder: Optional[bool] = False, +) -> List[paddle.Tensor]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + num_or_size_splits = x.shape[axis] + elif isinstance(num_or_size_splits, paddle.Tensor): + num_or_size_splits = num_or_size_splits.cast("int32") + num_or_size_splits = num_or_size_splits.tolist() + elif isinstance(num_or_size_splits, int): + num_chunks = x.shape[axis] // num_or_size_splits + remainder = x.shape[axis] % num_or_size_splits + if remainder != 0: + if with_remainder: + num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] + else: + raise ivy.utils.exceptions.IvyException( + "Split size is not compatible with input shape" + ) + + if isinstance(num_or_size_splits, (list, tuple)): + if sum(num_or_size_splits) < x.shape[axis]: + num_or_size_splits + type(num_or_size_splits)([-1]) + elif sum(num_or_size_splits) > x.shape[axis]: + raise ivy.utils.exceptions.IvyException( + "total split size is not compatible with input shape," + f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" + ) + + if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + imag_list = paddle.split(x.imag(), num_or_size_splits, axis) + real_list = paddle.split(x.real(), num_or_size_splits, axis) + return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] + ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) + return [tensor.cast(x.dtype) for tensor in ret] + return paddle.split(x, num_or_size_splits, axis) + + def squeeze( x: paddle.Tensor, /, @@ -241,106 +391,18 @@ def stack( return paddle.stack(arrays, axis=axis) -# Extra # -# ------# - - -def split( +def swapaxes( x: paddle.Tensor, + axis0: int, + axis1: int, /, *, copy: Optional[bool] = None, - num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, - axis: Optional[int] = 0, - with_remainder: Optional[bool] = False, -) -> List[paddle.Tensor]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - num_or_size_splits = x.shape[axis] - elif isinstance(num_or_size_splits, paddle.Tensor): - num_or_size_splits = num_or_size_splits.cast("int32") - num_or_size_splits = num_or_size_splits.tolist() - elif isinstance(num_or_size_splits, int): - num_chunks = x.shape[axis] // num_or_size_splits - remainder = x.shape[axis] % num_or_size_splits - if remainder != 0: - if with_remainder: - num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] - else: - raise ivy.utils.exceptions.IvyException( - "Split size is not compatible with input shape" - ) - - if isinstance(num_or_size_splits, (list, tuple)): - if sum(num_or_size_splits) < x.shape[axis]: - num_or_size_splits + type(num_or_size_splits)([-1]) - elif sum(num_or_size_splits) > x.shape[axis]: - raise ivy.utils.exceptions.IvyException( - "total split size is not compatible with input shape," - f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" - ) - - if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x): - imag_list = paddle.split(x.imag(), num_or_size_splits, axis) - real_list = paddle.split(x.real(), num_or_size_splits, axis) - return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] - ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) - return [tensor.cast(x.dtype) for tensor in ret] - return paddle.split(x, num_or_size_splits, axis) - - -def repeat( - x: paddle.Tensor, - /, - repeats: Union[int, Iterable[int]], - *, - axis: int = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - # handle the case when repeats contains 0 as paddle doesn't support it - if (isinstance(repeats, Number) and repeats == 0) or ( - isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 - ): - if axis is None: - return paddle.to_tensor([], dtype=x.dtype) - else: - shape = x.shape - shape[axis] = 0 - return paddle.zeros(shape=shape).cast(x.dtype) - - if isinstance(repeats, paddle.Tensor) and repeats.size == 1: - repeats = repeats.item() - - if axis is not None: - axis = axis % x.ndim - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), - paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), - ) - - return paddle.repeat_interleave( - x.cast("float32"), repeats=repeats, axis=axis - ).cast(x.dtype) - - return paddle.repeat_interleave(x, repeats=repeats, axis=axis) + axes = [x for x in range(x.ndim)] + axes[axis0], axes[axis1] = axes[axis1], axes[axis0] + return paddle_backend.permute_dims(x, axes) def tile( @@ -383,70 +445,6 @@ def tile( return paddle.tile(x, repeats) -def constant_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - paddings = [] - pad_width = list(pad_width) - for item in pad_width: - if len(item) != 2: - raise ivy.utils.exceptions.IvyException("Length of each item should be 2") - else: - paddings.append(item[0]) - paddings.append(item[1]) - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.nn.functional.pad( - x.cast("float32"), pad=paddings, value=value - ).cast(x.dtype) - return paddle.nn.functional.pad(x=x, pad=paddings, value=value) - - -def zero_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[paddle.Tensor] = None, -): - return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) - - -def swapaxes( - x: paddle.Tensor, - axis0: int, - axis1: int, - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axes = [x for x in range(x.ndim)] - axes[axis0], axes[axis1] = axes[axis1], axes[axis0] - return paddle_backend.permute_dims(x, axes) - - -def clip( - x: paddle.Tensor, - x_min: Union[Number, paddle.Tensor], - x_max: Union[Number, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) - - def unstack( x: paddle.Tensor, /, @@ -482,3 +480,13 @@ def unstack( if keepdims: return [paddle_backend.expand_dims(r, axis=axis) for r in ret] return ret + + +def zero_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[paddle.Tensor] = None, +): + return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) diff --git a/ivy/functional/backends/paddle/random.py b/ivy/functional/backends/paddle/random.py index 3ca640ccf531b..85d02e5b69146 100644 --- a/ivy/functional/backends/paddle/random.py +++ b/ivy/functional/backends/paddle/random.py @@ -19,39 +19,69 @@ ) from . import backend_version -# Extra # -# ------# + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "float32", + "float64", + ) + } + }, + backend_version, +) +def multinomial( + population_size: int, + num_samples: int, + /, + *, + batch_size: int = 1, + probs: Optional[paddle.Tensor] = None, + replace: bool = True, + device: Place, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if probs is None: + probs = paddle.ones((batch_size, num_samples)) / population_size + probs = paddle.cast(probs, paddle.float32) + if seed: + paddle.seed(seed) + x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) + return x @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def random_uniform( +def randint( + low: Union[int, paddle.Tensor], + high: Union[int, paddle.Tensor], + /, *, - low: Union[float, paddle.Tensor] = 0.0, - high: Union[float, paddle.Tensor] = 1.0, - shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: paddle.dtype, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, device: Place, - seed=None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + seed: Optional[int] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) + _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - # Set range and seed - rng = high - low + range = high - low if seed: _ = paddle.seed(seed) - random_base = paddle.uniform(shape, min=0.0, max=1.0) - return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( - dtype + _retval = paddle.cast( + paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype ) + return _retval if shape else _retval.squeeze(axis=0) @with_unsupported_device_and_dtypes( @@ -80,68 +110,39 @@ def random_normal( return paddle.normal(mean, std).cast(dtype) -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "float32", - "float64", - ) - } - }, - backend_version, -) -def multinomial( - population_size: int, - num_samples: int, - /, - *, - batch_size: int = 1, - probs: Optional[paddle.Tensor] = None, - replace: bool = True, - device: Place, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if probs is None: - probs = paddle.ones((batch_size, num_samples)) / population_size - probs = paddle.cast(probs, paddle.float32) - if seed: - paddle.seed(seed) - x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) - return x +# Extra # +# ------# @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def randint( - low: Union[int, paddle.Tensor], - high: Union[int, paddle.Tensor], - /, +def random_uniform( *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + low: Union[float, paddle.Tensor] = 0.0, + high: Union[float, paddle.Tensor] = 1.0, + shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: paddle.dtype, device: Place, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - seed: Optional[int] = None, + seed=None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) - _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - range = high - low + # Set range and seed + rng = high - low if seed: _ = paddle.seed(seed) + random_base = paddle.uniform(shape, min=0.0, max=1.0) - _retval = paddle.cast( - paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype + return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( + dtype ) - return _retval if shape else _retval.squeeze(axis=0) def seed(*, seed_value: int = 0) -> None: diff --git a/ivy/functional/backends/paddle/searching.py b/ivy/functional/backends/paddle/searching.py index 9519e03c23916..36fd90f771b40 100644 --- a/ivy/functional/backends/paddle/searching.py +++ b/ivy/functional/backends/paddle/searching.py @@ -94,6 +94,31 @@ def argmin( return ret.astype(dtype) +# Extra # +# ----- # + + +def argwhere( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.ndim == 0: + return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") + if x.dtype in [ + paddle.int8, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + ]: + if paddle.is_complex(x): + real_idx = paddle.nonzero(x.real()) + imag_idx = paddle.nonzero(x.imag()) + idx = paddle.concat([real_idx, imag_idx], axis=0) + return paddle.unique(idx, axis=0) + return paddle.nonzero(x.cast("float32")) + return paddle.nonzero(x) + + def nonzero( x: paddle.Tensor, /, @@ -175,28 +200,3 @@ def where( result = paddle.where(condition, x1, x2) return result.squeeze().cast(ret_dtype) if scalar_out else result.cast(ret_dtype) - - -# Extra # -# ----- # - - -def argwhere( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.ndim == 0: - return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") - if x.dtype in [ - paddle.int8, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - ]: - if paddle.is_complex(x): - real_idx = paddle.nonzero(x.real()) - imag_idx = paddle.nonzero(x.imag()) - idx = paddle.concat([real_idx, imag_idx], axis=0) - return paddle.unique(idx, axis=0) - return paddle.nonzero(x.cast("float32")) - return paddle.nonzero(x) diff --git a/ivy/functional/backends/paddle/sorting.py b/ivy/functional/backends/paddle/sorting.py index 6a1d8f38f10d8..b12e326d19c8c 100644 --- a/ivy/functional/backends/paddle/sorting.py +++ b/ivy/functional/backends/paddle/sorting.py @@ -33,29 +33,13 @@ def argsort( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, backend_version, ) -def sort( - x: paddle.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[paddle.Tensor] = None, +def msort( + a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( - x.dtype - ) - return paddle.sort(x, axis=axis, descending=descending) + return paddle.sort(a, axis=0) @with_unsupported_device_and_dtypes( @@ -115,10 +99,26 @@ def searchsorted( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def msort( - a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None +def sort( + x: paddle.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.sort(a, axis=0) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( + x.dtype + ) + return paddle.sort(x, axis=axis, descending=descending) diff --git a/ivy/functional/backends/paddle/statistical.py b/ivy/functional/backends/paddle/statistical.py index 4d0297b1b7aac..b03ff9d943fcd 100644 --- a/ivy/functional/backends/paddle/statistical.py +++ b/ivy/functional/backends/paddle/statistical.py @@ -1,7 +1,3 @@ -# global - -torch_scatter = None - from typing import Union, Optional, Sequence import paddle @@ -16,45 +12,163 @@ # local from . import backend_version -# Array API Standard # -# -------------------# +# global +torch_scatter = None -def min( + +# --- Helpers --- # +# --------------- # + + +def _std(x, axis, correction, keepdim): + u = paddle_backend.mean(x, axis=axis, keepdims=True) + out = paddle_backend.sum( + paddle_backend.pow(paddle_backend.subtract(x, u), 2), + axis=axis, + keepdims=keepdim, + ) + num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() + num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() + n = num_elm_out / num_elm_in + out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) + if correction: + n = paddle_backend.sqrt( + paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) + ) + out = paddle_backend.multiply(out, n) + return out + + +# --- Main --- # +# ------------ # + + +# Extra # +# ----- # +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") + } + }, + backend_version, +) +def cumprod( x: paddle.Tensor, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - ret_dtype = x.dtype - if x.dtype in [ + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ + paddle.uint8, paddle.int8, paddle.int16, + paddle.float16, + ]: + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumprod(x, dim=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumprod(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) + else: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def cumsum( + x: paddle.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ paddle.uint8, + paddle.int8, paddle.float16, - paddle.bfloat16, - paddle.complex64, - paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x): - real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) - imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) - ret = paddle.complex(real, imag) - else: - ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumsum(x, axis=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumsum(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) else: - ret = paddle.amin(x, axis=axis, keepdim=keepdims) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == x.ndim: - axis = None - if (x.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +def einsum( + equation: str, + *operands: paddle.Tensor, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() def max( @@ -128,6 +242,47 @@ def mean( return ret.astype(ret_dtype) +# Array API Standard # +# -------------------# + + +def min( + x: paddle.Tensor, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + ret_dtype = x.dtype + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bfloat16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) + imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) + ret = paddle.complex(real, imag) + else: + ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) + else: + ret = paddle.amin(x, axis=axis, keepdim=keepdims) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == x.ndim: + axis = None + if (x.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + def prod( x: paddle.Tensor, /, @@ -150,25 +305,6 @@ def prod( return ret -def _std(x, axis, correction, keepdim): - u = paddle_backend.mean(x, axis=axis, keepdims=True) - out = paddle_backend.sum( - paddle_backend.pow(paddle_backend.subtract(x, u), 2), - axis=axis, - keepdims=keepdim, - ) - num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() - num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() - n = num_elm_out / num_elm_in - out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) - if correction: - n = paddle_backend.sqrt( - paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) - ) - out = paddle_backend.multiply(out, n) - return out - - def std( x: paddle.Tensor, /, @@ -217,130 +353,3 @@ def var( ) -> paddle.Tensor: ret = paddle_backend.pow(_std(x, axis, correction, keepdims), 2).cast(x.dtype) return ret - - -# Extra # -# ----- # -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") - } - }, - backend_version, -) -def cumprod( - x: paddle.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ - paddle.uint8, - paddle.int8, - paddle.int16, - paddle.float16, - ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumprod(x, dim=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumprod(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) - else: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def cumsum( - x: paddle.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ - paddle.uint8, - paddle.int8, - paddle.float16, - paddle.bool, - ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumsum(x, axis=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumsum(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) - else: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -def einsum( - equation: str, - *operands: paddle.Tensor, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/tensorflow/activations.py b/ivy/functional/backends/tensorflow/activations.py index de43f95f9a8f7..719d6b8d6cccf 100644 --- a/ivy/functional/backends/tensorflow/activations.py +++ b/ivy/functional/backends/tensorflow/activations.py @@ -30,6 +30,11 @@ def gelu( return tf.nn.gelu(x, approximate) +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return x * tf.nn.relu6(x + 3) / 6 + + def leaky_relu( x: Tensor, /, @@ -41,6 +46,23 @@ def leaky_relu( return tf.nn.leaky_relu(x, alpha) +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def log_softmax( + x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None +): + return tf.nn.log_softmax(x, axis) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def mish( + x: Tensor, + /, + *, + out: Optional[Tensor] = None, +) -> Tensor: + return x * tf.math.tanh(tf.math.softplus(x)) + + def relu(x: Tensor, /, *, complex_mode="jax", out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu(x) @@ -89,25 +111,3 @@ def softplus( if threshold is not None: return tf.where(x_beta > threshold, x, res) return res - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def log_softmax( - x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None -): - return tf.nn.log_softmax(x, axis) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def mish( - x: Tensor, - /, - *, - out: Optional[Tensor] = None, -) -> Tensor: - return x * tf.math.tanh(tf.math.softplus(x)) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return x * tf.nn.relu6(x + 3) / 6 diff --git a/ivy/functional/backends/tensorflow/control_flow_ops.py b/ivy/functional/backends/tensorflow/control_flow_ops.py index 9b6f42c456d75..1f600eebbe861 100644 --- a/ivy/functional/backends/tensorflow/control_flow_ops.py +++ b/ivy/functional/backends/tensorflow/control_flow_ops.py @@ -1,42 +1,20 @@ import tensorflow as tf -# def if_exp(cond, if_true, if_false, expr_repr): -# def true_fn(): -# return if_true() -# -# def false_fn(): -# return if_false() -# -# return tf.cond(cond, true_fn, false_fn) +# --- Helpers --- # +# --------------- # -def if_else(cond, body_fn, orelse_fn, vars): - # back-compatibility - if isinstance(cond, bool): - v = cond - cond = lambda *_: v - cond = bool(cond(*vars)) - # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) - # use pythonic placeholder until the graph compiler supports callable arguments - - if cond: - return body_fn(*vars) - else: - return orelse_fn(*vars) +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) -def while_loop(test_fn, body_fn, vars): - def body_fn_wrapper(*loop_vars): - return body_fn(*loop_vars) - - def test_fn_wrapper(*loop_vars): - return test_fn(*loop_vars) +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} - if not vars: - vars = (0,) - return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) +# --- Main --- # +# ------------ # def for_loop( @@ -70,9 +48,40 @@ def empty_function(*args): return _dict_to_tuple(vars_dict) -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} +# def if_exp(cond, if_true, if_false, expr_repr): +# def true_fn(): +# return if_true() +# +# def false_fn(): +# return if_false() +# +# return tf.cond(cond, true_fn, false_fn) -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) +def if_else(cond, body_fn, orelse_fn, vars): + # back-compatibility + if isinstance(cond, bool): + v = cond + cond = lambda *_: v + cond = bool(cond(*vars)) + # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) + + # use pythonic placeholder until the graph compiler supports callable arguments + + if cond: + return body_fn(*vars) + else: + return orelse_fn(*vars) + + +def while_loop(test_fn, body_fn, vars): + def body_fn_wrapper(*loop_vars): + return body_fn(*loop_vars) + + def test_fn_wrapper(*loop_vars): + return test_fn(*loop_vars) + + if not vars: + vars = (0,) + + return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) diff --git a/ivy/functional/backends/tensorflow/creation.py b/ivy/functional/backends/tensorflow/creation.py index 25ce085eb766d..9db7dc926d4a9 100644 --- a/ivy/functional/backends/tensorflow/creation.py +++ b/ivy/functional/backends/tensorflow/creation.py @@ -20,6 +20,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +# --- Main --- # +# ------------ # + + # Array API Standard # # -------------------# @@ -97,6 +109,17 @@ def asarray( return tf.identity(ret) if copy else ret +def copy_array( + x: Union[tf.Tensor, tf.Variable], + *, + to_ivy_array: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if to_ivy_array: + return ivy.to_ivy(tf.identity(x)) + return tf.identity(x) + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -181,6 +204,27 @@ def from_dlpack( return tf.experimental.dlpack.from_dlpack(dlcapsule) +@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) +def frombuffer( + buffer: bytes, + dtype: Optional[tf.DType] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(buffer, bytearray): + buffer = bytes(buffer) + ret = tf.io.decode_raw(buffer, dtype) + dtype = tf.dtypes.as_dtype(dtype) + if offset > 0: + offset = int(offset / dtype.size) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + + return ret + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -210,10 +254,6 @@ def full_like( return tf.experimental.numpy.full_like(x, fill_value, dtype=dtype) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - def linspace( start: Union[tf.Tensor, tf.Variable, float], stop: Union[tf.Tensor, tf.Variable, float], @@ -265,6 +305,23 @@ def meshgrid( return res +def one_hot( + indices: Union[tf.Tensor, tf.Variable], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[tf.DType] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.one_hot( + indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype + ) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -310,6 +367,29 @@ def triu( return tf.experimental.numpy.triu(x, k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[Union[tf.Tensor, tf.Variable]]: + n_cols = n_rows if n_cols is None else n_cols + + if n_rows < 0 or n_cols < 0: + n_rows, n_cols = 0, 0 + + ret = [[], []] + + for i in range(0, min(n_rows, n_cols - k), 1): + for j in range(max(0, k + i), n_cols, 1): + ret[0].append(i) + ret[1].append(j) + + return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -336,75 +416,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: Union[tf.Tensor, tf.Variable], - *, - to_ivy_array: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if to_ivy_array: - return ivy.to_ivy(tf.identity(x)) - return tf.identity(x) - - -def one_hot( - indices: Union[tf.Tensor, tf.Variable], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[tf.DType] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.one_hot( - indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype - ) - - -@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) -def frombuffer( - buffer: bytes, - dtype: Optional[tf.DType] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(buffer, bytearray): - buffer = bytes(buffer) - ret = tf.io.decode_raw(buffer, dtype) - dtype = tf.dtypes.as_dtype(dtype) - if offset > 0: - offset = int(offset / dtype.size) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - - return ret - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: str, -) -> Tuple[Union[tf.Tensor, tf.Variable]]: - n_cols = n_rows if n_cols is None else n_cols - - if n_rows < 0 or n_cols < 0: - n_rows, n_cols = 0, 0 - - ret = [[], []] - - for i in range(0, min(n_rows, n_cols - k), 1): - for j in range(max(0, k + i), n_cols, 1): - ret[0].append(i) - ret[1].append(j) - - return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) diff --git a/ivy/functional/backends/tensorflow/data_type.py b/ivy/functional/backends/tensorflow/data_type.py index 819a8fc381d75..d6100609a38d8 100644 --- a/ivy/functional/backends/tensorflow/data_type.py +++ b/ivy/functional/backends/tensorflow/data_type.py @@ -28,7 +28,6 @@ tf.complex128: "complex128", tf.bool: "bool", } - native_dtype_dict = { "int8": tf.int8, "int16": tf.int16, @@ -91,93 +90,6 @@ def __repr__(self): ) -# Array API Standard # -# -------------------# - - -def astype( - x: Union[tf.Tensor, tf.Variable], - dtype: Union[DType, str], - /, - *, - copy: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return tf.experimental.numpy.copy(x) if copy else x - return tf.cast(x, dtype) - - -def broadcast_arrays( - *arrays: Union[tf.Tensor, tf.Variable], -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(arrays) > 1: - try: - desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - if len(arrays) > 2: - for i in range(2, len(arrays)): - try: - desired_shape = tf.broadcast_dynamic_shape( - desired_shape, arrays[i].shape - ) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - else: - return [arrays[0]] - result = [] - for tensor in arrays: - result.append(tf.broadcast_to(tensor, desired_shape)) - - return result - - -def broadcast_to( - x: Union[tf.Tensor, tf.Variable], - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if tf.rank(x) > len(shape): - return tf.broadcast_to(tf.reshape(x, -1), shape) - return tf.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - if ivy.as_native_dtype(type) == tf.bfloat16: - return Finfo(Bfloat16Finfo()) - return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def result_type( - *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], -) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return tf.experimental.numpy.result_type(arrays_and_dtypes) - - result = tf.experimental.numpy.result_type( - arrays_and_dtypes[0], arrays_and_dtypes[1] - ) - for i in range(2, len(arrays_and_dtypes)): - result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -247,6 +159,62 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: Union[tf.Tensor, tf.Variable], + dtype: Union[DType, str], + /, + *, + copy: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return tf.experimental.numpy.copy(x) if copy else x + return tf.cast(x, dtype) + + +def broadcast_arrays( + *arrays: Union[tf.Tensor, tf.Variable], +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(arrays) > 1: + try: + desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + if len(arrays) > 2: + for i in range(2, len(arrays)): + try: + desired_shape = tf.broadcast_dynamic_shape( + desired_shape, arrays[i].shape + ) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + else: + return [arrays[0]] + result = [] + for tensor in arrays: + result.append(tf.broadcast_to(tensor, desired_shape)) + + return result + + +def broadcast_to( + x: Union[tf.Tensor, tf.Variable], + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if tf.rank(x) > len(shape): + return tf.broadcast_to(tf.reshape(x, -1), shape) + return tf.broadcast_to(x, shape) + + def dtype( x: Union[tf.Tensor, tf.Variable, np.ndarray], *, as_native: bool = False ) -> ivy.Dtype: @@ -269,6 +237,22 @@ def dtype_bits(dtype_in: Union[tf.DType, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + if ivy.as_native_dtype(type) == tf.bfloat16: + return Finfo(Bfloat16Finfo()) + return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) + + def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -278,6 +262,16 @@ def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: return False -# ToDo: -# 1. result_type: Add support for bfloat16 with int16 -# 2. can_cast : Add support for complex64, complex128 +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def result_type( + *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], +) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return tf.experimental.numpy.result_type(arrays_and_dtypes) + + result = tf.experimental.numpy.result_type( + arrays_and_dtypes[0], arrays_and_dtypes[1] + ) + for i in range(2, len(arrays_and_dtypes)): + result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/tensorflow/device.py b/ivy/functional/backends/tensorflow/device.py index d0872e5e2da9d..ca7c545ecba6e 100644 --- a/ivy/functional/backends/tensorflow/device.py +++ b/ivy/functional/backends/tensorflow/device.py @@ -4,9 +4,6 @@ Collection of TensorFlow general functions, wrapped to fit Ivy syntax and signature. """ - -# global -_round = round import tensorflow as tf from typing import Union, Optional @@ -14,6 +11,33 @@ import ivy from ivy.functional.ivy.device import Profiler as BaseProfiler +# global +_round = round + + +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._options = tf.profiler.experimental.ProfilerOptions( + host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 + ) + + def start(self): + tf.profiler.experimental.start(self._save_dir, options=self._options) + + def stop(self): + tf.profiler.experimental.stop() + + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() + + +# --- Helpers --- # +# --------------- # + def _same_device(dev_a, dev_b): if dev_a is None or dev_b is None: @@ -23,34 +47,8 @@ def _same_device(dev_a, dev_b): ) -def dev( - x: Union[tf.Tensor, tf.Variable], - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, str]: - dv = x.device - if as_native: - return dv - return as_ivy_dev(dv) - - -def to_device( - x: Union[tf.Tensor, tf.Variable], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if device is None: - return x - device = as_native_dev(device) - current_dev = dev(x) - if not _same_device(current_dev, device): - with tf.device("/" + device.upper()): - return tf.identity(x) - return x +# --- Main --- # +# ------------ # def as_ivy_dev(device: str, /): @@ -79,47 +77,57 @@ def clear_cached_mem_on_dev(device: str, /): return None -def num_gpus() -> int: - return len(tf.config.list_physical_devices("GPU")) +def dev( + x: Union[tf.Tensor, tf.Variable], + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, str]: + dv = x.device + if as_native: + return dv + return as_ivy_dev(dv) def gpu_is_available() -> bool: return len(tf.config.list_physical_devices("GPU")) > 0 -def tpu_is_available() -> bool: - try: - resolver = tf.distribute.cluster_resolver.TPUClusterResolver() - tf.config.experimental_connect_to_cluster(resolver) - tf.tpu.experimental.initialize_tpu_system(resolver) - tf.config.list_logical_devices("TPU") - tf.distribute.experimental.TPUStrategy(resolver) - return True - except ValueError: - return False - - def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): default_device = ivy.default_device(device_shifting_dev, as_native=True) with tf.device(default_device): return fn(*args, **kwargs) -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._options = tf.profiler.experimental.ProfilerOptions( - host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 - ) +def num_gpus() -> int: + return len(tf.config.list_physical_devices("GPU")) - def start(self): - tf.profiler.experimental.start(self._save_dir, options=self._options) - def stop(self): - tf.profiler.experimental.stop() +def to_device( + x: Union[tf.Tensor, tf.Variable], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if device is None: + return x + device = as_native_dev(device) + current_dev = dev(x) + if not _same_device(current_dev, device): + with tf.device("/" + device.upper()): + return tf.identity(x) + return x - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + try: + resolver = tf.distribute.cluster_resolver.TPUClusterResolver() + tf.config.experimental_connect_to_cluster(resolver) + tf.tpu.experimental.initialize_tpu_system(resolver) + tf.config.list_logical_devices("TPU") + tf.distribute.experimental.TPUStrategy(resolver) + return True + except ValueError: + return False diff --git a/ivy/functional/backends/tensorflow/elementwise.py b/ivy/functional/backends/tensorflow/elementwise.py index 1e060e8e8cae3..38c717e0b4ce2 100644 --- a/ivy/functional/backends/tensorflow/elementwise.py +++ b/ivy/functional/backends/tensorflow/elementwise.py @@ -10,6 +10,9 @@ from . import backend_version +gcd.support_native_out = False + + def abs( x: Union[float, tf.Tensor, tf.Variable], /, @@ -56,6 +59,32 @@ def add( return tf.add(x1, x2) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def angle( + input: Union[tf.Tensor, tf.Variable], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if deg: + return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) + else: + return tf.math.angle(input, name=None) + + def asin( x: Union[tf.Tensor, tf.Variable], /, @@ -218,6 +247,16 @@ def cosh( return tf.cosh(x) +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def deg2rad( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.deg2rad(x) + + def divide( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -245,6 +284,20 @@ def equal( return tf.math.equal(x1, x2) +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def erf( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.erf(x) + + def exp( x: Union[tf.Tensor, tf.Variable], /, @@ -313,6 +366,37 @@ def fmin( return ret +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, + backend_version, +) +def fmod( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + # tf.math.floormod returns wrong results + res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) + return tf.where(x1 < 0, -res, res) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def gcd( + x1: Union[tf.Tensor, tf.Variable, int, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + return tf.experimental.numpy.gcd(x1, x2) + + @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def greater( x1: Union[float, tf.Tensor, tf.Variable], @@ -337,6 +421,28 @@ def greater_equal( return tf.math.greater_equal(x1, x2) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def imag( + val: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.imag(val, name=None) + + def isfinite( x: Union[tf.Tensor, tf.Variable], /, @@ -387,6 +493,15 @@ def isnan( return tf.math.is_nan(x) +def isreal( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.isreal(x) + + @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) def lcm( x1: Union[tf.Tensor, tf.Variable], @@ -471,16 +586,6 @@ def logaddexp( return tf.experimental.numpy.logaddexp(x1, x2) -@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) -def real( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.real(x) - - @with_unsupported_dtypes( { "2.13.0 and below": ( @@ -555,6 +660,64 @@ def logical_xor( return tf.math.logical_xor(tf.cast(x1, tf.bool), tf.cast(x2, tf.bool)) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def maximum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.maximum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def minimum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.minimum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) + + def multiply( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -566,6 +729,31 @@ def multiply( return tf.math.multiply(x1, x2) +def nan_to_num( + x: Union[tf.Tensor, tf.Variable], + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + posinf = posinf if posinf is not None else x.dtype.max + neginf = neginf if neginf is not None else x.dtype.min + posinf = tf.constant(posinf, x.dtype) + neginf = tf.constant(neginf, x.dtype) + nan = tf.constant(nan, x.dtype) + ret = tf.where(tf.math.is_nan(x), nan, x) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) + if copy: + return ret + else: + x = ret + return x + + @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) @@ -623,6 +811,49 @@ def pow( return tf.experimental.numpy.power(x1, x2) +def rad2deg( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.rad2deg(x) + + +@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) +def real( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.real(x) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "int8", + "int16", + "int32", + "int64", + ) + }, + backend_version, +) +def reciprocal( + x: Union[float, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.reciprocal(x) + + @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def remainder( x1: Union[float, tf.Tensor, tf.Variable], @@ -783,234 +1014,3 @@ def trunc( else: ret = (tf.math.floor if ret >= 0 else tf.math.ceil)(ret) return ret - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def erf( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.erf(x) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def maximum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.maximum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def minimum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.minimum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "int8", - "int16", - "int32", - "int64", - ) - }, - backend_version, -) -def reciprocal( - x: Union[float, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.reciprocal(x) - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def deg2rad( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.deg2rad(x) - - -def rad2deg( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rad2deg(x) - - -def isreal( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isreal(x) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, - backend_version, -) -def fmod( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - # tf.math.floormod returns wrong results - res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) - return tf.where(x1 < 0, -res, res) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def gcd( - x1: Union[tf.Tensor, tf.Variable, int, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - return tf.experimental.numpy.gcd(x1, x2) - - -gcd.support_native_out = False - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def angle( - input: Union[tf.Tensor, tf.Variable], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if deg: - return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) - else: - return tf.math.angle(input, name=None) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def imag( - val: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.imag(val, name=None) - - -def nan_to_num( - x: Union[tf.Tensor, tf.Variable], - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - posinf = posinf if posinf is not None else x.dtype.max - neginf = neginf if neginf is not None else x.dtype.min - posinf = tf.constant(posinf, x.dtype) - neginf = tf.constant(neginf, x.dtype) - nan = tf.constant(nan, x.dtype) - ret = tf.where(tf.math.is_nan(x), nan, x) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) - if copy: - return ret - else: - x = ret - return x diff --git a/ivy/functional/backends/tensorflow/experimental/activations.py b/ivy/functional/backends/tensorflow/experimental/activations.py index f798d7b907d98..0f5126a6678ff 100644 --- a/ivy/functional/backends/tensorflow/experimental/activations.py +++ b/ivy/functional/backends/tensorflow/experimental/activations.py @@ -10,6 +10,15 @@ from . import backend_version +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: + alpha = tf.cast(alpha, x.dtype) + ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) + + def logit( x: Union[tf.Tensor, tf.Variable], /, @@ -25,16 +34,9 @@ def logit( return tf.cast(tf.math.log(x / (1 - x)), x_dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) -def thresholded_relu( - x: Tensor, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[Tensor] = None, -) -> Tensor: - threshold = tf.cast(threshold, x.dtype) - return tf.cast(tf.where(x > threshold, x, 0), x.dtype) +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return tf.math.log_sigmoid(input) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -42,11 +44,6 @@ def relu6(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu6(x) -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return tf.math.log_sigmoid(input) - - @with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) def selu(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: ret = tf.nn.selu(x) @@ -68,10 +65,13 @@ def silu( return ivy.astype(ret, x.dtype) -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: - alpha = tf.cast(alpha, x.dtype) - ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) +@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) +def thresholded_relu( + x: Tensor, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[Tensor] = None, +) -> Tensor: + threshold = tf.cast(threshold, x.dtype) + return tf.cast(tf.where(x > threshold, x, 0), x.dtype) diff --git a/ivy/functional/backends/tensorflow/experimental/creation.py b/ivy/functional/backends/tensorflow/experimental/creation.py index 96cd0c15b84be..4462e9d61977e 100644 --- a/ivy/functional/backends/tensorflow/experimental/creation.py +++ b/ivy/functional/backends/tensorflow/experimental/creation.py @@ -8,28 +8,40 @@ from .. import backend_version -# Array API Standard # -# -------------------# +def blackman_window( + size: int, + /, + *, + periodic: bool = True, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if size < 2: + return tf.ones([size], dtype=tf.result_type(size, 0.0)) + if periodic: + count = tf.arange(size) / size + else: + count = tf.linspace(start=0, stop=size, num=size) + return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( + 0.08 * tf.cos(2 * tf.pi * 2 * count) + ) -@with_unsupported_device_and_dtypes( - {"2.13.0 and below": {"cpu": ("bfloat16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, + +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if window_length < 2: - return tf.ones([window_length], dtype=dtype) - if periodic is False: - return tf.signal.kaiser_window(window_length, beta, dtype=dtype) + if size < 2: + return tf.ones([size], dtype=dtype) + if periodic: + return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] else: - return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] + return tf.signal.hann_window(size, periodic=False, dtype=dtype) def kaiser_bessel_derived_window( @@ -42,29 +54,28 @@ def kaiser_bessel_derived_window( return tf.signal.kaiser_bessel_derived_window(window_length, beta, dtype) -def vorbis_window( - window_length: Union[tf.Tensor, tf.Variable], - *, - dtype: tf.DType = tf.dtypes.float32, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) +# Array API Standard # +# -------------------# -def hann_window( - size: int, - /, - *, +@with_unsupported_device_and_dtypes( + {"2.13.0 and below": {"cpu": ("bfloat16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=dtype) - if periodic: - return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] + if window_length < 2: + return tf.ones([window_length], dtype=dtype) + if periodic is False: + return tf.signal.kaiser_window(window_length, beta, dtype=dtype) else: - return tf.signal.hann_window(size, periodic=False, dtype=dtype) + return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] def tril_indices( @@ -90,6 +101,20 @@ def tril_indices( return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) +@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) +def trilu( + x: Union[tf.Tensor, tf.Variable], + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if upper: + return tf.experimental.numpy.triu(x, k) + return tf.experimental.numpy.tril(x, k) + + def unsorted_segment_min( data: tf.Tensor, segment_ids: tf.Tensor, @@ -98,26 +123,6 @@ def unsorted_segment_min( return tf.math.unsorted_segment_min(data, segment_ids, num_segments) -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=tf.result_type(size, 0.0)) - if periodic: - count = tf.arange(size) / size - else: - count = tf.linspace(start=0, stop=size, num=size) - - return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( - 0.08 * tf.cos(2 * tf.pi * 2 * count) - ) - - def unsorted_segment_sum( data: tf.Tensor, segment_ids: tf.Tensor, @@ -126,15 +131,10 @@ def unsorted_segment_sum( return tf.math.unsorted_segment_sum(data, segment_ids, num_segments) -@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) -def trilu( - x: Union[tf.Tensor, tf.Variable], - /, +def vorbis_window( + window_length: Union[tf.Tensor, tf.Variable], *, - k: int = 0, - upper: bool = True, + dtype: tf.DType = tf.dtypes.float32, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if upper: - return tf.experimental.numpy.triu(x, k) - return tf.experimental.numpy.tril(x, k) + return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) diff --git a/ivy/functional/backends/tensorflow/experimental/elementwise.py b/ivy/functional/backends/tensorflow/experimental/elementwise.py index 96dba1c88a4dd..487ed04c624b4 100644 --- a/ivy/functional/backends/tensorflow/experimental/elementwise.py +++ b/ivy/functional/backends/tensorflow/experimental/elementwise.py @@ -11,61 +11,54 @@ from .. import backend_version -@with_supported_dtypes( - {"2.13.0 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.lgamma(x) +# --- Helpers --- # +# --------------- # -def sinc( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x = ivy.pi * x - return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim -@with_supported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version -) -def fmax( +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis + + +# --- Main --- # +# ------------ # + + +def allclose( x1: Union[tf.Tensor, tf.Variable], x2: Union[tf.Tensor, tf.Variable], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - x1 = tf.where(tf.math.is_nan(x1), x2, x1) - x2 = tf.where(tf.math.is_nan(x2), x1, x2) - return tf.experimental.numpy.maximum(x1, x2) +) -> bool: + return tf.experimental.numpy.allclose( + x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan + ) -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def float_power( - x1: Union[tf.Tensor, tf.Variable, float, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], +def conj( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): - out_dtype = tf.complex128 - else: - out_dtype = tf.float64 - return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) + return tf.math.conj(x) def copysign( @@ -103,67 +96,35 @@ def count_nonzero( ) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def nansum( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[tf.DType] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) - - -def isclose( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan - ) - - -def signbit( - x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def diff( + x: Union[tf.Tensor, tf.Variable, list, tuple], /, *, + n: int = 1, + axis: int = -1, + prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, + append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.signbit(x) + if n == 0: + return x + if prepend is not None: + x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) + if append is not None: + x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) + return tf.experimental.numpy.diff(x, n=n, axis=axis) -def hypot( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], +def digamma( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) - - -def allclose( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> bool: - return tf.experimental.numpy.allclose( - x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan - ) + return tf.math.digamma(x) @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) @@ -176,74 +137,56 @@ def fix( return tf.cast(tf.where(x > 0, tf.math.floor(x), tf.math.ceil(x)), x.dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) -def nextafter( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.nextafter(x1, x2) - - @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) -def diff( - x: Union[tf.Tensor, tf.Variable, list, tuple], +def float_power( + x1: Union[tf.Tensor, tf.Variable, float, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], /, *, - n: int = 1, - axis: int = -1, - prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, - append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if n == 0: - return x - if prepend is not None: - x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) - if append is not None: - x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) - return tf.experimental.numpy.diff(x, n=n, axis=axis) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): + out_dtype = tf.complex128 + else: + out_dtype = tf.float64 + return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) @with_supported_dtypes( - { - "2.13.0 and below": ( - "float32", - "float64", - ) - }, - backend_version, + {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version ) -def zeta( - x: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, tf.Variable], +def fmax( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.zeta(x, q) - - -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim + x1, x2 = promote_types_of_inputs(x1, x2) + x1 = tf.where(tf.math.is_nan(x1), x2, x1) + x2 = tf.where(tf.math.is_nan(x2), x1, x2) + return tf.experimental.numpy.maximum(x1, x2) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) +def frexp( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[ + Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] + ] = None, +) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: + e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) + e = tf.cast(e, x.dtype) + while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): + e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) + m = x / tf.math.pow(2, e) + e = tf.cast(e, tf.int32) + return m, e def gradient( @@ -428,36 +371,29 @@ def gradient( return outvals -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float16", - "float32", - "float64", - "complex64", - "complex128", - ) - }, - backend_version, -) -def xlogy( - x: Union[tf.Tensor, tf.Variable], - y: Union[tf.Tensor, tf.Variable], +def hypot( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x, y = promote_types_of_inputs(x, y) - return tf.math.xlogy(x, y) + return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) -def conj( - x: Union[tf.Tensor, tf.Variable], +def isclose( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.conj(x) + return tf.experimental.numpy.isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan + ) @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) @@ -481,22 +417,17 @@ def ldexp( return tf.cast(ret, out_dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) -def frexp( +@with_supported_dtypes( + {"2.13.0 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( x: Union[tf.Tensor, tf.Variable], /, *, - out: Optional[ - Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] - ] = None, -) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: - e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) - e = tf.cast(e, x.dtype) - while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): - e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) - m = x / tf.math.pow(2, e) - e = tf.cast(e, tf.int32) - return m, e + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.lgamma(x) def modf( @@ -508,10 +439,87 @@ def modf( return tf.math.modf(x) -def digamma( +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def nansum( x: Union[tf.Tensor, tf.Variable], /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[tf.DType] = None, + keepdims: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.digamma(x) + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) + + +@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) +def nextafter( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.nextafter(x1, x2) + + +def signbit( + x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.signbit(x) + + +def sinc( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x = ivy.pi * x + return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float16", + "float32", + "float64", + "complex64", + "complex128", + ) + }, + backend_version, +) +def xlogy( + x: Union[tf.Tensor, tf.Variable], + y: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x, y = promote_types_of_inputs(x, y) + return tf.math.xlogy(x, y) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float32", + "float64", + ) + }, + backend_version, +) +def zeta( + x: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.zeta(x, q) diff --git a/ivy/functional/backends/tensorflow/experimental/layers.py b/ivy/functional/backends/tensorflow/experimental/layers.py index 94c60963871b5..2a528a74cdf3f 100644 --- a/ivy/functional/backends/tensorflow/experimental/layers.py +++ b/ivy/functional/backends/tensorflow/experimental/layers.py @@ -20,6 +20,10 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode, _get_size +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # Determine depth pooling kernel, strides, depth_pooling = _depth_max_pooling_helper( @@ -30,229 +34,71 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def max_pool1d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) +def _fft2_helper(x, shape, axes): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) - if data_format == "NCW": - x = tf.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor ) - if not depth_pooling: - new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - padding = [(pad_w // 2, pad_w - pad_w // 2)] + shape = shape_initialization(shape, axes, x) - if ceil_mode: - padding[0] = _padding_ceil_mode( - x.shape[1], new_kernel[0], padding[0], strides[0] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + rank = rank_initialization(axes) - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - if depth_pooling: - res = tf.transpose(res, (0, 2, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCW": - return tf.transpose(res, (0, 2, 1)) - return res + perm = get_perm(input_rank_tensor, axes) + x = transpose_x(x, perm, perform_transpose) -def max_pool2d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + x = fft2_operations(x, rank) - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) + x = transpose_x(x, tf.argsort(perm), perform_transpose) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - if not depth_pooling: - new_kernel = [ - kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) - ] - if isinstance(padding, str): - pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) - padding = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] + return x - x_shape = x.shape[1:-1] - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) +def _fft2_norm( + x: Union[tf.Tensor, tf.Variable], + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", +): + n = tf.constant(s[0] * s[1], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.sqrt(n) + elif norm == "forward": + return x / n else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def max_pool3d( +def _fft_norm( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + dim: int, /, *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCDHW": - x = tf.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - if not depth_pooling: - x_shape = x.shape[1:-1] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - padding = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) + norm: str = "backward", +): + n = tf.constant(x.shape[dim], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) + elif norm == "forward": + return x / tf.cast(n, x.dtype) else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 4, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCDHW": - return tf.transpose(res, (0, 4, 1, 2, 3)) - return res + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): @@ -280,13 +126,123 @@ def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): return padding, pad_specific, c -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) -def avg_pool1d( +def _ifft_norm( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, + dim: int, + *, + norm: str = "backward", +): + n = x.shape[dim] + if norm == "backward": + return x + elif norm == "ortho": + return x * math.sqrt(n) + elif norm == "forward": + return x * n + else: + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + + +def _ifftn_helper(x, shape, axes, norm): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + norm_factor = norm_initialization(norm, shape, x) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = ifft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +def _rfftn_helper(x, shape, axes, norm): + x = rfft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + norm_factor = norm_initialization(norm, shape, x) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = rfft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +def _right_pad_or_crop(tensor, shape): + input_shape = tf.shape(tensor) + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + with tf.control_dependencies( + [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] + ): + shape = tf.identity(shape) + shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) + + pad_sizes = tf.math.maximum(shape - input_shape, 0) + pad_sizes = tf.expand_dims(pad_sizes, -1) + pad_sizes = tf.concat( + [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 + ) + tensor = tf.pad(tensor, pad_sizes, constant_values=0) + + crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) + tensor = tf.slice(tensor, crop_tensor, shape) + return tensor + + +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) +def avg_pool1d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, *, data_format: str = "NWC", count_include_pad: bool = False, @@ -547,31 +503,15 @@ def avg_pool3d( return res -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def pool( - x: Union[tf.Tensor, tf.Variable], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, - *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.nn.pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ) +def axes_initialization(shape, axes, input_shape, input_rank_tensor): + if axes is None: + axes = ( + tf.range(-tf.size(input_shape), 0) + if shape is None + else tf.range(-tf.size(shape), 0) + ) + axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) + return axes @with_supported_dtypes({"2.13.0 and below": ("float32", "float64")}, backend_version) @@ -599,112 +539,6 @@ def dct( return dct_out -def idct( - x: Union[tf.Tensor, tf.Variable], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - -def _fft_norm( - x: Union[tf.Tensor, tf.Variable], - dim: int, - /, - *, - norm: str = "backward", -): - n = tf.constant(x.shape[dim], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) - elif norm == "forward": - return x / tf.cast(n, x.dtype) - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -def _ifft_norm( - x: Union[tf.Tensor, tf.Variable], - dim: int, - *, - norm: str = "backward", -): - n = x.shape[dim] - if norm == "backward": - return x - elif norm == "ortho": - return x * math.sqrt(n) - elif norm == "forward": - return x * n - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft( - x: Union[tf.Tensor, tf.Variable], - dim: int, - /, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.shape[dim] != n: - s = list(x.shape) - if s[dim] > n: - index = [slice(None)] * len(s) - index[dim] = slice(0, n) - x = x[tuple(index)] - del index - else: - s[dim] = n - s[dim] - z = tf.zeros(s, x.dtype) - x = tf.concat([x, z], dim) - del s - operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" - if dim != -1 or dim != len(x.shape) - 1: - permute = [i for i in range(len(x.shape))] - permute[dim], permute[-1] = permute[-1], permute[dim] - x = tf.transpose(x, permute) - ret = tf.signal.fft(x, operation_name) - ret = tf.transpose(ret, permute) - del permute - else: - ret = tf.signal.fft(x, operation_name) - ret = _fft_norm(ret, dim, norm=norm) - return ret - - def dropout( x: Union[tf.Tensor, tf.Variable], prob: float, @@ -791,10 +625,181 @@ def dropout3d( return res -def ifft( - x: Union[tf.Tensor, tf.Variable], - dim: int, - *, +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def embedding( + weights: Union[tf.Tensor, tf.Variable], + indices: Union[tf.Tensor, tf.Variable], + /, + *, + max_norm: Optional[float] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft( + x: Union[tf.Tensor, tf.Variable], + dim: int, + /, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if x.shape[dim] != n: + s = list(x.shape) + if s[dim] > n: + index = [slice(None)] * len(s) + index[dim] = slice(0, n) + x = x[tuple(index)] + del index + else: + s[dim] = n - s[dim] + z = tf.zeros(s, x.dtype) + x = tf.concat([x, z], dim) + del s + operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" + if dim != -1 or dim != len(x.shape) - 1: + permute = [i for i in range(len(x.shape))] + permute[dim], permute[-1] = permute[-1], permute[dim] + x = tf.transpose(x, permute) + ret = tf.signal.fft(x, operation_name) + ret = tf.transpose(ret, permute) + del permute + else: + ret = tf.signal.fft(x, operation_name) + ret = _fft_norm(ret, dim, norm=norm) + return ret + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft2( + x: Union[tf.Tensor, tf.Variable], + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if len(x.shape) > 2: + result = _fft2_helper(x, s, dim) + else: + x_new = trans_x_to_s(x, s, dim) + x_complex = tf.cast(x_new, tf.complex128) + result = tf.signal.fft2d(x_complex) + + result = _fft2_norm(result, s, dim, norm) + if x.dtype == tf.complex64: + result = tf.cast(result, dtype=tf.complex128) + return result + + +def fft2_operations(x, rank): + if x.shape.rank == 1: + x = tf.signal.fft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} + ) + else: + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.fft(x), + 1: lambda: tf.signal.fft2d(x), + 2: lambda: tf.signal.fft3d(x), + }, + ) + return x + + +# --- IFFTN --- # +def fft_input_validation(x): + if not x.dtype.is_complex: + raise TypeError( + "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( + x.dtype + ) + ) + return x + + +def get_perm(input_rank_tensor, axes): + all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) + perm = tf.concat( + [ + tf.boolean_mask( + all_dims, + tf.foldl( + lambda acc, elem: tf.math.logical_and( + acc, tf.math.not_equal(all_dims, elem) + ), + axes, + initializer=tf.fill(all_dims.shape, True), + ), + ), + axes, + ], + 0, + ) + return perm + + +def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): + if perform_padding: + pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) + pad_shape = tf.tensor_scatter_nd_update( + pad_shape, tf.expand_dims(axes, -1), shape + ) + x = _right_pad_or_crop(x, pad_shape) + return x + + +def idct( + x: Union[tf.Tensor, tf.Variable], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> tf.Tensor: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def ifft( + x: Union[tf.Tensor, tf.Variable], + dim: int, + *, norm: str = "backward", n: Union[int, Tuple[int]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, @@ -846,19 +851,41 @@ def ifft( return ret -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def embedding( - weights: Union[tf.Tensor, tf.Variable], - indices: Union[tf.Tensor, tf.Variable], - /, +def ifft_operations(x, rank, norm_factor): + if x.shape.rank == 1: + x = tf.signal.ifft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} + ) + else: + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.ifft(x), + 1: lambda: tf.signal.ifft2d(x), + 2: lambda: tf.signal.ifft3d(x), + }, + ) + x = x * norm_factor + return x + + +def ifftn( + x: Union[tf.Tensor, tf.Variable], + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - max_norm: Optional[float] = None, + norm: Optional[str] = "backward", out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) + result = _ifftn_helper(x, s, axes, norm) + + if out is not None: + out = result + return out + else: + return result def interpolate( @@ -919,417 +946,292 @@ def interpolate( return ret -interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (not align_corners and (len(x.shape) - 2) < 2) - and mode not in ["nearest", "area", "bicubic", "nd"] -) - - -def _fft2_norm( - x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", -): - n = tf.constant(s[0] * s[1], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.sqrt(n) - elif norm == "forward": - return x / n - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -def trans_x_to_s( +def max_pool1d( x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - """Change the shape of the input array x to the desired output shape s.""" - if x.dtype != tf.complex128 and x.dtype != tf.complex64: - x = tf.cast(x, tf.float32) - x_shape = x.shape - if dim == (-1, -2) or dim == (1, 0): - s = (s[1], s[0]) - if s[0] >= x_shape[0] and s[1] >= x_shape[1]: - paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) - x_new = tf.pad(x, paddings=paddings) - elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): - x_new = x[: s[0], : s[1]] - if s[0] != x_new.shape[0]: - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif s[1] != x_new.shape[1]: - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], 1) - elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], axis=1) - else: - x_new = x[: s[0], : s[1]] - return x_new - + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -def fft2_operations(x, rank): - if x.shape.rank == 1: - x = tf.signal.fft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} + if data_format == "NCW": + x = tf.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.fft(x), - 1: lambda: tf.signal.fft2d(x), - 2: lambda: tf.signal.fft3d(x), - }, + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - return x - -def _fft2_helper(x, shape, axes): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = fft2_operations(x, rank) + if not depth_pooling: + new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + padding = [(pad_w // 2, pad_w - pad_w // 2)] - x = transpose_x(x, tf.argsort(perm), perform_transpose) + if ceil_mode: + padding[0] = _padding_ceil_mode( + x.shape[1], new_kernel[0], padding[0], strides[0] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - return x + if depth_pooling: + res = tf.transpose(res, (0, 2, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCW": + return tf.transpose(res, (0, 2, 1)) + return res -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft2( +def max_pool2d( x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if len(x.shape) > 2: - result = _fft2_helper(x, s, dim) - else: - x_new = trans_x_to_s(x, s, dim) - x_complex = tf.cast(x_new, tf.complex128) - result = tf.signal.fft2d(x_complex) - - result = _fft2_norm(result, s, dim, norm) - if x.dtype == tf.complex64: - result = tf.cast(result, dtype=tf.complex128) - return result - + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -# --- IFFTN --- # -def fft_input_validation(x): - if not x.dtype.is_complex: - raise TypeError( - "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( - x.dtype - ) + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel ) - return x - - -def shape_and_axes_validation(shape, axes, input_rank_tensor): - if shape is not None: - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - checks_shape = [ - tf.debugging.assert_less_equal( - tf.size(shape), - input_rank_tensor, - message=( - "Argument `shape` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(shape), - ) - ] - with tf.control_dependencies(checks_shape): - shape = tf.identity(shape) - - if axes is not None: - axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) - checks_axes = [ - tf.debugging.assert_less_equal( - tf.size(axes), - input_rank_tensor, - message=( - "Argument `axes` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(axes), - ), - tf.debugging.assert_less( - axes, - input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - tf.debugging.assert_greater_equal( - axes, - -input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - ] - with tf.control_dependencies(checks_axes): - axes = tf.identity(axes) - - if shape is not None and axes is not None: - checks_shape_axes = [ - tf.debugging.assert_equal( - tf.size(shape), - tf.size(axes), - message=( - "Arguments `shape` and `axes` must have equal length. " - "Received: {}, {}" - ).format(shape, axes), - ) - ] - with tf.control_dependencies(checks_shape_axes): - shape, axes = tf.identity_n([shape, axes]) - - return shape, axes - - -def axes_initialization(shape, axes, input_shape, input_rank_tensor): - if axes is None: - axes = ( - tf.range(-tf.size(input_shape), 0) - if shape is None - else tf.range(-tf.size(shape), 0) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides ) - axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) - return axes - - -def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): - perform_padding = shape is not None - perform_transpose = tf.math.logical_not( - tf.math.reduce_all( - tf.math.equal( - axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) - ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - ) - return perform_padding, perform_transpose - - -def shape_initialization(shape, axes, x): - if shape is None: - shape = tf.gather(tf.shape(x), axes, axis=0) - return shape + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) -def rank_initialization(axes): - rank = tf.size(axes) - with tf.control_dependencies( - [ - tf.debugging.assert_less_equal( - rank, 3, message="N-D FFT supported only up to 3-D." - ) + if not depth_pooling: + new_kernel = [ + kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) ] - ): - rank = tf.identity(rank) - - return rank + if isinstance(padding, str): + pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) + padding = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + x_shape = x.shape[1:-1] -def norm_initialization(norm, shape, x): - if norm == "backward": - norm_factor = tf.constant(1, x.dtype) - elif norm == "forward" or norm == "ortho": - norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) - if norm == "ortho": - norm_factor = tf.math.sqrt(norm_factor) - return norm_factor + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) -def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): - if perform_padding: - pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) - pad_shape = tf.tensor_scatter_nd_update( - pad_shape, tf.expand_dims(axes, -1), shape - ) - x = _right_pad_or_crop(x, pad_shape) - return x + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res -def get_perm(input_rank_tensor, axes): - all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) - perm = tf.concat( - [ - tf.boolean_mask( - all_dims, - tf.foldl( - lambda acc, elem: tf.math.logical_and( - acc, tf.math.not_equal(all_dims, elem) - ), - axes, - initializer=tf.fill(all_dims.shape, True), - ), - ), - axes, - ], - 0, +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def max_pool3d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims ) - return perm - -def ifft_operations(x, rank, norm_factor): - if x.shape.rank == 1: - x = tf.signal.ifft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} + if data_format == "NCDHW": + x = tf.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.ifft(x), - 1: lambda: tf.signal.ifft2d(x), - 2: lambda: tf.signal.ifft3d(x), - }, + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - x = x * norm_factor - return x - - -def transpose_x(x, perm, perform_transpose): - x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) - return x - - -def static_output_shape(input_shape, shape, axes): - output_shape = input_shape.as_list() - if shape is not None: - if axes is None: - axes = list(range(-len(shape), 0)) - if isinstance(shape, tf.Tensor): - if isinstance(axes, tf.Tensor): - output_shape = [None] * len(output_shape) - else: - for ax in axes: - output_shape[ax] = None - else: - for idx, ax in enumerate(axes): - output_shape[ax] = shape[idx] - return tf.TensorShape(output_shape) - - -def _right_pad_or_crop(tensor, shape): - input_shape = tf.shape(tensor) - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - with tf.control_dependencies( - [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] - ): - shape = tf.identity(shape) - shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) - - pad_sizes = tf.math.maximum(shape - input_shape, 0) - pad_sizes = tf.expand_dims(pad_sizes, -1) - pad_sizes = tf.concat( - [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 - ) - tensor = tf.pad(tensor, pad_sizes, constant_values=0) - - crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) - tensor = tf.slice(tensor, crop_tensor, shape) - return tensor - - -def _ifftn_helper(x, shape, axes, norm): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) + if not depth_pooling: + x_shape = x.shape[1:-1] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + padding = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) - perm = get_perm(input_rank_tensor, axes) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - x = transpose_x(x, perm, perform_transpose) + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 4, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCDHW": + return tf.transpose(res, (0, 4, 1, 2, 3)) + return res - x = ifft_operations(x, rank, norm_factor) - x = transpose_x(x, tf.argsort(perm), perform_transpose) +def norm_initialization(norm, shape, x): + if norm == "backward": + norm_factor = tf.constant(1, x.dtype) + elif norm == "forward" or norm == "ortho": + norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) + if norm == "ortho": + norm_factor = tf.math.sqrt(norm_factor) + return norm_factor - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - return x +def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): + perform_padding = shape is not None + perform_transpose = tf.math.logical_not( + tf.math.reduce_all( + tf.math.equal( + axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) + ) + ) + ) + return perform_padding, perform_transpose -def ifftn( +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def pool( x: Union[tf.Tensor, tf.Variable], - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, *, - norm: Optional[str] = "backward", + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - result = _ifftn_helper(x, s, axes, norm) + return tf.nn.pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ) - if out is not None: - out = result - return out - else: - return result +def rank_initialization(axes): + rank = tf.size(axes) + with tf.control_dependencies( + [ + tf.debugging.assert_less_equal( + rank, 3, message="N-D FFT supported only up to 3-D." + ) + ] + ): + rank = tf.identity(rank) -""" -RFFTN Function -""" + return rank def rfft_input_validation(x): @@ -1364,40 +1266,6 @@ def rfft_operations(x, rank, norm_factor): return x -def _rfftn_helper(x, shape, axes, norm): - x = rfft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor - ) - - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = rfft_operations(x, rank, norm_factor) - - x = transpose_x(x, tf.argsort(perm), perform_transpose) - - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - - return x - - @with_supported_device_and_dtypes( { "2.5.0 and above": { @@ -1427,3 +1295,139 @@ def rfftn( else: # return result return tf.cast(result, tf.complex128) + + +def shape_and_axes_validation(shape, axes, input_rank_tensor): + if shape is not None: + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + checks_shape = [ + tf.debugging.assert_less_equal( + tf.size(shape), + input_rank_tensor, + message=( + "Argument `shape` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(shape), + ) + ] + with tf.control_dependencies(checks_shape): + shape = tf.identity(shape) + + if axes is not None: + axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) + checks_axes = [ + tf.debugging.assert_less_equal( + tf.size(axes), + input_rank_tensor, + message=( + "Argument `axes` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(axes), + ), + tf.debugging.assert_less( + axes, + input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + tf.debugging.assert_greater_equal( + axes, + -input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + ] + with tf.control_dependencies(checks_axes): + axes = tf.identity(axes) + + if shape is not None and axes is not None: + checks_shape_axes = [ + tf.debugging.assert_equal( + tf.size(shape), + tf.size(axes), + message=( + "Arguments `shape` and `axes` must have equal length. " + "Received: {}, {}" + ).format(shape, axes), + ) + ] + with tf.control_dependencies(checks_shape_axes): + shape, axes = tf.identity_n([shape, axes]) + + return shape, axes + + +def shape_initialization(shape, axes, x): + if shape is None: + shape = tf.gather(tf.shape(x), axes, axis=0) + return shape + + +def static_output_shape(input_shape, shape, axes): + output_shape = input_shape.as_list() + if shape is not None: + if axes is None: + axes = list(range(-len(shape), 0)) + if isinstance(shape, tf.Tensor): + if isinstance(axes, tf.Tensor): + output_shape = [None] * len(output_shape) + else: + for ax in axes: + output_shape[ax] = None + else: + for idx, ax in enumerate(axes): + output_shape[ax] = shape[idx] + return tf.TensorShape(output_shape) + + +def trans_x_to_s( + x: Union[tf.Tensor, tf.Variable], + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), +) -> Union[tf.Tensor, tf.Variable]: + """Change the shape of the input array x to the desired output shape s.""" + if x.dtype != tf.complex128 and x.dtype != tf.complex64: + x = tf.cast(x, tf.float32) + x_shape = x.shape + if dim == (-1, -2) or dim == (1, 0): + s = (s[1], s[0]) + if s[0] >= x_shape[0] and s[1] >= x_shape[1]: + paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) + x_new = tf.pad(x, paddings=paddings) + elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): + x_new = x[: s[0], : s[1]] + if s[0] != x_new.shape[0]: + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif s[1] != x_new.shape[1]: + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], 1) + elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], axis=1) + else: + x_new = x[: s[0], : s[1]] + return x_new + + +def transpose_x(x, perm, perform_transpose): + x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) + return x + + +interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (not align_corners and (len(x.shape) - 2) < 2) + and mode not in ["nearest", "area", "bicubic", "nd"] +) +"""RFFTN Function.""" diff --git a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py index c457e4ceb0ac0..e59d5c46a79f9 100644 --- a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py @@ -11,33 +11,59 @@ from .. import backend_version -@with_unsupported_dtypes( - {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version -) -def eigh_tridiagonal( - alpha: Union[tf.Tensor, tf.Variable], - beta: Union[tf.Tensor, tf.Variable], +dot.support_native_out = True + + +def adjoint( + x: Union[tf.Tensor, tf.Variable], /, *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] - ] = None, - tol: Optional[float] = None, -) -> Union[ - tf.Tensor, - tf.Variable, - Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], -]: - return tf.linalg.eigh_tridiagonal( - alpha, - beta, - eigvals_only=eigvals_only, - select=select, - select_range=select_range, - tol=tol, - ) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + _check_valid_dimension_size(x) + return tf.linalg.adjoint(x) + + +@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( + x: Union[tf.Tensor, tf.Variable], + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + svd = tf.linalg.svd(x, compute_uv=False) + if len(x.shape) >= 3: + ax = len(x.shape) // 2 + elif len(x.shape) >= 3 and p == -1: + ax = [-1, -2] + else: + ax = None + if p is None or p == 2: + k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) + elif p == "nuc": + svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) + k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) + elif p == "fro": + k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] + ) + elif p < 0: + if p == -1: + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=0), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) + elif p == -2: + k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) + elif p == -float("inf"): + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=1), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) + else: + k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord=p, axis=[-2, -1] + ) + return k def diagflat( @@ -75,23 +101,14 @@ def diagflat( return ret -def kron( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.kron(a, b) - - -def matrix_exp( - x: Union[tf.Tensor, tf.Variable], +def dot( + a: tf.Tensor, + b: tf.Tensor, /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.linalg.expm(x) + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + return tf.tensordot(a, b, axes=1) def eig( @@ -105,6 +122,35 @@ def eig( return tf.linalg.eig(x) +@with_unsupported_dtypes( + {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version +) +def eigh_tridiagonal( + alpha: Union[tf.Tensor, tf.Variable], + beta: Union[tf.Tensor, tf.Variable], + /, + *, + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] + ] = None, + tol: Optional[float] = None, +) -> Union[ + tf.Tensor, + tf.Variable, + Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], +]: + return tf.linalg.eigh_tridiagonal( + alpha, + beta, + eigvals_only=eigvals_only, + select=select, + select_range=select_range, + tol=tol, + ) + + def eigvals( x: Union[tf.Tensor, tf.Variable], /, @@ -114,14 +160,33 @@ def eigvals( return tf.linalg.eigvals(x) -def adjoint( +def kron( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.kron(a, b) + + +def lu_factor( + x: Union[tf.Tensor, tf.Variable], + /, + *, + pivot: Optional[bool] = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + raise IvyNotImplementedException() + + +def matrix_exp( x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - _check_valid_dimension_size(x) - return tf.linalg.adjoint(x) + return tf.linalg.expm(x) @with_supported_dtypes( @@ -149,68 +214,3 @@ def multi_dot( raise ValueError("Expecting at least two tensors.") dot_out = _reduce(tf.matmul, x) return dot_out - - -@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( - x: Union[tf.Tensor, tf.Variable], - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - svd = tf.linalg.svd(x, compute_uv=False) - if len(x.shape) >= 3: - ax = len(x.shape) // 2 - elif len(x.shape) >= 3 and p == -1: - ax = [-1, -2] - else: - ax = None - if p is None or p == 2: - k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) - elif p == "nuc": - svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) - k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) - elif p == "fro": - k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] - ) - elif p < 0: - if p == -1: - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=0), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) - elif p == -2: - k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) - elif p == -float("inf"): - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=1), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) - else: - k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord=p, axis=[-2, -1] - ) - return k - - -def lu_factor( - x: Union[tf.Tensor, tf.Variable], - /, - *, - pivot: Optional[bool] = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - raise IvyNotImplementedException() - - -def dot( - a: tf.Tensor, - b: tf.Tensor, - /, - *, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - return tf.tensordot(a, b, axes=1) - - -dot.support_native_out = True diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index 7e0a806ee1e19..bf0be94c43b24 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -10,99 +10,126 @@ import ivy -def moveaxis( - a: Union[tf.Tensor, tf.Variable], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, +def atleast_1d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.moveaxis(a, source, destination) +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_1d(*arys) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def heaviside( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) +def atleast_2d( + *arys: Union[tf.Tensor, tf.Variable], + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_2d(*arys) -def flipud( - m: Union[tf.Tensor, tf.Variable], +def atleast_3d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_3d(*arys) + + +def broadcast_shapes( + *shapes: Union[List[int], List[Tuple]], +) -> Tuple[int, ...]: + if len(shapes) > 1: + desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) + if len(shapes) > 2: + for i in range(2, len(shapes)): + desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) + else: + return [shapes[0]] + return tuple(desired_shape.numpy().tolist()) + + +def concat_from_sequence( + input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.flipud(m) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + highest_dtype = input_sequence[0].dtype + for i in input_sequence: + highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) + if new_axis == 0: + ret = tf.concat(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = tf.stack(input_sequence, axis=axis) + return ret -def vstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + +def dsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vstack(arrays) + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def dstack( arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.hstack(arrays) + return tf.experimental.numpy.dstack(arrays) -def rot90( - m: Union[tf.Tensor, tf.Variable], +def expand( + x: Union[tf.Tensor, tf.Variable], + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Union[tf.Tensor, tf.Variable] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rot90(m, k, axes) + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape.num_elements() / tf.reduce_prod( + [s for s in shape if s > 0] + ) + return tf.broadcast_to(x, shape) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) -def top_k( - x: tf.Tensor, - k: int, +def fill_diagonal( + a: tf.Tensor, + v: Union[int, float], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, -) -> Tuple[tf.Tensor, tf.Tensor]: - k = min(k, x.shape[axis]) - if not largest: - indices = tf.experimental.numpy.argsort(x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) + wrap: bool = False, +): + shape = tf.shape(a) + max_end = tf.math.reduce_prod(shape) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] else: - indices = tf.experimental.numpy.argsort(-x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) - if not sorted: - indices = tf.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) - val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) - indices = tf.dtypes.cast(indices, tf.int64) - return topk_res(val, indices) + step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) + a = tf.reshape(a, (-1,)) + end = min(end, max_end) + indices = [[i] for i in range(0, end, step)] + ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) + a = tf.tensor_scatter_nd_update(a, indices, ups) + a = tf.reshape(a, shape) + return a def fliplr( @@ -115,72 +142,80 @@ def fliplr( return tf.experimental.numpy.fliplr(m) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def i0( - x: Union[tf.Tensor, tf.Variable], +def flipud( + m: Union[tf.Tensor, tf.Variable], /, *, + copy: Optional[bool] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.bessel_i0(x, name=None) + return tf.experimental.numpy.flipud(m) -def vsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def heaviside( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) -def dsplit( +def hsplit( ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, ) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def atleast_1d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_1d(*arys) +def hstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.hstack(arrays) -def dstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def i0( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.dstack(arrays) + return tf.math.bessel_i0(x, name=None) -def atleast_2d( - *arys: Union[tf.Tensor, tf.Variable], +def moveaxis( + a: Union[tf.Tensor, tf.Variable], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_2d(*arys) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.moveaxis(a, source, destination) -def atleast_3d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], +def rot90( + m: Union[tf.Tensor, tf.Variable], + /, + *, copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_3d(*arys) + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Union[tf.Tensor, tf.Variable] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.rot90(m, k, axes) def take_along_axis( @@ -228,69 +263,36 @@ def take_along_axis( return tf.experimental.numpy.take_along_axis(arr, indices, axis) -def hsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Tuple[int, ...]], +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) +def top_k( + x: tf.Tensor, + k: int, /, *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes( - *shapes: Union[List[int], List[Tuple]], -) -> Tuple[int, ...]: - if len(shapes) > 1: - desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) - if len(shapes) > 2: - for i in range(2, len(shapes)): - desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, +) -> Tuple[tf.Tensor, tf.Tensor]: + k = min(k, x.shape[axis]) + if not largest: + indices = tf.experimental.numpy.argsort(x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) else: - return [shapes[0]] - return tuple(desired_shape.numpy().tolist()) - - -def expand( - x: Union[tf.Tensor, tf.Variable], - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape.num_elements() / tf.reduce_prod( - [s for s in shape if s > 0] - ) - return tf.broadcast_to(x, shape) - - -def concat_from_sequence( - input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], - /, - *, - new_axis: int = 0, - axis: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - highest_dtype = input_sequence[0].dtype - for i in input_sequence: - highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) - - if new_axis == 0: - ret = tf.concat(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = tf.stack(input_sequence, axis=axis) - return ret + indices = tf.experimental.numpy.argsort(-x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) + if not sorted: + indices = tf.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) + val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) + indices = tf.dtypes.cast(indices, tf.int64) + return topk_res(val, indices) def unique_consecutive( @@ -346,26 +348,24 @@ def unique_consecutive( ) -def fill_diagonal( - a: tf.Tensor, - v: Union[int, float], +def vsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, - wrap: bool = False, -): - shape = tf.shape(a) - max_end = tf.math.reduce_prod(shape) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) - a = tf.reshape(a, (-1,)) - end = min(end, max_end) - indices = [[i] for i in range(0, end, step)] - ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) - a = tf.tensor_scatter_nd_update(a, indices, ups) - a = tf.reshape(a, shape) - return a + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vstack(arrays) diff --git a/ivy/functional/backends/tensorflow/experimental/norms.py b/ivy/functional/backends/tensorflow/experimental/norms.py index 5fc38b60283af..22f00921bcf96 100644 --- a/ivy/functional/backends/tensorflow/experimental/norms.py +++ b/ivy/functional/backends/tensorflow/experimental/norms.py @@ -4,30 +4,6 @@ from . import backend_version -def l1_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - -def l2_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) def batch_norm( x: Union[tf.Tensor, tf.Variable], @@ -153,6 +129,30 @@ def instance_norm( ) +def l1_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + +def l2_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + def lp_normalize( x: Union[tf.Tensor, tf.Variable], /, diff --git a/ivy/functional/backends/tensorflow/experimental/random.py b/ivy/functional/backends/tensorflow/experimental/random.py index a5bbf97481be0..fc31d6a1c8d3f 100644 --- a/ivy/functional/backends/tensorflow/experimental/random.py +++ b/ivy/functional/backends/tensorflow/experimental/random.py @@ -15,6 +15,51 @@ ) +def bernoulli( + probs: Union[float, tf.Tensor, tf.Variable], + *, + logits: Union[float, tf.Tensor, tf.Variable] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: DType, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if seed is not None: + tf.random.set_seed(seed) + if logits is not None: + logits = tf.cast(logits, dtype) + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + probs = tf.cast(probs, dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return tfp.distributions.Bernoulli( + logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True + ).sample(shape, seed) + + +def beta( + alpha: Union[float, tf.Tensor, tf.Variable], + beta: Union[float, tf.Tensor, tf.Variable], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, + dtype: Optional[Union[DType, ivy.Dtype]] = None, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not dtype: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + shape = _check_bounds_and_get_shape(alpha, beta, shape).shape + alpha = tf.cast(alpha, dtype) + beta = tf.cast(beta, dtype) + return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) + + # dirichlet @with_unsupported_dtypes( { @@ -54,26 +99,6 @@ def dirichlet( ) -def beta( - alpha: Union[float, tf.Tensor, tf.Variable], - beta: Union[float, tf.Tensor, tf.Variable], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, - dtype: Optional[Union[DType, ivy.Dtype]] = None, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not dtype: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - shape = _check_bounds_and_get_shape(alpha, beta, shape).shape - alpha = tf.cast(alpha, dtype) - beta = tf.cast(beta, dtype) - return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) - - def gamma( alpha: Union[float, tf.Tensor, tf.Variable], beta: Union[float, tf.Tensor, tf.Variable], @@ -117,28 +142,3 @@ def poisson( if tf.reduce_any(lam < 0): return tf.where(lam < 0, fill_value, ret) return ret - - -def bernoulli( - probs: Union[float, tf.Tensor, tf.Variable], - *, - logits: Union[float, tf.Tensor, tf.Variable] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: DType, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if seed is not None: - tf.random.set_seed(seed) - if logits is not None: - logits = tf.cast(logits, dtype) - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - probs = tf.cast(probs, dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return tfp.distributions.Bernoulli( - logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True - ).sample(shape, seed) diff --git a/ivy/functional/backends/tensorflow/experimental/statistical.py b/ivy/functional/backends/tensorflow/experimental/statistical.py index 9cd91437a5f4a..2a5feef312527 100644 --- a/ivy/functional/backends/tensorflow/experimental/statistical.py +++ b/ivy/functional/backends/tensorflow/experimental/statistical.py @@ -14,202 +14,79 @@ from copy import deepcopy -def histogram( - a: tf.Tensor, - /, - *, - bins: Optional[Union[int, tf.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[tf.DType] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[tf.Tensor] = None, - density: Optional[bool] = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - min_a = tf.reduce_min(a) - max_a = tf.reduce_max(a) - if isinstance(bins, tf.Tensor) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - if isinstance(bins, int): - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - elif isinstance(bins, int): - range = (min_a, max_a) - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - if tf.shape(bins)[0] < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - if min_a < bins[0] and not extend_lower_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_lower_interval to deal with this." - ) - if max_a > bins[-1] and not extend_upper_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_upper_interval to deal with this." - ) - ret = tfp.stats.histogram( - x=a, - edges=bins, - axis=axis, - weights=weights, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - name="histogram", - ) - if density: - pass - # TODO: Tensorflow native dtype argument is not working - if dtype: - ret = tf.cast(ret, dtype) - bins = tf.cast(bins, dtype) - # TODO: weird error when returning bins: return ret, bins - return ret +# --- Helpers --- # +# --------------- # -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float", - "complex", +def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: + values, indices = [], [] + if ( + isinstance(x[0], tf.Tensor) + and isinstance(x[0].numpy().tolist(), list) + and len(x[0].numpy().tolist()) >= 1 + ): + if axis >= 1: + for ret1 in x: + value, indice = __find_cummax(ret1, axis=axis - 1) + indices.append(indice) + values.append(value) + else: + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + else: + x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) + values, indices = tf.scan( + lambda a, b: ( + a + if a > b + or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 + else b + ), + (x, x_indices), ) - }, - backend_version, -) -def median( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tfp.stats.percentile( - input, - 50.0, - axis=axis, - interpolation="midpoint", - keepdims=keepdims, + + return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( + tf.convert_to_tensor(indices), dtype=tf.int64 ) -def nanmean( - a: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) - +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] -def _validate_quantile(q): - if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: - for i in range(tf.size(q)): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) else: - if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): - return False - return True - - -def to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis - - -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = tf.experimental.numpy.ndim(a) - axis_arg = deepcopy(axis) - if axis is not None: - axis = to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) - - for i, s in enumerate(sorted(keep)): - a = tf.experimental.numpy.moveaxis(a, s, i) - a = tf.reshape( - a, - [ - *a.shape[:nkeep], - -1, - ], - ) - axis_arg = -1 - - ret = fn(a, q, axis=axis_arg) - - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] - - return ret - - -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if tf.experimental.numpy.ndim(q) > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = tf.reshape(a, [-1]) - elif axis != 0: - a = tf.experimental.numpy.moveaxis(a, axis, 0) - axis = 0 - - n = a.shape[axis] - - indices = q * (n - 1) - - a = tf.sort(a, axis) - - indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) - indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) - - weights = indices - tf.cast(indices_below, dtype=ret_dtype) - - indices_below = tf.clip_by_value(indices_below, 0, n - 1) - indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) - tensor_upper = tf.gather(a, indices_upper, axis=axis) - tensor_below = tf.gather(a, indices_below, axis=axis) - - pred = weights <= 0.5 - out = tf.where(pred, tensor_below, tensor_upper) - - return tf.cast(out, ret_dtype) + indices.append((lst, tuple(prefix))) + return indices def _compute_quantile_wrapper( @@ -247,50 +124,39 @@ def _compute_quantile_wrapper( ) -def quantile( - a: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, float], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = tf.experimental.numpy.ndim(a) + axis_arg = deepcopy(axis) + if axis is not None: + axis = to_positive_axis(axis, nd) + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) -def corrcoef( - x: tf.Tensor, - /, - *, - y: tf.Tensor, - rowvar: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = tf.concat([x, y], axis=axis) + for i, s in enumerate(sorted(keep)): + a = tf.experimental.numpy.moveaxis(a, s, i) + a = tf.reshape( + a, + [ + *a.shape[:nkeep], + -1, + ], + ) + axis_arg = -1 - if rowvar: - mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) - cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) - else: - mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) - cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) + ret = fn(a, q, axis=axis_arg) - cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) - cor = cov2_t @ cov_t @ cov2_t - return cor + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret def _nanmedian_helper(input, axis=None, keepdims=False): @@ -452,22 +318,52 @@ def _nanmedian_helper(input, axis=None, keepdims=False): return result -def nanmedian( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if overwrite_input: - copied_input = tf.identity(input) - return _nanmedian_helper(copied_input, axis, keepdims) +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if tf.experimental.numpy.ndim(q) > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = tf.reshape(a, [-1]) + elif axis != 0: + a = tf.experimental.numpy.moveaxis(a, axis, 0) + axis = 0 + + n = a.shape[axis] + + indices = q * (n - 1) + a = tf.sort(a, axis) + + indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) + indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) + + weights = indices - tf.cast(indices_below, dtype=ret_dtype) + + indices_below = tf.clip_by_value(indices_below, 0, n - 1) + indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) + tensor_upper = tf.gather(a, indices_upper, axis=axis) + tensor_below = tf.gather(a, indices_below, axis=axis) + + pred = weights <= 0.5 + out = tf.where(pred, tensor_below, tensor_upper) + + return tf.cast(out, ret_dtype) + + +def _validate_quantile(q): + if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: + for i in range(tf.size(q)): + if not (0.0 <= q[i] <= 1.0): + return False else: - result = _nanmedian_helper(input, axis, keepdims) - return result + if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): + return False + return True + + +# --- Main --- # +# ------------ # @with_supported_device_and_dtypes( @@ -505,19 +401,30 @@ def bincount( ) -@with_supported_device_and_dtypes( - { - "2.13.0 and below": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def igamma( - a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None +def corrcoef( + x: tf.Tensor, + /, + *, + y: tf.Tensor, + rowvar: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> tf.Tensor: - return tf.math.igamma(a, x) + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = tf.concat([x, y], axis=axis) + + if rowvar: + mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) + cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) + else: + mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) + cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) + + cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) + cor = cov2_t @ cov_t @ cov2_t + return cor @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) @@ -680,77 +587,6 @@ def cummax( return __find_cummax(x, axis=axis) -def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: - values, indices = [], [] - if ( - isinstance(x[0], tf.Tensor) - and isinstance(x[0].numpy().tolist(), list) - and len(x[0].numpy().tolist()) >= 1 - ): - if axis >= 1: - for ret1 in x: - value, indice = __find_cummax(ret1, axis=axis - 1) - indices.append(indice) - values.append(value) - else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - else: - x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) - values, indices = tf.scan( - lambda a, b: ( - a - if a > b - or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 - else b - ), - (x, x_indices), - ) - - return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( - tf.convert_to_tensor(indices), dtype=tf.int64 - ) - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "complex")}, backend_version, @@ -781,3 +617,175 @@ def cummin( return cummin_x else: return tf.cast(cummin_x, dtype) + + +def histogram( + a: tf.Tensor, + /, + *, + bins: Optional[Union[int, tf.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[tf.DType] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[tf.Tensor] = None, + density: Optional[bool] = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + min_a = tf.reduce_min(a) + max_a = tf.reduce_max(a) + if isinstance(bins, tf.Tensor) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + if isinstance(bins, int): + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + elif isinstance(bins, int): + range = (min_a, max_a) + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + if tf.shape(bins)[0] < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + if min_a < bins[0] and not extend_lower_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_lower_interval to deal with this." + ) + if max_a > bins[-1] and not extend_upper_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_upper_interval to deal with this." + ) + ret = tfp.stats.histogram( + x=a, + edges=bins, + axis=axis, + weights=weights, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + name="histogram", + ) + if density: + pass + # TODO: Tensorflow native dtype argument is not working + if dtype: + ret = tf.cast(ret, dtype) + bins = tf.cast(bins, dtype) + # TODO: weird error when returning bins: return ret, bins + return ret + + +@with_supported_device_and_dtypes( + { + "2.13.0 and below": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def igamma( + a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None +) -> tf.Tensor: + return tf.math.igamma(a, x) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float", + "complex", + ) + }, + backend_version, +) +def median( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tfp.stats.percentile( + input, + 50.0, + axis=axis, + interpolation="midpoint", + keepdims=keepdims, + ) + + +def nanmean( + a: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) + + +def nanmedian( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if overwrite_input: + copied_input = tf.identity(input) + return _nanmedian_helper(copied_input, axis, keepdims) + + else: + result = _nanmedian_helper(input, axis, keepdims) + return result + + +def quantile( + a: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, float], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) + + +def to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis diff --git a/ivy/functional/backends/tensorflow/general.py b/ivy/functional/backends/tensorflow/general.py index ba239549e3fdd..e6052c0b5e49e 100644 --- a/ivy/functional/backends/tensorflow/general.py +++ b/ivy/functional/backends/tensorflow/general.py @@ -23,12 +23,26 @@ _round = round -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, (tf.Tensor, tf.Variable)): - if exclusive and isinstance(x, tf.Variable): - return False - return True - return False +# --- Helpers --- # +# --------------- # + + +def _check_query(query): + return not isinstance(query, list) and ( + not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) + ) + + +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view + + +# --- Main --- # +# ------------ # def array_equal( @@ -48,58 +62,6 @@ def current_backend_str() -> str: return "tensorflow" -def _check_query(query): - return not isinstance(query, list) and ( - not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) - ) - - -def get_item( - x: Union[tf.Tensor, tf.Variable], - /, - query: Union[tf.Tensor, tf.Variable, Tuple], - *, - copy: bool = None, -) -> Union[tf.Tensor, tf.Variable]: - return x.__getitem__(query) - - -get_item.partial_mixed_handler = lambda x, query, **kwargs: ( - all(_check_query(i) for i in query) - if isinstance(query, tuple) - else _check_query(query) -) - - -def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: - # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions - if ( - ivy.is_array(x) - and get_num_dims(x) == 0 - and ivy.as_native_dtype(x.dtype) is tf.bfloat16 - ): - x = tf.expand_dims(x, 0) - if copy: - return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) - else: - return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) - if copy: - return np.array(tf.convert_to_tensor(x)) - else: - return np.asarray(tf.convert_to_tensor(x)) - - -def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: - ret = to_numpy(x).item() - if x.dtype == tf.bfloat16: - return float(ret) - return ret - - -def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: - return x.numpy().tolist() - - def gather( params: Union[tf.Tensor, tf.Variable], indices: Union[tf.Tensor, tf.Variable], @@ -115,38 +77,6 @@ def gather( return tf.gather(params, indices, axis=axis, batch_dims=batch_dims) -def gather_nd_helper(params, indices): - indices_shape = tf.shape(indices) - params_shape = tf.shape(params) - num_index_dims = indices_shape[-1] - result_dim_sizes_list = [ - tf.math.reduce_prod(params_shape[i + 1 :]) for i in range(len(params_shape) - 1) - ] + [1] - result_dim_sizes = tf.convert_to_tensor(result_dim_sizes_list, dtype=indices.dtype) - implicit_indices_factor = result_dim_sizes[num_index_dims - 1] - flat_params = tf.reshape(params, (-1,)) - new_shape = [1] * (len(indices_shape) - 1) + [num_index_dims] - indices_scales = tf.reshape(result_dim_sizes[0:num_index_dims], new_shape) - indices_for_flat_tiled = tf.reshape( - tf.reduce_sum(indices * indices_scales, -1, keepdims=True), (-1, 1) - ) - indices_for_flat_tiled = tf.repeat( - indices_for_flat_tiled, implicit_indices_factor, axis=1 - ) - implicit_indices = tf.repeat( - tf.expand_dims(tf.range(implicit_indices_factor), 0), - indices_for_flat_tiled.shape[0], - axis=0, - ) - indices_for_flat = indices_for_flat_tiled + implicit_indices - flat_indices_for_flat = tf.reshape(indices_for_flat, (-1,)) - flat_gather = tf.gather(flat_params, flat_indices_for_flat) - res = tf.reshape( - flat_gather, tf.concat([indices_shape[:-1], params_shape[num_index_dims:]], 0) - ) - return res - - def gather_nd( params: Union[tf.Tensor, tf.Variable], indices: Union[tf.Tensor, tf.Variable], @@ -184,6 +114,48 @@ def gather_nd( return result +def gather_nd_helper(params, indices): + indices_shape = tf.shape(indices) + params_shape = tf.shape(params) + num_index_dims = indices_shape[-1] + result_dim_sizes_list = [ + tf.math.reduce_prod(params_shape[i + 1 :]) for i in range(len(params_shape) - 1) + ] + [1] + result_dim_sizes = tf.convert_to_tensor(result_dim_sizes_list, dtype=indices.dtype) + implicit_indices_factor = result_dim_sizes[num_index_dims - 1] + flat_params = tf.reshape(params, (-1,)) + new_shape = [1] * (len(indices_shape) - 1) + [num_index_dims] + indices_scales = tf.reshape(result_dim_sizes[0:num_index_dims], new_shape) + indices_for_flat_tiled = tf.reshape( + tf.reduce_sum(indices * indices_scales, -1, keepdims=True), (-1, 1) + ) + indices_for_flat_tiled = tf.repeat( + indices_for_flat_tiled, implicit_indices_factor, axis=1 + ) + implicit_indices = tf.repeat( + tf.expand_dims(tf.range(implicit_indices_factor), 0), + indices_for_flat_tiled.shape[0], + axis=0, + ) + indices_for_flat = indices_for_flat_tiled + implicit_indices + flat_indices_for_flat = tf.reshape(indices_for_flat, (-1,)) + flat_gather = tf.gather(flat_params, flat_indices_for_flat) + res = tf.reshape( + flat_gather, tf.concat([indices_shape[:-1], params_shape[num_index_dims:]], 0) + ) + return res + + +def get_item( + x: Union[tf.Tensor, tf.Variable], + /, + query: Union[tf.Tensor, tf.Variable, Tuple], + *, + copy: bool = None, +) -> Union[tf.Tensor, tf.Variable]: + return x.__getitem__(query) + + def get_num_dims(x, /, *, as_array=False): return ( tf.cast(tf.shape(tf.shape(x))[0], tf.int64) @@ -286,18 +258,49 @@ def inplace_update( return val -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view - - def inplace_variables_supported(): return True +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, (tf.Tensor, tf.Variable)): + if exclusive and isinstance(x, tf.Variable): + return False + return True + return False + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def isin( + elements: tf.Tensor, + test_elements: tf.Tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> tf.Tensor: + input_shape = elements.shape + + if tf.rank(elements) == 0: + elements = tf.reshape(elements, [1]) + if tf.rank(test_elements) == 0: + test_elements = tf.reshape(test_elements, [1]) + if not assume_unique: + test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] + + elements = tf.reshape(elements, [-1]) + test_elements = tf.reshape(test_elements, [-1]) + + output = tf.reduce_any( + tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 + ) + return tf.reshape(output, input_shape) ^ invert + + +def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: + return x.dtype.size + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -347,9 +350,6 @@ def scatter_flat( return res -scatter_flat.support_native_out = True - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def scatter_nd( indices: Union[tf.Tensor, tf.Variable], @@ -406,9 +406,6 @@ def scatter_nd( return res -scatter_nd.support_native_out = True - - def shape( x: Union[tf.Tensor, tf.Variable], /, @@ -421,6 +418,35 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: + return x.numpy().tolist() + + +def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: + # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions + if ( + ivy.is_array(x) + and get_num_dims(x) == 0 + and ivy.as_native_dtype(x.dtype) is tf.bfloat16 + ): + x = tf.expand_dims(x, 0) + if copy: + return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) + else: + return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) + if copy: + return np.array(tf.convert_to_tensor(x)) + else: + return np.asarray(tf.convert_to_tensor(x)) + + +def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: + ret = to_numpy(x).item() + if x.dtype == tf.bfloat16: + return float(ret) + return ret + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -507,32 +533,10 @@ def _vmap(*args, **kwargs): return _vmap -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def isin( - elements: tf.Tensor, - test_elements: tf.Tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> tf.Tensor: - input_shape = elements.shape - - if tf.rank(elements) == 0: - elements = tf.reshape(elements, [1]) - if tf.rank(test_elements) == 0: - test_elements = tf.reshape(test_elements, [1]) - if not assume_unique: - test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] - - elements = tf.reshape(elements, [-1]) - test_elements = tf.reshape(test_elements, [-1]) - - output = tf.reduce_any( - tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 - ) - return tf.reshape(output, input_shape) ^ invert - - -def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: - return x.dtype.size +get_item.partial_mixed_handler = lambda x, query, **kwargs: ( + all(_check_query(i) for i in query) + if isinstance(query, tuple) + else _check_query(query) +) +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True diff --git a/ivy/functional/backends/tensorflow/gradients.py b/ivy/functional/backends/tensorflow/gradients.py index 3eb4207d2df36..6968f2f245ac8 100644 --- a/ivy/functional/backends/tensorflow/gradients.py +++ b/ivy/functional/backends/tensorflow/gradients.py @@ -21,17 +21,8 @@ ) -def variable(x, /): - with tf.device(ivy.dev(x, as_native=True)): - return tf.Variable(x, trainable=True) - - -def is_variable(x, /, *, exclusive=False): - return isinstance(x, tf.Variable) - - -def variable_data(x: tf.Variable, /) -> tf.Variable: - return x.value() +# --- Helpers --- # +# --------------- # def _grad_func(y, xs, xs_required, tape): @@ -63,6 +54,10 @@ def _grad_func(y, xs, xs_required, tape): return grads +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: Union[tf.Tensor, tf.Variable], @@ -117,81 +112,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - def grad_fn(xs): - grads = ivy.nested_map( - xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False - ) - with tf.GradientTape(watch_accessed_variables=False) as tape: - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - tape.watch(xs) - y = func(xs) - y = y.to_native(y) - grads_ = tape.gradient(y, xs) - grads_ = ivy.nested_map( - grads_, - lambda x: ivy.to_ivy(x), - include_derived=True, - ) - grads_ = ivy.to_ivy(grads_) - grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) - grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) - xs = ivy.to_ivy(xs) - if isinstance(xs, ivy.Array): - grads = grads_ - else: - ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) - y = ivy.to_ivy(y) - return y, grads - - return grad_fn - - -def stop_gradient( - x: Union[tf.Tensor, tf.Variable], - /, - *, - preserve_type: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - is_var = is_variable(x) - x = tf.stop_gradient(x) - if is_var and preserve_type: - return variable(x) - return x - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(x_in): - with tf.GradientTape(persistent=True) as tape: - ivy.nested_map(x_in, ivy.copy_array) - x_in = ivy.to_native(x_in, nested=True) - tape.watch(x_in) - y = grad_fn(x_in) - - # Deal with multiple outputs - if not isinstance(y, ivy.NativeArray): - jacobian = ivy.nested_map( - y, - lambda yi: ivy.to_ivy( - tape.jacobian(yi, x_in, unconnected_gradients="zero"), - nested=True, - ), - include_derived=True, - ) - else: - jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) - return jacobian - - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -240,5 +160,93 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive=False): + return isinstance(x, tf.Variable) + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(x_in): + with tf.GradientTape(persistent=True) as tape: + ivy.nested_map(x_in, ivy.copy_array) + x_in = ivy.to_native(x_in, nested=True) + tape.watch(x_in) + y = grad_fn(x_in) + + # Deal with multiple outputs + if not isinstance(y, ivy.NativeArray): + jacobian = ivy.nested_map( + y, + lambda yi: ivy.to_ivy( + tape.jacobian(yi, x_in, unconnected_gradients="zero"), + nested=True, + ), + include_derived=True, + ) + else: + jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) + return jacobian + + return callback_fn + + +def stop_gradient( + x: Union[tf.Tensor, tf.Variable], + /, + *, + preserve_type: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + is_var = is_variable(x) + x = tf.stop_gradient(x) + if is_var and preserve_type: + return variable(x) + return x + + +def value_and_grad(func): + def grad_fn(xs): + grads = ivy.nested_map( + xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False + ) + with tf.GradientTape(watch_accessed_variables=False) as tape: + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + tape.watch(xs) + y = func(xs) + y = y.to_native(y) + grads_ = tape.gradient(y, xs) + grads_ = ivy.nested_map( + grads_, + lambda x: ivy.to_ivy(x), + include_derived=True, + ) + grads_ = ivy.to_ivy(grads_) + grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) + grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) + xs = ivy.to_ivy(xs) + if isinstance(xs, ivy.Array): + grads = grads_ + else: + ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) + y = ivy.to_ivy(y) + return y, grads + + return grad_fn + + +def variable(x, /): + with tf.device(ivy.dev(x, as_native=True)): + return tf.Variable(x, trainable=True) + + +def variable_data(x: tf.Variable, /) -> tf.Variable: + return x.value() + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/tensorflow/layers.py b/ivy/functional/backends/tensorflow/layers.py index c3070846e21cb..d361fdd9827dc 100644 --- a/ivy/functional/backends/tensorflow/layers.py +++ b/ivy/functional/backends/tensorflow/layers.py @@ -17,6 +17,10 @@ ) +# --- Helpers --- # +# --------------- # + + def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): if filter_format == "channel_first": filters = tf.transpose(filters, (*range(2, dims + 2), 1, 0)) @@ -33,6 +37,24 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters +def _output_shape( + x_shape, filter_shape, output_shape, strides, padding, dims, dilations +): + dilations = [dilations] * dims if isinstance(dilations, int) else dilations + strides = [strides] * dims if isinstance(strides, int) else strides + if output_shape is None: + out_shape = [ + _deconv_length( + x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] + ) + for i in range(dims) + ] + output_shape = [x_shape[0], *out_shape, filter_shape[-2]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] + return output_shape + + def _pad_before_conv(x, filters, strides, padding, dims, dilations): dilations = [dilations] * dims if isinstance(dilations, int) else dilations strides = [strides] * dims if isinstance(strides, int) else strides @@ -64,22 +86,8 @@ def _pad_before_conv(x, filters, strides, padding, dims, dilations): ) -def _output_shape( - x_shape, filter_shape, output_shape, strides, padding, dims, dilations -): - dilations = [dilations] * dims if isinstance(dilations, int) else dilations - strides = [strides] * dims if isinstance(strides, int) else strides - if output_shape is None: - out_shape = [ - _deconv_length( - x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] - ) - for i in range(dims) - ] - output_shape = [x_shape[0], *out_shape, filter_shape[-2]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] - return output_shape +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) @@ -205,32 +213,6 @@ def conv2d_transpose( return res -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def depthwise_conv2d( - x: Union[tf.Tensor, tf.Variable], - filters: Union[tf.Tensor, tf.Variable], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - if tf.rank(filters) == 3: - filters = tf.expand_dims(filters, -1) - x = _pad_before_conv(x, filters, strides, padding, 2, dilations) - strides = [1, strides[0], strides[1], 1] - res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def conv3d( x: Union[tf.Tensor, tf.Variable], @@ -491,3 +473,29 @@ def conv_general_transpose( if data_format == "channel_first": res = tf.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def depthwise_conv2d( + x: Union[tf.Tensor, tf.Variable], + filters: Union[tf.Tensor, tf.Variable], + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + if tf.rank(filters) == 3: + filters = tf.expand_dims(filters, -1) + x = _pad_before_conv(x, filters, strides, padding, 2, dilations) + strides = [1, strides[0], strides[1], 1] + res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res diff --git a/ivy/functional/backends/tensorflow/linear_algebra.py b/ivy/functional/backends/tensorflow/linear_algebra.py index 5281feade12b5..28c91cb2bcac4 100644 --- a/ivy/functional/backends/tensorflow/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/linear_algebra.py @@ -79,6 +79,21 @@ def det( return tf.linalg.det(x) +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def diag( + x: Union[tf.Tensor, tf.Variable], + /, + *, + k: int = 0, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.diag(x, k=k) + + @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def diagonal( x: Union[tf.Tensor, tf.Variable], @@ -614,6 +629,21 @@ def trace( return tf.experimental.numpy.trace(x, offset=offset, axis1=axis1, axis2=axis2) +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, + backend_version, +) +def vander( + x: Union[tf.Tensor, tf.Variable], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -667,36 +697,6 @@ def vector_norm( return tf.reduce_sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def diag( - x: Union[tf.Tensor, tf.Variable], - /, - *, - k: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.diag(x, k=k) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, - backend_version, -) -def vander( - x: Union[tf.Tensor, tf.Variable], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes( { "2.13.0 and below": ( diff --git a/ivy/functional/backends/tensorflow/manipulation.py b/ivy/functional/backends/tensorflow/manipulation.py index 41800bca81d4f..a26c08e2c8abe 100644 --- a/ivy/functional/backends/tensorflow/manipulation.py +++ b/ivy/functional/backends/tensorflow/manipulation.py @@ -15,12 +15,47 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _reshape_fortran_tf(x, shape): if len(x.shape) > 0: x = tf.transpose(x) return tf.transpose(tf.reshape(x, shape[::-1])) +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def clip( + x: Union[tf.Tensor, tf.Variable], + x_min: Union[Number, tf.Tensor, tf.Variable], + x_max: Union[Number, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): + promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) + promoted_type = ivy.as_native_dtype( + ivy.promote_types(promoted_type, x_max.dtype) + ) + x = tf.cast(x, promoted_type) + x_min = tf.cast(x_min, promoted_type) + x_max = tf.cast(x_max, promoted_type) + if tf.size(x) == 0: + ret = x + elif x.dtype == tf.bool: + ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) + ret = tf.cast(ret, x.dtype) + else: + ret = tf.clip_by_value(x, x_min, x_max) + return ret + + # Array API Standard # # -------------------# @@ -54,6 +89,14 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None +): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width, constant_values=value) + + def expand_dims( x: Union[tf.Tensor, tf.Variable], /, @@ -106,6 +149,18 @@ def permute_dims( return tf.transpose(x, perm=axes) +@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) +def repeat( + x: Union[tf.Tensor, tf.Variable], + /, + repeats: Union[int, List[int]], + *, + axis: int = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.repeat(x, repeats, axis) + + def reshape( x: Union[tf.Tensor, tf.Variable], /, @@ -148,6 +203,47 @@ def roll( return ret +# Extra # +# ------# + + +def split( + x: Union[tf.Tensor, tf.Variable], + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> Union[tf.Tensor, tf.Variable]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + dim_size = tf.shape(x)[axis] + num_or_size_splits = int(dim_size) + if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): + num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) + num_or_size_splits = num_or_size_splits.numpy().tolist() + elif isinstance(num_or_size_splits, int) and with_remainder: + num_chunks = x.shape[axis] / num_or_size_splits + num_chunks_int = math.floor(num_chunks) + remainder = num_chunks - num_chunks_int + if remainder != 0: + num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ + int(remainder * num_or_size_splits) + ] + + return tf.split(x, num_or_size_splits, axis) + + def squeeze( x: Union[tf.Tensor, tf.Variable], /, @@ -202,57 +298,25 @@ def stack( raise ivy.utils.exceptions.IvyIndexError(e) -# Extra # -# ------# - - -def split( - x: Union[tf.Tensor, tf.Variable], +def swapaxes( + x, + axis0, + axis1, /, *, copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> Union[tf.Tensor, tf.Variable]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - dim_size = tf.shape(x)[axis] - num_or_size_splits = int(dim_size) - if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): - num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) - num_or_size_splits = num_or_size_splits.numpy().tolist() - elif isinstance(num_or_size_splits, int) and with_remainder: - num_chunks = x.shape[axis] / num_or_size_splits - num_chunks_int = math.floor(num_chunks) - remainder = num_chunks - num_chunks_int - if remainder != 0: - num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ - int(remainder * num_or_size_splits) - ] - - return tf.split(x, num_or_size_splits, axis) - - -@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) -def repeat( - x: Union[tf.Tensor, tf.Variable], - /, - repeats: Union[int, List[int]], - *, - axis: int = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.repeat(x, repeats, axis) +): + x_shape = x.shape + num_dims = len(x_shape) + axis0 %= num_dims + axis1 %= num_dims + config = list(range(num_dims)) + config.pop(axis0) + config.insert(axis0, axis1) + config.pop(axis1) + config.insert(axis1, axis0) + return tf.transpose(x, config) @with_unsupported_dtypes( @@ -293,68 +357,6 @@ def tile( return tf.tile(x, repeats) -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None -): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width, constant_values=value) - - -def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width) - - -def swapaxes( - x, - axis0, - axis1, - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -): - x_shape = x.shape - num_dims = len(x_shape) - axis0 %= num_dims - axis1 %= num_dims - config = list(range(num_dims)) - config.pop(axis0) - config.insert(axis0, axis1) - config.pop(axis1) - config.insert(axis1, axis0) - return tf.transpose(x, config) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def clip( - x: Union[tf.Tensor, tf.Variable], - x_min: Union[Number, tf.Tensor, tf.Variable], - x_max: Union[Number, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): - promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) - promoted_type = ivy.as_native_dtype( - ivy.promote_types(promoted_type, x_max.dtype) - ) - x = tf.cast(x, promoted_type) - x_min = tf.cast(x_min, promoted_type) - x_max = tf.cast(x_max, promoted_type) - if tf.size(x) == 0: - ret = x - elif x.dtype == tf.bool: - ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) - ret = tf.cast(ret, x.dtype) - else: - ret = tf.clip_by_value(x, x_min, x_max) - return ret - - def unstack( x: Union[tf.Tensor, tf.Variable], /, @@ -369,3 +371,9 @@ def unstack( if keepdims: return [tf.expand_dims(r, axis) for r in ret] return ret + + +def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width) diff --git a/ivy/functional/backends/tensorflow/random.py b/ivy/functional/backends/tensorflow/random.py index b0f3b97a3d877..ae0f6b4a26c31 100644 --- a/ivy/functional/backends/tensorflow/random.py +++ b/ivy/functional/backends/tensorflow/random.py @@ -22,47 +22,6 @@ from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, tf.Tensor, tf.Variable] = 0.0, - high: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, - dtype: DType, - device: str, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - shape = _check_bounds_and_get_shape(low, high, shape).shape - low = tf.cast(low, dtype) - high = tf.cast(high, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) - - -def random_normal( - *, - mean: Union[float, tf.Tensor, tf.Variable] = 0.0, - std: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: DType, - seed: Optional[int] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - mean = tf.cast(mean, dtype) - std = tf.cast(std, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) def multinomial( population_size: int, @@ -142,6 +101,47 @@ def randint( return tf.cast(tf.random.uniform(shape, low, high, "float32", seed=seed), dtype) +def random_normal( + *, + mean: Union[float, tf.Tensor, tf.Variable] = 0.0, + std: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: DType, + seed: Optional[int] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + mean = tf.cast(mean, dtype) + std = tf.cast(std, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, tf.Tensor, tf.Variable] = 0.0, + high: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, + dtype: DType, + device: str, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + shape = _check_bounds_and_get_shape(low, high, shape).shape + low = tf.cast(low, dtype) + high = tf.cast(high, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) + + def seed(*, seed_value: int = 0) -> None: tf.random.set_seed(seed_value) return diff --git a/ivy/functional/backends/tensorflow/searching.py b/ivy/functional/backends/tensorflow/searching.py index 047441cff9b08..13d461c8c54c5 100644 --- a/ivy/functional/backends/tensorflow/searching.py +++ b/ivy/functional/backends/tensorflow/searching.py @@ -77,6 +77,29 @@ def argmin( return tf.cast(ret, dtype) if dtype is not None else ret +# Extra # +# ----- # + + +def argwhere( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(x, tf.Variable): + x_ndim = x.shape.rank + else: + x_ndim = x.ndim + if x_ndim == 0: + return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") + where_x = tf.experimental.numpy.nonzero(x) + res = tf.experimental.numpy.concatenate( + [tf.expand_dims(item, -1) for item in where_x], -1 + ) + return res + + def nonzero( x: Union[tf.Tensor, tf.Variable], /, @@ -114,26 +137,3 @@ def where( ) -> Union[tf.Tensor, tf.Variable]: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return tf.cast(tf.experimental.numpy.where(condition, x1, x2), x1.dtype) - - -# Extra # -# ----- # - - -def argwhere( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(x, tf.Variable): - x_ndim = x.shape.rank - else: - x_ndim = x.ndim - if x_ndim == 0: - return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") - where_x = tf.experimental.numpy.nonzero(x) - res = tf.experimental.numpy.concatenate( - [tf.expand_dims(item, -1) for item in where_x], -1 - ) - return res diff --git a/ivy/functional/backends/tensorflow/sorting.py b/ivy/functional/backends/tensorflow/sorting.py index 06403a1d358ca..2e1202a943ee0 100644 --- a/ivy/functional/backends/tensorflow/sorting.py +++ b/ivy/functional/backends/tensorflow/sorting.py @@ -27,29 +27,6 @@ def argsort( return tf.cast(ret, dtype=tf.int64) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def sort( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # TODO: handle stable sort when it's supported in tensorflow - # currently it supports only quicksort (unstable) - direction = "DESCENDING" if descending else "ASCENDING" - x = tf.convert_to_tensor(x) - is_bool = x.dtype.is_bool - if is_bool: - x = tf.cast(x, tf.int32) - ret = tf.sort(x, axis=axis, direction=direction) - if is_bool: - ret = tf.cast(ret, dtype=tf.bool) - return ret - - # msort @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def msort( @@ -103,3 +80,26 @@ def searchsorted( if is_supported_int_ret_dtype: return tf.searchsorted(x, v, side=side, out_type=ret_dtype) return tf.cast(tf.searchsorted(x, v, side=side), ret_dtype) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def sort( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # TODO: handle stable sort when it's supported in tensorflow + # currently it supports only quicksort (unstable) + direction = "DESCENDING" if descending else "ASCENDING" + x = tf.convert_to_tensor(x) + is_bool = x.dtype.is_bool + if is_bool: + x = tf.cast(x, tf.int32) + ret = tf.sort(x, axis=axis, direction=direction) + if is_bool: + ret = tf.cast(ret, dtype=tf.bool) + return ret diff --git a/ivy/functional/backends/tensorflow/statistical.py b/ivy/functional/backends/tensorflow/statistical.py index 53fc8b5bb57e6..4ab04d803cd90 100644 --- a/ivy/functional/backends/tensorflow/statistical.py +++ b/ivy/functional/backends/tensorflow/statistical.py @@ -9,21 +9,82 @@ from . import backend_version from ivy.utils.einsum_parser import legalise_einsum_expr -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def min( + +def _infer_dtype(dtype: tf.DType): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) +def cumprod( x: Union[tf.Tensor, tf.Variable], /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - axis = tuple(axis) if isinstance(axis, list) else axis - return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumprod(x, axis, exclusive, reverse) + + +def cumsum( + x: Union[tf.Tensor, tf.Variable], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumsum(x, axis, exclusive, reverse) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("unsigned", "int8", "int16")}, + backend_version, +) +def einsum( + equation: str, + *operands: Union[tf.Tensor, tf.Variable], + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = _get_promoted_type_of_operands(operands) + equation = legalise_einsum_expr(*[equation, *operands]) + return tf.cast(tf.einsum(equation, *operands), dtype) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -51,11 +112,21 @@ def mean( return tf.math.reduce_mean(x, axis=axis, keepdims=keepdims) -def _infer_dtype(dtype: tf.DType): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +# Array API Standard # +# -------------------# + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def min( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + axis = tuple(axis) if isinstance(axis, list) else axis + return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) def prod( @@ -144,65 +215,3 @@ def var( * size / (size - correction) ) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) -def cumprod( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumprod(x, axis, exclusive, reverse) - - -def cumsum( - x: Union[tf.Tensor, tf.Variable], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumsum(x, axis, exclusive, reverse) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("unsigned", "int8", "int16")}, - backend_version, -) -def einsum( - equation: str, - *operands: Union[tf.Tensor, tf.Variable], - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = _get_promoted_type_of_operands(operands) - equation = legalise_einsum_expr(*[equation, *operands]) - return tf.cast(tf.einsum(equation, *operands), dtype) diff --git a/ivy/functional/backends/torch/activations.py b/ivy/functional/backends/torch/activations.py index afb85ebe96f5d..f40e405bf579b 100644 --- a/ivy/functional/backends/torch/activations.py +++ b/ivy/functional/backends/torch/activations.py @@ -17,23 +17,7 @@ from . import backend_version -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def relu( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.relu(x) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def leaky_relu( - x: torch.Tensor, - /, - *, - alpha: float = 0.2, - complex_mode="jax", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.leaky_relu(x, alpha) +sigmoid.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -52,43 +36,31 @@ def gelu( return torch.nn.functional.gelu(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - if not ivy.is_array(x): - x = torch.tensor(x) - return torch.sigmoid(x, out=out) - - -sigmoid.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def softmax( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "complex", + "float16", + ) + }, + backend_version, +) +def hardswish( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if axis is None: - axis = -1 - return torch.nn.functional.softmax(x, axis) + return torch.nn.functional.hardswish(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def softplus( +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def leaky_relu( x: torch.Tensor, /, *, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, + alpha: float = 0.2, complex_mode="jax", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - kwargs = { - k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None - } - return torch.nn.functional.softplus(x, **kwargs) + return torch.nn.functional.leaky_relu(x, alpha) @with_unsupported_dtypes( @@ -125,16 +97,44 @@ def mish(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.mish(x) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "complex", - "float16", - ) - }, - backend_version, -) -def hardswish( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def relu( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.nn.functional.hardswish(x) + return torch.relu(x) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + if not ivy.is_array(x): + x = torch.tensor(x) + return torch.sigmoid(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def softmax( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if axis is None: + axis = -1 + return torch.nn.functional.softmax(x, axis) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def softplus( + x: torch.Tensor, + /, + *, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, + complex_mode="jax", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + kwargs = { + k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None + } + return torch.nn.functional.softplus(x, **kwargs) diff --git a/ivy/functional/backends/torch/creation.py b/ivy/functional/backends/torch/creation.py index 7c5a41482e44c..8f1603b9c8f79 100644 --- a/ivy/functional/backends/torch/creation.py +++ b/ivy/functional/backends/torch/creation.py @@ -25,6 +25,10 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + # noinspection PyProtectedMember @@ -47,6 +51,28 @@ def _differentiable_linspace(start, stop, num, *, device, dtype=None): return res +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +def _stack_tensors(x, dtype): + if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): + for i, item in enumerate(x): + x[i] = _stack_tensors(item, dtype) + x = torch.stack(x) + else: + if isinstance(x, (list, tuple)): + if isinstance(x[0], torch.Tensor): + x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) + else: + x = torch.as_tensor(x, dtype=dtype) + return x + + +# --- Main --- # +# ------------ # + + @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def arange( start: float, @@ -78,23 +104,6 @@ def arange( return torch.arange(start, stop, step, dtype=dtype, device=device) -arange.support_native_out = True - - -def _stack_tensors(x, dtype): - if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): - for i, item in enumerate(x): - x[i] = _stack_tensors(item, dtype) - x = torch.stack(x) - else: - if isinstance(x, (list, tuple)): - if isinstance(x[0], torch.Tensor): - x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) - else: - x = torch.as_tensor(x, dtype=dtype) - return x - - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) @asarray_to_native_arrays_and_back @asarray_infer_device @@ -137,6 +146,17 @@ def asarray( return ret.clone().detach() if copy else ret +def copy_array( + x: torch.Tensor, + *, + to_ivy_array: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if to_ivy_array: + return ivy.to_ivy(x.clone()) + return x.clone() + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -152,9 +172,6 @@ def empty( ) -empty.support_native_out = True - - def empty_like( x: torch.Tensor, /, @@ -225,14 +242,23 @@ def eye( return ret -eye.support_native_out = True - - def from_dlpack(x, /, *, out: Optional[torch.Tensor] = None): x = x.detach() if x.requires_grad else x return torch.utils.dlpack.from_dlpack(x) +def frombuffer( + buffer: bytes, + dtype: Optional[torch.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> torch.Tensor: + buffer_copy = copy.deepcopy(buffer) + dtype = ivy.as_native_dtype(dtype) + + return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -254,9 +280,6 @@ def full( ) -full.support_native_out = True - - def full_like( x: torch.Tensor, /, @@ -270,10 +293,6 @@ def full_like( return torch.full_like(x, fill_value, dtype=dtype, device=device) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - @with_unsupported_device_and_dtypes( {"2.0.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -325,9 +344,6 @@ def linspace( return ans.to(dtype) -linspace.support_native_out = True - - def linspace_helper(start, stop, num, axis=None, *, dtype=None, device): num = num.detach().numpy().item() if isinstance(num, torch.Tensor) else num start_is_array = isinstance(start, torch.Tensor) @@ -428,6 +444,46 @@ def meshgrid( return res +def one_hot( + indices: torch.Tensor, + depth: int, + /, + *, + on_value: Optional[torch.Tensor] = None, + off_value: Optional[torch.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = torch.float32 + else: + if not on_none: + dtype = torch.tensor(on_value).dtype + elif not off_none: + dtype = torch.tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) + off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) + + res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) + + if not on_none or not off_none: + res = torch.where(res == 1, on_value, off_value) + + if axis is not None: + res = torch.moveaxis(res, -1, axis) + + return res.to(device, dtype) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -438,18 +494,21 @@ def ones( return torch.ones(shape, dtype=dtype, device=device, out=out) -ones.support_native_out = True - - -def ones_like_v_0p4p0_and_above( +def ones_like_v_0p1p12_to_0p2p0( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.ones_like(x, dtype=dtype, device=device) +): + if len(x.shape) == 1: + for i in range(x.shape[0]): + x[i] = 1 + return x + for i in range(x.shape[0]): + x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) + return x def ones_like_v_0p3p0_to_0p3p1( @@ -463,21 +522,15 @@ def ones_like_v_0p3p0_to_0p3p1( return torch.ones_like(x, out=out) -def ones_like_v_0p1p12_to_0p2p0( +def ones_like_v_0p4p0_and_above( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -): - if len(x.shape) == 1: - for i in range(x.shape[0]): - x[i] = 1 - return x - for i in range(x.shape[0]): - x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) - return x +) -> torch.Tensor: + return torch.ones_like(x, dtype=dtype, device=device) def tril( @@ -486,16 +539,26 @@ def tril( return torch.tril(x, diagonal=k, out=out) -tril.support_native_out = True - - def triu( x: torch.Tensor, /, *, k: int = 0, out: Optional[torch.Tensor] = None ) -> torch.Tensor: return torch.triu(x, diagonal=k, out=out) -triu.support_native_out = True +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: torch.device, +) -> Tuple[torch.Tensor]: + n_cols = n_rows if n_cols is None else n_cols + return tuple( + torch.triu_indices( + row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device + ) + ) def zeros( @@ -508,9 +571,6 @@ def zeros( return torch.zeros(shape, dtype=dtype, device=device, out=out) -zeros.support_native_out = True - - def zeros_like( x: torch.Tensor, /, @@ -527,82 +587,12 @@ def zeros_like( array = asarray - - -def copy_array( - x: torch.Tensor, - *, - to_ivy_array: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if to_ivy_array: - return ivy.to_ivy(x.clone()) - return x.clone() - - -def one_hot( - indices: torch.Tensor, - depth: int, - /, - *, - on_value: Optional[torch.Tensor] = None, - off_value: Optional[torch.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = torch.float32 - else: - if not on_none: - dtype = torch.tensor(on_value).dtype - elif not off_none: - dtype = torch.tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) - off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) - - res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) - - if not on_none or not off_none: - res = torch.where(res == 1, on_value, off_value) - - if axis is not None: - res = torch.moveaxis(res, -1, axis) - - return res.to(device, dtype) - - -def frombuffer( - buffer: bytes, - dtype: Optional[torch.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> torch.Tensor: - buffer_copy = copy.deepcopy(buffer) - dtype = ivy.as_native_dtype(dtype) - - return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: torch.device, -) -> Tuple[torch.Tensor]: - n_cols = n_rows if n_cols is None else n_cols - return tuple( - torch.triu_indices( - row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device - ) - ) +arange.support_native_out = True +empty.support_native_out = True +eye.support_native_out = True +full.support_native_out = True +linspace.support_native_out = True +ones.support_native_out = True +tril.support_native_out = True +triu.support_native_out = True +zeros.support_native_out = True diff --git a/ivy/functional/backends/torch/data_type.py b/ivy/functional/backends/torch/data_type.py index 05b4ccb37a4c4..67f5aeabdaf84 100644 --- a/ivy/functional/backends/torch/data_type.py +++ b/ivy/functional/backends/torch/data_type.py @@ -23,7 +23,6 @@ torch.complex128: "complex128", torch.bool: "bool", } - native_dtype_dict = { "int8": torch.int8, "int16": torch.int16, @@ -68,73 +67,6 @@ def smallest_normal(self): return self._torch_finfo.tiny -# Array API Standard # -# -------------------# - - -def astype( - x: torch.Tensor, - dtype: torch.dtype, - /, - *, - copy: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return x.clone() if copy else x - return x.to(dtype) - - -def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: - try: - return list(torch.broadcast_tensors(*arrays)) - except RuntimeError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: torch.Tensor, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return torch.broadcast_to(x.reshape(-1), shape) - return torch.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return Finfo(torch.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return torch.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: - input = [] - for val in arrays_and_dtypes: - torch_val = as_native_dtype(val) - if isinstance(torch_val, torch.dtype): - torch_val = torch.tensor(1, dtype=torch_val) - input.append(torch_val) - - result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) - - for i in range(2, len(input)): - result = torch.tensor(1, dtype=torch.result_type(result, input[i])) - return as_ivy_dtype(result.dtype) - - # Extra # # ------# @@ -204,6 +136,44 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: torch.Tensor, + dtype: torch.dtype, + /, + *, + copy: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return x.clone() if copy else x + return x.to(dtype) + + +def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: + try: + return list(torch.broadcast_tensors(*arrays)) + except RuntimeError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: torch.Tensor, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return torch.broadcast_to(x.reshape(-1), shape) + return torch.broadcast_to(x, shape) + + def dtype(x: Union[torch.tensor, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -224,6 +194,20 @@ def dtype_bits(dtype_in: Union[torch.dtype, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return Finfo(torch.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return torch.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -231,3 +215,18 @@ def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: + input = [] + for val in arrays_and_dtypes: + torch_val = as_native_dtype(val) + if isinstance(torch_val, torch.dtype): + torch_val = torch.tensor(1, dtype=torch_val) + input.append(torch_val) + + result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) + + for i in range(2, len(input)): + result = torch.tensor(1, dtype=torch.result_type(result, input[i])) + return as_ivy_dtype(result.dtype) diff --git a/ivy/functional/backends/torch/device.py b/ivy/functional/backends/torch/device.py index e7f768a7f5fe1..4a86dcf02f129 100644 --- a/ivy/functional/backends/torch/device.py +++ b/ivy/functional/backends/torch/device.py @@ -18,37 +18,26 @@ torch_scatter = None -# API # -# ----# +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._prof = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True + ) -def dev( - x: torch.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, torch.device]: - dv = x.device - if as_native: - if isinstance(dv, torch.device): - dv = dv.type - elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return torch.device(dv.replace("gpu", "mps")) - return torch.device(dv.replace("gpu", "cuda")) - return as_ivy_dev(dv) + def start(self): + self._prof.__enter__() + def stop(self): + self._prof.__exit__(None, None, None) + self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) -def to_device( - x: torch.Tensor, - device: torch.device, - /, - *, - stream: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if device is None: - return x - ret = x.to(as_native_dev(device)) - if isinstance(x, torch.nn.Parameter): - return torch.nn.Parameter(ret) - return ret + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() def as_ivy_dev(device: torch.device, /): @@ -89,10 +78,21 @@ def clear_cached_mem_on_dev(device: Union[ivy.Device, torch.device], /) -> None: mps.empty_cache() -def num_gpus() -> int: - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return 1 - return torch.cuda.device_count() +# API # +# ----# + + +def dev( + x: torch.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, torch.device]: + dv = x.device + if as_native: + if isinstance(dv, torch.device): + dv = dv.type + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return torch.device(dv.replace("gpu", "mps")) + return torch.device(dv.replace("gpu", "cuda")) + return as_ivy_dev(dv) def gpu_is_available() -> bool: @@ -103,13 +103,6 @@ def gpu_is_available() -> bool: return False -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - if importlib.util.find_spec("torch_xla") is not None: - return True - return False - - def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( *args, device_shifting_dev=device_shifting_dev, **kwargs @@ -121,22 +114,30 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return fn(*args, **kwargs) -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._prof = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True - ) +def num_gpus() -> int: + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return 1 + return torch.cuda.device_count() - def start(self): - self._prof.__enter__() - def stop(self): - self._prof.__exit__(None, None, None) - self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) +def to_device( + x: torch.Tensor, + device: torch.device, + /, + *, + stream: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if device is None: + return x + ret = x.to(as_native_dev(device)) + if isinstance(x, torch.nn.Parameter): + return torch.nn.Parameter(ret) + return ret - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + if importlib.util.find_spec("torch_xla") is not None: + return True + return False diff --git a/ivy/functional/backends/torch/elementwise.py b/ivy/functional/backends/torch/elementwise.py index a0a3f5c11e053..1d03e1b29cfcc 100644 --- a/ivy/functional/backends/torch/elementwise.py +++ b/ivy/functional/backends/torch/elementwise.py @@ -13,12 +13,49 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _cast_for_unary_op(x): if not isinstance(x, torch.Tensor): x = torch.tensor(x) return x +# --- Main --- # +# ------------ # + + +@handle_numpy_arrays_in_specific_backend +def abs( + x: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = _cast_for_unary_op(x) + if x.dtype is torch.bool: + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.abs(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.acos(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.acosh(x, out=out) + + @handle_numpy_arrays_in_specific_backend def add( x1: Union[float, torch.Tensor], @@ -34,48 +71,69 @@ def add( return torch.add(x1, x2, out=out) -add.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_xor( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def angle( + input: torch.Tensor, /, *, + deg: Optional[bool] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_xor(x1, x2, out=out) + if deg: + return torch.angle(input, out=out) * (180 / pi) + else: + return torch.angle(input, out=out) -bitwise_xor.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.asin(x, out=out) -def imag( - val: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if val.dtype not in (torch.complex64, torch.complex128): - ret = torch.imag(val.to(torch.complex64)) - return ret.to(val.dtype) - return torch.imag(val) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.asinh(x, out=out) -imag.support_native_out = False +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.atan(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version +) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) @handle_numpy_arrays_in_specific_backend -def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def atan2( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.atan2(x1, x2, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.expm1(x, out=out) + return torch.atanh(x, out=out) -expm1.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def bitwise_and( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_and(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -87,68 +145,49 @@ def bitwise_invert( return torch.bitwise_not(x, out=out) -bitwise_invert.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.isfinite(x) - - +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def isinf( - x: torch.Tensor, +def bitwise_left_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, - detect_positive: bool = True, - detect_negative: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - if detect_negative and detect_positive: - return torch.isinf(x) - elif detect_negative: - return torch.isneginf(x) - elif detect_positive: - return torch.isposinf(x) - return torch.full_like(x, False, dtype=torch.bool) + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_left_shift(x1, x2, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def bitwise_or( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.eq(x1, x2, out=out) - - -equal.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_or(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less_equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def bitwise_right_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.less_equal(x1, x2, out=out) - - -less_equal.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) + return torch.bitwise_right_shift(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_and( +def bitwise_xor( x1: Union[int, bool, torch.Tensor], x2: Union[int, bool, torch.Tensor], /, @@ -156,10 +195,7 @@ def bitwise_and( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_and(x1, x2, out=out) - - -bitwise_and.support_native_out = True + return torch.bitwise_xor(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @@ -173,135 +209,162 @@ def ceil(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.ceil(x, out=out) -ceil.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.cos(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.floor(x, out=out) + return torch.cosh(x, out=out) -floor.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.deg2rad(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmin( - x1: torch.Tensor, - x2: torch.Tensor, +@handle_numpy_arrays_in_specific_backend +def divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.fmin(x1, x2, out=None) - - -fmin.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + ret = torch.div(x1, x2) + if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): + ret = ivy.astype(ret, x1.dtype, copy=False) + else: + ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) + return ret -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.asin(x, out=out) +def equal( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.eq(x1, x2, out=out) -asin.support_native_out = True +# Extra # +# ------# -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.asinh(x, out=out) + return torch.erf(x, out=out) -asinh.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.exp(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def sign( - x: torch.Tensor, +def exp2( + x: Union[torch.Tensor, float, list, tuple], /, *, - np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - if "complex" in str(x.dtype): - if np_variant: - return torch.where( - x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j - ) - return torch.sgn(x, out=out) - return torch.sign(x, out=out) - - -sign.support_native_out = True + return torch.exp2(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.sqrt(x, out=out) - - -sqrt.support_native_out = True + return torch.expm1(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.cosh(x, out=out) - - -cosh.support_native_out = True + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.floor(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log10(x, out=out) - - -log10.support_native_out = True +def floor_divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if ivy.exists(out): + if not ivy.is_float_dtype(out): + return ivy.inplace_update( + out, torch.floor(torch.div(x1, x2)).type(out.dtype) + ) + return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log2(x, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmin( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.fmin(x1, x2, out=None) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "complex")}, + backend_version, +) @handle_numpy_arrays_in_specific_backend -def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log1p(x, out=out) - - -log1p.support_native_out = True +def fmod( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmod(x1, x2, out=None) -@handle_numpy_arrays_in_specific_backend -def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.isnan(x) +def gcd( + x1: Union[torch.Tensor, int, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.gcd(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less( +def greater( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -309,14 +372,12 @@ def less( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.lt(x1, x2, out=out) - - -less.support_native_out = True + return torch.greater(x1, x2, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def multiply( +def greater_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -324,56 +385,73 @@ def multiply( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.multiply(x1, x2, out=out) + return torch.greater_equal(x1, x2, out=out) -multiply.support_native_out = True +def imag( + val: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if val.dtype not in (torch.complex64, torch.complex128): + ret = torch.imag(val.to(torch.complex64)) + return ret.to(val.dtype) + return torch.imag(val) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.cos(x, out=out) + return torch.isfinite(x) -cos.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def isinf( + x: torch.Tensor, + /, + *, + detect_positive: bool = True, + detect_negative: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = _cast_for_unary_op(x) + if detect_negative and detect_positive: + return torch.isinf(x) + elif detect_negative: + return torch.isneginf(x) + elif detect_positive: + return torch.isposinf(x) + return torch.full_like(x, False, dtype=torch.bool) @handle_numpy_arrays_in_specific_backend -def logical_not( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: +def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.logical_not(x.type(torch.bool), out=out) + return torch.isnan(x) -logical_not.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.isreal(x) +@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def lcm( + x1: torch.Tensor, + x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2) - if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): - ret = ivy.astype(ret, x1.dtype, copy=False) - else: - ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) - return ret - - -divide.support_native_out = True + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.lcm(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater( +def less( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -381,15 +459,12 @@ def greater( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater(x1, x2, out=out) - - -greater.support_native_out = True + return torch.lt(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater_equal( +def less_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -397,46 +472,60 @@ def greater_equal( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater_equal(x1, x2, out=out) - - -greater_equal.support_native_out = True + return torch.less_equal(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.acos(x, out=out) + return torch.log(x, out=out) -acos.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log10(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def lcm( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.lcm(x1, x2, out=out) +def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log1p(x, out=out) -lcm.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log2(x, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def logical_xor( +def logaddexp( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.logaddexp(x1, x2, out=out) -logical_xor.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def logaddexp2( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + if not ivy.is_float_dtype(x1): + x1 = x1.type(ivy.default_float_dtype(as_native=True)) + x2 = x2.type(ivy.default_float_dtype(as_native=True)) + return torch.logaddexp2(x1, x2, out=out) @handle_numpy_arrays_in_specific_backend @@ -446,7 +535,12 @@ def logical_and( return torch.logical_and(x1.type(torch.bool), x2.type(torch.bool), out=out) -logical_and.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def logical_not( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.logical_not(x.type(torch.bool), out=out) @handle_numpy_arrays_in_specific_backend @@ -456,27 +550,72 @@ def logical_or( return torch.logical_or(x1.type(torch.bool), x2.type(torch.bool), out=out) -logical_or.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.acosh(x, out=out) - - -acosh.support_native_out = True +def logical_xor( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sin(x, out=out) +def maximum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 >= x2, x1, x2, out=out) + return torch.maximum(x1, x2, out=out) -sin.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def minimum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 <= x2, x1, x2, out=out) + return torch.minimum(x1, x2, out=out) + + +@handle_numpy_arrays_in_specific_backend +def multiply( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.multiply(x1, x2, out=out) + + +def nan_to_num( + x: torch.Tensor, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if copy: + return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) + else: + x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) + return x @handle_numpy_arrays_in_specific_backend @@ -487,9 +626,6 @@ def negative( return torch.neg(x, out=out) -negative.support_native_out = True - - @handle_numpy_arrays_in_specific_backend def not_equal( x1: Union[float, torch.Tensor], @@ -502,56 +638,108 @@ def not_equal( return torch.not_equal(x1, x2, out=out) -not_equal.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def positive( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.positive(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def tanh( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None +def pow( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tanh(x, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.pow(x1, x2, out=out) -tanh.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.rad2deg(x, out=out) + + +def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.real(x) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def reciprocal( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.reciprocal(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor_divide( +def remainder( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, + modulus: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.exists(out): - if not ivy.is_float_dtype(out): - return ivy.inplace_update( - out, torch.floor(torch.div(x1, x2)).type(out.dtype) - ) - return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) + if not modulus: + res = x1 / x2 + res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) + diff = res - res_floored + diff, x2 = ivy.promote_types_of_inputs(diff, x2) + if ivy.exists(out): + if out.dtype != x2.dtype: + return ivy.inplace_update( + out, torch.round(torch.mul(diff, x2)).to(out.dtype) + ) + return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) + return torch.remainder(x1, x2, out=out).to(x1.dtype) -floor_divide.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def round( + x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.round(x, decimals=decimals, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_or( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def sign( + x: torch.Tensor, /, *, + np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_or(x1, x2, out=out) + x = _cast_for_unary_op(x) + if "complex" in str(x.dtype): + if np_variant: + return torch.where( + x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j + ) + return torch.sgn(x, out=out) + return torch.sign(x, out=out) -bitwise_or.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sin(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -561,15 +749,11 @@ def sinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.sinh(x, out=out) -sinh.support_native_out = True - - +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def positive( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: +def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.positive(x) + return torch.sqrt(x, out=out) @handle_numpy_arrays_in_specific_backend @@ -578,37 +762,35 @@ def square(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.T return torch.square(x, out=out) -square.support_native_out = True - - @handle_numpy_arrays_in_specific_backend -def pow( +def subtract( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.pow(x1, x2, out=out) + if alpha not in (1, None): + return torch.subtract(x1, x2, alpha=alpha, out=out) + return torch.subtract(x1, x2, out=out) -pow.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.tan(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def round( - x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None +def tanh( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.round(x, decimals=decimals, out=out) - - -round.support_native_out = True + x = _cast_for_unary_op(x) + return torch.tanh(x, out=out) def trapz( @@ -634,9 +816,6 @@ def trapz( return torch.trapezoid(y, x=x, dim=axis) -trapz.support_native_out = False - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: @@ -649,388 +828,87 @@ def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return ret -trunc.support_native_out = True - - +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def abs( - x: Union[float, torch.Tensor], +def trunc_divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x = _cast_for_unary_op(x) - if x.dtype is torch.bool: - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.abs(x, out=out) - - -abs.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def logaddexp( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.logaddexp(x1, x2, out=out) - - -logaddexp.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def logaddexp2( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - if not ivy.is_float_dtype(x1): - x1 = x1.type(ivy.default_float_dtype(as_native=True)) - x2 = x2.type(ivy.default_float_dtype(as_native=True)) - return torch.logaddexp2(x1, x2, out=out) - - -logaddexp2.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tan(x, out=out) + ret = torch.div(x1, x2, rounding_mode="trunc") + if ivy.is_float_dtype(x1.dtype): + ret = ret.to(x1.dtype) + else: + ret = ret.to(ivy.default_float_dtype(as_native=True)) + return ret +add.support_native_out = True +bitwise_xor.support_native_out = True +imag.support_native_out = False +expm1.support_native_out = True +bitwise_invert.support_native_out = True +equal.support_native_out = True +less_equal.support_native_out = True +bitwise_and.support_native_out = True +ceil.support_native_out = True +floor.support_native_out = True +fmin.support_native_out = True +asin.support_native_out = True +asinh.support_native_out = True +sign.support_native_out = True +sqrt.support_native_out = True +cosh.support_native_out = True +log10.support_native_out = True +log1p.support_native_out = True +less.support_native_out = True +multiply.support_native_out = True +cos.support_native_out = True +logical_not.support_native_out = True +divide.support_native_out = True +greater.support_native_out = True +greater_equal.support_native_out = True +acos.support_native_out = True +lcm.support_native_out = True +logical_xor.support_native_out = True +logical_and.support_native_out = True +logical_or.support_native_out = True +acosh.support_native_out = True +sin.support_native_out = True +negative.support_native_out = True +not_equal.support_native_out = True +tanh.support_native_out = True +floor_divide.support_native_out = True +bitwise_or.support_native_out = True +sinh.support_native_out = True +square.support_native_out = True +pow.support_native_out = True +round.support_native_out = True +trapz.support_native_out = False +trunc.support_native_out = True +abs.support_native_out = True +logaddexp.support_native_out = True +logaddexp2.support_native_out = True tan.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.atan(x, out=out) - - atan.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version -) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) -@handle_numpy_arrays_in_specific_backend -def atan2( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.atan2(x1, x2, out=out) - - atan2.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log(x, out=out) - - log.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.exp(x, out=out) - - exp.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def exp2( - x: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.exp2(x, out=out) - - exp2.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def subtract( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - alpha: Optional[Union[int, float]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if alpha not in (1, None): - return torch.subtract(x1, x2, alpha=alpha, out=out) - return torch.subtract(x1, x2, out=out) - - subtract.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def remainder( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - modulus: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if not modulus: - res = x1 / x2 - res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) - diff = res - res_floored - diff, x2 = ivy.promote_types_of_inputs(diff, x2) - if ivy.exists(out): - if out.dtype != x2.dtype: - return ivy.inplace_update( - out, torch.round(torch.mul(diff, x2)).to(out.dtype) - ) - return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) - return torch.remainder(x1, x2, out=out).to(x1.dtype) - - remainder.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.atanh(x, out=out) - - atanh.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_right_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) - return torch.bitwise_right_shift(x1, x2, out=out) - - bitwise_right_shift.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_left_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_left_shift(x1, x2, out=out) - - bitwise_left_shift.support_native_out = True - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.erf(x, out=out) - - erf.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def minimum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 <= x2, x1, x2, out=out) - return torch.minimum(x1, x2, out=out) - - minimum.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def maximum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 >= x2, x1, x2, out=out) - return torch.maximum(x1, x2, out=out) - - maximum.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def reciprocal( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.reciprocal(x, out=out) - - reciprocal.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.deg2rad(x, out=out) - - deg2rad.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.rad2deg(x, out=out) - - rad2deg.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def trunc_divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2, rounding_mode="trunc") - if ivy.is_float_dtype(x1.dtype): - ret = ret.to(x1.dtype) - else: - ret = ret.to(ivy.default_float_dtype(as_native=True)) - return ret - - -@handle_numpy_arrays_in_specific_backend -def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.isreal(x) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "complex")}, - backend_version, -) -@handle_numpy_arrays_in_specific_backend -def fmod( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmod(x1, x2, out=None) - - fmod.support_native_out = True - - -def gcd( - x1: Union[torch.Tensor, int, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.gcd(x1, x2, out=out) - - gcd.support_native_out = True - - -def angle( - input: torch.Tensor, - /, - *, - deg: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if deg: - return torch.angle(input, out=out) * (180 / pi) - else: - return torch.angle(input, out=out) - - angle.support_native_out = True - - -def nan_to_num( - x: torch.Tensor, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if copy: - return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) - else: - x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) - return x - - -def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.real(x) diff --git a/ivy/functional/backends/torch/experimental/activations.py b/ivy/functional/backends/torch/experimental/activations.py index 98d72ac526c45..008312f394795 100644 --- a/ivy/functional/backends/torch/experimental/activations.py +++ b/ivy/functional/backends/torch/experimental/activations.py @@ -10,6 +10,16 @@ from . import backend_version +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def elu( + x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + ret = torch.nn.functional.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def logit( x: torch.Tensor, @@ -21,15 +31,11 @@ def logit( return torch.logit(x, eps=eps, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def thresholded_relu( - x: torch.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = None, - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def logsigmoid( + input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.threshold(x, threshold=threshold, value=0) + return torch.nn.functional.logsigmoid(input) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -37,13 +43,6 @@ def relu6(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return torch.nn.functional.relu6(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def logsigmoid( - input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.nn.functional.logsigmoid(input) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def selu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: ret = torch.nn.functional.selu(x) @@ -57,11 +56,12 @@ def silu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.silu(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def elu( - x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def thresholded_relu( + x: torch.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - ret = torch.nn.functional.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) + return torch.threshold(x, threshold=threshold, value=0) diff --git a/ivy/functional/backends/torch/experimental/creation.py b/ivy/functional/backends/torch/experimental/creation.py index a1b94a8a894cf..14dd6d3c12b86 100644 --- a/ivy/functional/backends/torch/experimental/creation.py +++ b/ivy/functional/backends/torch/experimental/creation.py @@ -12,33 +12,26 @@ ) from .. import backend_version -# noinspection PyProtectedMember - -# Array API Standard # -# -------------------# +vorbis_window.support_native_out = False +hann_window.support_native_out = False +blackman_window.support_native_out = False +trilu.support_native_out = True -@with_unsupported_device_and_dtypes( - {"2.0.1 and below": {"cpu": ("float16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def blackman_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.kaiser_window( - window_length, - periodic, - beta, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.blackman_window( + size, + periodic=periodic, dtype=dtype, - layout=torch.strided, - device=None, - requires_grad=False, ) @@ -64,29 +57,6 @@ def hamming_window( ) -def vorbis_window( - window_length: torch.tensor, - *, - dtype: torch.dtype = torch.float32, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.tensor( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, - ) - - -vorbis_window.support_native_out = False - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def hann_window( size: int, @@ -103,7 +73,34 @@ def hann_window( ) -hann_window.support_native_out = False +# noinspection PyProtectedMember + + +# Array API Standard # +# -------------------# + + +@with_unsupported_device_and_dtypes( + {"2.0.1 and below": {"cpu": ("float16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.kaiser_window( + window_length, + periodic, + beta, + dtype=dtype, + layout=torch.strided, + device=None, + requires_grad=False, + ) def tril_indices( @@ -126,6 +123,19 @@ def tril_indices( ) +def trilu( + x: torch.Tensor, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if upper: + return torch.triu(x, diagonal=k, out=out) + return torch.tril(x, diagonal=k, out=out) + + def unsorted_segment_min( data: torch.Tensor, segment_ids: torch.Tensor, @@ -152,25 +162,6 @@ def unsorted_segment_min( return res -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.blackman_window( - size, - periodic=periodic, - dtype=dtype, - ) - - -blackman_window.support_native_out = False - - def unsorted_segment_sum( data: torch.Tensor, segment_ids: torch.Tensor, @@ -196,17 +187,21 @@ def unsorted_segment_sum( return res -def trilu( - x: torch.Tensor, - /, +def vorbis_window( + window_length: torch.tensor, *, - k: int = 0, - upper: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if upper: - return torch.triu(x, diagonal=k, out=out) - return torch.tril(x, diagonal=k, out=out) - - -trilu.support_native_out = True + dtype: torch.dtype = torch.float32, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.tensor( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, + ) diff --git a/ivy/functional/backends/torch/experimental/elementwise.py b/ivy/functional/backends/torch/experimental/elementwise.py index dcb4e30dcf1dc..ab2d5a0c2125d 100644 --- a/ivy/functional/backends/torch/experimental/elementwise.py +++ b/ivy/functional/backends/torch/experimental/elementwise.py @@ -14,49 +14,69 @@ from .. import backend_version -@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) -def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.lgamma(x, out=out) +# --- Helpers --- # +# --------------- # -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmax( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmax(x1, x2, out=None) - +def _are_suitable_types_for_torch_lerp(input, end, weight): + suitable_types = [ + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ] -fmax.support_native_out = True + if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): + return False + else: + if input.dtype not in suitable_types or end.dtype not in suitable_types: + return False + if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): + return False + else: + if isinstance(weight, torch.Tensor): + if weight.dtype not in [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ]: + return False -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sinc(x, out=out) + return True -sinc.support_native_out = True +# --- Main --- # +# ------------ # -def float_power( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], +def allclose( + x1: torch.Tensor, + x2: torch.Tensor, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - # Native out is supported but with restrictions leading - # to failures hence letting ivy handle it. - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.float_power(x1, x2, out=out) +) -> bool: + ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) + return torch.tensor(ret) -float_power.support_native_out = True +def conj( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + conj_x = torch.conj(x) + return torch.resolve_conj(input=conj_x) def copysign( @@ -73,9 +93,6 @@ def copysign( return torch.copysign(torch.as_tensor(x1), x2, out=out) -copysign.support_native_out = True - - def count_nonzero( a: torch.Tensor, /, @@ -107,42 +124,6 @@ def count_nonzero( return x -count_nonzero.support_native_out = False - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def nansum( - x: torch.Tensor, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[torch.dtype] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) - - -nansum.support_native_out = False - - -def isclose( - a: torch.Tensor, - b: torch.Tensor, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -isclose.support_native_out = False - - def diff( x: Union[torch.Tensor, list, tuple], /, @@ -167,41 +148,14 @@ def diff( return torch.diff(x, n=n, dim=axis, prepend=prepend, append=append) -def signbit( - x: Union[torch.Tensor, float, int, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.signbit(x, out=out) - - -signbit.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def hypot( - x1: torch.Tensor, - x2: torch.Tensor, +def digamma( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.hypot(x1, x2) - - -def allclose( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[torch.Tensor] = None, -) -> bool: - ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - return torch.tensor(ret) + return torch.special.digamma(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -214,39 +168,39 @@ def fix( return torch.fix(x, out=out) -fix.support_native_out = True +def float_power( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # Native out is supported but with restrictions leading + # to failures hence letting ivy handle it. + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.float_power(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def nextafter( +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmax( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nextafter(x1, x2) - - -nextafter.support_native_out = True + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmax(x1, x2, out=None) -def zeta( +def frexp( x: torch.Tensor, - q: torch.Tensor, /, *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) - temp = torch.logical_and(temp, torch.le(q, 0)) - nan_indices = torch.logical_or(temp, torch.lt(x, 1)) - result = torch.special.zeta(x, q) - result.masked_fill_(nan_indices, float("nan")) - return result - - -zeta.support_native_out = False + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + mantissa, exponent = torch.frexp(x, out=out) + return mantissa, exponent def gradient( @@ -270,25 +224,28 @@ def gradient( return grad -@with_supported_dtypes( - {"2.0.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def xlogy( - x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None -) -> torch.tensor: - x, y = promote_types_of_inputs(x, y) - return torch.xlogy(x, y, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def hypot( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.hypot(x1, x2) -def conj( - x: torch.Tensor, +def isclose( + a: torch.Tensor, + b: torch.Tensor, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - conj_x = torch.conj(x) - return torch.resolve_conj(input=conj_x) + return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) def ldexp( @@ -301,39 +258,6 @@ def ldexp( return torch.ldexp(x1, x2, out=out) -def _are_suitable_types_for_torch_lerp(input, end, weight): - suitable_types = [ - torch.int8, - torch.int16, - torch.int32, - torch.int64, - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ] - - if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): - return False - else: - if input.dtype not in suitable_types or end.dtype not in suitable_types: - return False - - if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): - return False - else: - if isinstance(weight, torch.Tensor): - if weight.dtype not in [ - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ]: - return False - - return True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def lerp( input: torch.Tensor, @@ -346,40 +270,100 @@ def lerp( return torch.lerp(input, end, weight, out=out) -lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( - _are_suitable_types_for_torch_lerp(input, end, weight) -) -lerp.support_native_out = True +@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) +def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.lgamma(x, out=out) -def frexp( +def modf( x: torch.Tensor, /, *, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - mantissa, exponent = torch.frexp(x, out=out) - return mantissa, exponent + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + modf_x = torch.modf(x) + return torch.resolve_modf(input=modf_x) -def modf( +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def nansum( x: torch.Tensor, /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[torch.dtype] = None, + keepdims: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - modf_x = torch.modf(x) - return torch.resolve_modf(input=modf_x) + dtype = ivy.as_native_dtype(dtype) + return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def digamma( +def nextafter( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nextafter(x1, x2) + + +def signbit( + x: Union[torch.Tensor, float, int, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.signbit(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sinc(x, out=out) + + +@with_supported_dtypes( + {"2.0.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def xlogy( + x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None +) -> torch.tensor: + x, y = promote_types_of_inputs(x, y) + return torch.xlogy(x, y, out=out) + + +def zeta( x: torch.Tensor, + q: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.digamma(x, out=out) + temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) + temp = torch.logical_and(temp, torch.le(q, 0)) + nan_indices = torch.logical_or(temp, torch.lt(x, 1)) + result = torch.special.zeta(x, q) + result.masked_fill_(nan_indices, float("nan")) + return result +fmax.support_native_out = True +sinc.support_native_out = True +float_power.support_native_out = True +copysign.support_native_out = True +count_nonzero.support_native_out = False +nansum.support_native_out = False +isclose.support_native_out = False +signbit.support_native_out = True +fix.support_native_out = True +nextafter.support_native_out = True +zeta.support_native_out = False +lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( + _are_suitable_types_for_torch_lerp(input, end, weight) +) +lerp.support_native_out = True digamma.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/layers.py b/ivy/functional/backends/torch/experimental/layers.py index 9d35ad4f31232..0c26508af9aca 100644 --- a/ivy/functional/backends/torch/experimental/layers.py +++ b/ivy/functional/backends/torch/experimental/layers.py @@ -16,257 +16,34 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode -def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_first"): - # Determine depth pooling - kernel, strides, depth_pooling = _depth_max_pooling_helper( - x.shape, kernel, strides, dims=dims, data_format=data_format - ) - if depth_pooling: - x = torch.permute(x, (0, 2, 1, *range(3, dims + 2))) - return x, kernel, strides, depth_pooling - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def max_pool1d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NWC": - x = x.permute((0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) - - if not depth_pooling: - new_kernel = [dilation[0] * (kernel[0] - 1) + 1] - - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2] - else: - pad_list = [item for sublist in padding for item in sublist] - - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) - - if depth_pooling: - res = torch.permute(res, (0, 2, 1)) - if data_format == "NWC": - res = res.permute((0, 2, 1)) - return res - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def max_pool2d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - kernel = ( - [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 3, 1, 2]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 3, 1, 2]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) - - x_shape = list(x.shape[2:]) - if not depth_pooling: - new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] - - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] - else: - # torch pad takes width padding first, then height padding - padding = (padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] +# --- Helpers --- # +# --------------- # - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = torch.permute(res, (0, 2, 1, 3)) - if data_format == "NHWC": - return res.permute(0, 2, 3, 1) - return res +def _add_ceil_pad_to_pad_list(num_pad, k, c): + return num_pad + (num_pad - ((k * num_pad) / (k - c))) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", +def _adjust_num_padded_values_to_ceil( + pad_specific, num_padded_values, x_shape, kernel, strides, dims +): + for i in range(dims): + pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] + _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) + num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( + num_padded_values[i][-1], kernel[i], c ) - }, - backend_version, -) -def max_pool3d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + return num_padded_values - if data_format == "NDHWC": - x = x.permute(0, 4, 1, 2, 3) - kernel = ( - [kernel[i] for i in [0, 4, 1, 2, 3]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 4, 1, 2, 3]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 4, 1, 2, 3]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" +def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_first"): + # Determine depth pooling + kernel, strides, depth_pooling = _depth_max_pooling_helper( + x.shape, kernel, strides, dims=dims, data_format=data_format ) - - if not depth_pooling: - x_shape = x.shape[2:] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - pad_list = [ - pad_w // 2, - pad_w - pad_w // 2, - pad_h // 2, - pad_h - pad_h // 2, - pad_d // 2, - pad_d - pad_d // 2, - ] - else: - # torch pad takes width padding first, then height, then depth - padding = (padding[2], padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] - - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = res.permute(0, 2, 1, 3, 4) - if data_format == "NDHWC": - res = res.permute(0, 2, 3, 4, 1) - return res - - -def _add_ceil_pad_to_pad_list(num_pad, k, c): - return num_pad + (num_pad - ((k * num_pad) / (k - c))) + x = torch.permute(x, (0, 2, 1, *range(3, dims + 2))) + return x, kernel, strides, depth_pooling def _get_specific_pad(x_shape, kernel, strides, padding, dims): @@ -293,6 +70,27 @@ def _get_specific_pad(x_shape, kernel, strides, padding, dims): return padding, pad_specific +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool1d(input, output_size): + return torch.nn.functional.adaptive_avg_pool1d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool2d(input, output_size): + return torch.nn.functional.adaptive_avg_pool2d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_max_pool2d( + input: torch.Tensor, output_size: Union[Sequence[int], int] +) -> torch.Tensor: + return torch.nn.functional.adaptive_max_pool2d(input, output_size) + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) def avg_pool1d( x: torch.Tensor, @@ -363,18 +161,6 @@ def avg_pool1d( return res -def _adjust_num_padded_values_to_ceil( - pad_specific, num_padded_values, x_shape, kernel, strides, dims -): - for i in range(dims): - pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] - _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) - num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( - num_padded_values[i][-1], kernel[i], c - ) - return num_padded_values - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -630,23 +416,222 @@ def dct( else: x = x * axis_dim_float - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(None, axis_dim) - dct_out = torch.real( - torch.fft.irfft( - scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis - ) - )[axis_idx] - return dct_out + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(None, axis_dim) + dct_out = torch.real( + torch.fft.irfft( + scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis + ) + )[axis_idx] + return dct_out + + elif type == 4: + dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(1, None, 2) + dct_out = dct_2[axis_idx] + if norm == "ortho": + dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) + return dct_out + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def dropout( + x: torch.Tensor, + prob: float, + /, + *, + scale: bool = True, + dtype: torch.dtype = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = ivy.astype(x, dtype) if dtype else x + res = torch.nn.functional.dropout(x, prob, training=training) + res = torch.multiply(res, (1.0 - prob)) if not scale else res + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout1d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 3 + if data_format == "NWC": + perm = (0, 2, 1) if is_batched else (1, 0) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout1d(x, prob, training=training) + if data_format == "NWC": + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout2d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 4 + if data_format == "NHWC": + perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout2d(x, prob, training=training) + if data_format == "NHWC": + perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def dropout3d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 5 + if data_format == "NDHWC": + perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout3d(x, prob, training=training) + if data_format == "NDHWC": + perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def embedding( + weights: torch.Tensor, + indices: torch.Tensor, + /, + *, + max_norm: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def fft( + x: torch.Tensor, + dim: int, + /, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if x.dtype in [torch.int64, torch.float64, torch.complex128]: + out_dtype = torch.complex128 + else: + out_dtype = torch.complex64 + return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) + - elif type == 4: - dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(1, None, 2) - dct_out = dct_2[axis_idx] - if norm == "ortho": - dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) - return dct_out +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def fft2( + x: torch.Tensor, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " + ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return torch.tensor( + torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 + ) def idct( @@ -663,19 +648,9 @@ def idct( return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def fft( +def ifft( x: torch.Tensor, dim: int, - /, *, norm: str = "backward", n: Union[int, Tuple[int]] = None, @@ -702,89 +677,200 @@ def fft( ) if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.dtype in [torch.int64, torch.float64, torch.complex128]: - out_dtype = torch.complex128 - else: - out_dtype = torch.complex64 - return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) + return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def dropout( +def ifftn( x: torch.Tensor, - prob: float, - /, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - scale: bool = True, - dtype: torch.dtype = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, + norm: Optional[str] = "backward", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = ivy.astype(x, dtype) if dtype else x - res = torch.nn.functional.dropout(x, prob, training=training) - res = torch.multiply(res, (1.0 - prob)) if not scale else res - return res + return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) -dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( - kwargs.get("noise_shape") is None and kwargs.get("seed") is None -) +def interpolate( + x: torch.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Literal[ + "linear", + "bilinear", + "trilinear", + "nearest", + "area", + "nearest_exact", + "tf_area", + "bicubic", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", + ] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: bool = False, + out: Optional[torch.Tensor] = None, +): + return torch.nn.functional.interpolate( + x, + size=size, + mode=mode, + align_corners=align_corners, + antialias=antialias, + scale_factor=scale_factor, + recompute_scale_factor=recompute_scale_factor, + ) -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def dropout1d( +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def max_pool1d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 3 + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NWC": - perm = (0, 2, 1) if is_batched else (1, 0) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout1d(x, prob, training=training) + x = x.permute((0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + if not depth_pooling: + new_kernel = [dilation[0] * (kernel[0] - 1) + 1] + + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2] + else: + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) + + if depth_pooling: + res = torch.permute(res, (0, 2, 1)) if data_format == "NWC": - res = torch.permute(res, perm) + res = res.permute((0, 2, 1)) return res @with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, backend_version, ) -def dropout2d( +def max_pool2d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 4 + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NHWC": - perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout2d(x, prob, training=training) + x = x.permute(0, 3, 1, 2) + kernel = ( + [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 3, 1, 2]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 3, 1, 2]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + x_shape = list(x.shape[2:]) + if not depth_pooling: + new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] + + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] + else: + # torch pad takes width padding first, then height padding + padding = (padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) + if depth_pooling: + res = torch.permute(res, (0, 2, 1, 3)) if data_format == "NHWC": - perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) - res = torch.permute(res, perm) + return res.permute(0, 2, 3, 1) return res @@ -797,184 +883,87 @@ def dropout2d( }, backend_version, ) -def dropout3d( +def max_pool3d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 5 - if data_format == "NDHWC": - perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout3d(x, prob, training=training) - if data_format == "NDHWC": - perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) - res = torch.permute(res, perm) - return res - + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -def ifft( - x: torch.Tensor, - dim: int, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " + if data_format == "NDHWC": + x = x.permute(0, 4, 1, 2, 3) + kernel = ( + [kernel[i] for i in [0, 4, 1, 2, 3]] + if len(kernel) == (dims + 2) + else kernel ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" + strides = ( + [strides[i] for i in [0, 4, 1, 2, 3]] + if len(strides) == (dims + 2) + else strides ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" + padding = ( + [padding[i] for i in [0, 4, 1, 2, 3]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def embedding( - weights: torch.Tensor, - indices: torch.Tensor, - /, - *, - max_norm: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) - - -embedding.support_native_out = False - -def interpolate( - x: torch.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Literal[ - "linear", - "bilinear", - "trilinear", - "nearest", - "area", - "nearest_exact", - "tf_area", - "bicubic", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", - ] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: bool = False, - out: Optional[torch.Tensor] = None, -): - return torch.nn.functional.interpolate( - x, - size=size, - mode=mode, - align_corners=align_corners, - antialias=antialias, - scale_factor=scale_factor, - recompute_scale_factor=recompute_scale_factor, + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" ) + if not depth_pooling: + x_shape = x.shape[2:] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] -interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ - "tf_area", - "nd", - "bicubic_tensorflow", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", -] - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_max_pool2d( - input: torch.Tensor, output_size: Union[Sequence[int], int] -) -> torch.Tensor: - return torch.nn.functional.adaptive_max_pool2d(input, output_size) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool1d(input, output_size): - return torch.nn.functional.adaptive_avg_pool1d(input, output_size) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool2d(input, output_size): - return torch.nn.functional.adaptive_avg_pool2d(input, output_size) - + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + pad_list = [ + pad_w // 2, + pad_w - pad_w // 2, + pad_h // 2, + pad_h - pad_h // 2, + pad_d // 2, + pad_d - pad_d // 2, + ] + else: + # torch pad takes width padding first, then height, then depth + padding = (padding[2], padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def fft2( - x: torch.Tensor, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.tensor( - torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 - ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) -def ifftn( - x: torch.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) + if depth_pooling: + res = res.permute(0, 2, 1, 3, 4) + if data_format == "NDHWC": + res = res.permute(0, 2, 3, 4, 1) + return res @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -1010,3 +999,18 @@ def rfftn( return torch.tensor( torch.fft.rfftn(x, s, axes, norm=norm, out=out), dtype=torch.complex128 ) + + +dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( + kwargs.get("noise_shape") is None and kwargs.get("seed") is None +) +embedding.support_native_out = False +interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ + "tf_area", + "nd", + "bicubic_tensorflow", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", +] diff --git a/ivy/functional/backends/torch/experimental/linear_algebra.py b/ivy/functional/backends/torch/experimental/linear_algebra.py index 68e98bff30a60..1c3a668fe152c 100644 --- a/ivy/functional/backends/torch/experimental/linear_algebra.py +++ b/ivy/functional/backends/torch/experimental/linear_algebra.py @@ -12,6 +12,37 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size +diagflat.support_native_out = False +kron.support_native_out = True +matrix_exp.support_native_out = True +eig.support_native_out = False +eigvals.support_native_out = False +multi_dot.support_native_out = True +cond.support_native_out = False +dot.support_native_out = True + + +def adjoint( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_valid_dimension_size(x) + return torch.adjoint(x).resolve_conj() + + +@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( + x: torch.Tensor, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.cond(x, p=p, out=out) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def diagflat( x: torch.Tensor, @@ -97,32 +128,14 @@ def diagflat( return ret -diagflat.support_native_out = False - - -def kron( +def dot( a: torch.Tensor, b: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.kron(a, b, out=out) - - -kron.support_native_out = True - - -def matrix_exp( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.linalg.matrix_exp(x) - - -matrix_exp.support_native_out = True + return torch.matmul(a, b) def eig( @@ -133,53 +146,20 @@ def eig( return torch.linalg.eig(x) -eig.support_native_out = False - - def eigvals(x: torch.Tensor, /) -> torch.Tensor: if not torch.is_complex(x): x = x.to(torch.complex128) return torch.linalg.eigvals(x) -eigvals.support_native_out = False - - -def adjoint( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_dimension_size(x) - return torch.adjoint(x).resolve_conj() - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def multi_dot( - x: Sequence[torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.multi_dot(x, out=out) - - -multi_dot.support_native_out = True - - -@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( - x: torch.Tensor, +def kron( + a: torch.Tensor, + b: torch.Tensor, /, *, - p: Optional[Union[None, int, str]] = None, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.cond(x, p=p, out=out) - - -cond.support_native_out = False +) -> torch.tensor: + return torch.kron(a, b, out=out) def lu_factor( @@ -192,14 +172,20 @@ def lu_factor( raise IvyNotImplementedException() -def dot( - a: torch.Tensor, - b: torch.Tensor, +def matrix_exp( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.matmul(a, b) + return torch.linalg.matrix_exp(x) -dot.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def multi_dot( + x: Sequence[torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.multi_dot(x, out=out) diff --git a/ivy/functional/backends/torch/experimental/losses.py b/ivy/functional/backends/torch/experimental/losses.py index db3442685a678..e3b6e9485e543 100644 --- a/ivy/functional/backends/torch/experimental/losses.py +++ b/ivy/functional/backends/torch/experimental/losses.py @@ -3,6 +3,24 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, + backend_version, +) +def huber_loss( + input: torch.Tensor, + target: torch.Tensor, + /, + *, + reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, +) -> torch.Tensor: + return torch.nn.functional.huber_loss( + input, target, reduction=reduction, delta=delta + ) + + # Assuming ivy and backend_version are imported and defined properly @@ -52,20 +70,3 @@ def smooth_l1_loss( beta=beta, reduction=reduction, ) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, - backend_version, -) -def huber_loss( - input: torch.Tensor, - target: torch.Tensor, - /, - *, - reduction: Optional[str] = "mean", - delta: Optional[float] = 1.0, -) -> torch.Tensor: - return torch.nn.functional.huber_loss( - input, target, reduction=reduction, delta=delta - ) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index 8683957a72c97..42793dacfbbf2 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -10,135 +10,130 @@ import ivy -def moveaxis( - a: torch.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.moveaxis(a, source, destination) - - moveaxis.support_native_out = False +heaviside.support_native_out = True +flipud.support_native_out = False +fliplr.support_native_out = False +i0.support_native_out = True +flatten.partial_mixed_handler = ( + lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" +) +take_along_axis.support_native_out = True +broadcast_shapes.support_native_out = False +expand.support_native_out = False -def heaviside( - x1: torch.tensor, - x2: torch.tensor, - /, - *, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.heaviside( - x1, - x2, - out=out, - ) +def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_1d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -heaviside.support_native_out = True +def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_2d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -def flipud( - m: torch.Tensor, - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.flipud(m) +def atleast_3d( + *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[torch.Tensor]: + transformed = torch.atleast_3d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -flipud.support_native_out = False +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return tuple(torch.broadcast_shapes(*shapes)) -def vstack( - arrays: Sequence[torch.Tensor], +def concat_from_sequence( + input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if not isinstance(arrays, tuple): - arrays = tuple(arrays) - return torch.vstack(arrays, out=None) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = torch.cat(input_sequence, dim=axis) + return ret + elif new_axis == 1: + ret = torch.stack(input_sequence, dim=axis) + return ret -def hstack( - arrays: Sequence[torch.Tensor], +def dsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(arrays, tuple): - arrays = tuple(arrays) - return torch.hstack(arrays, out=None) + copy: Optional[bool] = None, +) -> List[torch.Tensor]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def rot90( - m: torch.Tensor, +def dstack( + arrays: Sequence[torch.Tensor], /, *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.rot90(m, k, axes) + if not isinstance(arrays, tuple): + arrays = tuple(arrays) + return torch.dstack(arrays, out=out) -def top_k( +def expand( x: torch.Tensor, - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] - ) - if not largest: - indices = torch.argsort(x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - else: - indices = torch.argsort(-x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - if not sorted: - indices = torch.sort(indices, dim=axis)[0] - val = torch.gather(x, axis, indices) - return topk_res(val, indices) - - -def fliplr( - m: torch.Tensor, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.fliplr(m) - - -fliplr.support_native_out = False + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return x.expand(shape) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def i0( - x: torch.Tensor, +def fill_diagonal( + a: torch.Tensor, + v: Union[int, float], /, *, - out: Optional[torch.Tensor] = None, + wrap: bool = False, ) -> torch.Tensor: - return torch.i0(x, out=out) + shape = a.shape + max_end = torch.prod(torch.tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + end = max_end if end > max_end else end + a = torch.reshape(a, (-1,)) + w = torch.zeros(a.shape, dtype=bool).to(a.device) + ins = torch.arange(0, max_end).to(a.device) + steps = torch.arange(0, end, step).to(a.device) -i0.support_native_out = True + for i in steps: + i = ins == i + w = torch.logical_or(w, i) + a = torch.where(w, v, a) + a = torch.reshape(a, shape) + return a def flatten( @@ -154,47 +149,53 @@ def flatten( return torch.flatten(x, start_dim=start_dim, end_dim=end_dim) -flatten.partial_mixed_handler = ( - lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" -) +def fliplr( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.fliplr(m) -def vsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], +def flipud( + m: torch.Tensor, /, *, copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.flipud(m) -def dsplit( +def heaviside( + x1: torch.tensor, + x2: torch.tensor, + /, + *, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.heaviside( + x1, + x2, + out=out, + ) + + +def hsplit( ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, ) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_1d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def dstack( +def hstack( arrays: Sequence[torch.Tensor], /, *, @@ -202,23 +203,41 @@ def dstack( ) -> torch.Tensor: if not isinstance(arrays, tuple): arrays = tuple(arrays) - return torch.dstack(arrays, out=out) + return torch.hstack(arrays, out=None) -def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_2d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def i0( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.i0(x, out=out) -def atleast_3d( - *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[torch.Tensor]: - transformed = torch.atleast_3d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +def moveaxis( + a: torch.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.moveaxis(a, source, destination) + + +def rot90( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.rot90(m, k, axes) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -263,59 +282,30 @@ def take_along_axis( return torch.take_along_dim(arr, indices, axis, out=out) -def hsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -take_along_axis.support_native_out = True - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return tuple(torch.broadcast_shapes(*shapes)) - - -broadcast_shapes.support_native_out = False - - -def expand( +def top_k( x: torch.Tensor, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return x.expand(shape) - - -expand.support_native_out = False - - -def concat_from_sequence( - input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = torch.cat(input_sequence, dim=axis) - return ret - elif new_axis == 1: - ret = torch.stack(input_sequence, dim=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] + ) + if not largest: + indices = torch.argsort(x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) + else: + indices = torch.argsort(-x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) + if not sorted: + indices = torch.sort(indices, dim=axis)[0] + val = torch.gather(x, axis, indices) + return topk_res(val, indices) @with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) @@ -342,32 +332,26 @@ def unique_consecutive( ) -def fill_diagonal( - a: torch.Tensor, - v: Union[int, float], +def vsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, - wrap: bool = False, -) -> torch.Tensor: - shape = a.shape - max_end = torch.prod(torch.tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + copy: Optional[bool] = None, +) -> List[torch.Tensor]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - end = max_end if end > max_end else end - a = torch.reshape(a, (-1,)) - w = torch.zeros(a.shape, dtype=bool).to(a.device) - ins = torch.arange(0, max_end).to(a.device) - steps = torch.arange(0, end, step).to(a.device) - for i in steps: - i = ins == i - w = torch.logical_or(w, i) - a = torch.where(w, v, a) - a = torch.reshape(a, shape) - return a +def vstack( + arrays: Sequence[torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not isinstance(arrays, tuple): + arrays = tuple(arrays) + return torch.vstack(arrays, out=None) diff --git a/ivy/functional/backends/torch/experimental/norms.py b/ivy/functional/backends/torch/experimental/norms.py index 5a3e0c31f764b..9524918277dc9 100644 --- a/ivy/functional/backends/torch/experimental/norms.py +++ b/ivy/functional/backends/torch/experimental/norms.py @@ -5,31 +5,27 @@ from .. import backend_version -def l1_normalize( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) - - l1_normalize.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def l2_normalize( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) - - l2_normalize.support_native_out = True +batch_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) + ) +) +instance_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) + ) +) +lp_normalize.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -73,15 +69,29 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def group_norm( + x: torch.Tensor, + num_groups: int = 1, + /, + *, + offset: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", + out: Optional[torch.Tensor] = None, +): + xdims = x.ndim + if data_format == "NSC": + x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) + xnormalized = torch.nn.functional.group_norm( + x, num_groups, weight=scale, bias=offset, eps=eps ) -) + + if data_format == "NSC": + xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) + + return xnormalized @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -127,40 +137,25 @@ def instance_norm( return xnormalized, runningmean, runningvariance -instance_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) - ) -) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def group_norm( +def l1_normalize( x: torch.Tensor, - num_groups: int = 1, /, *, - offset: Optional[torch.Tensor] = None, - scale: Optional[torch.Tensor] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", + axis: Optional[int] = None, out: Optional[torch.Tensor] = None, -): - xdims = x.ndim - if data_format == "NSC": - x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) - xnormalized = torch.nn.functional.group_norm( - x, num_groups, weight=scale, bias=offset, eps=eps - ) +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) - if data_format == "NSC": - xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) - return xnormalized +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def l2_normalize( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -173,6 +168,3 @@ def lp_normalize( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: return torch.nn.functional.normalize(x, p=p, dim=axis, out=out) - - -lp_normalize.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/random.py b/ivy/functional/backends/torch/experimental/random.py index a65cee4f57431..b8e58426ed24b 100644 --- a/ivy/functional/backends/torch/experimental/random.py +++ b/ivy/functional/backends/torch/experimental/random.py @@ -12,23 +12,46 @@ ) -# dirichlet -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def dirichlet( - alpha: Union[torch.tensor, float, Sequence[float]], - /, +# --- Helpers --- # +# --------------- # + + +def _poisson_with_neg_lam(lam, fill_value, device, dtype): + if torch.any(lam < 0): + pos_lam = torch.where(lam < 0, 0, lam) + ret = torch.poisson(pos_lam).type(dtype).to(device) + ret = torch.where(lam < 0, fill_value, ret) + else: + ret = torch.poisson(lam).type(dtype).to(device) + return ret + + +# --- Main --- # +# ------------ # + + +def bernoulli( + probs: Union[float, torch.Tensor], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - out: Optional[torch.Tensor] = None, + logits: Union[float, torch.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: torch.device, + dtype: torch.dtype, seed: Optional[int] = None, - dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - size = size if size is not None else len(alpha) - if seed is not None: + if seed: torch.manual_seed(seed) - return torch.tensor( - torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), - dtype=dtype, + if logits is not None: + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return ( + torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) + .sample(shape) + .to(device, dtype) ) @@ -53,6 +76,26 @@ def beta( return ret +# dirichlet +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def dirichlet( + alpha: Union[torch.tensor, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + out: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + size = size if size is not None else len(alpha) + if seed is not None: + torch.manual_seed(seed) + return torch.tensor( + torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), + dtype=dtype, + ) + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, torch.Tensor], @@ -74,16 +117,6 @@ def gamma( return ret -def _poisson_with_neg_lam(lam, fill_value, device, dtype): - if torch.any(lam < 0): - pos_lam = torch.where(lam < 0, 0, lam) - ret = torch.poisson(pos_lam).type(dtype).to(device) - ret = torch.where(lam < 0, fill_value, ret) - else: - ret = torch.poisson(lam).type(dtype).to(device) - return ret - - def poisson( lam: Union[float, torch.Tensor], *, @@ -104,28 +137,3 @@ def poisson( _check_shapes_broadcastable(lam.shape, list_shape) lam = torch.broadcast_to(lam, list_shape) return _poisson_with_neg_lam(lam, fill_value, device, dtype) - - -def bernoulli( - probs: Union[float, torch.Tensor], - *, - logits: Union[float, torch.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: torch.device, - dtype: torch.dtype, - seed: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if seed: - torch.manual_seed(seed) - if logits is not None: - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return ( - torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) - .sample(shape) - .to(device, dtype) - ) diff --git a/ivy/functional/backends/torch/experimental/searching.py b/ivy/functional/backends/torch/experimental/searching.py index 456d6204f8056..2c2a8040bdf21 100644 --- a/ivy/functional/backends/torch/experimental/searching.py +++ b/ivy/functional/backends/torch/experimental/searching.py @@ -3,6 +3,9 @@ import torch +unravel_index.support_native_out = False + + def unravel_index( indices: torch.Tensor, shape: Tuple[int], @@ -16,6 +19,3 @@ def unravel_index( output.append(temp % dim) temp = temp // dim return tuple(reversed(output)) - - -unravel_index.support_native_out = False diff --git a/ivy/functional/backends/torch/experimental/sorting.py b/ivy/functional/backends/torch/experimental/sorting.py index c6b7ee0f5bb06..db91a013dde4e 100644 --- a/ivy/functional/backends/torch/experimental/sorting.py +++ b/ivy/functional/backends/torch/experimental/sorting.py @@ -6,6 +6,9 @@ import ivy +lexsort.support_native_out = False + + # invert_permutation def invert_permutation( x: Union[torch.Tensor, list, tuple], @@ -44,6 +47,3 @@ def lexsort( # only valid for torch > 2.0.1 result = result[temp] return result - - -lexsort.support_native_out = False diff --git a/ivy/functional/backends/torch/experimental/statistical.py b/ivy/functional/backends/torch/experimental/statistical.py index cc1d4d5fa05ce..78b4e43e4db7e 100644 --- a/ivy/functional/backends/torch/experimental/statistical.py +++ b/ivy/functional/backends/torch/experimental/statistical.py @@ -10,6 +10,333 @@ from copy import deepcopy +# --- Helpers --- # +# --------------- # + + +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + return torch.quantile( + x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out + ) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + + +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = a.ndim + axis_arg = deepcopy(axis) + if axis is not None: + axis = _to_positive_axis(axis, nd) + + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) + + for i, s in enumerate(sorted(keep)): + a = torch.moveaxis(a, s, i) + a = a.view( + [ + *a.shape[:nkeep], + -1, + ] + ) + axis_arg = -1 + + ret = fn(a, q, axis=axis_arg) + + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret + + +def _nanmedian(input, axis, keepdims): + dtype = input.dtype + temp = input.to(torch.float64) + num_dim = len(temp.size()) + keepdim_shape = list(temp.size()) + q = 0.5 + + axis = [axis] if isinstance(axis, int) else list(axis) + + for i in axis: + keepdim_shape[i] = 1 + axis = [num_dim + x if x < 0 else x for x in axis] + axis.sort() + dimension = len(temp.size()) + while len(axis) > 0: + axis1 = axis[0] + for axis2 in range(axis1 + 1, dimension): + temp = torch.transpose(temp, axis1, axis2) + axis1 = axis2 + axis = [x - 1 for x in axis] + axis.pop(0) + dimension = dimension - 1 + temp = torch.flatten(temp, start_dim=dimension - len(axis)) + ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") + if keepdims: + keepdim_shape = tuple(keepdim_shape) + ret = ret.reshape(keepdim_shape) + + if dtype in [torch.int32, torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=torch.float16) + else: + ret = torch.asarray(ret, dtype=torch.float32) + + return ret + + +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if isinstance(q, float): + q = torch.as_tensor(q) + if isinstance(q, torch.Tensor) and q.ndim > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = a.flatten() + + n = a.shape[axis] + + indices = q * (n - 1) + + a = torch.sort(a, axis)[axis] + indices_below = torch.floor(indices).to(torch.int64) + indices_upper = torch.ceil(indices).to(torch.int64) + + weights = indices - indices_below.to(torch.float64) + + indices_below = torch.clip(indices_below, 0, n - 1) + indices_upper = torch.clip(indices_upper, 0, n - 1) + tensor_upper = torch.index_select(a, 0, indices_upper) + tensor_below = torch.index_select(a, 0, indices_below) + + pred = weights <= 0.5 + out = torch.where(pred, tensor_below, tensor_upper) + return out.to(ret_dtype) + + +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis + + +def _validate_quantile(q): + if isinstance(q, float): + q = torch.as_tensor(q) + if q.ndim == 1 and torch.numel(q) < 10: + for i in range(torch.numel(q)): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (torch.all(0 <= q) and torch.all(q <= 1)): + return False + return True + + +# --- Main --- # +# ------------ # + + +def bincount( + x: torch.Tensor, + /, + *, + weights: Optional[torch.Tensor] = None, + minlength: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if weights is None: + ret = torch.bincount(x, minlength=minlength) + ret = ret.to(x.dtype) + else: + ret = torch.bincount(x, weights=weights, minlength=minlength) + ret = ret.to(weights.dtype) + return ret + + +def corrcoef( + x: torch.Tensor, + /, + *, + y: Optional[torch.Tensor] = None, + rowvar: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = torch.concat([x, y], dim=axis) + xarr = xarr.T if not rowvar else xarr + + return torch.corrcoef(xarr) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def cov( + x1: torch.Tensor, + x2: torch.Tensor = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[torch.Tensor] = None, + aweights: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + # dtype casts separately + if fweights is not None: + fweights = fweights.type(torch.int64) + if aweights is not None: + aweights = aweights.type(torch.float64) + + if x1.dim() > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if x2.dim() > 2: + raise ValueError("x2 has more than 2 dimensions") + + if ddof is None: + if bias == 0: + ddof = 1 + else: + ddof = 0 + + if dtype is None: + x1 = x1.type(torch.float64) + if x2 is not None: + x2 = x2.type(torch.float64) + else: + x1 = x1.type(dtype) + if x2 is not None: + x2 = x2.type(dtype) + + X = x1 + if not rowVar and len(x1.shape) != 1: + X = torch.t(x1) + + if x2 is not None: + if not rowVar and len(x2.shape) != 1: + x2 = torch.t(x2) + X = torch.vstack((X, x2)) + + return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, + backend_version, +) +def cummax( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if x.dtype in (torch.bool, torch.float16): + x = x.to(dtype=torch.float64) + elif x.dtype in (torch.int16, torch.int8, torch.uint8): + x = x.to(dtype=torch.int64) + elif x.dtype in (torch.complex64, torch.complex128): + x = x.real.to(dtype=torch.float64) + + if exclusive or reverse: + if exclusive and reverse: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + x1, x2 = torch.concat( + (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 + ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x1, x2 = torch.cummax(x, -1) + res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + else: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + return res1, res2 + + return torch.cummax(x, axis, out=out) + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cummin( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + if not (reverse): + ret = torch.cummin(x, axis)[0] + else: + ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] + ret = torch.flip(ret, (axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret.to(dtype)) + return ret.to(dtype) + + @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -136,254 +463,54 @@ def histogram( return ret -histogram.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) -def median( - input: torch.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if isinstance(axis, tuple): - if len(axis) == 1: - axis = axis[0] - ret = quantile( - input, - 0.5, - axis=axis, - keepdims=keepdims, - interpolation="midpoint", - ) - if input.dtype in [torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif input.dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=input.dtype) - else: - ret = torch.asarray(ret, dtype=torch.float32) - return ret - - -median.support_native_out = False - - -def nanmean( +def igamma( a: torch.Tensor, /, *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[torch.dtype] = None, + x: torch.Tensor, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) - - -nanmean.support_native_out = True - - -def _validate_quantile(q): - if isinstance(q, float): - q = torch.as_tensor(q) - if q.ndim == 1 and torch.numel(q) < 10: - for i in range(torch.numel(q)): - if not (0.0 <= q[i] <= 1.0): - return False - else: - if not (torch.all(0 <= q) and torch.all(q <= 1)): - return False - return True - - -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis - - -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = a.ndim - axis_arg = deepcopy(axis) - if axis is not None: - axis = _to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) - - for i, s in enumerate(sorted(keep)): - a = torch.moveaxis(a, s, i) - a = a.view( - [ - *a.shape[:nkeep], - -1, - ] - ) - axis_arg = -1 - - ret = fn(a, q, axis=axis_arg) - - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] - - return ret - - -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if isinstance(q, float): - q = torch.as_tensor(q) - if isinstance(q, torch.Tensor) and q.ndim > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = a.flatten() - - n = a.shape[axis] - - indices = q * (n - 1) - - a = torch.sort(a, axis)[axis] - indices_below = torch.floor(indices).to(torch.int64) - indices_upper = torch.ceil(indices).to(torch.int64) - - weights = indices - indices_below.to(torch.float64) - - indices_below = torch.clip(indices_below, 0, n - 1) - indices_upper = torch.clip(indices_upper, 0, n - 1) - tensor_upper = torch.index_select(a, 0, indices_upper) - tensor_below = torch.index_select(a, 0, indices_below) - - pred = weights <= 0.5 - out = torch.where(pred, tensor_below, tensor_upper) - return out.to(ret_dtype) - - -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - return torch.quantile( - x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out - ) - else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) + return torch.special.gammainc(a, x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def quantile( - a: torch.Tensor, - q: Union[torch.Tensor, float], +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) +def median( + input: torch.Tensor, /, *, - axis: Optional[Union[Sequence[int], int]] = None, + axis: Optional[Union[Tuple[int], int]] = None, keepdims: bool = False, - interpolation: str = "linear", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, + if isinstance(axis, tuple): + if len(axis) == 1: + axis = axis[0] + ret = quantile( + input, + 0.5, axis=axis, keepdims=keepdims, - interpolation=interpolation, - out=out, + interpolation="midpoint", ) + if input.dtype in [torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif input.dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=input.dtype) + else: + ret = torch.asarray(ret, dtype=torch.float32) + return ret -quantile.support_native_out = True - - -def corrcoef( - x: torch.Tensor, +def nanmean( + a: torch.Tensor, /, *, - y: Optional[torch.Tensor] = None, - rowvar: bool = True, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = torch.concat([x, y], dim=axis) - xarr = xarr.T if not rowvar else xarr - - return torch.corrcoef(xarr) - - -def _nanmedian(input, axis, keepdims): - dtype = input.dtype - temp = input.to(torch.float64) - num_dim = len(temp.size()) - keepdim_shape = list(temp.size()) - q = 0.5 - - axis = [axis] if isinstance(axis, int) else list(axis) - - for i in axis: - keepdim_shape[i] = 1 - axis = [num_dim + x if x < 0 else x for x in axis] - axis.sort() - dimension = len(temp.size()) - while len(axis) > 0: - axis1 = axis[0] - for axis2 in range(axis1 + 1, dimension): - temp = torch.transpose(temp, axis1, axis2) - axis1 = axis2 - axis = [x - 1 for x in axis] - axis.pop(0) - dimension = dimension - 1 - temp = torch.flatten(temp, start_dim=dimension - len(axis)) - ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") - if keepdims: - keepdim_shape = tuple(keepdim_shape) - ret = ret.reshape(keepdim_shape) - - if dtype in [torch.int32, torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=torch.float16) - else: - ret = torch.asarray(ret, dtype=torch.float32) - - return ret + return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -438,166 +565,33 @@ def nanmedian( return _nanmedian(input, axis, keepdims) -nanmedian.support_native_out = True - - -def bincount( - x: torch.Tensor, - /, - *, - weights: Optional[torch.Tensor] = None, - minlength: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if weights is None: - ret = torch.bincount(x, minlength=minlength) - ret = ret.to(x.dtype) - else: - ret = torch.bincount(x, weights=weights, minlength=minlength) - ret = ret.to(weights.dtype) - return ret - - -bincount.support_native_out = False - - -def igamma( +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def quantile( a: torch.Tensor, + q: Union[torch.Tensor, float], /, *, - x: torch.Tensor, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: bool = False, + interpolation: str = "linear", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.gammainc(a, x, out=out) + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + out=out, + ) +histogram.support_native_out = True +median.support_native_out = False +nanmean.support_native_out = True +quantile.support_native_out = True +nanmedian.support_native_out = True +bincount.support_native_out = False igamma.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def cov( - x1: torch.Tensor, - x2: torch.Tensor = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[torch.Tensor] = None, - aweights: Optional[torch.Tensor] = None, - dtype: Optional[torch.dtype] = None, -) -> torch.Tensor: - # dtype casts separately - if fweights is not None: - fweights = fweights.type(torch.int64) - if aweights is not None: - aweights = aweights.type(torch.float64) - - if x1.dim() > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if x2.dim() > 2: - raise ValueError("x2 has more than 2 dimensions") - - if ddof is None: - if bias == 0: - ddof = 1 - else: - ddof = 0 - - if dtype is None: - x1 = x1.type(torch.float64) - if x2 is not None: - x2 = x2.type(torch.float64) - else: - x1 = x1.type(dtype) - if x2 is not None: - x2 = x2.type(dtype) - - X = x1 - if not rowVar and len(x1.shape) != 1: - X = torch.t(x1) - - if x2 is not None: - if not rowVar and len(x2.shape) != 1: - x2 = torch.t(x2) - X = torch.vstack((X, x2)) - - return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) - - cov.support_native_out = False - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, - backend_version, -) -def cummax( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - if x.dtype in (torch.bool, torch.float16): - x = x.to(dtype=torch.float64) - elif x.dtype in (torch.int16, torch.int8, torch.uint8): - x = x.to(dtype=torch.int64) - elif x.dtype in (torch.complex64, torch.complex128): - x = x.real.to(dtype=torch.float64) - - if exclusive or reverse: - if exclusive and reverse: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - x1, x2 = torch.concat( - (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 - ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x1, x2 = torch.cummax(x, -1) - res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - else: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - return res1, res2 - - return torch.cummax(x, axis, out=out) - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cummin( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - if not (reverse): - ret = torch.cummin(x, axis)[0] - else: - ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] - ret = torch.flip(ret, (axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret.to(dtype)) - return ret.to(dtype) diff --git a/ivy/functional/backends/torch/general.py b/ivy/functional/backends/torch/general.py index f8cd14960cddb..c75cc289b5bc7 100644 --- a/ivy/functional/backends/torch/general.py +++ b/ivy/functional/backends/torch/general.py @@ -4,11 +4,6 @@ from numbers import Number from operator import mul from typing import Optional, Union, Sequence, Callable, List, Tuple - -try: - import functorch -except ImportError: - functorch = () # for torch 1.10.1 import numpy as np import torch @@ -18,9 +13,18 @@ from . import backend_version, is_variable from ...ivy.general import _broadcast_to +try: + import functorch +except ImportError: + functorch = () # for torch 1.10.1 + torch_scatter = None +# --- Helpers --- # +# --------------- # + + def _parse_index(indices, ndims): ind = list() for so in indices: @@ -44,12 +48,8 @@ def _parse_index(indices, ndims): return ind -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, torch.Tensor): - if exclusive and x.requires_grad: - return False - return True - return False +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"2.0.1 and below": ("complex", "bfloat16")}, backend_version) @@ -66,111 +66,6 @@ def current_backend_str() -> str: return "torch" -def neg_step(query): - return ( - not isinstance(query, (int, bool)) - and not ivy.is_array(query) - and query is not None - and query is not Ellipsis - and ( - (isinstance(query, slice) and query.step is not None and query.step < 0) - or ( - not isinstance(query, slice) - and any( - isinstance(q, slice) and q.step is not None and q.step < 0 - for q in query - ) - ) - ) - ) - - -def get_item( - x: torch.Tensor, - /, - query: Union[torch.Tensor, Tuple], - *, - copy: bool = None, -) -> torch.Tensor: - return x.__getitem__(query) - - -get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) - - -def set_item( - x: torch.Tensor, - query: Union[torch.Tensor, Tuple], - val: torch.Tensor, - /, - *, - copy: Optional[bool] = False, -) -> torch.Tensor: - if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: - val = val.to(x.dtype) - if copy: - x = x.clone() - x.__setitem__(query, val) - return x - - -set_item.partial_mixed_handler = ( - lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad -) - - -def to_numpy( - x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif torch.is_tensor(x): - x = x.resolve_neg().resolve_conj() - if copy: - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16") - return x.detach().cpu().numpy() - else: - raise ivy.utils.exceptions.IvyException( - "Overwriting the same address is not supported for torch." - ) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - -def to_scalar(x: torch.Tensor, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - -def to_list(x: torch.Tensor, /) -> list: - if isinstance(x, np.ndarray): - return x.tolist() - elif torch.is_tensor(x): - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16").tolist() - else: - return x.detach().cpu().numpy().tolist() - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - def gather( params: torch.Tensor, indices: torch.Tensor, @@ -207,6 +102,36 @@ def gather( return result +def gather_nd( + params: torch.Tensor, + indices: torch.Tensor, + /, + *, + batch_dims: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = torch.stack(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -237,34 +162,14 @@ def gather_nd_helper(params, indices): return res -def gather_nd( - params: torch.Tensor, - indices: torch.Tensor, +def get_item( + x: torch.Tensor, /, + query: Union[torch.Tensor, Tuple], *, - batch_dims: int = 0, - out: Optional[torch.Tensor] = None, + copy: bool = None, ) -> torch.Tensor: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = torch.stack(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result + return x.__getitem__(query) def get_num_dims( @@ -338,6 +243,37 @@ def inplace_variables_supported(): return True +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, torch.Tensor): + if exclusive and x.requires_grad: + return False + return True + return False + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version +) +def isin( + elements: torch.tensor, + test_elements: torch.tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> torch.tensor: + return torch.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + +def itemsize(x: torch.tensor) -> int: + return x.element_size() + + def multiprocessing(context: Optional[str] = None): import torch.multiprocessing @@ -346,6 +282,25 @@ def multiprocessing(context: Optional[str] = None): return torch.multiprocessing.get_context(context) +def neg_step(query): + return ( + not isinstance(query, (int, bool)) + and not ivy.is_array(query) + and query is not None + and query is not Ellipsis + and ( + (isinstance(query, slice) and query.step is not None and query.step < 0) + or ( + not isinstance(query, slice) + and any( + isinstance(q, slice) and q.step is not None and q.step < 0 + for q in query + ) + ) + ) + ) + + @with_unsupported_dtypes( { "2.0.1 and below": ("bfloat16",), @@ -395,9 +350,6 @@ def scatter_flat( return res -scatter_flat.support_native_out = True - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -492,7 +444,20 @@ def scatter_nd( return res -scatter_nd.support_native_out = True +def set_item( + x: torch.Tensor, + query: Union[torch.Tensor, Tuple], + val: torch.Tensor, + /, + *, + copy: Optional[bool] = False, +) -> torch.Tensor: + if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: + val = val.to(x.dtype) + if copy: + x = x.clone() + x.__setitem__(query, val) + return x def shape( @@ -507,6 +472,58 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: torch.Tensor, /) -> list: + if isinstance(x, np.ndarray): + return x.tolist() + elif torch.is_tensor(x): + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16").tolist() + else: + return x.detach().cpu().numpy().tolist() + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + +def to_numpy( + x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif torch.is_tensor(x): + x = x.resolve_neg().resolve_conj() + if copy: + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16") + return x.detach().cpu().numpy() + else: + raise ivy.utils.exceptions.IvyException( + "Overwriting the same address is not supported for torch." + ) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + +def to_scalar(x: torch.Tensor, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def vmap( func: Callable, @@ -523,27 +540,10 @@ def _vmap(*args): return _vmap -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version +get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) +set_item.partial_mixed_handler = ( + lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad ) -def isin( - elements: torch.tensor, - test_elements: torch.tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> torch.tensor: - return torch.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True isin.support_native_out = True - - -def itemsize(x: torch.tensor) -> int: - return x.element_size() diff --git a/ivy/functional/backends/torch/gradients.py b/ivy/functional/backends/torch/gradients.py index 012e8aafe4a17..ad88bc84f0d3d 100644 --- a/ivy/functional/backends/torch/gradients.py +++ b/ivy/functional/backends/torch/gradients.py @@ -19,20 +19,8 @@ ) -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = ivy.astype(x, ivy.default_float_dtype()).to_native() - if not x.is_leaf: - return x.detach().requires_grad_() - return x.clone().requires_grad_() - - -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, torch.Tensor) and x.requires_grad - - -def variable_data(x: torch.Tensor, /) -> torch.Tensor: - return x.data +# --- Helpers --- # +# --------------- # def _grad_func(y, xs, retain_grads): @@ -92,6 +80,10 @@ def grad_(x): return grads +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: torch.Tensor, @@ -140,61 +132,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = torch.autograd.grad(y, x, allow_unused=True)[0] - grad = ( - grad - if grad is not None - else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) - ) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def stop_gradient( - x: Optional[torch.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[torch.Tensor] = None, -): - if is_variable(x) and preserve_type: - if x.grad_fn: - x = x.detach() - x.requires_grad = True - elif x.grad: - x.grad.data.zero_() - return x - return x.detach() - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - callback_fn = lambda x_in: ivy.to_ivy( - torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), - nested=True, - include_derived=True, - ) - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -257,5 +194,76 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, torch.Tensor) and x.requires_grad + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + callback_fn = lambda x_in: ivy.to_ivy( + torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), + nested=True, + include_derived=True, + ) + return callback_fn + + +def stop_gradient( + x: Optional[torch.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[torch.Tensor] = None, +): + if is_variable(x) and preserve_type: + if x.grad_fn: + x = x.detach() + x.requires_grad = True + elif x.grad: + x.grad.data.zero_() + return x + return x.detach() + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = torch.autograd.grad(y, x, allow_unused=True)[0] + grad = ( + grad + if grad is not None + else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) + ) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = ivy.astype(x, ivy.default_float_dtype()).to_native() + if not x.is_leaf: + return x.detach().requires_grad_() + return x.clone().requires_grad_() + + +def variable_data(x: torch.Tensor, /) -> torch.Tensor: + return x.data + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/torch/layers.py b/ivy/functional/backends/torch/layers.py index 0ae11f08e2e9b..ecb3f3f29d9bd 100644 --- a/ivy/functional/backends/torch/layers.py +++ b/ivy/functional/backends/torch/layers.py @@ -11,22 +11,8 @@ from ivy.functional.ivy.layers import _handle_padding, _deconv_length -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, - backend_version, -) -def linear( - x: torch.Tensor, - weight: torch.Tensor, - /, - *, - bias: Optional[torch.Tensor] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.linear(x, weight, bias) - - -linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 +# --- Helpers --- # +# --------------- # def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): @@ -126,6 +112,10 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding +# --- Main --- # +# ------------ # + + @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version, @@ -288,48 +278,6 @@ def conv2d_transpose( return res -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - "complex", - ) - }, - backend_version, -) -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: torch.Tensor, - filters: torch.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters - filters = torch.unsqueeze(filters, -1) - dims_in = filters.shape[-2] - filters = filters.permute(2, 3, 0, 1) - x, padding = _pad_before_conv( - x, filters, strides, padding, 2, dilations, "channel_first" - ) - # noinspection PyArgumentEqualDefault - res = torch.nn.functional.conv2d( - x, filters, None, strides, padding, dilations, dims_in - ) - if data_format == "NHWC": - return res.permute(0, 2, 3, 1) - return res - - @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version ) @@ -544,3 +492,63 @@ def conv_general_transpose( if data_format == "channel_last": res = res.permute(0, *range(2, dims + 2), 1) return res + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + "complex", + ) + }, + backend_version, +) +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: torch.Tensor, + filters: torch.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NHWC": + x = x.permute(0, 3, 1, 2) + filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters + filters = torch.unsqueeze(filters, -1) + dims_in = filters.shape[-2] + filters = filters.permute(2, 3, 0, 1) + x, padding = _pad_before_conv( + x, filters, strides, padding, 2, dilations, "channel_first" + ) + # noinspection PyArgumentEqualDefault + res = torch.nn.functional.conv2d( + x, filters, None, strides, padding, dilations, dims_in + ) + if data_format == "NHWC": + return res.permute(0, 2, 3, 1) + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, + backend_version, +) +def linear( + x: torch.Tensor, + weight: torch.Tensor, + /, + *, + bias: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.linear(x, weight, bias) + + +linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 diff --git a/ivy/functional/backends/torch/linear_algebra.py b/ivy/functional/backends/torch/linear_algebra.py index 950c2ff6fd295..341eef4f46c7f 100644 --- a/ivy/functional/backends/torch/linear_algebra.py +++ b/ivy/functional/backends/torch/linear_algebra.py @@ -14,6 +14,27 @@ from .elementwise import _cast_for_unary_op +cholesky.support_native_out = True +cross.support_native_out = True +det.support_native_out = True +eigh.support_native_out = True +eigvalsh.support_native_out = True +inner.support_native_out = True +inv.support_native_out = True +matmul.support_native_out = True +matrix_norm.support_native_out = True +eig.support_native_out = True +matrix_power.support_native_out = True +matrix_rank.support_native_out = True +outer.support_native_out = True +pinv.support_native_out = True +slogdet.support_native_out = True +svdvals.support_native_out = True +vecdot.support_native_out = True +vector_norm.support_native_out = True +vector_to_skew_symmetric_matrix.support_native_out = True + + # Array API Standard # # -------------------# @@ -40,9 +61,6 @@ def cholesky( return ret -cholesky.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) def cross( x1: torch.Tensor, @@ -68,15 +86,24 @@ def cross( ) -cross.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def det(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: return torch.linalg.det(x, out=out) -det.support_native_out = True +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def diag( + x: torch.Tensor, + /, + *, + k: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.tensor: + return torch.diag(x, diagonal=k) def diagonal( @@ -91,6 +118,17 @@ def diagonal( return torch.diagonal(x, offset=offset, dim1=axis1, dim2=axis2) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def eig( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] + ) + eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) + return result_tuple(eigenvalues, eigenvectors) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -102,9 +140,6 @@ def eigh( return result_tuple(eigenvalues, eigenvectors) -eigh.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigvalsh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -112,9 +147,6 @@ def eigvalsh( return torch.linalg.eigvalsh(x, UPLO=UPLO, out=out) -eigvalsh.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def inner( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None @@ -132,9 +164,6 @@ def inner( return torch.inner(x1, x2, out=out) -inner.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def inv( x: torch.Tensor, @@ -157,9 +186,6 @@ def inv( return ret -inv.support_native_out = True - - @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "bool")}, backend_version ) @@ -191,9 +217,6 @@ def matmul( return torch.matmul(x1, x2, out=out) -matmul.support_native_out = True - - @with_supported_dtypes({"2.0.1 and below": ("float", "complex")}, backend_version) def matrix_norm( x: torch.Tensor, @@ -207,23 +230,6 @@ def matrix_norm( return torch.linalg.matrix_norm(x, ord=ord, dim=axis, keepdim=keepdims, out=out) -matrix_norm.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def eig( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] - ) - eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) - return result_tuple(eigenvalues, eigenvectors) - - -eig.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_power( x: torch.Tensor, n: int, /, *, out: Optional[torch.Tensor] = None @@ -231,9 +237,6 @@ def matrix_power( return torch.linalg.matrix_power(x, n, out=out) -matrix_power.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_rank( x: torch.Tensor, @@ -271,9 +274,6 @@ def matrix_rank( return ret -matrix_rank.support_native_out = True - - def matrix_transpose( x: torch.Tensor, /, *, conjugate: bool = False, out: Optional[torch.Tensor] = None ) -> torch.Tensor: @@ -294,9 +294,6 @@ def outer( return torch.outer(x1, x2, out=out) -outer.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def pinv( x: torch.Tensor, @@ -310,21 +307,6 @@ def pinv( return torch.linalg.pinv(x, rtol, out=out) -pinv.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def tensorsolve( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.tensorsolve(x1, x2, dims=axes) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def qr( x: torch.Tensor, @@ -359,9 +341,6 @@ def slogdet( return results(sign, logabsdet) -slogdet.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def solve( x1: torch.Tensor, @@ -420,9 +399,6 @@ def svdvals(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch. return torch.linalg.svdvals(x, out=out) -svdvals.support_native_out = True - - # ToDo: re-add int32 support once # (https://github.com/pytorch/pytorch/issues/84530) is fixed @with_unsupported_dtypes({"2.0.1 and below": ("int32",)}, backend_version) @@ -448,6 +424,18 @@ def tensordot( return ret +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def tensorsolve( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.tensorsolve(x1, x2, dims=axes) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def trace( x: torch.Tensor, @@ -465,6 +453,30 @@ def trace( return ret +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def vander( + x: torch.tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + # torch.vander hasn't been used as it produces 0 gradients + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + ret = torch.pow( + torch.transpose(torch.unsqueeze(x, 0), 0, 1), + torch.arange(start, stop, step), + out=out, + ) + if ret.dtype != x.dtype: + return ret.to(x.dtype) + return ret + + def vecdot( x1: torch.Tensor, x2: torch.Tensor, @@ -485,9 +497,6 @@ def vecdot( return torch.tensordot(x1, x2, dims=([axis], [axis])).to(dtype) -vecdot.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("integer",)}, backend_version) def vector_norm( x: torch.Tensor, @@ -507,48 +516,6 @@ def vector_norm( return torch.linalg.vector_norm(x, ord, axis, keepdims, out=out) -vector_norm.support_native_out = True - - -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def diag( - x: torch.Tensor, - /, - *, - k: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.diag(x, diagonal=k) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def vander( - x: torch.tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - # torch.vander hasn't been used as it produces 0 gradients - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - ret = torch.pow( - torch.transpose(torch.unsqueeze(x, 0), 0, 1), - torch.arange(start, stop, step), - out=out, - ) - if ret.dtype != x.dtype: - return ret.to(x.dtype) - return ret - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -576,6 +543,3 @@ def vector_to_skew_symmetric_matrix( row3 = torch.cat((-a2s, a1s, zs), -1) # BS x 3 x 3 return torch.cat((row1, row2, row3), -2, out=out) - - -vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/torch/manipulation.py b/ivy/functional/backends/torch/manipulation.py index fe8b37b4dcafd..1b5efcec1cf06 100644 --- a/ivy/functional/backends/torch/manipulation.py +++ b/ivy/functional/backends/torch/manipulation.py @@ -14,12 +14,42 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _reshape_fortran_torch(x, shape): if len(x.shape) > 0: x = x.permute(*reversed(range(len(x.shape)))) return x.reshape(shape[::-1]).permute(list(range(len(shape)))[::-1]) +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version +) +def clip( + x: torch.Tensor, + x_min: Union[Number, torch.Tensor], + x_max: Union[Number, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if hasattr(x_min, "dtype"): + x_min = torch.asarray(x_min, device=x.device) + x_max = torch.asarray(x_max, device=x.device) + promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) + promoted_type = torch.promote_types(promoted_type, x.dtype) + x_min = x_min.to(promoted_type) + x_max = x_max.to(promoted_type) + x = x.to(promoted_type) + return torch.clamp(x, x_min, x_max, out=out) + + # Array API Standard # # -------------------# @@ -43,7 +73,26 @@ def concat( return torch.cat(xs, dim=axis, out=out) -concat.support_native_out = True +def constant_pad( + x: torch.Tensor, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if 0 in x.shape: + new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] + return torch.ones(new_shape, dtype=x.dtype) * value + if x.shape == (): + x = x.unsqueeze(0) + if isinstance(pad_width, torch.Tensor): + pad_width = pad_width.detach().cpu().numpy().tolist() + pad_width_flat: List[int] = list() + for pad_width_sec in reversed(pad_width): + for item in pad_width_sec: + pad_width_flat.append(item) + return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) def expand_dims( @@ -93,6 +142,23 @@ def permute_dims( return torch.permute(x, axes) +@with_unsupported_dtypes( + {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version +) +def repeat( + x: torch.Tensor, + /, + repeats: Union[int, Iterable[int]], + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if len(x.shape) == 0 and axis in [0, -1]: + axis = None + repeats = torch.tensor(repeats) + return torch.repeat_interleave(x, repeats, axis) + + def reshape( x: torch.Tensor, /, @@ -132,69 +198,6 @@ def roll( return torch.roll(x, shift, axis) -def squeeze( - x: torch.Tensor, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if isinstance(axis, int): - if x.size(dim=axis) > 1: - raise ValueError( - "Expected dimension of size [{}, {}], but found " - "dimension size {}".format(-x.dim(), x.dim(), axis) - ) - if x.shape[axis] != 1: - raise ivy.utils.exceptions.IvyException( - f"Expected size of axis to be 1 but was {x.shape[axis]}" - ) - return torch.squeeze(x, axis) - if axis is None: - if copy: - newarr = torch.clone(x) - return torch.squeeze(newarr) - return torch.squeeze(x) - newarr = torch.clone(x) - if isinstance(axis, tuple): - axis = list(axis) - normalise_axis = [ - (len(x.shape) - abs(element)) if element < 0 else element for element in axis - ] - normalise_axis.sort() - axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] - dim = x.dim() - for i in axis_updated_after_squeeze: - shape = x.shape[i] - if shape > 1 and (shape < -dim or dim <= shape): - raise ValueError( - "Expected dimension of size [{}, {}], " - "but found dimension size {}".format(-dim, dim, shape) - ) - else: - if copy: - newarr = torch.squeeze(newarr, i) - else: - x = torch.squeeze(x, i) - if copy: - return newarr - return x - - -def stack( - arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], - /, - *, - axis: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.stack(arrays, axis, out=out) - - -stack.support_native_out = True - - # Extra # # ------# @@ -244,61 +247,64 @@ def split( return list(torch.split(x, num_or_size_splits, axis)) -@with_unsupported_dtypes( - {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version -) -def repeat( +def squeeze( x: torch.Tensor, /, - repeats: Union[int, Iterable[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if len(x.shape) == 0 and axis in [0, -1]: - axis = None - repeats = torch.tensor(repeats) - return torch.repeat_interleave(x, repeats, axis) - - -def tile( - x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - if isinstance(repeats, torch.Tensor): - repeats = repeats.detach().cpu().numpy().tolist() - return x.repeat(repeats) + if isinstance(axis, int): + if x.size(dim=axis) > 1: + raise ValueError( + "Expected dimension of size [{}, {}], but found " + "dimension size {}".format(-x.dim(), x.dim(), axis) + ) + if x.shape[axis] != 1: + raise ivy.utils.exceptions.IvyException( + f"Expected size of axis to be 1 but was {x.shape[axis]}" + ) + return torch.squeeze(x, axis) + if axis is None: + if copy: + newarr = torch.clone(x) + return torch.squeeze(newarr) + return torch.squeeze(x) + newarr = torch.clone(x) + if isinstance(axis, tuple): + axis = list(axis) + normalise_axis = [ + (len(x.shape) - abs(element)) if element < 0 else element for element in axis + ] + normalise_axis.sort() + axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] + dim = x.dim() + for i in axis_updated_after_squeeze: + shape = x.shape[i] + if shape > 1 and (shape < -dim or dim <= shape): + raise ValueError( + "Expected dimension of size [{}, {}], " + "but found dimension size {}".format(-dim, dim, shape) + ) + else: + if copy: + newarr = torch.squeeze(newarr, i) + else: + x = torch.squeeze(x, i) + if copy: + return newarr + return x -def constant_pad( - x: torch.Tensor, +def stack( + arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], /, - pad_width: List[List[int]], *, - value: Number = 0.0, + axis: int = 0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if 0 in x.shape: - new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] - return torch.ones(new_shape, dtype=x.dtype) * value - if x.shape == (): - x = x.unsqueeze(0) - if isinstance(pad_width, torch.Tensor): - pad_width = pad_width.detach().cpu().numpy().tolist() - pad_width_flat: List[int] = list() - for pad_width_sec in reversed(pad_width): - for item in pad_width_sec: - pad_width_flat.append(item) - return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) - - -def zero_pad( - x: torch.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[torch.Tensor] = None, -): - return constant_pad(x, pad_width, value=0.0) + return torch.stack(arrays, axis, out=out) def swapaxes( @@ -313,29 +319,12 @@ def swapaxes( return torch.transpose(x, axis0, axis1) -@with_unsupported_dtypes( - {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version -) -def clip( - x: torch.Tensor, - x_min: Union[Number, torch.Tensor], - x_max: Union[Number, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, +def tile( + x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if hasattr(x_min, "dtype"): - x_min = torch.asarray(x_min, device=x.device) - x_max = torch.asarray(x_max, device=x.device) - promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) - promoted_type = torch.promote_types(promoted_type, x.dtype) - x_min = x_min.to(promoted_type) - x_max = x_max.to(promoted_type) - x = x.to(promoted_type) - return torch.clamp(x, x_min, x_max, out=out) - - -clip.support_native_out = True + if isinstance(repeats, torch.Tensor): + repeats = repeats.detach().cpu().numpy().tolist() + return x.repeat(repeats) def unstack( @@ -355,3 +344,18 @@ def unstack( if keepdims: return [r.unsqueeze(axis) for r in ret] return ret + + +def zero_pad( + x: torch.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[torch.Tensor] = None, +): + return constant_pad(x, pad_width, value=0.0) + + +concat.support_native_out = True +stack.support_native_out = True +clip.support_native_out = True diff --git a/ivy/functional/backends/torch/random.py b/ivy/functional/backends/torch/random.py index 55495e074b067..12f7a19686be8 100644 --- a/ivy/functional/backends/torch/random.py +++ b/ivy/functional/backends/torch/random.py @@ -14,52 +14,10 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, torch.Tensor] = 0.0, - high: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - device: torch.device, - seed: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - shape = _check_bounds_and_get_shape(low, high, shape).shape - rand_range = high - low - if seed: - torch.manual_seed(seed) - if torch.is_tensor(shape): - shape = shape.tolist() - return ( - torch.rand(shape, device=device, dtype=torch.float) * rand_range + low - ).type(dtype) - - -def random_normal( - *, - mean: Union[float, torch.Tensor] = 0.0, - std: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - seed: Optional[int] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - dtype = ivy.as_native_dtype(dtype) - if seed: - torch.manual_seed(seed) - if isinstance(mean, (int, float)) and isinstance(std, (int, float)): - return torch.normal(mean, std, shape, out=out).type(dtype).to(device) - return torch.normal(mean, std, out=out).type(dtype).to(device) - random_normal.support_native_out = True +multinomial.support_native_out = True +shuffle.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) @@ -90,9 +48,6 @@ def multinomial( return torch.multinomial(probs.float(), num_samples, replace, out=out).to(device) -multinomial.support_native_out = True - - def randint( low: Union[int, torch.Tensor], high: Union[int, torch.Tensor], @@ -115,6 +70,51 @@ def randint( return (torch.rand(shape, device=device) * rand_range + low).to(dtype) +def random_normal( + *, + mean: Union[float, torch.Tensor] = 0.0, + std: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + seed: Optional[int] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + dtype = ivy.as_native_dtype(dtype) + if seed: + torch.manual_seed(seed) + if isinstance(mean, (int, float)) and isinstance(std, (int, float)): + return torch.normal(mean, std, shape, out=out).type(dtype).to(device) + return torch.normal(mean, std, out=out).type(dtype).to(device) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, torch.Tensor] = 0.0, + high: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + device: torch.device, + seed: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + shape = _check_bounds_and_get_shape(low, high, shape).shape + rand_range = high - low + if seed: + torch.manual_seed(seed) + if torch.is_tensor(shape): + shape = shape.tolist() + return ( + torch.rand(shape, device=device, dtype=torch.float) * rand_range + low + ).type(dtype) + + def seed(*, seed_value: int = 0) -> None: torch.manual_seed(seed_value) torch.cuda.manual_seed(seed_value) @@ -140,6 +140,3 @@ def shuffle( if seed: torch.manual_seed(seed) return torch.index_select(x, 0, torch.randperm(batch_size), out=out) - - -shuffle.support_native_out = True diff --git a/ivy/functional/backends/torch/searching.py b/ivy/functional/backends/torch/searching.py index 702100b3b5ab1..778a0a3b5d141 100644 --- a/ivy/functional/backends/torch/searching.py +++ b/ivy/functional/backends/torch/searching.py @@ -69,6 +69,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.argwhere(x) + + def nonzero( x: torch.Tensor, /, @@ -107,16 +120,3 @@ def where( if condition.dtype is not torch.bool: condition = condition == 1.0 return ivy.astype(torch.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.argwhere(x) diff --git a/ivy/functional/backends/torch/sorting.py b/ivy/functional/backends/torch/sorting.py index 6fea3ca9794c3..f5e085b0088a3 100644 --- a/ivy/functional/backends/torch/sorting.py +++ b/ivy/functional/backends/torch/sorting.py @@ -8,6 +8,12 @@ from . import backend_version +argsort.support_native_out = True +sort.support_native_out = True +msort.support_native_out = True +searchsorted.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def argsort( x: torch.Tensor, @@ -26,30 +32,6 @@ def argsort( return sorted_indices -argsort.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def sort( - x: torch.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if out is not None: - out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) - sorted_tensor, _ = torch.sort( - x, dim=axis, descending=descending, stable=stable, out=out - ) - return sorted_tensor - - -sort.support_native_out = True - - # msort @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def msort( @@ -58,9 +40,6 @@ def msort( return torch.msort(a, out=out) -msort.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def searchsorted( x: torch.Tensor, @@ -113,4 +92,19 @@ def searchsorted( return torch.searchsorted(x, v, sorter=sorter, side=side).to(ret_dtype) -searchsorted.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def sort( + x: torch.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if out is not None: + out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) + sorted_tensor, _ = torch.sort( + x, dim=axis, descending=descending, stable=stable, out=out + ) + return sorted_tensor diff --git a/ivy/functional/backends/torch/statistical.py b/ivy/functional/backends/torch/statistical.py index 04694a0bdca3f..930858211703a 100644 --- a/ivy/functional/backends/torch/statistical.py +++ b/ivy/functional/backends/torch/statistical.py @@ -1,5 +1,3 @@ -# global -torch_scatter = None from typing import Union, Optional, Sequence import torch @@ -10,30 +8,130 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Array API Standard # -# -------------------# +# global +torch_scatter = None -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def min( +# --- Helpers --- # +# --------------- # + + +def _infer_dtype(dtype: torch.dtype) -> torch.dtype: + default_dtype = ivy.infer_default_dtype(dtype) + if default_dtype in ivy.valid_dtypes: + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return ivy.as_native_dtype(default_dtype) + return ivy.as_native_dtype(dtype) + + +# --- Main --- # +# ------------ # + + +# Extra # +# ----- # + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + }, + backend_version, +) +def cumprod( x: torch.Tensor, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if axis == (): - if ivy.exists(out): - return ivy.inplace_update(out, x) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + + if not (exclusive or reverse): + return torch.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + ret = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumprod(x, -1, dtype=dtype) + ret = torch.transpose(x, axis, -1) + else: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + ret = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "1.12.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cumsum( + x: torch.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + res = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumsum(x, -1, dtype=dtype) + res = torch.transpose(x, axis, -1) else: - return x - if not keepdims and not axis and axis != 0: - return torch.amin(input=x, out=out) - return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + res = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res + return torch.cumsum(x, axis, dtype=dtype, out=out) -min.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def einsum( + equation: str, + *operands: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = _get_promoted_type_of_operands(operands) + return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -55,9 +153,6 @@ def max( return torch.amax(input=x, dim=axis, keepdim=keepdims, out=out) -max.support_native_out = True - - def mean( x: torch.Tensor, /, @@ -81,15 +176,27 @@ def mean( return torch.mean(x, dim=axis, keepdim=keepdims, out=out) -mean.support_native_out = True +# Array API Standard # +# -------------------# -def _infer_dtype(dtype: torch.dtype) -> torch.dtype: - default_dtype = ivy.infer_default_dtype(dtype) - if default_dtype in ivy.valid_dtypes: - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return ivy.as_native_dtype(default_dtype) - return ivy.as_native_dtype(dtype) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def min( + x: torch.Tensor, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if axis == (): + if ivy.exists(out): + return ivy.inplace_update(out, x) + else: + return x + if not keepdims and not axis and axis != 0: + return torch.amin(input=x, out=out) + return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) # Function does support uint8, but allowing support for unsigned will cause @@ -214,113 +321,8 @@ def var( ).to(x.dtype) -# Extra # -# ----- # - - -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - }, - backend_version, -) -def cumprod( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - - if not (exclusive or reverse): - return torch.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - ret = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumprod(x, -1, dtype=dtype) - ret = torch.transpose(x, axis, -1) - else: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - ret = torch.flip(x, dims=(axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - +min.support_native_out = True +max.support_native_out = True +mean.support_native_out = True cumprod.support_native_out = True - - -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "1.12.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cumsum( - x: torch.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - res = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumsum(x, -1, dtype=dtype) - res = torch.transpose(x, axis, -1) - else: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - res = torch.flip(x, dims=(axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res - return torch.cumsum(x, axis, dtype=dtype, out=out) - - cumsum.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def einsum( - equation: str, - *operands: torch.Tensor, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = _get_promoted_type_of_operands(operands) - return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) diff --git a/ivy/functional/backends/torch/utility.py b/ivy/functional/backends/torch/utility.py index ca643ce0d4095..5d62f84cb7479 100644 --- a/ivy/functional/backends/torch/utility.py +++ b/ivy/functional/backends/torch/utility.py @@ -3,6 +3,10 @@ from typing import Union, Optional, Sequence +all.support_native_out = True +any.support_native_out = True + + def all( x: torch.Tensor, /, @@ -25,9 +29,6 @@ def all( return x -all.support_native_out = True - - def any( x: torch.Tensor, /, @@ -48,6 +49,3 @@ def any( for i, a in enumerate(axis): x = torch.any(x, dim=a if keepdims else a - i, keepdim=keepdims, out=out) return x - - -any.support_native_out = True diff --git a/ivy/functional/frontends/torch/comparison_ops.py b/ivy/functional/frontends/torch/comparison_ops.py index b743b38d135f0..eeaadadfbf864 100644 --- a/ivy/functional/frontends/torch/comparison_ops.py +++ b/ivy/functional/frontends/torch/comparison_ops.py @@ -290,7 +290,7 @@ def topk(input, k, dim=None, largest=True, sorted=True, *, out=None): gt = greater +ne = not_equal ge = greater_equal le = less_equal lt = less -ne = not_equal diff --git a/ivy/functional/ivy/activations.py b/ivy/functional/ivy/activations.py index f19180ab7a529..d55d9dfcde068 100644 --- a/ivy/functional/ivy/activations.py +++ b/ivy/functional/ivy/activations.py @@ -18,6 +18,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _gelu_jax_like( x: Union[ivy.Array, ivy.NativeArray], /, @@ -31,6 +35,87 @@ def _gelu_jax_like( return fn_original(x, approximate=True, out=out) +def _leaky_relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original: Optional[Callable] = None, + alpha: float = 0.2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.astype(x * alpha, x.dtype), + x, + ) + + +def _relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.array(0.0, dtype=x.dtype), + x, + ) + + +def _softplus_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, + out: Optional[ivy.Array] = None, +): + if beta is not None: + x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) + else: + x_beta = x + amax = ivy.relu(x_beta) + res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) + res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) + res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( + x.dtype + ) * ivy.astype(1j, x.dtype) + if beta is not None: + res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) + if threshold is not None: + res = ivy.where( + ivy.real(x_beta) < threshold, + res, + x, + ).astype(x.dtype) + return res + + +def _wrap_between(y, a): + """Wrap y between [-a, a]""" + a = ivy.array(a, dtype=y.dtype) + a2 = ivy.array(2 * a, dtype=y.dtype) + zero = ivy.array(0, dtype=y.dtype) + rem = ivy.remainder(ivy.add(y, a), a2) + rem = ivy.where(rem < zero, rem + a2, rem) - a + return rem + + +# --- Main --- # +# ------------ # + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -99,26 +184,52 @@ def gelu( return current_backend(x).gelu(x, approximate=approximate, out=out) -gelu.jax_like = _gelu_jax_like +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def hardswish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply the hardswish activation function element-wise. + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. -def _leaky_relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original: Optional[Callable] = None, - alpha: float = 0.2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.astype(x * alpha, x.dtype), - x, - ) + Returns + ------- + ret + an array containing the hardswish activation of each element in ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 0., 4.]) + >>> y = ivy.hardswish(x) + >>> y + ivy.array([0., 0., 4.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) + >>> x = ivy.hardswish(x, out=x) + >>> x + { + a: ivy.array([-0., 4., 5.]), + b: ivy.array([0., 5.]) + } + """ + return current_backend(x).hardswish(x, out=out) @handle_exceptions @@ -199,9 +310,6 @@ def leaky_relu( return current_backend(x).leaky_relu(x, alpha=alpha, out=out) -leaky_relu.jax_like = _leaky_relu_jax_like - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -276,22 +384,60 @@ def log_softmax( return current_backend(x).log_softmax(x, axis=axis, out=out) -def _relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - out: Optional[ivy.Array] = None, +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def mish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None ) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.array(0.0, dtype=x.dtype), - x, - ) + """ + Apply the mish activation function element-wise. + + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the mish activation of each element in + ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.mish(x) + >>> print(y) + ivy.array([-0.30340147, 0. , 0.86509842]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.mish(x, out = y) + >>> print(y) + ivy.array([ 1.40337825, 0.56114835, -0.20788449]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.mish(x) + >>> print(x) + { + a: ivy.array([0.86509842, -0.30883577]), + b: ivy.array([0.28903052, -0.10714479]) + } + """ + return current_backend(x).mish(x, out=out) @handle_exceptions @@ -363,9 +509,6 @@ def relu( return current_backend(x).relu(x, out=out) -relu.jax_like = _relu_jax_like - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -496,46 +639,6 @@ def softmax( return current_backend(x).softmax(x, axis=axis, out=out) -def _wrap_between(y, a): - """Wrap y between [-a, a]""" - a = ivy.array(a, dtype=y.dtype) - a2 = ivy.array(2 * a, dtype=y.dtype) - zero = ivy.array(0, dtype=y.dtype) - rem = ivy.remainder(ivy.add(y, a), a2) - rem = ivy.where(rem < zero, rem + a2, rem) - a - return rem - - -def _softplus_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, - out: Optional[ivy.Array] = None, -): - if beta is not None: - x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) - else: - x_beta = x - amax = ivy.relu(x_beta) - res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) - res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) - res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( - x.dtype - ) * ivy.astype(1j, x.dtype) - if beta is not None: - res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) - if threshold is not None: - res = ivy.where( - ivy.real(x_beta) < threshold, - res, - x, - ).astype(x.dtype) - return res - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -607,108 +710,7 @@ def softplus( return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out) +gelu.jax_like = _gelu_jax_like +leaky_relu.jax_like = _leaky_relu_jax_like +relu.jax_like = _relu_jax_like softplus.jax_like = _softplus_jax_like - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def mish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply the mish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the mish activation of each element in - ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.mish(x) - >>> print(y) - ivy.array([-0.30340147, 0. , 0.86509842]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.mish(x, out = y) - >>> print(y) - ivy.array([ 1.40337825, 0.56114835, -0.20788449]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.mish(x) - >>> print(x) - { - a: ivy.array([0.86509842, -0.30883577]), - b: ivy.array([0.28903052, -0.10714479]) - } - """ - return current_backend(x).mish(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -def hardswish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply the hardswish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the hardswish activation of each element in ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 0., 4.]) - >>> y = ivy.hardswish(x) - >>> y - ivy.array([0., 0., 4.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) - >>> x = ivy.hardswish(x, out=x) - >>> x - { - a: ivy.array([-0., 4., 5.]), - b: ivy.array([0., 5.]) - } - """ - return current_backend(x).hardswish(x, out=out) diff --git a/ivy/functional/ivy/constants.py b/ivy/functional/ivy/constants.py index a830e3df30d66..461877e75b3b2 100644 --- a/ivy/functional/ivy/constants.py +++ b/ivy/functional/ivy/constants.py @@ -1,64 +1,55 @@ # global import math - +atto = 1e-18 +centi = 1e-2 +deci = 1e-1 +deka = 1e1 # Array API Standard # # -------------------# e = math.e -"""IEEE 754 floating-point representation of Euler's constant.""" - -pi = math.pi -"""IEEE 754 floating-point representation of the mathematical constant π.""" - -nan = math.nan -"""IEEE 754 floating-point representation of Not a Number (NaN).""" - -inf = math.inf -"""IEEE 754 floating-point representation of (positive) infinity.""" - -newaxis = None -"""An alias for None which is useful for indexing arrays.""" - - +exa = 1e18 +exbi = 2**60 +femto = 1e-15 +gibi = 2**30 +giga = 1e9 # Mathematical constants # # ------# golden = golden_ratio = (1 + math.sqrt(5)) / 2 -quetta = 1e30 -ronna = 1e27 -yotta = 1e24 -zetta = 1e21 -exa = 1e18 -peta = 1e15 -tera = 1e12 -giga = 1e9 -mega = 1e6 -kilo = 1e3 hecto = 1e2 -deka = 1e1 -deci = 1e-1 -centi = 1e-2 -milli = 1e-3 -micro = 1e-6 -nano = 1e-9 -pico = 1e-12 -femto = 1e-15 -atto = 1e-18 -zepto = 1e-21 -yocto = 1e-24 -ronto = 1e-27 -quecto = 1e-30 - - +inf = math.inf # Binary prefixes # # ------# kibi = 2**10 +kilo = 1e3 mebi = 2**20 -gibi = 2**30 -tebi = 2**40 +mega = 1e6 +micro = 1e-6 +milli = 1e-3 +nan = math.nan +nano = 1e-9 +newaxis = None pebi = 2**50 -exbi = 2**60 -zebi = 2**70 +peta = 1e15 +pi = math.pi +pico = 1e-12 +quecto = 1e-30 +quetta = 1e30 +ronna = 1e27 +ronto = 1e-27 +tebi = 2**40 +tera = 1e12 yobi = 2**80 +yocto = 1e-24 +yotta = 1e24 +zebi = 2**70 +zepto = 1e-21 +zetta = 1e21 +"""IEEE 754 floating-point representation of Euler's constant.""" +"""IEEE 754 floating-point representation of the mathematical constant π.""" +"""IEEE 754 floating-point representation of Not a Number (NaN).""" +"""IEEE 754 floating-point representation of (positive) infinity.""" +"""An alias for None which is useful for indexing arrays.""" diff --git a/ivy/functional/ivy/control_flow_ops.py b/ivy/functional/ivy/control_flow_ops.py index 4c2d78dc13c39..e41773bbcf366 100644 --- a/ivy/functional/ivy/control_flow_ops.py +++ b/ivy/functional/ivy/control_flow_ops.py @@ -9,6 +9,100 @@ ) +# --- Helpers --- # +# --------------- # + + +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) + + +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} + + +# --- Main --- # +# ------------ # + + +def cast_bool(x): + return bool(x) + + +# todo (nightcrab) find a better place for these cmp functions + + +def cmp_is(left, right): + return left is right + + +def cmp_isnot(left, right): + return left is not right + + +def for_loop( + iterable: Iterable[Any], + body_fn: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + """ + Loops over an iterable, passing the current iteration along with a tuple of + variables into the provided body function. + + Parameters + ---------- + iterable + The iterable to loop over. + body_fn + A function to call each iteration, first taking the iterator value + and then a tuple of extra parameters. + vars + Extra parameters to be passed to body_fn. + + Returns + ------- + ret + The loop's return value (if any). + + Example + ---- + ``` + def body_fn(k, args): + print(k+1) + return args + + lst = [5,6] + + ivy.for_loop(lst, body_fn, ()) + >>> 5 + >>> 6 + ``` + """ + iterator = iterable.__iter__() + + vars_dict = _tuple_to_dict(vars) + + def test_fn(iterator, original_body, vars_dict): + try: + val = iterator.__next__() + except StopIteration: + return False + + vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) + + for k in range(len(vars_tuple)): + vars_dict[k] = vars_tuple[k] + + return True + + def empty_function(iterator, original_body, vars_dict): + return (iterator, original_body, vars_dict) + + packed_vars = (iterator, body_fn, vars_dict) + + return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) + + def if_else( cond: Callable, body_fn: Callable, @@ -70,6 +164,17 @@ def _if_else(cond, body_fn, orelse_fn, vars): return _if_else(cond, body_fn, orelse_fn, vars) +def try_except( + body1: Callable, + body2: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + try: + return body1(*vars) + except Exception as e: + return body2(*vars, e) + + def while_loop( test_fn: Callable, body_fn: Callable, @@ -124,100 +229,3 @@ def _while_loop(test_fn, body_fn, vars): body_fn = to_ivy_arrays_and_back(body_fn) return _while_loop(test_fn, body_fn, vars) - - -def for_loop( - iterable: Iterable[Any], - body_fn: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - """ - Loops over an iterable, passing the current iteration along with a tuple of - variables into the provided body function. - - Parameters - ---------- - iterable - The iterable to loop over. - body_fn - A function to call each iteration, first taking the iterator value - and then a tuple of extra parameters. - vars - Extra parameters to be passed to body_fn. - - Returns - ------- - ret - The loop's return value (if any). - - Example - ---- - ``` - def body_fn(k, args): - print(k+1) - return args - - lst = [5,6] - - ivy.for_loop(lst, body_fn, ()) - >>> 5 - >>> 6 - ``` - """ - iterator = iterable.__iter__() - - vars_dict = _tuple_to_dict(vars) - - def test_fn(iterator, original_body, vars_dict): - try: - val = iterator.__next__() - except StopIteration: - return False - - vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) - - for k in range(len(vars_tuple)): - vars_dict[k] = vars_tuple[k] - - return True - - def empty_function(iterator, original_body, vars_dict): - return (iterator, original_body, vars_dict) - - packed_vars = (iterator, body_fn, vars_dict) - - return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) - - -def try_except( - body1: Callable, - body2: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - try: - return body1(*vars) - except Exception as e: - return body2(*vars, e) - - -# todo (nightcrab) find a better place for these cmp functions - - -def cmp_is(left, right): - return left is right - - -def cmp_isnot(left, right): - return left is not right - - -def cast_bool(x): - return bool(x) - - -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} - - -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) diff --git a/ivy/functional/ivy/creation.py b/ivy/functional/ivy/creation.py index 0273b8f4b0e55..f12f82585431e 100644 --- a/ivy/functional/ivy/creation.py +++ b/ivy/functional/ivy/creation.py @@ -35,44 +35,32 @@ handle_backend_invalid, ) -# Helpers # -# --------# +# Type hints # +# -----------# -def asarray_handle_nestable(fn: Callable) -> Callable: - fn_name = fn.__name__ +SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") +_T_co = TypeVar("_T_co", covariant=True) - @functools.wraps(fn) - def _asarray_handle_nestable(*args, **kwargs): - """ - Call `fn` with the *nestable* property of the function correctly handled. This - means mapping the function to the container leaves if any containers are passed - in the input. - Parameters - ---------- - args - The arguments to be passed to the function. +class NestedSequence(Protocol[_T_co]): + def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: + ... - kwargs - The keyword arguments to be passed to the function. + def __len__(self, /) -> int: + ... - Returns - ------- - The return of the function, with the nestable property handled correctly. - """ - # This decorator should only be applied to ivy.asarray, so we know where - # the container must be if there is one. - cont_fn = getattr(ivy.Container, "static_" + fn_name) - if isinstance(args[0], ivy.Container): - return cont_fn(*args, **kwargs) - # if the passed arguments does not contain a container, the function using - # the passed arguments, returning an ivy or a native array. - return fn(*args, **kwargs) +# --- Helpers --- # +# --------------- # - _asarray_handle_nestable.handle_nestable = True - return _asarray_handle_nestable + +def _flatten_nest(xs): + for x in xs: + if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): + yield from _flatten_nest(x) + else: + yield x def _ivy_to_native(x): @@ -92,6 +80,15 @@ def _ivy_to_native(x): return x +def _remove_np_bfloat16(obj): + # unlike other frameworks, torch and paddle do not support creating tensors + # from numpy arrays that have bfloat16 dtype using any extension because + # bfloat16 in not supported natively by numpy (as of version <=1.25) + if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": + return obj.tolist() + return obj + + def _shape_to_native(x): # checks the first element of the leaf list and # converts it to a native array if it is an ivy array @@ -109,168 +106,8 @@ def _shape_to_native(x): return x -def _flatten_nest(xs): - for x in xs: - if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): - yield from _flatten_nest(x) - else: - yield x - - -def _remove_np_bfloat16(obj): - # unlike other frameworks, torch and paddle do not support creating tensors - # from numpy arrays that have bfloat16 dtype using any extension because - # bfloat16 in not supported natively by numpy (as of version <=1.25) - if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": - return obj.tolist() - return obj - - -def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): - """ - Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances - and return arrays are all converted to `ivy.Array` instances. - - This wrapper is specifically for the backend implementations of - asarray. - - It assumes either all the elements in a leaf list are ivy arrays - or none of them are. It checks the first element of all the leaf - list. If it is an ivy array, it converts all the elements in the - leaf list to native otherwise it skips that leaf list. - """ - new_arg = _ivy_to_native(args[0]) - new_args = (new_arg,) + args[1:] - if dtype is not None: - dtype = ivy.default_dtype(dtype=dtype, as_native=True) - return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) - - return _asarray_to_native_arrays_and_back - - -def asarray_infer_dtype(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_dtype(*args, dtype=None, **kwargs): - """ - Determine the correct `dtype`, and then calls the function with the `dtype` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. - - Parameters - ---------- - args - The arguments to be passed to the function. - - dtype - The dtype for the function. - - kwargs - The keyword arguments to be passed to the function. - - Returns - ------- - The return of the function, with `dtype` passed explicitly. - """ - - def _infer_dtype(obj): - if isinstance(obj, ivy.NativeShape): - obj = list(obj) - if hasattr(obj, "dtype"): - return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype - else: - return ivy.default_dtype(item=obj) - - if not ivy.exists(dtype): - arr = args[0] - # get default dtypes for all elements - dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] - # flatten the nested structure - dtype_list = _flatten_nest(dtype_list) - # keep unique dtypes - dtype_list = list(set(dtype_list)) - if len(dtype_list) != 0: # handle the case of empty input - # promote all dtypes to a single dtype - dtype = dtype_list[0] - # we disable precise mode to avoid wider than necessary casting - # that might result from the mixing of int32 and float32 - with ivy.PreciseMode(False): - for dt in dtype_list[1:]: - dtype = ivy.promote_types(dtype, dt) - else: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - # call the function with dtype provided explicitly - return fn(*args, dtype=dtype, **kwargs) - - _asarray_infer_dtype.infer_dtype = True - return _asarray_infer_dtype - - -def asarray_infer_device(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_device(*args, device=None, **kwargs): - """ - Determine the correct `device`, and then calls the function with the `device` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. - - Parameters - ---------- - args - The arguments to be passed to the function. - - device - The device for the function. - - kwargs - The keyword arguments to be passed to the function. - - Returns - ------- - The return of the function, with `device` passed explicitly. - """ - if isinstance(args[0], list): - return fn( - *args, device=ivy.default_device(device, as_native=True), **kwargs - ) - - # find the first array argument, if required - arr = None if ivy.exists(device) else args[0] - # infer the correct device - device = ivy.default_device(device, item=arr, as_native=True) - # call the function with device provided explicitly - return fn(*args, device=device, **kwargs) - - _asarray_infer_device.infer_device = True - return _asarray_infer_device - - -def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: - @functools.wraps(fn) - def _inputs_to_native_shapes(*args, **kwargs): - new_arg = _shape_to_native(args[0]) - new_args = (new_arg,) + args[1:] - return fn(*new_args, **kwargs) - - _inputs_to_native_shapes.inputs_to_native_shapes = True - return _inputs_to_native_shapes - - -# Type hints # -# -----------# - -SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") -_T_co = TypeVar("_T_co", covariant=True) - - -class NestedSequence(Protocol[_T_co]): - def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: - ... - - def __len__(self, /) -> int: - ... +# --- Main --- # +# ------------ # # Array API Standard # @@ -472,434 +309,199 @@ def asarray( ) -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def zeros( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with zeros. +def asarray_handle_nestable(fn: Callable) -> Callable: + fn_name = fn.__name__ - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type must - be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + @functools.wraps(fn) + def _asarray_handle_nestable(*args, **kwargs): + """ + Call `fn` with the *nestable* property of the function correctly handled. This + means mapping the function to the container leaves if any containers are passed + in the input. - Returns - ------- - ret - an array containing zeros. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.NativeShape` input: - >>> shape = (3, 5) - >>> x = ivy.zeros(shape) - >>> print(x) - ivy.array([[0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.]]) - - >>> x = ivy.zeros(5) - >>> print(x) - ivy.array([0., 0., 0., 0., 0.]) - """ - return current_backend().zeros(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with ones. - - .. note:: - - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal to - zero (i.e., ``1 + 0j``). - - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing ones. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Shape` input: - - >>> shape = (2,2) - >>> x = ivy.ones(shape) - >>> print(x) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Dtype` input: - - >>> shape = (3,2) - >>> d_type = ivy.int64 - >>> y = ivy.ones(shape, dtype=d_type) - >>> print(y) - ivy.array([[1, 1], - [1, 1], - [1, 1]]) - - With :class:`ivy.Device` input: - - >>> shape = (3,2) - >>> y = ivy.ones(shape, device="cpu") - >>> print(y) - ivy.array([[1., 1.], - [1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input: - - >>> shape = (1, 5, 2) - >>> x = ivy.zeros(shape) - >>> ivy.ones(shape, out=x) - >>> print(x) - ivy.array([[[1., 1.], - [1., 1.], - [1., 1.], - [1., 1.], - [1., 1.]]]) - """ - return current_backend().ones(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def full_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - fill_value: Number, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ``fill_value`` and having the same ``shape`` as an - input array ``x`` . - - Parameters - ---------- - x - input array from which to derive the output array shape. - fill_value - Scalar fill value - dtype - output array data type. If ``dtype`` is `None`, the output array data type must - be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array having the same shape as ``x`` and where every element is equal to - ``fill_value``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :code:`int` datatype: - - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> fill_value = 1 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) - - >>> fill_value = 0.000123 - >>> x = ivy.ones(5) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) - - With float datatype: + Parameters + ---------- + args + The arguments to be passed to the function. - >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + kwargs + The keyword arguments to be passed to the function. - With :class:`ivy.NativeArray` input: + Returns + ------- + The return of the function, with the nestable property handled correctly. + """ + # This decorator should only be applied to ivy.asarray, so we know where + # the container must be if there is one. + cont_fn = getattr(ivy.Container, "static_" + fn_name) + if isinstance(args[0], ivy.Container): + return cont_fn(*args, **kwargs) - >>> x = ivy.native_array([3.0, 8.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x,fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123]) + # if the passed arguments does not contain a container, the function using + # the passed arguments, returning an ivy or a native array. + return fn(*args, **kwargs) - >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([[0.000123, 0.000123, 0.000123], - [0.000123, 0.000123, 0.000123]]) + _asarray_handle_nestable.handle_nestable = True + return _asarray_handle_nestable - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), - ... b=ivy.array([4.123, 5.23, 6.23])) - >>> fill_value = 15.0 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - { - a: ivy.array([15., 15., 15.]), - b: ivy.array([15., 15., 15.]) - } - """ - return current_backend(x).full_like( - x, fill_value, dtype=dtype, device=device, out=out - ) +def asarray_infer_device(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_device(*args, device=None, **kwargs): + """ + Determine the correct `device`, and then calls the function with the `device` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. + Parameters + ---------- + args + The arguments to be passed to the function. -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ones and having the same shape as an input array - ``x``. + device + The device for the function. - .. note:: + kwargs + The keyword arguments to be passed to the function. - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal - to zero (i.e., ``1 + 0j``). + Returns + ------- + The return of the function, with `device` passed explicitly. + """ + if isinstance(args[0], list): + return fn( + *args, device=ivy.default_device(device, as_native=True), **kwargs + ) - Parameters - ---------- - x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default ``None``. - device - device on which to place the created array. If device is ``None``, the output - array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + # find the first array argument, if required + arr = None if ivy.exists(device) else args[0] + # infer the correct device + device = ivy.default_device(device, item=arr, as_native=True) + # call the function with device provided explicitly + return fn(*args, device=device, **kwargs) - Returns - ------- - ret - an array having the same shape as ``x`` and filled with ``ones``. + _asarray_infer_device.infer_device = True + return _asarray_infer_device - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. +def asarray_infer_dtype(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_dtype(*args, dtype=None, **kwargs): + """ + Determine the correct `dtype`, and then calls the function with the `dtype` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Parameters + ---------- + args + The arguments to be passed to the function. - Functional Examples - ------------------- + dtype + The dtype for the function. - With :class:`ivy.Array` input: + kwargs + The keyword arguments to be passed to the function. - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) + Returns + ------- + The return of the function, with `dtype` passed explicitly. + """ - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1., 1., 1.], - [1., 1., 1.]]) + def _infer_dtype(obj): + if isinstance(obj, ivy.NativeShape): + obj = list(obj) + if hasattr(obj, "dtype"): + return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype + else: + return ivy.default_dtype(item=obj) - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.zeros(3) - >>> ivy.ones_like(x, out=y) - >>> print(y) - ivy.array([1., 1., 1.]) + if not ivy.exists(dtype): + arr = args[0] + # get default dtypes for all elements + dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] + # flatten the nested structure + dtype_list = _flatten_nest(dtype_list) + # keep unique dtypes + dtype_list = list(set(dtype_list)) + if len(dtype_list) != 0: # handle the case of empty input + # promote all dtypes to a single dtype + dtype = dtype_list[0] + # we disable precise mode to avoid wider than necessary casting + # that might result from the mixing of int32 and float32 + with ivy.PreciseMode(False): + for dt in dtype_list[1:]: + dtype = ivy.promote_types(dtype, dt) + else: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + # call the function with dtype provided explicitly + return fn(*args, dtype=dtype, **kwargs) - With :class:`ivy.NativeArray` input: + _asarray_infer_dtype.infer_dtype = True + return _asarray_infer_dtype - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1, 1, 1], - [1, 1, 1]]) - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) +def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: + @functools.wraps(fn) + def _inputs_to_native_shapes(*args, **kwargs): + new_arg = _shape_to_native(args[0]) + new_args = (new_arg,) + args[1:] + return fn(*new_args, **kwargs) - With :class:`ivy.Container` input: + _inputs_to_native_shapes.inputs_to_native_shapes = True + return _inputs_to_native_shapes - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.ones_like(x) - >>> print(y) - { - a: ivy.array([1, 1, 1]), - b: ivy.array([1, 1, 1]) - } - With :class:`ivy.Array` input: +def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): + """ + Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances + and return arrays are all converted to `ivy.Array` instances. - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.ones_like() - >>> print(y) - ivy.array([1, 1, 1, 1, 1]) + This wrapper is specifically for the backend implementations of + asarray. - With :class:'ivy.Container' input: + It assumes either all the elements in a leaf list are ivy arrays + or none of them are. It checks the first element of all the leaf + list. If it is an ivy array, it converts all the elements in the + leaf list to native otherwise it skips that leaf list. + """ + new_arg = _ivy_to_native(args[0]) + new_args = (new_arg,) + args[1:] + if dtype is not None: + dtype = ivy.default_dtype(dtype=dtype, as_native=True) + return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.ones_like() - >>> print(y) - { - a: ivy.array([1., 1.]), - b: ivy.array([1., 1.]) - } - """ - return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) + return _asarray_to_native_arrays_and_back @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function -@infer_dtype @handle_device_shifting -@infer_device -def zeros_like( +def copy_array( x: Union[ivy.Array, ivy.NativeArray], /, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + to_ivy_array: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array filled with zeros and having the same ``shape`` as an input array - ``x``. + Copy an array. Parameters ---------- x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. + array, input array containing elements to copy. + to_ivy_array + boolean, if True the returned array will be an ivy.Array object otherwise + returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., + depending on the backend), defaults to True. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -907,195 +509,83 @@ def zeros_like( Returns ------- ret - an array having the same shape as ``x`` and filled with ``zeros``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: + a copy of the input array ``x``. - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.zeros_like(x) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) + Examples + -------- + With one :class:`ivy.Array` input: - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.zeros_like(x) + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.copy_array(x) >>> print(y) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) + ivy.array([-1, 0, 1]) - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.ones(3) - >>> ivy.zeros_like(x, out=y) + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = ivy.copy_array(x) >>> print(y) - ivy.array([0., 0., 0.]) - - With :class:`ivy.NativeArray` input: + ivy.array([1, 0, 1, 1]) - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.zeros_like(x) + >>> x = ivy.array([1, 0, 1, -1]) + >>> y = ivy.zeros((1, 4)) + >>> ivy.copy_array(x, out=y) >>> print(y) - ivy.array([[0, 0, 0],[0, 0, 0]]) - + ivy.array([1, 0, 1, -1]) - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) + >>> x = ivy.array([1, 0, 1, 1]) + >>> ivy.copy_array(x, out=x) + >>> print(x) + ivy.array([1, 0, 1, 1]) - With :class:`ivy.Container` input: + With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.zeros_like(x) + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.copy_array(x) >>> print(y) { - a: ivy.array([0, 0, 0]), - b: ivy.array([0, 0, 0]) + a: ivy.array([-1, 0, 1]) } - - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.zeros_like() - >>> print(y) - ivy.array([0, 0, 0, 0, 0]) - - With :class:'ivy.Container' input: - - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.zeros_like() + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.copy_array(x) >>> print(y) { - a: ivy.array([0., 0.]), - b: ivy.array([0., 0.]) + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) } - """ - return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def tril( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the lower triangular part of a matrix (or a stack of matrices) ``x``. - - .. note:: - - The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` - on the interval ``[0, min(M, N) - 1]``. - - Parameters - ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. - k - diagonal above which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the lower triangular part(s). The returned array must have - the same shape and data type as x. All elements above the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. + With one :class:`ivy.Container` static method: - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - return current_backend(x).tril(x, k=k, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def triu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the upper triangular part of a matrix (or a stack of matrices) ``x``. - - .. note:: - - The upper triangular part of the matrix is defined as the elements - on and above the specified diagonal ``k``. + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.Container.static_copy_array(x) + >>> print(y) + { + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) + } - Parameters - ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. *, - k - diagonal below which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + With one :class:`ivy.Array` instance method: - Returns - ------- - ret - an array containing the upper triangular part(s). The returned array must have - the same shape and data type as x. All elements below the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. + >>> x = ivy.array([-1, 0, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([-1, 0, 1]) + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([1, 0, 1, 1]) - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + With :class:`ivy.Container` instance method: - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) + >>> y = x.copy_array() + >>> print(y) + { + a: ivy.array([1, 0, 1]), + b: ivy.array([-1, 0, 1, 1]) + } """ - return current_backend(x).triu(x, k=k, out=out) + return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) @handle_backend_invalid @@ -1352,43 +842,18 @@ def eye( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@infer_dtype @handle_device_shifting -@infer_device -def linspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], - /, - num: int, - *, - axis: Optional[int] = None, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, +def from_dlpack( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None ) -> ivy.Array: """ - Generate a certain number of evenly-spaced values in an interval along a given axis. - - See :math:`arange` that allows to specify the step size of evenly spaced values in - an interval. + Return a new array containing the data from another (array) object with a + ``__dlpack__`` method. Parameters ---------- - start - First entry in the range. - stop - Final entry in the range. - num - Number of values to generate. - axis - Axis along which the operation is performed. - endpoint - If True, stop is the last sample. Otherwise, it is not included. - dtype - output array data type. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + x object + input (array) object. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1396,186 +861,87 @@ def linspace( Returns ------- ret - Tensor of evenly-spaced values. + an array containing the data in `x`. + + .. admonition:: Note + :class: note + + The returned array may be either a copy or a view. See + :ref:`data-interchange` for details. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.from_dlpack.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - - Functional Examples - ------------------- - - With float input: - - >>> x = ivy.linspace(1, 2, 3) - >>> print(x) - ivy.array([1. , 1.5, 2. ]) - - >>> x = ivy.linspace(1, 2, 4, endpoint=False) - >>> print(x) - ivy.array([1., 1.25, 1.5 , 1.75]) - - >>> x = ivy.linspace(1, 10, 4, dtype="int32") - >>> print(x) - ivy.array([ 1, 4, 7, 10]) - - >>> x = ivy.linspace(1, 2, 4, device= "cpu") - >>> print(x) - ivy.array([1., 1.33333337, 1.66666663, 2.]) - - >>> y = ivy.array([0,0,0,0]) - >>> ivy.linspace(1, 2, 4, out= y) - >>> print(y) - ivy.array([1, 1, 1, 2]) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2]) - >>> y = ivy.array([4,5]) - >>> z = ivy.linspace(x, y, 4, axis = 0) - >>> print(z) - ivy.array([[1, 2], - [2, 3], - [3, 4], - [4, 5]]) """ - return current_backend(start).linspace( - start, - stop, - num, - axis=axis, - endpoint=endpoint, - dtype=dtype, - device=device, - out=out, - ) + return current_backend(x).from_dlpack(x, out=out) -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def meshgrid( - *arrays: Union[ivy.Array, ivy.NativeArray], - sparse: bool = False, - indexing: str = "xy", - out: Optional[ivy.Array] = None, -) -> List[ivy.Array]: - """ - Return coordinate matrices from coordinate vectors. - - Parameters - ---------- - arrays - an arbitrary number of one-dimensional arrays representing grid coordinates. - Each array should have the same numeric data type. - sparse - if True, a sparse grid is returned in order to conserve memory. - Default: ``False``. - indexing - Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or - one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, - respectively), the ``indexing`` keyword has no effect and should be ignored. - Default: ``'xy'``. - - Returns - ------- - ret - list of N arrays, where ``N`` is the number of provided one-dimensional input - arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional - arrays having lengths ``Ni = len(xi)``, - - - if matrix indexing ``ij``, then each returned array must have the shape - ``(N1, N2, N3, ..., Nn)``. - - if Cartesian indexing ``xy``, then each returned array must have shape - ``(N2, N1, N3, ..., Nn)``. - - Accordingly, for the two-dimensional case with input one-dimensional arrays of - length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must - have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned - array must have shape ``(N, M)``. - - Similarly, for the three-dimensional case with input one-dimensional arrays of - length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned - array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then - each returned array must have shape ``(N, M, P)``. - - Each returned array should have the same data type as the input arrays. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of - the `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) - - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) +@outputs_to_ivy_arrays +def frombuffer( + buffer: bytes, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> ivy.Array: + r""" + Interpret a buffer as a 1-dimensional array. - >>> x = ivy.array([1, 2, 5]) - >>> y = ivy.array([4, 1]) - >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') - >>> print(xv) - ivy.array([[1, 1], - [2, 2], - [5, 5]]) + .. note:: + Note that either of the following must be true: + 1. count is a positive non-zero number, and the total number of bytes + in the buffer is equal or greater than offset plus count times the size + (in bytes) of dtype. + 2. count is negative, and the length (number of bytes) of the buffer + subtracted by the offset is a multiple of the size (in bytes) of dtype. - >>> print(yv) - ivy.array([[4, 1], - [4, 1], - [4, 1]]) + Parameters + ---------- + buffer + An object that exposes the buffer interface. + dtype + Data-type of the returned array; default: float. + count + Number of items to read. -1 means all data in the buffer. + offset + Start reading the buffer from this offset (in bytes); default: 0. - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> xv, yv = ivy.meshgrid(x, y, sparse=True) - >>> print(xv) - ivy.array([[1, 2, 3]]) + Returns + ------- + out + 1-dimensional array. - >>> print(yv) - ivy.array([[4], [5], [6]]) + Examples + -------- + With :class:`bytes` inputs: - With :class:`ivy.NativeArray` input: + >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' + >>> y = ivy.frombuffer(x, dtype=ivy.float64) + >>> print(y) + ivy.array([1., 2.]) - >>> x = ivy.native_array([1, 2]) - >>> y = ivy.native_array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) + >>> x = b'\x01\x02\x03\x04' + >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) + >>> print(y) + ivy.array([2, 3, 4]) - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) + >>> x = b'\x00<\x00@\x00B\x00D\x00E' + >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) + >>> print(y) + ivy.array([2., 3., 4., 5.]) """ - return current_backend().meshgrid( - *arrays, sparse=sparse, indexing=indexing, out=out + return current_backend().frombuffer( + buffer, + dtype=dtype, + count=count, + offset=offset, ) @@ -1686,8 +1052,230 @@ def full( [[False, False]]]) } """ - return current_backend().full( - shape, fill_value, dtype=dtype, device=device, out=out + return current_backend().full( + shape, fill_value, dtype=dtype, device=device, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def full_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + fill_value: Number, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ``fill_value`` and having the same ``shape`` as an + input array ``x`` . + + Parameters + ---------- + x + input array from which to derive the output array shape. + fill_value + Scalar fill value + dtype + output array data type. If ``dtype`` is `None`, the output array data type must + be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and where every element is equal to + ``fill_value``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :code:`int` datatype: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> fill_value = 1 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> fill_value = 0.000123 + >>> x = ivy.ones(5) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With float datatype: + + >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([3.0, 8.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x,fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123]) + + >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([[0.000123, 0.000123, 0.000123], + [0.000123, 0.000123, 0.000123]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), + ... b=ivy.array([4.123, 5.23, 6.23])) + >>> fill_value = 15.0 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + { + a: ivy.array([15., 15., 15.]), + b: ivy.array([15., 15., 15.]) + } + """ + return current_backend(x).full_like( + x, fill_value, dtype=dtype, device=device, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def linspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], + /, + num: int, + *, + axis: Optional[int] = None, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a certain number of evenly-spaced values in an interval along a given axis. + + See :math:`arange` that allows to specify the step size of evenly spaced values in + an interval. + + Parameters + ---------- + start + First entry in the range. + stop + Final entry in the range. + num + Number of values to generate. + axis + Axis along which the operation is performed. + endpoint + If True, stop is the last sample. Otherwise, it is not included. + dtype + output array data type. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Tensor of evenly-spaced values. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With float input: + + >>> x = ivy.linspace(1, 2, 3) + >>> print(x) + ivy.array([1. , 1.5, 2. ]) + + >>> x = ivy.linspace(1, 2, 4, endpoint=False) + >>> print(x) + ivy.array([1., 1.25, 1.5 , 1.75]) + + >>> x = ivy.linspace(1, 10, 4, dtype="int32") + >>> print(x) + ivy.array([ 1, 4, 7, 10]) + + >>> x = ivy.linspace(1, 2, 4, device= "cpu") + >>> print(x) + ivy.array([1., 1.33333337, 1.66666663, 2.]) + + >>> y = ivy.array([0,0,0,0]) + >>> ivy.linspace(1, 2, 4, out= y) + >>> print(y) + ivy.array([1, 1, 1, 2]) + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2]) + >>> y = ivy.array([4,5]) + >>> z = ivy.linspace(x, y, 4, axis = 0) + >>> print(z) + ivy.array([[1, 2], + [2, 3], + [3, 4], + [4, 5]]) + """ + return current_backend(start).linspace( + start, + stop, + num, + axis=axis, + endpoint=endpoint, + dtype=dtype, + device=device, + out=out, ) @@ -1697,163 +1285,237 @@ def full( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def from_dlpack( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +@infer_dtype +@infer_device +def logspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], + /, + num: int, + *, + base: float = 10.0, + axis: int = 0, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array containing the data from another (array) object with a - ``__dlpack__`` method. + Generate a certain number of evenly-spaced values in log space, in an interval along + a given axis. Parameters ---------- - x object - input (array) object. + start + First value in the range in log space. base ** start is the starting value in + the sequence. Can be an array or a float. + stop + Last value in the range in log space. base ** stop is the final value in the + sequence. Can be an array or a float. + num + Number of values to generate. + base + The base of the log space. Default is 10.0 + axis + Axis along which the operation is performed. Relevant only if start or stop are + array-like. Default is 0. + endpoint + If True, stop is the last sample. Otherwise, it is not included. Default is + True. + dtype + The data type of the output tensor. If None, the dtype of on_value is used or if + that is None, the dtype of off_value is used, or if that is None, defaults to + float32. Default is None. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is + None. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + inputs broadcast to. Default is None. Returns ------- ret - an array containing the data in `x`. + Tensor of evenly-spaced values in log space. - .. admonition:: Note - :class: note + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - The returned array may be either a copy or a view. See - :ref:`data-interchange` for details. + Functional Examples + ------------------- + With float input: + >>> print(ivy.logspace(1, 2, 4)) + ivy.array([ 10., 21.5443469, 46.41588834, 100.]) - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + >>> print(ivy.logspace(1, 2, 4, endpoint=False)) + ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - return current_backend(x).from_dlpack(x, out=out) + >>> print(ivy.logspace(1, 2, 4, dtype= int)) + ivy.array([ 10., 10., 10., 100.]) + >>> out = ivy.array([0,0,0,0]) + >>> ivy.logspace(1, 2, 4, out = out) + >>> print(out) + ivy.array([ 10, 21, 46, 100]) -# Extra # -# ------# + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[1.e+01, 1.e+02], + [1.e+02, 1.e+03], + [1.e+03, 1.e+04], + [1.e+04, 1.e+05]) + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4, axis = 1)) + ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], + [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) -array = asarray + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[ 10., 100.], + [ 100., 100.], + [ 1000., 1000.], + [10000., 10000.]]) + """ + result = base ** linspace( + start, + stop, + num, + endpoint=endpoint, + axis=axis, + dtype=dtype, + device=device, + ) + if ivy.exists(out): + return ivy.inplace_update(out, result) + return result @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def copy_array( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - to_ivy_array: bool = True, +def meshgrid( + *arrays: Union[ivy.Array, ivy.NativeArray], + sparse: bool = False, + indexing: str = "xy", out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> List[ivy.Array]: """ - Copy an array. + Return coordinate matrices from coordinate vectors. Parameters ---------- - x - array, input array containing elements to copy. - to_ivy_array - boolean, if True the returned array will be an ivy.Array object otherwise - returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., - depending on the backend), defaults to True. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + arrays + an arbitrary number of one-dimensional arrays representing grid coordinates. + Each array should have the same numeric data type. + sparse + if True, a sparse grid is returned in order to conserve memory. + Default: ``False``. + indexing + Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or + one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, + respectively), the ``indexing`` keyword has no effect and should be ignored. + Default: ``'xy'``. Returns ------- ret - a copy of the input array ``x``. + list of N arrays, where ``N`` is the number of provided one-dimensional input + arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional + arrays having lengths ``Ni = len(xi)``, - Examples - -------- - With one :class:`ivy.Array` input: + - if matrix indexing ``ij``, then each returned array must have the shape + ``(N1, N2, N3, ..., Nn)``. + - if Cartesian indexing ``xy``, then each returned array must have shape + ``(N2, N1, N3, ..., Nn)``. - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.copy_array(x) - >>> print(y) - ivy.array([-1, 0, 1]) + Accordingly, for the two-dimensional case with input one-dimensional arrays of + length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must + have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned + array must have shape ``(N, M)``. - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = ivy.copy_array(x) - >>> print(y) - ivy.array([1, 0, 1, 1]) + Similarly, for the three-dimensional case with input one-dimensional arrays of + length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned + array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then + each returned array must have shape ``(N, M, P)``. - >>> x = ivy.array([1, 0, 1, -1]) - >>> y = ivy.zeros((1, 4)) - >>> ivy.copy_array(x, out=y) - >>> print(y) - ivy.array([1, 0, 1, -1]) + Each returned array should have the same data type as the input arrays. - >>> x = ivy.array([1, 0, 1, 1]) - >>> ivy.copy_array(x, out=x) - >>> print(x) - ivy.array([1, 0, 1, 1]) - With one :class:`ivy.Container` input: + This function conforms to the `Array API Standard + `_. This docstring is an extension of + the `docstring `_ + in the standard. - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]) - } + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) - } + Functional Examples + ------------------- - With one :class:`ivy.Container` static method: + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) + + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) + + >>> x = ivy.array([1, 2, 5]) + >>> y = ivy.array([4, 1]) + >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') + >>> print(xv) + ivy.array([[1, 1], + [2, 2], + [5, 5]]) - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.Container.static_copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) - } + >>> print(yv) + ivy.array([[4, 1], + [4, 1], + [4, 1]]) - With one :class:`ivy.Array` instance method: + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> xv, yv = ivy.meshgrid(x, y, sparse=True) + >>> print(xv) + ivy.array([[1, 2, 3]]) - >>> x = ivy.array([-1, 0, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([-1, 0, 1]) + >>> print(yv) + ivy.array([[4], [5], [6]]) - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([1, 0, 1, 1]) + With :class:`ivy.NativeArray` input: - With :class:`ivy.Container` instance method: + >>> x = ivy.native_array([1, 2]) + >>> y = ivy.native_array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) - >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) - >>> y = x.copy_array() - >>> print(y) - { - a: ivy.array([1, 0, 1]), - b: ivy.array([-1, 0, 1, 1]) - } + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) """ - return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) + return current_backend().meshgrid( + *arrays, sparse=sparse, indexing=indexing, out=out + ) @handle_backend_invalid @@ -1961,69 +1623,290 @@ def one_hot( Returns ------- ret - Tensor of zeros with the same shape and type as a, unless dtype provided which - overrides. - - Examples - -------- - With :class:`ivy.Array` inputs: + Tensor of zeros with the same shape and type as a, unless dtype provided which + overrides. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([3, 1]) + >>> y = 5 + >>> z = x.one_hot(5) + >>> print(z) + ivy.array([[0., 0., 0., 1., 0.], + ... [0., 1., 0., 0., 0.]]) + + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, y) + ivy.array([[1., 0., 0., 0., 0.]]) + + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, 5, out=z) + ivy.array([[1., 0., 0., 0., 0.]]) + >>> print(z) + ivy.array([[1., 0., 0., 0., 0.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1, 2]), \ + b=ivy.array([3, 1]), c=ivy.array([2, 3])) + >>> y = 5 + >>> z = x.one_hot(y) + >>> print(z) + { + a: ivy.array([[0., 1., 0., 0., 0.], + [0., 0., 1., 0., 0.]]), + b: ivy.array([[0., 0., 0., 1., 0.], + [0., 1., 0., 0., 0.]]), + c: ivy.array([[0., 0., 1., 0., 0.], + [0., 0., 0., 1., 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([2]), \ + b=ivy.array([]), c=ivy.native_array([4])) + >>> y = 7 + >>> z = x.one_hot(y) + >>> print(z) + { + a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), + b: ivy.array([], shape=(0, 7)), + c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) + } + """ + return current_backend(indices).one_hot( + indices, + depth, + on_value=on_value, + off_value=off_value, + axis=axis, + dtype=dtype, + device=device, + out=out, + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with ones. + + .. note:: + + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal to + zero (i.e., ``1 + 0j``). + + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing ones. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Shape` input: + + >>> shape = (2,2) + >>> x = ivy.ones(shape) + >>> print(x) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Dtype` input: + + >>> shape = (3,2) + >>> d_type = ivy.int64 + >>> y = ivy.ones(shape, dtype=d_type) + >>> print(y) + ivy.array([[1, 1], + [1, 1], + [1, 1]]) + + With :class:`ivy.Device` input: + + >>> shape = (3,2) + >>> y = ivy.ones(shape, device="cpu") + >>> print(y) + ivy.array([[1., 1.], + [1., 1.], + [1., 1.]]) + + With :class:`ivy.Array` input: + + >>> shape = (1, 5, 2) + >>> x = ivy.zeros(shape) + >>> ivy.ones(shape, out=x) + >>> print(x) + ivy.array([[[1., 1.], + [1., 1.], + [1., 1.], + [1., 1.], + [1., 1.]]]) + """ + return current_backend().ones(shape, dtype=dtype, device=device, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ones and having the same shape as an input array + ``x``. + + .. note:: + + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal + to zero (i.e., ``1 + 0j``). + + Parameters + ---------- + x + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default ``None``. + device + device on which to place the created array. If device is ``None``, the output + array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and filled with ``ones``. - >>> x = ivy.array([3, 1]) - >>> y = 5 - >>> z = x.one_hot(5) - >>> print(z) - ivy.array([[0., 0., 0., 1., 0.], - ... [0., 1., 0., 0., 0.]]) - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, y) - ivy.array([[1., 0., 0., 0., 0.]]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, 5, out=z) - ivy.array([[1., 0., 0., 0., 0.]]) - >>> print(z) - ivy.array([[1., 0., 0., 0., 0.]]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1., 1., 1.], + [1., 1., 1.]]) + + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.zeros(3) + >>> ivy.ones_like(x, out=y) + >>> print(y) + ivy.array([1., 1., 1.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1, 1, 1], + [1, 1, 1]]) + + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 2]), \ - b=ivy.array([3, 1]), c=ivy.array([2, 3])) - >>> y = 5 - >>> z = x.one_hot(y) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.ones_like(x) + >>> print(y) { - a: ivy.array([[0., 1., 0., 0., 0.], - [0., 0., 1., 0., 0.]]), - b: ivy.array([[0., 0., 0., 1., 0.], - [0., 1., 0., 0., 0.]]), - c: ivy.array([[0., 0., 1., 0., 0.], - [0., 0., 0., 1., 0.]]) + a: ivy.array([1, 1, 1]), + b: ivy.array([1, 1, 1]) } - >>> x = ivy.Container(a=ivy.array([2]), \ - b=ivy.array([]), c=ivy.native_array([4])) - >>> y = 7 - >>> z = x.one_hot(y) - >>> print(z) + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.ones_like() + >>> print(y) + ivy.array([1, 1, 1, 1, 1]) + + With :class:'ivy.Container' input: + + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.ones_like() + >>> print(y) { - a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), - b: ivy.array([], shape=(0, 7)), - c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) + a: ivy.array([1., 1.]), + b: ivy.array([1., 1.]) } """ - return current_backend(indices).one_hot( - indices, - depth, - on_value=on_value, - off_value=off_value, - axis=axis, - dtype=dtype, - device=device, - out=out, - ) + return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) @handle_backend_invalid @@ -2032,178 +1915,110 @@ def one_hot( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@infer_dtype -@infer_device -def logspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], +@handle_device_shifting +def tril( + x: Union[ivy.Array, ivy.NativeArray], /, - num: int, *, - base: float = 10.0, - axis: int = 0, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + k: int = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Generate a certain number of evenly-spaced values in log space, in an interval along - a given axis. + Return the lower triangular part of a matrix (or a stack of matrices) ``x``. + + .. note:: + + The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` + on the interval ``[0, min(M, N) - 1]``. Parameters ---------- - start - First value in the range in log space. base ** start is the starting value in - the sequence. Can be an array or a float. - stop - Last value in the range in log space. base ** stop is the final value in the - sequence. Can be an array or a float. - num - Number of values to generate. - base - The base of the log space. Default is 10.0 - axis - Axis along which the operation is performed. Relevant only if start or stop are - array-like. Default is 0. - endpoint - If True, stop is the last sample. Otherwise, it is not included. Default is - True. - dtype - The data type of the output tensor. If None, the dtype of on_value is used or if - that is None, the dtype of off_value is used, or if that is None, defaults to - float32. Default is None. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is - None. + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. + k + diagonal above which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. Default is None. + inputs broadcast to. Returns ------- ret - Tensor of evenly-spaced values in log space. + an array containing the lower triangular part(s). The returned array must have + the same shape and data type as x. All elements above the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - - Functional Examples - ------------------- - With float input: - - >>> print(ivy.logspace(1, 2, 4)) - ivy.array([ 10., 21.5443469, 46.41588834, 100.]) - - >>> print(ivy.logspace(1, 2, 4, endpoint=False)) - ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) - - >>> print(ivy.logspace(1, 2, 4, dtype= int)) - ivy.array([ 10., 10., 10., 100.]) - - >>> out = ivy.array([0,0,0,0]) - >>> ivy.logspace(1, 2, 4, out = out) - >>> print(out) - ivy.array([ 10, 21, 46, 100]) - - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[1.e+01, 1.e+02], - [1.e+02, 1.e+03], - [1.e+03, 1.e+04], - [1.e+04, 1.e+05]) - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4, axis = 1)) - ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], - [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[ 10., 100.], - [ 100., 100.], - [ 1000., 1000.], - [10000., 10000.]]) """ - result = base ** linspace( - start, - stop, - num, - endpoint=endpoint, - axis=axis, - dtype=dtype, - device=device, - ) - if ivy.exists(out): - return ivy.inplace_update(out, result) - return result + return current_backend(x).tril(x, k=k, out=out) +@handle_backend_invalid @handle_nestable -@outputs_to_ivy_arrays -def frombuffer( - buffer: bytes, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - count: Optional[int] = -1, - offset: Optional[int] = 0, +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def triu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Interpret a buffer as a 1-dimensional array. + """ + Return the upper triangular part of a matrix (or a stack of matrices) ``x``. .. note:: - Note that either of the following must be true: - 1. count is a positive non-zero number, and the total number of bytes - in the buffer is equal or greater than offset plus count times the size - (in bytes) of dtype. - 2. count is negative, and the length (number of bytes) of the buffer - subtracted by the offset is a multiple of the size (in bytes) of dtype. + + The upper triangular part of the matrix is defined as the elements + on and above the specified diagonal ``k``. Parameters ---------- - buffer - An object that exposes the buffer interface. - dtype - Data-type of the returned array; default: float. - count - Number of items to read. -1 means all data in the buffer. - offset - Start reading the buffer from this offset (in bytes); default: 0. + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. *, + k + diagonal below which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - out - 1-dimensional array. - - Examples - -------- - With :class:`bytes` inputs: + ret + an array containing the upper triangular part(s). The returned array must have + the same shape and data type as x. All elements below the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. - >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' - >>> y = ivy.frombuffer(x, dtype=ivy.float64) - >>> print(y) - ivy.array([1., 2.]) - >>> x = b'\x01\x02\x03\x04' - >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) - >>> print(y) - ivy.array([2, 3, 4]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = b'\x00<\x00@\x00B\x00D\x00E' - >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) - >>> print(y) - ivy.array([2., 3., 4., 5.]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. """ - return current_backend().frombuffer( - buffer, - dtype=dtype, - count=count, - offset=offset, - ) + return current_backend(x).triu(x, k=k, out=out) @handle_exceptions @@ -2300,3 +2115,193 @@ def triu_indices( (ivy.array([0, 0, 0, 0, 1, 1, 1, 1]), ivy.array([0, 1, 2, 3, 0, 1, 2, 3])) """ return current_backend().triu_indices(n_rows, n_cols, k, device=device) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def zeros( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with zeros. + + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type must + be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing zeros. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.NativeShape` input: + >>> shape = (3, 5) + >>> x = ivy.zeros(shape) + >>> print(x) + ivy.array([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + + >>> x = ivy.zeros(5) + >>> print(x) + ivy.array([0., 0., 0., 0., 0.]) + """ + return current_backend().zeros(shape, dtype=dtype, device=device, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def zeros_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with zeros and having the same ``shape`` as an input array + ``x``. + + Parameters + ---------- + x + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and filled with ``zeros``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([0, 0, 0, 0, 0, 0]) + + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) + + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.ones(3) + >>> ivy.zeros_like(x, out=y) + >>> print(y) + ivy.array([0., 0., 0.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([[0, 0, 0],[0, 0, 0]]) + + + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) + >>> print(y) + ivy.array([0, 0, 0, 0, 0, 0]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.zeros_like(x) + >>> print(y) + { + a: ivy.array([0, 0, 0]), + b: ivy.array([0, 0, 0]) + } + + + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.zeros_like() + >>> print(y) + ivy.array([0, 0, 0, 0, 0]) + + With :class:'ivy.Container' input: + + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.zeros_like() + >>> print(y) + { + a: ivy.array([0., 0.]), + b: ivy.array([0., 0.]) + } + """ + return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) + + +# Extra # +# ------# + + +array = asarray diff --git a/ivy/functional/ivy/data_type.py b/ivy/functional/ivy/data_type.py index 4a084d541e98c..b5114cfda5ca2 100644 --- a/ivy/functional/ivy/data_type.py +++ b/ivy/functional/ivy/data_type.py @@ -26,52 +26,188 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# --------# +# Array API Standard # +# -------------------# +Finfo = None +Iinfo = None +default_complex_dtype_stack = list() +# Extra # +# ------# -def _is_valid_dtypes_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): - fn_supported_dtypes = fn.supported_dtypes - fn_unsupported_dtypes = fn.unsupported_dtypes - if isinstance(fn_supported_dtypes, dict): - if isinstance(fn_unsupported_dtypes, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_dtypes - and backend_str in fn_unsupported_dtypes - ): - return False - else: - if isinstance(fn_unsupported_dtypes, tuple): - return False - return True +default_dtype_stack = list() +default_float_dtype_stack = list() +default_int_dtype_stack = list() +default_uint_dtype_stack = list() -def _handle_nestable_dtype_info(fn): - def _handle_nestable_dtype_info_wrapper(type): - if isinstance(type, ivy.Container): - type = type.cont_map(lambda x, kc: fn(x)) - type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) - type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) - return type - return fn(type) +class DefaultDtype: + """Ivy's DefaultDtype class.""" - return _handle_nestable_dtype_info_wrapper + def __init__(self, dtype: ivy.Dtype): + self._dtype = dtype + def __enter__(self): + set_default_dtype(self._dtype) + return self -# Unindent every line in the source such that -# class methods can be compiled as normal methods -def _lstrip_lines(source: str) -> str: - # Separate all lines - source = source.split("\n") - # Check amount of indent before first character - indent = len(source[0]) - len(source[0].lstrip()) - # Remove same spaces from all lines - for i in range(len(source)): - source[i] = source[i][indent:] - source = "\n".join(source) - return source + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultFloatDtype: + """Ivy's DefaultFloatDtype class.""" + + def __init__(self, float_dtype: ivy.Dtype): + self._float_dtype = float_dtype + + def __enter__(self): + set_default_float_dtype(self._float_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_float_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultIntDtype: + """Ivy's DefaultIntDtype class.""" + + def __init__(self, int_dtype: ivy.Dtype): + self._int_dtype = int_dtype + + def __enter__(self): + set_default_int_dtype(self._int_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_int_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultUintDtype: + """Ivy's DefaultUintDtype class.""" + + def __init__(self, uint_dtype: ivy.UintDtype): + self._uint_dtype = uint_dtype + + def __enter__(self): + set_default_uint_dtype(self._uint_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_uint_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultComplexDtype: + """Ivy's DefaultComplexDtype class.""" + + def __init__(self, complex_dtype: ivy.Dtype): + self._complex_dtype = complex_dtype + + def __enter__(self): + set_default_complex_dtype(self._complex_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_complex_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +# --- Helpers --- # +# --------------- # + + +def _check_complex128(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "complex128" + elif isinstance(input, np.ndarray): + return str(input.dtype) == "complex128" + if hasattr(input, "real") and hasattr(input, "imag"): + return _check_float64(input.real) and _check_float64(input.imag) + return False + + +def _check_float64(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "float64" + if math.isfinite(input): + m, e = math.frexp(input) + return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) + return False + + +# Get the list of dtypes supported by the function +# by default returns the supported dtypes +def _get_dtypes(fn, complement=True): + supported = set(ivy.valid_dtypes) + + # We only care about getting dtype info from the base function + # if we do need to at some point use dtype information from the parent function + # we can comment out the following condition + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") + if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: + if complement: + supported = set(ivy.all_dtypes).difference(supported) + return supported + + # Their values are formatted like either + # 1. fn.supported_dtypes = ("float16",) + # Could also have the "all" value for the framework + basic = [ + ("supported_dtypes", set.intersection, ivy.valid_dtypes), + ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), + ] + + # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. + typesets = { + "valid": ivy.valid_dtypes, + "numeric": ivy.valid_numeric_dtypes, + "float": ivy.valid_float_dtypes, + "integer": ivy.valid_int_dtypes, + "unsigned": ivy.valid_uint_dtypes, + "complex": ivy.valid_complex_dtypes, + } + + for key, merge_fn, base in basic: + if hasattr(fn, key): + dtypes = getattr(fn, key) + # only einops allowed to be a dictionary + if isinstance(dtypes, dict): + dtypes = dtypes.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(dtypes, tuple) + dtypes = list(dtypes) + typeset_list = [] + for i, dtype in reversed(list(enumerate(dtypes))): + if dtype in typesets: + typeset_list.extend(typesets[dtype]) + dtypes.pop(i) + dtypes = dtypes + typeset_list + supported = merge_fn(supported, set(dtypes)) + + if complement: + supported = set(ivy.all_dtypes).difference(supported) + + return tuple(supported) # Get the list of function used the function @@ -126,6 +262,50 @@ def _get_functions_from_string(func_names, module): return ret +def _handle_nestable_dtype_info(fn): + def _handle_nestable_dtype_info_wrapper(type): + if isinstance(type, ivy.Container): + type = type.cont_map(lambda x, kc: fn(x)) + type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) + type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) + return type + return fn(type) + + return _handle_nestable_dtype_info_wrapper + + +def _is_valid_dtypes_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): + fn_supported_dtypes = fn.supported_dtypes + fn_unsupported_dtypes = fn.unsupported_dtypes + if isinstance(fn_supported_dtypes, dict): + if isinstance(fn_unsupported_dtypes, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_dtypes + and backend_str in fn_unsupported_dtypes + ): + return False + else: + if isinstance(fn_unsupported_dtypes, tuple): + return False + return True + + +# Unindent every line in the source such that +# class methods can be compiled as normal methods +def _lstrip_lines(source: str) -> str: + # Separate all lines + source = source.split("\n") + # Check amount of indent before first character + indent = len(source[0]) - len(source[0].lstrip()) + # Remove same spaces from all lines + for i in range(len(source)): + source[i] = source[i][indent:] + source = "\n".join(source) + return source + + # Get dtypes/device of nested functions, used for unsupported and supported dtypes # IMPORTANT: a few caveats: # 1. The base functions must be defined in ivy or the same module @@ -181,67 +361,44 @@ def _nested_get(f, base_set, merge_fn, get_fn, wrapper=set): return out -# Get the list of dtypes supported by the function -# by default returns the supported dtypes -def _get_dtypes(fn, complement=True): - supported = set(ivy.valid_dtypes) +# --- Main --- # +# ------------ # - # We only care about getting dtype info from the base function - # if we do need to at some point use dtype information from the parent function - # we can comment out the following condition - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") - if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: - if complement: - supported = set(ivy.all_dtypes).difference(supported) - return supported - # Their values are formatted like either - # 1. fn.supported_dtypes = ("float16",) - # Could also have the "all" value for the framework - basic = [ - ("supported_dtypes", set.intersection, ivy.valid_dtypes), - ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), - ] - - # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. - typesets = { - "valid": ivy.valid_dtypes, - "numeric": ivy.valid_numeric_dtypes, - "float": ivy.valid_float_dtypes, - "integer": ivy.valid_int_dtypes, - "unsigned": ivy.valid_uint_dtypes, - "complex": ivy.valid_complex_dtypes, - } +@handle_exceptions +def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: + """ + Convert native data type to string representation. - for key, merge_fn, base in basic: - if hasattr(fn, key): - dtypes = getattr(fn, key) - # only einops allowed to be a dictionary - if isinstance(dtypes, dict): - dtypes = dtypes.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(dtypes, tuple) - dtypes = list(dtypes) - typeset_list = [] - for i, dtype in reversed(list(enumerate(dtypes))): - if dtype in typesets: - typeset_list.extend(typesets[dtype]) - dtypes.pop(i) - dtypes = dtypes + typeset_list - supported = merge_fn(supported, set(dtypes)) + Parameters + ---------- + dtype_in + The data type to convert to string. - if complement: - supported = set(ivy.all_dtypes).difference(supported) + Returns + ------- + ret + data type string 'float32' + """ + return current_backend(None).as_ivy_dtype(dtype_in) - return tuple(supported) +@handle_exceptions +def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: + """ + Convert data type string representation to native data type. -# Array API Standard # -# -------------------# + Parameters + ---------- + dtype_in + The data type string to convert to native data type. -Finfo = None -Iinfo = None + Returns + ------- + ret + data type e.g. ivy.float32. + """ + return current_backend(None).as_native_dtype(dtype_in) @handle_exceptions @@ -560,475 +717,49 @@ def can_cast( >>> print(ivy.can_cast(x, ivy.float64)) True - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([[-1, -1, -1], - ... [1, 1, 1]], - ... dtype='int16') - >>> print(ivy.can_cast(x, 'uint8')) - False - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3, 4, 5])) - >>> print(ivy.can_cast(x, 'int64')) - { - a: False, - b: True - } - """ - if isinstance(from_, (ivy.Array, ivy.NativeArray)): - from_ = from_.dtype - try: - dtype = ivy.promote_types(from_, to) - return dtype == to - except KeyError: - return False - - -@handle_exceptions -@handle_backend_invalid -@inputs_to_native_arrays -@handle_device_shifting -def finfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Finfo: - """ - Machine limits for floating-point data types. - - Parameters - ---------- - type - the kind of floating-point data-type about which to get information. - - Returns - ------- - ret - an object having the followng attributes: - - - **bits**: *int* - - number of bits occupied by the floating-point data type. - - - **eps**: *float* - - difference between 1.0 and the next smallest representable floating-point - number larger than 1.0 according to the IEEE-754 standard. - - - **max**: *float* - - largest representable number. - - - **min**: *float* - - smallest representable number. - - - **smallest_normal**: *float* - - smallest positive floating-point number with full precision. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Dtype` input: - - >>> y = ivy.finfo(ivy.float32) - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - - With :code:`str` input: - - >>> y = ivy.finfo('float32') - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) - >>> print(ivy.finfo(x)) - finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) - - >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) - >>> print(ivy.finfo(x)) - finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) - - With :class:`ivy.Container` input: - - >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), - ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) - >>> print(ivy.finfo(c)) - { - x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), - y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) - } - """ - return current_backend(None).finfo(type) - - -@handle_exceptions -@handle_backend_invalid -@inputs_to_native_arrays -@handle_device_shifting -def iinfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Iinfo: - """ - Machine limits for integer data types. - - Parameters - ---------- - type - the kind of integer data-type about which to get information. - - Returns - ------- - ret - a class with that encapsules the following attributes: - - - **bits**: *int* - - number of bits occupied by the type. - - - **max**: *int* - - largest representable number. - - - **min**: *int* - - smallest representable number. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Dtype` input: - - >>> ivy.iinfo(ivy.int32) - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :code:`str` input: - - >>> ivy.iinfo('int32') - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([13,21,34], dtype=ivy.int8) - >>> ivy.iinfo(x) - iinfo(min=-128, max=127, dtype=int8) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) - >>> ivy.iinfo(x) - iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) - - With :class:`ivy.Container` input: - - >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), - ... y=ivy.array([76,81,16], dtype=ivy.uint32)) - >>> ivy.iinfo(c) - { - x: iinfo(min=0, max=65535, dtype=uint16), - y: iinfo(min=0, max=4294967295, dtype=uint32) - } - """ - return current_backend(None).iinfo(type) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def result_type( - *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] -) -> ivy.Dtype: - """ - Return the dtype that results from applying the type promotion rules (see - :ref:`type-promotion`) to the arguments. - - .. note:: - If provided mixed dtypes (e.g., integer and floating-point), the returned dtype - will be implementation-specific. - - Parameters - ---------- - arrays_and_dtypes - an arbitrary number of input arrays and/or dtypes. - - Returns - ------- - ret - the dtype resulting from an operation involving the input arrays and dtypes. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.array([3., 4., 5.]) - >>> d = ivy.result_type(x, y) - >>> print(d) - float32 - - With :class:`ivy.Dtype` input: - - >>> d = ivy.result_type(ivy.uint8, ivy.uint64) - >>> print(d) - uint64 - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = x.a.dtype - >>> print(d) - int32 - - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = ivy.result_type(x, ivy.float64) - >>> print(d) - { - a: float64 - } - """ - return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) - - -# Extra # -# ------# - -default_dtype_stack = list() -default_float_dtype_stack = list() -default_int_dtype_stack = list() -default_uint_dtype_stack = list() -default_complex_dtype_stack = list() - - -class DefaultDtype: - """Ivy's DefaultDtype class.""" - - def __init__(self, dtype: ivy.Dtype): - self._dtype = dtype - - def __enter__(self): - set_default_dtype(self._dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultFloatDtype: - """Ivy's DefaultFloatDtype class.""" - - def __init__(self, float_dtype: ivy.Dtype): - self._float_dtype = float_dtype - - def __enter__(self): - set_default_float_dtype(self._float_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_float_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultIntDtype: - """Ivy's DefaultIntDtype class.""" - - def __init__(self, int_dtype: ivy.Dtype): - self._int_dtype = int_dtype - - def __enter__(self): - set_default_int_dtype(self._int_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_int_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultUintDtype: - """Ivy's DefaultUintDtype class.""" - - def __init__(self, uint_dtype: ivy.UintDtype): - self._uint_dtype = uint_dtype - - def __enter__(self): - set_default_uint_dtype(self._uint_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_uint_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultComplexDtype: - """Ivy's DefaultComplexDtype class.""" - - def __init__(self, complex_dtype: ivy.Dtype): - self._complex_dtype = complex_dtype - - def __enter__(self): - set_default_complex_dtype(self._complex_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_complex_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -@handle_exceptions -def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: - """ - Get the number of bits used for representing the input data type. - - Parameters - ---------- - dtype_in - The data type to determine the number of bits for. - - Returns - ------- - ret - The number of bits used to represent the data type. - - Examples - -------- - With :class:`ivy.Dtype` inputs: - - >>> x = ivy.dtype_bits(ivy.float32) - >>> print(x) - 32 - - >>> x = ivy.dtype_bits('int64') - >>> print(x) - 64 - - With :class:`ivy.NativeDtype` inputs: - - >>> x = ivy.dtype_bits(ivy.native_bool) - >>> print(x) - 1 - """ - return current_backend(dtype_in).dtype_bits(dtype_in) - - -@handle_exceptions -def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: - """ - Check if the given data type is hashable or not. - - Parameters - ---------- - dtype_in - The data type to check. - - Returns - ------- - ret - True if data type is hashable else False - """ - try: - hash(dtype_in) - return True - except TypeError: - return False - + With :class:`ivy.NativeArray` input: -@handle_exceptions -def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: - """ - Convert native data type to string representation. + >>> x = ivy.native_array([[-1, -1, -1], + ... [1, 1, 1]], + ... dtype='int16') + >>> print(ivy.can_cast(x, 'uint8')) + False - Parameters - ---------- - dtype_in - The data type to convert to string. + With :class:`ivy.Container` input: - Returns - ------- - ret - data type string 'float32' + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3, 4, 5])) + >>> print(ivy.can_cast(x, 'int64')) + { + a: False, + b: True + } """ - return current_backend(None).as_ivy_dtype(dtype_in) + if isinstance(from_, (ivy.Array, ivy.NativeArray)): + from_ = from_.dtype + try: + dtype = ivy.promote_types(from_, to) + return dtype == to + except KeyError: + return False @handle_exceptions -def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: +def check_float(x: Any) -> bool: """ - Convert data type string representation to native data type. + Check if the input is a float or a float-like object. Parameters ---------- - dtype_in - The data type string to convert to native data type. + x + Input to check. Returns ------- ret - data type e.g. ivy.float32. + "True" if the input is a float or a float-like object, otherwise "False". """ - return current_backend(None).as_native_dtype(dtype_in) - - -def _check_float64(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "float64" - if math.isfinite(input): - m, e = math.frexp(input) - return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) - return False - - -def _check_complex128(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "complex128" - elif isinstance(input, np.ndarray): - return str(input.dtype) == "complex128" - if hasattr(input, "real") and hasattr(input, "imag"): - return _check_float64(input.real) and _check_float64(input.imag) - return False + return isinstance(x, (int, float)) and x is not bool @handle_exceptions @@ -1075,52 +806,53 @@ def closest_valid_dtype(type: Union[ivy.Dtype, str, None], /) -> Union[ivy.Dtype @handle_exceptions @handle_nestable @inputs_to_ivy_arrays -def default_float_dtype( +@handle_device_shifting +def default_complex_dtype( *, input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - float_dtype: Optional[Union[ivy.FloatDtype, ivy.NativeDtype]] = None, + complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, as_native: bool = False, ) -> Union[ivy.Dtype, str, ivy.NativeDtype]: """ Parameters ---------- input - Number or array for inferring the float dtype. - float_dtype - The float dtype to be returned. + Number or array for inferring the complex dtype. + complex_dtype + The complex dtype to be returned. as_native - Whether to return the float dtype as native dtype. + Whether to return the complex dtype as native dtype. Returns ------- - Return ``float_dtype`` as native or ivy dtype if provided, else - if ``input`` is given, return its float dtype, otherwise return the - global default float dtype. + Return ``complex_dtype`` as native or ivy dtype if provided, else + if ``input`` is given, return its complex dtype, otherwise return the + global default complex dtype. Examples -------- - >>> ivy.default_float_dtype() - 'float32' + >>> ivy.default_complex_dtype() + 'complex64' - >>> ivy.set_default_float_dtype(ivy.FloatDtype("float64")) - >>> ivy.default_float_dtype() - 'float64' + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' - >>> ivy.default_float_dtype(float_dtype=ivy.FloatDtype("float16")) - 'float16' + >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) + 'complex128' - >>> ivy.default_float_dtype(input=4294.967346) - 'float32' + >>> ivy.default_complex_dtype(input=4294.967346) + 'complex64' - >>> x = ivy.array([9.8,8.9], dtype="float16") - >>> ivy.default_float_dtype(input=x) - 'float16' + >>> x = ivy.array([9.8,8.9], dtype="complex128") + >>> ivy.default_complex_dtype(input=x) + 'complex128' """ - global default_float_dtype_stack - if ivy.exists(float_dtype): + global default_complex_dtype_stack + if ivy.exists(complex_dtype): if as_native is True: - return ivy.as_native_dtype(float_dtype) - return ivy.FloatDtype(ivy.as_ivy_dtype(float_dtype)) + return ivy.as_native_dtype(complex_dtype) + return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) as_native = ivy.default(as_native, False) if ivy.exists(input): if ivy.is_array(input): @@ -1129,90 +861,42 @@ def default_float_dtype( ret = str(input.dtype) elif isinstance(input, (list, tuple, dict)): if ivy.nested_argwhere( - input, lambda x: _check_float64(x), stop_after_n_found=1 + input, lambda x: _check_complex128(x), stop_after_n_found=1 ): - ret = ivy.float64 + ret = ivy.complex128 else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] elif isinstance(input, Number): - if _check_float64(input): - ret = ivy.float64 + if _check_complex128(input): + ret = ivy.complex128 else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] if as_native: return ivy.as_native_dtype(ret) - return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) - - -@handle_exceptions -def infer_default_dtype( - dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: - """ - Summary. - - Parameters - ---------- - dtype - - as_native - (Default value = False) - - Returns - ------- - Return the default data type for the “kind” (integer or floating-point) of dtype - - Examples - -------- - >>> ivy.set_default_int_dtype("int32") - >>> ivy.infer_default_dtype("int8") - 'int8' - - >>> ivy.set_default_float_dtype("float64") - >>> ivy.infer_default_dtype("float32") - 'float64' - - >>> ivy.set_default_uint_dtype("uint32") - >>> x = ivy.array([0], dtype="uint64") - >>> ivy.infer_default_dtype(x.dtype) - 'uint32' - """ - if ivy.is_complex_dtype(dtype): - default_dtype = ivy.default_complex_dtype(as_native=as_native) - elif ivy.is_float_dtype(dtype): - default_dtype = ivy.default_float_dtype(as_native=as_native) - elif ivy.is_uint_dtype(dtype): - default_dtype = ivy.default_uint_dtype(as_native=as_native) - elif ivy.is_int_dtype(dtype): - default_dtype = ivy.default_int_dtype(as_native=as_native) - elif as_native: - default_dtype = ivy.as_native_dtype("bool") - else: - default_dtype = ivy.as_ivy_dtype("bool") - return default_dtype + return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) @handle_exceptions @@ -1293,19 +977,114 @@ def default_dtype( elif as_native: return ivy.as_native_dtype("bool") else: - return "bool" - global default_dtype_stack - if not default_dtype_stack: - global default_float_dtype_stack - if default_float_dtype_stack: + return "bool" + global default_dtype_stack + if not default_dtype_stack: + global default_float_dtype_stack + if default_float_dtype_stack: + ret = default_float_dtype_stack[-1] + else: + ret = "float32" + else: + ret = default_dtype_stack[-1] + if as_native: + return ivy.as_native_dtype(ret) + return ivy.as_ivy_dtype(ret) + + +@handle_exceptions +@handle_nestable +@inputs_to_ivy_arrays +def default_float_dtype( + *, + input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + float_dtype: Optional[Union[ivy.FloatDtype, ivy.NativeDtype]] = None, + as_native: bool = False, +) -> Union[ivy.Dtype, str, ivy.NativeDtype]: + """ + Parameters + ---------- + input + Number or array for inferring the float dtype. + float_dtype + The float dtype to be returned. + as_native + Whether to return the float dtype as native dtype. + + Returns + ------- + Return ``float_dtype`` as native or ivy dtype if provided, else + if ``input`` is given, return its float dtype, otherwise return the + global default float dtype. + + Examples + -------- + >>> ivy.default_float_dtype() + 'float32' + + >>> ivy.set_default_float_dtype(ivy.FloatDtype("float64")) + >>> ivy.default_float_dtype() + 'float64' + + >>> ivy.default_float_dtype(float_dtype=ivy.FloatDtype("float16")) + 'float16' + + >>> ivy.default_float_dtype(input=4294.967346) + 'float32' + + >>> x = ivy.array([9.8,8.9], dtype="float16") + >>> ivy.default_float_dtype(input=x) + 'float16' + """ + global default_float_dtype_stack + if ivy.exists(float_dtype): + if as_native is True: + return ivy.as_native_dtype(float_dtype) + return ivy.FloatDtype(ivy.as_ivy_dtype(float_dtype)) + as_native = ivy.default(as_native, False) + if ivy.exists(input): + if ivy.is_array(input): + ret = ivy.dtype(input) + elif isinstance(input, np.ndarray): + ret = str(input.dtype) + elif isinstance(input, (list, tuple, dict)): + if ivy.nested_argwhere( + input, lambda x: _check_float64(x), stop_after_n_found=1 + ): + ret = ivy.float64 + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: + ret = default_float_dtype_stack[-1] + elif isinstance(input, Number): + if _check_float64(input): + ret = ivy.float64 + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: + ret = default_float_dtype_stack[-1] + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: ret = default_float_dtype_stack[-1] - else: - ret = "float32" - else: - ret = default_dtype_stack[-1] if as_native: return ivy.as_native_dtype(ret) - return ivy.as_ivy_dtype(ret) + return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) @handle_exceptions @@ -1522,149 +1301,177 @@ def default_uint_dtype( @handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_device_shifting -def default_complex_dtype( - *, - input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, - as_native: bool = False, -) -> Union[ivy.Dtype, str, ivy.NativeDtype]: +def dtype( + x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: """ + Get the data type for input array x. + Parameters ---------- - input - Number or array for inferring the complex dtype. - complex_dtype - The complex dtype to be returned. + x + Tensor for which to get the data type. as_native - Whether to return the complex dtype as native dtype. + Whether or not to return the dtype in string format. Default is ``False``. Returns ------- - Return ``complex_dtype`` as native or ivy dtype if provided, else - if ``input`` is given, return its complex dtype, otherwise return the - global default complex dtype. + ret + Data type of the array. Examples -------- - >>> ivy.default_complex_dtype() - 'complex64' + With :class:`ivy.Array` inputs: - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' + >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) + >>> y = ivy.dtype(x1) + >>> print(y) + float32 - >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) - 'complex128' + With :class:`ivy.NativeArray` inputs: - >>> ivy.default_complex_dtype(input=4294.967346) - 'complex64' + >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) + >>> y = ivy.dtype(x1) + >>> print(y) + int32 - >>> x = ivy.array([9.8,8.9], dtype="complex128") - >>> ivy.default_complex_dtype(input=x) - 'complex128' + With :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), + ... b=ivy.native_array([1, 0, 0, 0, 1])) + >>> y = ivy.dtype(x.a) + >>> print(y) + float32 """ - global default_complex_dtype_stack - if ivy.exists(complex_dtype): - if as_native is True: - return ivy.as_native_dtype(complex_dtype) - return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - as_native = ivy.default(as_native, False) - if ivy.exists(input): - if ivy.is_array(input): - ret = ivy.dtype(input) - elif isinstance(input, np.ndarray): - ret = str(input.dtype) - elif isinstance(input, (list, tuple, dict)): - if ivy.nested_argwhere( - input, lambda x: _check_complex128(x), stop_after_n_found=1 - ): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - elif isinstance(input, Number): - if _check_complex128(input): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - if as_native: - return ivy.as_native_dtype(ret) - return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) + return current_backend(x).dtype(x, as_native=as_native) + + +@handle_exceptions +def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: + """ + Get the number of bits used for representing the input data type. + + Parameters + ---------- + dtype_in + The data type to determine the number of bits for. + + Returns + ------- + ret + The number of bits used to represent the data type. + + Examples + -------- + With :class:`ivy.Dtype` inputs: + + >>> x = ivy.dtype_bits(ivy.float32) + >>> print(x) + 32 + + >>> x = ivy.dtype_bits('int64') + >>> print(x) + 64 + + With :class:`ivy.NativeDtype` inputs: + + >>> x = ivy.dtype_bits(ivy.native_bool) + >>> print(x) + 1 + """ + return current_backend(dtype_in).dtype_bits(dtype_in) @handle_exceptions @handle_backend_invalid -@handle_nestable @inputs_to_native_arrays @handle_device_shifting -def dtype( - x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: +def finfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Finfo: """ - Get the data type for input array x. + Machine limits for floating-point data types. Parameters ---------- - x - Tensor for which to get the data type. - as_native - Whether or not to return the dtype in string format. Default is ``False``. + type + the kind of floating-point data-type about which to get information. Returns ------- ret - Data type of the array. + an object having the followng attributes: + + - **bits**: *int* + + number of bits occupied by the floating-point data type. + + - **eps**: *float* + + difference between 1.0 and the next smallest representable floating-point + number larger than 1.0 according to the IEEE-754 standard. + + - **max**: *float* + + largest representable number. + + - **min**: *float* + + smallest representable number. + + - **smallest_normal**: *float* + + smallest positive floating-point number with full precision. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Dtype` input: - >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) - >>> y = ivy.dtype(x1) + >>> y = ivy.finfo(ivy.float32) + >>> print(y) + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) + + With :code:`str` input: + + >>> y = ivy.finfo('float32') >>> print(y) - float32 + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - With :class:`ivy.NativeArray` inputs: + With :class:`ivy.Array` input: - >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) - >>> y = ivy.dtype(x1) - >>> print(y) - int32 + >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) + >>> print(ivy.finfo(x)) + finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) - With :class:`ivy.Container` inputs: + >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) + >>> print(ivy.finfo(x)) + finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) - >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), - ... b=ivy.native_array([1, 0, 0, 0, 1])) - >>> y = ivy.dtype(x.a) - >>> print(y) - float32 + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), + ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) + >>> print(ivy.finfo(c)) + { + x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), + y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) + } """ - return current_backend(x).dtype(x, as_native=as_native) + return current_backend(None).finfo(type) @handle_exceptions @@ -1768,6 +1575,131 @@ def function_unsupported_dtypes( ) +@handle_exceptions +@handle_backend_invalid +@inputs_to_native_arrays +@handle_device_shifting +def iinfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Iinfo: + """ + Machine limits for integer data types. + + Parameters + ---------- + type + the kind of integer data-type about which to get information. + + Returns + ------- + ret + a class with that encapsules the following attributes: + + - **bits**: *int* + + number of bits occupied by the type. + + - **max**: *int* + + largest representable number. + + - **min**: *int* + + smallest representable number. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Examples + -------- + With :class:`ivy.Dtype` input: + + >>> ivy.iinfo(ivy.int32) + iinfo(min=-2147483648, max=2147483647, dtype=int32) + + With :code:`str` input: + + >>> ivy.iinfo('int32') + iinfo(min=-2147483648, max=2147483647, dtype=int32) + + With :class:`ivy.Array` input: + + >>> x = ivy.array([13,21,34], dtype=ivy.int8) + >>> ivy.iinfo(x) + iinfo(min=-128, max=127, dtype=int8) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) + >>> ivy.iinfo(x) + iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) + + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), + ... y=ivy.array([76,81,16], dtype=ivy.uint32)) + >>> ivy.iinfo(c) + { + x: iinfo(min=0, max=65535, dtype=uint16), + y: iinfo(min=0, max=4294967295, dtype=uint32) + } + """ + return current_backend(None).iinfo(type) + + +@handle_exceptions +def infer_default_dtype( + dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: + """ + Summary. + + Parameters + ---------- + dtype + + as_native + (Default value = False) + + Returns + ------- + Return the default data type for the “kind” (integer or floating-point) of dtype + + Examples + -------- + >>> ivy.set_default_int_dtype("int32") + >>> ivy.infer_default_dtype("int8") + 'int8' + + >>> ivy.set_default_float_dtype("float64") + >>> ivy.infer_default_dtype("float32") + 'float64' + + >>> ivy.set_default_uint_dtype("uint32") + >>> x = ivy.array([0], dtype="uint64") + >>> ivy.infer_default_dtype(x.dtype) + 'uint32' + """ + if ivy.is_complex_dtype(dtype): + default_dtype = ivy.default_complex_dtype(as_native=as_native) + elif ivy.is_float_dtype(dtype): + default_dtype = ivy.default_float_dtype(as_native=as_native) + elif ivy.is_uint_dtype(dtype): + default_dtype = ivy.default_uint_dtype(as_native=as_native) + elif ivy.is_int_dtype(dtype): + default_dtype = ivy.default_int_dtype(as_native=as_native) + elif as_native: + default_dtype = ivy.as_native_dtype("bool") + else: + default_dtype = ivy.as_ivy_dtype("bool") + return default_dtype + + @handle_exceptions def invalid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -1822,30 +1754,151 @@ def is_bool_dtype( ret "True" if the input data type is a bool, otherwise "False". - Both the description and the type hints above assumes an array input for - simplicity but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Both the description and the type hints above assumes an array input for + simplicity but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. + """ + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, np.ndarray): + return "bool" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return ( + True + if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) + else False + ) + elif isinstance(dtype_in, (list, tuple, dict)): + return ( + True + if ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (bool, np.bool)) and x is not int, + ) + else False + ) + return "bool" in ivy.as_ivy_dtype(dtype_in) + + +@handle_exceptions +@handle_nestable +@inputs_to_ivy_arrays +def is_complex_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: + """ + Determine whether the input data type is a complex dtype. + + Parameters + ---------- + dtype_in + The array or data type to check + + Returns + ------- + ret + Whether or not the array or data type is of a complex dtype + + Examples + -------- + >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) + True + + >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) + True + + >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) + False + """ + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() + elif isinstance(dtype_in, np.ndarray): + return "complex" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return isinstance(dtype_in, (complex, np.complexfloating)) + elif isinstance(dtype_in, (list, tuple, dict)): + return ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (complex, np.complexfloating)) + or (ivy.is_array(x) and "complex" in ivy.dtype(x)), + ) + return "complex" in as_ivy_dtype(dtype_in) + + +@handle_exceptions +@handle_nestable +@inputs_to_native_arrays +def is_float_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: + """ + Determine whether the input data type is a float dtype. + + Parameters + ---------- + dtype_in + The array or data type to check + + Returns + ------- + ret + Whether or not the array or data type is of a floating point dtype + + Examples + -------- + >>> x = ivy.is_float_dtype(ivy.float32) + >>> print(x) + True + + >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) + >>> print(ivy.is_float_dtype(arr)) + True """ if ivy.is_array(dtype_in): dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() elif isinstance(dtype_in, np.ndarray): - return "bool" in dtype_in.dtype.name + return "float" in dtype_in.dtype.name elif isinstance(dtype_in, Number): - return ( - True - if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) - else False - ) + return True if isinstance(dtype_in, (float, np.floating)) else False elif isinstance(dtype_in, (list, tuple, dict)): return ( True if ivy.nested_argwhere( dtype_in, - lambda x: isinstance(x, (bool, np.bool)) and x is not int, + lambda x: isinstance(x, (float, np.floating)) + or (ivy.is_array(x) and "float" in ivy.dtype(x)), ) else False ) - return "bool" in ivy.as_ivy_dtype(dtype_in) + return "float" in as_ivy_dtype(dtype_in) + + +@handle_exceptions +def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: + """ + Check if the given data type is hashable or not. + + Parameters + ---------- + dtype_in + The data type to check. + + Returns + ------- + ret + True if data type is hashable else False + """ + try: + hash(dtype_in) + return True + except TypeError: + return False @handle_exceptions @@ -1937,72 +1990,34 @@ def is_int_dtype( @handle_exceptions -def check_float(x: Any) -> bool: - """ - Check if the input is a float or a float-like object. - - Parameters - ---------- - x - Input to check. - - Returns - ------- - ret - "True" if the input is a float or a float-like object, otherwise "False". - """ - return isinstance(x, (int, float)) and x is not bool - - -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -def is_float_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], - /, -) -> bool: +def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: """ - Determine whether the input data type is a float dtype. + Determine whether the input dtype is a Native dtype. Parameters ---------- dtype_in - The array or data type to check + Determine whether the input data type is a native data type object. Returns ------- ret - Whether or not the array or data type is of a floating point dtype + Boolean, whether or not dtype_in is a native data type. Examples -------- - >>> x = ivy.is_float_dtype(ivy.float32) - >>> print(x) + >>> ivy.set_backend('numpy') + >>> ivy.is_native_dtype(np.int32) True - >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) - >>> print(ivy.is_float_dtype(arr)) - True + >>> ivy.set_backend('numpy') + >>> ivy.is_native_array(ivy.float64) + False """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "float" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return True if isinstance(dtype_in, (float, np.floating)) else False - elif isinstance(dtype_in, (list, tuple, dict)): - return ( - True - if ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (float, np.floating)) - or (ivy.is_array(x) and "float" in ivy.dtype(x)), - ) - else False - ) - return "float" in as_ivy_dtype(dtype_in) + try: + return current_backend(None).is_native_dtype(dtype_in) + except ValueError: + return False @handle_exceptions @@ -2054,101 +2069,203 @@ def is_uint_dtype( @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -def is_complex_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], +def promote_types( + type1: Union[ivy.Dtype, ivy.NativeDtype], + type2: Union[ivy.Dtype, ivy.NativeDtype], /, -) -> bool: + *, + array_api_promotion: bool = False, +) -> ivy.Dtype: """ - Determine whether the input data type is a complex dtype. + Promote the datatypes type1 and type2, returning the data type they promote to. Parameters ---------- - dtype_in - The array or data type to check + type1 + the first of the two types to promote + type2 + the second of the two types to promote + array_api_promotion + whether to only use the array api promotion rules Returns ------- ret - Whether or not the array or data type is of a complex dtype + The type that both input types promote to + """ + query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] + query.sort(key=lambda x: str(x)) + query = tuple(query) + + def _promote(query): + if array_api_promotion: + return ivy.array_api_promotion_table[query] + return ivy.promotion_table[query] + + try: + ret = _promote(query) + except KeyError: + # try again with the dtypes swapped + query = (query[1], query[0]) + try: + ret = _promote(query) + except KeyError: + raise ivy.utils.exceptions.IvyDtypePromotionError( + "these dtypes ({} and {}) are not type promotable, ".format( + type1, type2 + ) + ) + return ret + + +@handle_exceptions +def promote_types_of_inputs( + x1: Union[ivy.NativeArray, Number, Iterable[Number]], + x2: Union[ivy.NativeArray, Number, Iterable[Number]], + /, + *, + array_api_promotion: bool = False, +) -> Tuple[ivy.NativeArray, ivy.NativeArray]: + """ + Promote the dtype of the given native array inputs to a common dtype based on type + promotion rules. + + While passing float or integer values or any other non-array input + to this function, it should be noted that the return will be an + array-like object. Therefore, outputs from this function should be + used as inputs only for those functions that expect an array-like or + tensor-like objects, otherwise it might give unexpected results. + """ + + def _special_case(a1, a2): + # check for float number and integer array case + return isinstance(a1, float) and "int" in str(a2.dtype) + + def _get_target_dtype(scalar, arr): + # identify a good dtype to give the scalar value, + # based on it's own type and that of the arr value + if _special_case(scalar, arr): + return "float64" + elif arr.dtype == bool and not isinstance(scalar, bool): + return None # let ivy infer a dtype + elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): + return "complex128" + else: + return arr.dtype + + if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): + device = ivy.default_device(item=x1, as_native=True) + x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) + elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): + device = ivy.default_device(item=x2, as_native=True) + x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) + elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): + x1 = ivy.asarray(x1) + x2 = ivy.asarray(x2) + + if x1.dtype != x2.dtype: + promoted = promote_types( + x1.dtype, x2.dtype, array_api_promotion=array_api_promotion + ) + x1 = ivy.astype(x1, promoted, copy=False) + x2 = ivy.astype(x2, promoted, copy=False) + + ivy.utils.assertions._check_jax_x64_flag(x1.dtype) + return ivy.to_native(x1), ivy.to_native(x2) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def result_type( + *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] +) -> ivy.Dtype: + """ + Return the dtype that results from applying the type promotion rules (see + :ref:`type-promotion`) to the arguments. + + .. note:: + If provided mixed dtypes (e.g., integer and floating-point), the returned dtype + will be implementation-specific. + + Parameters + ---------- + arrays_and_dtypes + an arbitrary number of input arrays and/or dtypes. + + Returns + ------- + ret + the dtype resulting from an operation involving the input arrays and dtypes. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.array([3., 4., 5.]) + >>> d = ivy.result_type(x, y) + >>> print(d) + float32 - Examples - -------- - >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) - True + With :class:`ivy.Dtype` input: - >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) - True + >>> d = ivy.result_type(ivy.uint8, ivy.uint64) + >>> print(d) + uint64 - >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) - False + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = x.a.dtype + >>> print(d) + int32 + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = ivy.result_type(x, ivy.float64) + >>> print(d) + { + a: float64 + } """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "complex" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return isinstance(dtype_in, (complex, np.complexfloating)) - elif isinstance(dtype_in, (list, tuple, dict)): - return ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (complex, np.complexfloating)) - or (ivy.is_array(x) and "complex" in ivy.dtype(x)), - ) - return "complex" in as_ivy_dtype(dtype_in) + return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) @handle_exceptions -def promote_types( - type1: Union[ivy.Dtype, ivy.NativeDtype], - type2: Union[ivy.Dtype, ivy.NativeDtype], - /, - *, - array_api_promotion: bool = False, -) -> ivy.Dtype: +def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): """ - Promote the datatypes type1 and type2, returning the data type they promote to. + Set the 'complex_dtype' as the default data type. Parameters ---------- - type1 - the first of the two types to promote - type2 - the second of the two types to promote - array_api_promotion - whether to only use the array api promotion rules + complex_dtype + The complex data type to be set as the default. - Returns - ------- - ret - The type that both input types promote to - """ - query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] - query.sort(key=lambda x: str(x)) - query = tuple(query) + Examples + -------- + With :class: `ivy.Dtype` input: - def _promote(query): - if array_api_promotion: - return ivy.array_api_promotion_table[query] - return ivy.promotion_table[query] + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' - try: - ret = _promote(query) - except KeyError: - # try again with the dtypes swapped - query = (query[1], query[0]) - try: - ret = _promote(query) - except KeyError: - raise ivy.utils.exceptions.IvyDtypePromotionError( - "these dtypes ({} and {}) are not type promotable, ".format( - type1, type2 - ) - ) - return ret + >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) + >>> ivy.default_complex_dtype() + 'complex128' + """ + complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) + ivy.utils.assertions._check_jax_x64_flag(complex_dtype) + global default_complex_dtype_stack + default_complex_dtype_stack.append(complex_dtype) @handle_exceptions @@ -2269,34 +2386,6 @@ def set_default_uint_dtype(uint_dtype: Union[ivy.Dtype, str], /): default_uint_dtype_stack.append(uint_dtype) -@handle_exceptions -def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): - """ - Set the 'complex_dtype' as the default data type. - - Parameters - ---------- - complex_dtype - The complex data type to be set as the default. - - Examples - -------- - With :class: `ivy.Dtype` input: - - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' - - >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) - >>> ivy.default_complex_dtype() - 'complex128' - """ - complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - ivy.utils.assertions._check_jax_x64_flag(complex_dtype) - global default_complex_dtype_stack - default_complex_dtype_stack.append(complex_dtype) - - @handle_exceptions def type_promote_arrays( x1: Union[ivy.Array, ivy.NativeArray], @@ -2323,6 +2412,27 @@ def type_promote_arrays( return ivy.astype(x1, new_type), ivy.astype(x2, new_type) +@handle_exceptions +def unset_default_complex_dtype(): + """ + Reset the current default complex dtype to the previous state. + + Examples + -------- + >>> ivy.set_default_complex_dtype(ivy.complex64) + >>> ivy.set_default_complex_dtype(ivy.complex128) + >>> ivy.default_complex_dtype_stack + ['complex64','complex128'] + + >>> ivy.unset_default_complex_dtype() + >>> ivy.default_complex_dtype_stack + ['complex64'] + """ + global default_complex_dtype_stack + if default_complex_dtype_stack: + default_complex_dtype_stack.pop(-1) + + @handle_exceptions def unset_default_dtype(): """ @@ -2413,27 +2523,6 @@ def unset_default_uint_dtype(): default_uint_dtype_stack.pop(-1) -@handle_exceptions -def unset_default_complex_dtype(): - """ - Reset the current default complex dtype to the previous state. - - Examples - -------- - >>> ivy.set_default_complex_dtype(ivy.complex64) - >>> ivy.set_default_complex_dtype(ivy.complex128) - >>> ivy.default_complex_dtype_stack - ['complex64','complex128'] - - >>> ivy.unset_default_complex_dtype() - >>> ivy.default_complex_dtype_stack - ['complex64'] - """ - global default_complex_dtype_stack - if default_complex_dtype_stack: - default_complex_dtype_stack.pop(-1) - - @handle_exceptions def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -2466,90 +2555,3 @@ def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bo if dtype_in is None: return True return ivy.as_ivy_dtype(dtype_in) in ivy.valid_dtypes - - -@handle_exceptions -def promote_types_of_inputs( - x1: Union[ivy.NativeArray, Number, Iterable[Number]], - x2: Union[ivy.NativeArray, Number, Iterable[Number]], - /, - *, - array_api_promotion: bool = False, -) -> Tuple[ivy.NativeArray, ivy.NativeArray]: - """ - Promote the dtype of the given native array inputs to a common dtype based on type - promotion rules. - - While passing float or integer values or any other non-array input - to this function, it should be noted that the return will be an - array-like object. Therefore, outputs from this function should be - used as inputs only for those functions that expect an array-like or - tensor-like objects, otherwise it might give unexpected results. - """ - - def _special_case(a1, a2): - # check for float number and integer array case - return isinstance(a1, float) and "int" in str(a2.dtype) - - def _get_target_dtype(scalar, arr): - # identify a good dtype to give the scalar value, - # based on it's own type and that of the arr value - if _special_case(scalar, arr): - return "float64" - elif arr.dtype == bool and not isinstance(scalar, bool): - return None # let ivy infer a dtype - elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): - return "complex128" - else: - return arr.dtype - - if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): - device = ivy.default_device(item=x1, as_native=True) - x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) - elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): - device = ivy.default_device(item=x2, as_native=True) - x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) - elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): - x1 = ivy.asarray(x1) - x2 = ivy.asarray(x2) - - if x1.dtype != x2.dtype: - promoted = promote_types( - x1.dtype, x2.dtype, array_api_promotion=array_api_promotion - ) - x1 = ivy.astype(x1, promoted, copy=False) - x2 = ivy.astype(x2, promoted, copy=False) - - ivy.utils.assertions._check_jax_x64_flag(x1.dtype) - return ivy.to_native(x1), ivy.to_native(x2) - - -@handle_exceptions -def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: - """ - Determine whether the input dtype is a Native dtype. - - Parameters - ---------- - dtype_in - Determine whether the input data type is a native data type object. - - Returns - ------- - ret - Boolean, whether or not dtype_in is a native data type. - - Examples - -------- - >>> ivy.set_backend('numpy') - >>> ivy.is_native_dtype(np.int32) - True - - >>> ivy.set_backend('numpy') - >>> ivy.is_native_array(ivy.float64) - False - """ - try: - return current_backend(None).is_native_dtype(dtype_in) - except ValueError: - return False diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py index 1fee8c85ac25b..058b9dd513163 100644 --- a/ivy/functional/ivy/device.py +++ b/ivy/functional/ivy/device.py @@ -10,6 +10,21 @@ import types from typing import Type, Optional, Tuple +# nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. + +from typing import Union, Callable, Iterable, Any + +# local +import ivy +from ivy.func_wrapper import ( + handle_out_argument, + to_native_arrays_and_back, + handle_nestable, + handle_array_like_without_promotion, + handle_backend_invalid, +) +from ivy.utils.exceptions import handle_exceptions + # noinspection PyUnresolvedReferences try: import pynvml @@ -24,26 +39,11 @@ " of the Ivy's device module will be limited. Please install pynvml if" " you wish to use GPUs with Ivy." ) - # nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. - -from typing import Union, Callable, Iterable, Any - -# local -import ivy -from ivy.func_wrapper import ( - handle_out_argument, - to_native_arrays_and_back, - handle_nestable, - handle_array_like_without_promotion, - handle_backend_invalid, -) -from ivy.utils.exceptions import handle_exceptions default_device_stack = list() +max_chunk_sizes = dict() soft_device_mode_stack = list() -dev_handles = dict() split_factors = dict() -max_chunk_sizes = dict() # Extra # @@ -136,10 +136,87 @@ def __exit__( return self -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return ivy.current_backend().handle_soft_device_variable( - *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs - ) +# Profiler # + + +class Profiler(abc.ABC): + """ + The profiler class is used to profile the execution of some code. + + Parameters + ---------- + save_dir + The directory to save the profile data to. + """ + + def __init__(self, save_dir: str): + self._save_dir = save_dir + + @abc.abstractmethod + def start(self): + """ + Start the profiler. + + This should be called before the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def stop(self): + """ + Stop the profiler. + + This should be called after the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def __enter__(self): + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def __exit__(self, exc_type, exc_val, exc_tb): + raise ivy.utils.exceptions.IvyNotImplementedException + + +# --- Helpers --- # +# --------------- # + + +def _get_devices(fn: Callable, complement: bool = True) -> Tuple: + valid_devices = ivy.valid_devices + invalid_devices = ivy.invalid_devices + all_devices = ivy.all_devices + + supported = set(ivy.valid_devices) + + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + supported = set(all_devices).difference(supported) + return supported + + # Their values are formated like either + # 1. fn.supported_devices = ("cpu",) + # Could also have the "all" value for the framework + basic = [ + ("supported_devices", set.intersection, valid_devices), + ("unsupported_devices", set.difference, invalid_devices), + ] + for key, merge_fn, base in basic: + if hasattr(fn, key): + v = getattr(fn, key) + if "einops" in fn.__name__ and isinstance(v, dict): + v = v.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(v, tuple) + supported = merge_fn(supported, set(v)) + + if complement: + supported = set(all_devices).difference(supported) + + return tuple(supported) # Helpers # @@ -155,6 +232,24 @@ def _get_nvml_gpu_handle(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: return handle +def _is_valid_devices_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): + fn_supported_devices = fn.supported_devices + fn_unsupported_devices = fn.unsupported_devices + if isinstance(fn_supported_devices, dict): + if isinstance(fn_unsupported_devices, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_devices + and backend_str in fn_unsupported_devices + ): + return False + else: + if isinstance(fn_unsupported_devices, tuple): + return False + return True + + def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kwargs): with ivy.ArrayMode(False): default_device = ivy.default_device(device_shifting_dev, as_native=True) @@ -169,176 +264,177 @@ def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kw return args, kwargs, default_device -# Device Queries # +# --- Main --- # +# ------------ # -# Array Printing + +# Conversions @handle_exceptions -def get_all_ivy_arrays_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, -) -> ivy.Container: +def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: """ - Get all ivy arrays which are currently alive on the specified device. + Convert device to string representation. Parameters ---------- device - The device handle from which to get the arrays + The device handle to convert to string. Returns ------- ret - Container with the arrays found for the specified device [identity, array] + Device string e.g. 'cuda:0'. Examples -------- - >>> x = ivy.array([1,0,2]) - >>> y = ivy.dev(x) - >>> z = ivy.get_all_ivy_arrays_on_dev(y) - >>> print(z) - {139740789224448:ivy.array([1,0,2])}, + >>> y = ivy.as_ivy_dev('cpu') + >>> print(y) + cpu """ - device = ivy.as_ivy_dev(device) - all_arrays = list() - for obj in gc.get_objects(): - if ( - obj is ivy.data_classes.array.array.Array - and ivy.is_ivy_array(obj) - and ivy.dev(obj) == device - ): - all_arrays.append(obj) - - return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) + return ivy.current_backend().as_ivy_dev(device) @handle_exceptions -def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: +def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: """ - Return the number of arrays which are currently alive on the specified device. + Convert device string representation to native device type. Parameters ---------- device - The device handle from which to count the arrays + The device string to convert to native device handle. + A native device handle can be passed in instead - in this case + the unmodified parameter is returned. Returns ------- ret - Number of arrays on the specified device + Native device handle. Examples -------- - >>> x1 = ivy.array([-1, 0, 5.2]) - >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 2 + With :class:`ivy.Device` input: - >>> x1 = ivy.native_array([-1, 0, 5.2]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 0 + >>> ivy.set_backend("numpy") + >>> ivy.as_native_dev("cpu") + 'cpu' - >>> x = ivy.Container(x1=ivy.array([-1]), - ... x2=ivy.native_array([-1])) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 1 + >>> ivy.set_backend("tensorflow") + >>> ivy.as_native_dev("tpu:3") + '/TPU:3' + + With :class:`ivy.NativeDevice` input: + + >>> import torch + >>> device = torch.device("cuda") + >>> device + device(type='cuda') + + >>> ivy.as_native_dev(device) + device(type='cuda') """ - return len(ivy.get_all_ivy_arrays_on_dev(device)) + return ivy.current_backend().as_native_dev(device) + + +# Memory @handle_exceptions -@handle_nestable -def print_all_ivy_arrays_on_dev( - *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - attr_only: bool = True, -) -> None: +def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: """ - Print the shape and dtype for all ivy arrays which are currently alive on the - specified device. + Clear memory cache on target device. Parameters ---------- device - The device on which to print the arrays - - attr_only - Whether or not to only print the `shape` and `dtype` attributes of the array + The device string to convert to native device handle or native device handle. Examples -------- - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y) - ((3,), 'int32') - ((3,), 'int32') - - - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) - [1,0,2] - [3,2,1] + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.clear_cached_mem_on_dev(device) """ - arrs = ivy.get_all_ivy_arrays_on_dev(device).values() - if attr_only: - [print((arr.shape, arr.dtype)) for arr in arrs] - else: - [print(arr) for arr in arrs] + ivy.current_backend().clear_cached_mem_on_dev(device) -ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False +# Default Device # +# noinspection PyShadowingNames @handle_exceptions -def set_soft_device_mode(mode: bool) -> None: - """ - Set the mode of whether to move input arrays to `ivy.default_device()` before - performing an operation. - - Parameter - --------- - mode - boolean whether to move input arrays - Examples - -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode - False - >>> ivy.set_soft_device_mode(True) - >>> ivy.soft_device_mode - True +def default_device( + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + /, + *, + item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, + as_native: bool = None, +) -> Union[ivy.Device, ivy.NativeDevice]: """ - global soft_device_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - soft_device_mode_stack.append(mode) - ivy.__setattr__("soft_device_mode", mode, True) + Return the input device or the default device. If the as_native flag is set, the + device will be converted to a native device. If the item is provided, the item's + device is returned. If the device is not provided, the last default device is + returned. If a default device has not been set, the first gpu is returned if + available, otherwise the cpu is returned. + Parameters + ---------- + device + The device to be returned or converted. + item + The item to get the device from. + as_native + Whether to convert the device to a native device. -@handle_exceptions -def unset_soft_device_mode() -> None: - """ - Reset the mode of moving input arrays to `ivy.default_device()` before performing an - operation. + Returns + ------- + ret + Device handle or string. Examples -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode - False - >>> ivy.unset_soft_device_mode() - >>> ivy.soft_device_mode - True + >>> ivy.default_device() + device(type='cpu') + + >>> ivy.default_device("gpu:0") + 'gpu:0' + + >>> ivy.default_device(item=[], as_native=False) + 'cpu' + + >>> ivy.default_device(item=(), as_native=True) + device(type='cpu') + + >>> ivy.default_device(item={"a": 1}, as_native=True) + device(type='cpu') + + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'gpu:0') + >>> ivy.default_device(item=x, as_native=True) + device(type='gpu', id=0) """ - global soft_device_mode_stack - if soft_device_mode_stack: - soft_device_mode_stack.pop(-1) - mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False - ivy.__setattr__("soft_device_mode", mode, True) + if ivy.exists(device): + if as_native is True: + return ivy.as_native_dev(device) + elif as_native is False: + return ivy.as_ivy_dev(device) + return device + as_native = ivy.default(as_native, False) + if ivy.exists(item): + if isinstance(item, (list, tuple, dict)) and len(item) == 0: + pass + elif ivy.is_array(item): + return ivy.dev(item, as_native=as_native) + global default_device_stack + if not default_device_stack: + ret = "gpu:0" if ivy.gpu_is_available() else "cpu" + else: + ret = default_device_stack[-1] + if as_native: + return ivy.as_native_dev(ret) + return ivy.as_ivy_dev(ret) # Retrieval @@ -385,195 +481,301 @@ def dev( return ivy.current_backend(x).dev(x, as_native=as_native) -# Conversions +# Utilization @handle_exceptions -def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: +def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Convert device to string representation. + Get the current utilization (%) for a given device. Parameters ---------- device - The device handle to convert to string. + The device string of the device to query utilization for. Returns ------- ret - Device string e.g. 'cuda:0'. + The device utilization (%) - Examples - -------- - >>> y = ivy.as_ivy_dev('cpu') - >>> print(y) - cpu + Example + ------- + >>> ivy.dev_util('cpu') + 13.4 + >>> ivy.dev_util('gpu:0') + 7.8 + >>> ivy.dev_util('cpu') + 93.4 + >>> ivy.dev_util('gpu:2') + 57.4 + >>> ivy.dev_util('cpu') + 84.2 """ - return ivy.current_backend().as_ivy_dev(device) + if device == "cpu": + return psutil.cpu_percent() + elif "gpu" in device: + handle = _get_nvml_gpu_handle(device) + return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions -def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: +@handle_nestable +def function_supported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: """ - Convert device string representation to native device type. + Return the supported devices of the current backend's function. The function returns + a dict containing the supported devices for the compositional and primary + implementations in case of partial mixed functions. Parameters ---------- - device - The device string to convert to native device handle. - A native device handle can be passed in instead - in this case - the unmodified parameter is returned. + fn + The function to check for the supported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. Returns ------- ret - Native device handle. + Tuple or dict containing the supported devices of the function Examples -------- - With :class:`ivy.Device` input: + >>> import ivy + >>> print(ivy.function_supported_devices(ivy.ones)) + ('cpu', 'gpu') + """ + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=False), + } + else: + supported_devices = set(_get_devices(fn, complement=False)) + if recurse: + supported_devices = ivy.functional.data_type._nested_get( + fn, supported_devices, set.intersection, function_supported_devices + ) - >>> ivy.set_backend("numpy") - >>> ivy.as_native_dev("cpu") - 'cpu' + return ( + supported_devices + if isinstance(supported_devices, dict) + else tuple(supported_devices) + ) - >>> ivy.set_backend("tensorflow") - >>> ivy.as_native_dev("tpu:3") - '/TPU:3' - With :class:`ivy.NativeDevice` input: +@handle_exceptions +@handle_nestable +def function_unsupported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: + """ + Return the unsupported devices of the current backend's function. The function + returns a dict containing the unsupported devices for the compositional and primary + implementations in case of partial mixed functions. - >>> import torch - >>> device = torch.device("cuda") - >>> device - device(type='cuda') + Parameters + ---------- + fn + The function to check for the unsupported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. - >>> ivy.as_native_dev(device) - device(type='cuda') + Returns + ------- + ret + Tuple or dict containing the unsupported devices of the function + + Examples + -------- + >>> print(ivy.function_unsupported_devices(ivy.ones)) + ('tpu',) """ - return ivy.current_backend().as_native_dev(device) + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=True), + } + else: + unsupported_devices = set(_get_devices(fn, complement=True)) + if recurse: + unsupported_devices = ivy.functional.data_type._nested_get( + fn, unsupported_devices, set.union, function_unsupported_devices + ) + return ( + unsupported_devices + if isinstance(unsupported_devices, dict) + else tuple(unsupported_devices) + ) -# Memory +# Device Queries # + +# Array Printing @handle_exceptions -def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +def get_all_ivy_arrays_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, +) -> ivy.Container: """ - Clear memory cache on target device. + Get all ivy arrays which are currently alive on the specified device. Parameters ---------- device - The device string to convert to native device handle or native device handle. + The device handle from which to get the arrays + + Returns + ------- + ret + Container with the arrays found for the specified device [identity, array] Examples -------- - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.clear_cached_mem_on_dev(device) + >>> x = ivy.array([1,0,2]) + >>> y = ivy.dev(x) + >>> z = ivy.get_all_ivy_arrays_on_dev(y) + >>> print(z) + {139740789224448:ivy.array([1,0,2])}, """ - ivy.current_backend().clear_cached_mem_on_dev(device) + device = ivy.as_ivy_dev(device) + all_arrays = list() + for obj in gc.get_objects(): + if ( + obj is ivy.data_classes.array.array.Array + and ivy.is_ivy_array(obj) + and ivy.dev(obj) == device + ): + all_arrays.append(obj) + + return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) + + +# Availability @handle_exceptions -def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +def gpu_is_available() -> bool: """ - Get the total amount of memory (in GB) for a given device string. In case of CPU, - the total RAM is returned. + Determine whether a GPU is available to use, with the backend framework. + + Returns + ------- + ret + Boolean, as to whether a gpu is available. + + Examples + -------- + >>> print(ivy.gpu_is_available()) + False + """ + return ivy.current_backend().gpu_is_available() + + +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return ivy.current_backend().handle_soft_device_variable( + *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs + ) + + +@handle_exceptions +def num_cpu_cores(*, logical: bool = True) -> int: + """ + Determine the number of cores available in the cpu. Parameters ---------- - device - The device string to convert to native device handle. + logical + Whether request is for number of physical or logical cores available in CPU Returns ------- ret - The total memory on the device in GB. + Number of cores available in CPU Examples -------- - >>> x = ivy.total_mem_on_dev("cpu") - >>> print(x) - 53.66700032 - - >>> x = ivy.total_mem_on_dev("gpu:0") - >>> print(x) - 8.589934592 + >>> print(ivy.num_cpu_cores(logical=False)) + 2 """ - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.total / 1e9 - elif device == "cpu": - return psutil.virtual_memory().total / 1e9 + if logical: + return psutil.cpu_count(logical=logical) else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + return psutil.cpu_count(logical=False) @handle_exceptions -def used_mem_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - process_specific: bool = False, -) -> float: +def num_gpus() -> int: """ - Get the used memory (in GB) for a given device string. In case of CPU, the used RAM - is returned. + Determine the number of available GPUs, with the backend framework. + + Returns + ------- + ret + Number of available GPUs. + + Examples + -------- + >>> print(ivy.num_gpus()) + 1 + """ + return ivy.current_backend().num_gpus() + + +@handle_exceptions +def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: + """ + Return the number of arrays which are currently alive on the specified device. Parameters ---------- device - The device string to convert to native device handle. - process_specific - Whether to check the memory used by this python process alone. Default is - False. + The device handle from which to count the arrays Returns ------- ret - The used memory on the device in GB. + Number of arrays on the specified device Examples -------- - >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) - >>> print(x) - 6.219563008 + >>> x1 = ivy.array([-1, 0, 5.2]) + >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 2 - >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) - >>> print(x) - 0.902400346 + >>> x1 = ivy.native_array([-1, 0, 5.2]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 0 - >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) + >>> x = ivy.Container(x1=ivy.array([-1]), + ... x2=ivy.native_array([-1])) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) >>> print(y) - 0.525205504 + 1 """ - ivy.clear_cached_mem_on_dev(device) - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - if process_specific: - pid = os.getpid() - for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): - if process.pid == pid: - return process.usedGpuMemory / 1e9 - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.used / 1e9 - elif device == "cpu": - if process_specific: - return psutil.Process(os.getpid()).memory_info().rss / 1e9 - vm = psutil.virtual_memory() - return (vm.total - vm.available) / 1e9 - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + return len(ivy.get_all_ivy_arrays_on_dev(device)) @handle_exceptions @@ -637,308 +839,151 @@ def percent_used_mem_on_dev( ) -# Utilization - - @handle_exceptions -def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +@handle_nestable +def print_all_ivy_arrays_on_dev( + *, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + attr_only: bool = True, +) -> None: """ - Get the current utilization (%) for a given device. + Print the shape and dtype for all ivy arrays which are currently alive on the + specified device. Parameters ---------- device - The device string of the device to query utilization for. + The device on which to print the arrays - Returns - ------- - ret - The device utilization (%) + attr_only + Whether or not to only print the `shape` and `dtype` attributes of the array - Example - ------- - >>> ivy.dev_util('cpu') - 13.4 - >>> ivy.dev_util('gpu:0') - 7.8 - >>> ivy.dev_util('cpu') - 93.4 - >>> ivy.dev_util('gpu:2') - 57.4 - >>> ivy.dev_util('cpu') - 84.2 - """ - if device == "cpu": - return psutil.cpu_percent() - elif "gpu" in device: - handle = _get_nvml_gpu_handle(device) - return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + Examples + -------- + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y) + ((3,), 'int32') + ((3,), 'int32') -# Availability + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) + [1,0,2] + [3,2,1] + """ + arrs = ivy.get_all_ivy_arrays_on_dev(device).values() + if attr_only: + [print((arr.shape, arr.dtype)) for arr in arrs] + else: + [print(arr) for arr in arrs] @handle_exceptions -def gpu_is_available() -> bool: +def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: """ - Determine whether a GPU is available to use, with the backend framework. + Set the default device to given device instance. - Returns - ------- - ret - Boolean, as to whether a gpu is available. + Parameters + ---------- + device + The device to set as the default device Examples -------- - >>> print(ivy.gpu_is_available()) - False + >>> ivy.set_default_device("cpu") + >>> ivy.default_device() + 'cpu' + + >>> ivy.set_backend("torch") + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device(as_native=True) + device(type='cuda', index=0) + + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_default_device(device) + >>> ivy.default_device(as_native=True) + device(type='cuda') """ - return ivy.current_backend().gpu_is_available() + global default_device_stack + default_device_stack.append(device) @handle_exceptions -def num_cpu_cores(*, logical: bool = True) -> int: +def set_soft_device_mode(mode: bool) -> None: """ - Determine the number of cores available in the cpu. - - Parameters - ---------- - logical - Whether request is for number of physical or logical cores available in CPU - - Returns - ------- - ret - Number of cores available in CPU - - Examples - -------- - >>> print(ivy.num_cpu_cores(logical=False)) - 2 - """ - if logical: - return psutil.cpu_count(logical=logical) - else: - return psutil.cpu_count(logical=False) - - -@handle_exceptions -def num_gpus() -> int: - """ - Determine the number of available GPUs, with the backend framework. - - Returns - ------- - ret - Number of available GPUs. - - Examples - -------- - >>> print(ivy.num_gpus()) - 1 - """ - return ivy.current_backend().num_gpus() - - -@handle_exceptions -def tpu_is_available() -> bool: - """ - Determine whether a TPU is available to use, with the backend framework. - - Returns - ------- - ret - Boolean, as to whether a tpu is available. + Set the mode of whether to move input arrays to `ivy.default_device()` before + performing an operation. + Parameter + --------- + mode + boolean whether to move input arrays Examples -------- - >>> print(ivy.tpu_is_available()) + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode False + >>> ivy.set_soft_device_mode(True) + >>> ivy.soft_device_mode + True """ - return ivy.current_backend().tpu_is_available() - - -# Default Device # - - -# noinspection PyShadowingNames -@handle_exceptions -def default_device( - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - /, - *, - item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, - as_native: bool = None, -) -> Union[ivy.Device, ivy.NativeDevice]: - """ - Return the input device or the default device. If the as_native flag is set, the - device will be converted to a native device. If the item is provided, the item's - device is returned. If the device is not provided, the last default device is - returned. If a default device has not been set, the first gpu is returned if - available, otherwise the cpu is returned. - - Parameters - ---------- - device - The device to be returned or converted. - item - The item to get the device from. - as_native - Whether to convert the device to a native device. - - Returns - ------- - ret - Device handle or string. - - Examples - -------- - >>> ivy.default_device() - device(type='cpu') - - >>> ivy.default_device("gpu:0") - 'gpu:0' - - >>> ivy.default_device(item=[], as_native=False) - 'cpu' - - >>> ivy.default_device(item=(), as_native=True) - device(type='cpu') - - >>> ivy.default_device(item={"a": 1}, as_native=True) - device(type='cpu') - - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'gpu:0') - >>> ivy.default_device(item=x, as_native=True) - device(type='gpu', id=0) - """ - if ivy.exists(device): - if as_native is True: - return ivy.as_native_dev(device) - elif as_native is False: - return ivy.as_ivy_dev(device) - return device - as_native = ivy.default(as_native, False) - if ivy.exists(item): - if isinstance(item, (list, tuple, dict)) and len(item) == 0: - pass - elif ivy.is_array(item): - return ivy.dev(item, as_native=as_native) - global default_device_stack - if not default_device_stack: - ret = "gpu:0" if ivy.gpu_is_available() else "cpu" - else: - ret = default_device_stack[-1] - if as_native: - return ivy.as_native_dev(ret) - return ivy.as_ivy_dev(ret) + global soft_device_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + soft_device_mode_stack.append(mode) + ivy.__setattr__("soft_device_mode", mode, True) @handle_exceptions -def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +def set_split_factor( + factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None +) -> None: """ - Set the default device to given device instance. + Set the global split factor for a given device, which can be used to scale batch + splitting chunk sizes for the device across the codebase. Parameters ---------- + factor + The factor to set the device-specific split factor to. device - The device to set as the default device + The device to set the split factor for. Sets the default device by default. Examples -------- - >>> ivy.set_default_device("cpu") - >>> ivy.default_device() - 'cpu' + >>> print(ivy.default_device()) + cpu - >>> ivy.set_backend("torch") - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device(as_native=True) - device(type='cuda', index=0) + >>> ivy.set_split_factor(0.5) + >>> print(ivy.split_factors) + {'cpu': 0.5} >>> import torch >>> ivy.set_backend("torch") >>> device = torch.device("cuda") - >>> ivy.set_default_device(device) - >>> ivy.default_device(as_native=True) - device(type='cuda') - """ - global default_device_stack - default_device_stack.append(device) - - -@handle_exceptions -def unset_default_device() -> None: - """ - Reset the default device to "cpu". - - Examples - -------- - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device() - "gpu:0" - >>> ivy.unset_default_device() - >>> ivy.default_device() - "cpu" - """ - global default_device_stack - if default_device_stack: - default_device_stack.pop(-1) - - -# Device Allocation # - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def to_device( - x: Union[ivy.Array, ivy.NativeArray], - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - stream: Optional[Union[int, Any]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Move the input array x to the desired device, specified by device string. - - Parameters - ---------- - x - input array to be moved to the desired device - device - device to move the input array `x` to - stream - stream object to use during copy. In addition to the types supported in - array.__dlpack__(), implementations may choose to support any library-specific - stream object with the caveat that any code using such an object would not be - portable. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> ivy.set_split_factor(0.3, device=device) + >>> print(ivy.split_factors) + {device(type='cuda'): 0.3} - Returns - ------- - ret - input array x placed on the desired device + >>> ivy.set_split_factor(0.4, device="tpu") + >>> print(ivy.split_factors) + {'tpu': 0.4} - Examples - -------- - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'cpu') - >>> print(x.device) - cpu + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_split_factor(0.2) + >>> ivy.set_split_factor(0.3, device='gpu') + >>> print(ivy.split_factors) + {'cpu': 0.2, 'gpu': 0.3} """ - return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) + ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) + global split_factors + device = ivy.default(device, default_device()) + split_factors[device] = factor # Function Splitting # @@ -983,55 +1028,6 @@ def split_factor( return split_factors.setdefault(device, 0.0) -@handle_exceptions -def set_split_factor( - factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None -) -> None: - """ - Set the global split factor for a given device, which can be used to scale batch - splitting chunk sizes for the device across the codebase. - - Parameters - ---------- - factor - The factor to set the device-specific split factor to. - device - The device to set the split factor for. Sets the default device by default. - - Examples - -------- - >>> print(ivy.default_device()) - cpu - - >>> ivy.set_split_factor(0.5) - >>> print(ivy.split_factors) - {'cpu': 0.5} - - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.3, device=device) - >>> print(ivy.split_factors) - {device(type='cuda'): 0.3} - - >>> ivy.set_split_factor(0.4, device="tpu") - >>> print(ivy.split_factors) - {'tpu': 0.4} - - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.2) - >>> ivy.set_split_factor(0.3, device='gpu') - >>> print(ivy.split_factors) - {'cpu': 0.2, 'gpu': 0.3} - """ - ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) - global split_factors - device = ivy.default(device, default_device()) - split_factors[device] = factor - - @handle_exceptions def split_func_call( func: Callable, @@ -1166,200 +1162,213 @@ def split_func_call( return ret[0] if len(ret) == 1 else ret -def _is_valid_devices_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): - fn_supported_devices = fn.supported_devices - fn_unsupported_devices = fn.unsupported_devices - if isinstance(fn_supported_devices, dict): - if isinstance(fn_unsupported_devices, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_devices - and backend_str in fn_unsupported_devices - ): - return False - else: - if isinstance(fn_unsupported_devices, tuple): - return False - return True - - -def _get_devices(fn: Callable, complement: bool = True) -> Tuple: - valid_devices = ivy.valid_devices - invalid_devices = ivy.invalid_devices - all_devices = ivy.all_devices +# Device Allocation # - supported = set(ivy.valid_devices) - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - supported = set(all_devices).difference(supported) - return supported +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def to_device( + x: Union[ivy.Array, ivy.NativeArray], + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + stream: Optional[Union[int, Any]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Move the input array x to the desired device, specified by device string. - # Their values are formated like either - # 1. fn.supported_devices = ("cpu",) - # Could also have the "all" value for the framework - basic = [ - ("supported_devices", set.intersection, valid_devices), - ("unsupported_devices", set.difference, invalid_devices), - ] - for key, merge_fn, base in basic: - if hasattr(fn, key): - v = getattr(fn, key) - if "einops" in fn.__name__ and isinstance(v, dict): - v = v.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(v, tuple) - supported = merge_fn(supported, set(v)) + Parameters + ---------- + x + input array to be moved to the desired device + device + device to move the input array `x` to + stream + stream object to use during copy. In addition to the types supported in + array.__dlpack__(), implementations may choose to support any library-specific + stream object with the caveat that any code using such an object would not be + portable. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - if complement: - supported = set(all_devices).difference(supported) + Returns + ------- + ret + input array x placed on the desired device - return tuple(supported) + Examples + -------- + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'cpu') + >>> print(x.device) + cpu + """ + return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) @handle_exceptions -@handle_nestable -def function_supported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: +def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Return the supported devices of the current backend's function. The function returns - a dict containing the supported devices for the compositional and primary - implementations in case of partial mixed functions. + Get the total amount of memory (in GB) for a given device string. In case of CPU, + the total RAM is returned. Parameters ---------- - fn - The function to check for the supported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + device + The device string to convert to native device handle. Returns ------- ret - Tuple or dict containing the supported devices of the function + The total memory on the device in GB. Examples -------- - >>> import ivy - >>> print(ivy.function_supported_devices(ivy.ones)) - ('cpu', 'gpu') + >>> x = ivy.total_mem_on_dev("cpu") + >>> print(x) + 53.66700032 + + >>> x = ivy.total_mem_on_dev("gpu:0") + >>> print(x) + 8.589934592 """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=False), - } + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.total / 1e9 + elif device == "cpu": + return psutil.virtual_memory().total / 1e9 else: - supported_devices = set(_get_devices(fn, complement=False)) - if recurse: - supported_devices = ivy.functional.data_type._nested_get( - fn, supported_devices, set.intersection, function_supported_devices - ) - - return ( - supported_devices - if isinstance(supported_devices, dict) - else tuple(supported_devices) - ) + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions -@handle_nestable -def function_unsupported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: +def tpu_is_available() -> bool: """ - Return the unsupported devices of the current backend's function. The function - returns a dict containing the unsupported devices for the compositional and primary - implementations in case of partial mixed functions. - - Parameters - ---------- - fn - The function to check for the unsupported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + Determine whether a TPU is available to use, with the backend framework. Returns ------- ret - Tuple or dict containing the unsupported devices of the function + Boolean, as to whether a tpu is available. Examples -------- - >>> print(ivy.function_unsupported_devices(ivy.ones)) - ('tpu',) + >>> print(ivy.tpu_is_available()) + False """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=True), - } - else: - unsupported_devices = set(_get_devices(fn, complement=True)) - if recurse: - unsupported_devices = ivy.functional.data_type._nested_get( - fn, unsupported_devices, set.union, function_unsupported_devices - ) - return ( - unsupported_devices - if isinstance(unsupported_devices, dict) - else tuple(unsupported_devices) - ) + return ivy.current_backend().tpu_is_available() -# Profiler # +@handle_exceptions +def unset_default_device() -> None: + """ + Reset the default device to "cpu". + Examples + -------- + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device() + "gpu:0" + >>> ivy.unset_default_device() + >>> ivy.default_device() + "cpu" + """ + global default_device_stack + if default_device_stack: + default_device_stack.pop(-1) -class Profiler(abc.ABC): + +@handle_exceptions +def unset_soft_device_mode() -> None: """ - The profiler class is used to profile the execution of some code. + Reset the mode of moving input arrays to `ivy.default_device()` before performing an + operation. - Parameters - ---------- - save_dir - The directory to save the profile data to. + Examples + -------- + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode + False + >>> ivy.unset_soft_device_mode() + >>> ivy.soft_device_mode + True """ + global soft_device_mode_stack + if soft_device_mode_stack: + soft_device_mode_stack.pop(-1) + mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False + ivy.__setattr__("soft_device_mode", mode, True) - def __init__(self, save_dir: str): - self._save_dir = save_dir - @abc.abstractmethod - def start(self): - """ - Start the profiler. +@handle_exceptions +def used_mem_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + process_specific: bool = False, +) -> float: + """ + Get the used memory (in GB) for a given device string. In case of CPU, the used RAM + is returned. - This should be called before the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException + Parameters + ---------- + device + The device string to convert to native device handle. + process_specific + Whether to check the memory used by this python process alone. Default is + False. - @abc.abstractmethod - def stop(self): - """ - Stop the profiler. + Returns + ------- + ret + The used memory on the device in GB. - This should be called after the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException + Examples + -------- + >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) + >>> print(x) + 6.219563008 - @abc.abstractmethod - def __enter__(self): - raise ivy.utils.exceptions.IvyNotImplementedException + >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) + >>> print(x) + 0.902400346 - @abc.abstractmethod - def __exit__(self, exc_type, exc_val, exc_tb): - raise ivy.utils.exceptions.IvyNotImplementedException + >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) + >>> print(y) + 0.525205504 + """ + ivy.clear_cached_mem_on_dev(device) + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + if process_specific: + pid = os.getpid() + for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): + if process.pid == pid: + return process.usedGpuMemory / 1e9 + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.used / 1e9 + elif device == "cpu": + if process_specific: + return psutil.Process(os.getpid()).memory_info().rss / 1e9 + vm = psutil.virtual_memory() + return (vm.total - vm.available) / 1e9 + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) + + +ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False +dev_handles = dict() diff --git a/ivy/functional/ivy/elementwise.py b/ivy/functional/ivy/elementwise.py index 5fd532ccfb9c9..e5280d85a6f3e 100644 --- a/ivy/functional/ivy/elementwise.py +++ b/ivy/functional/ivy/elementwise.py @@ -18,6 +18,20 @@ from ivy.utils.exceptions import handle_exceptions +pow.unsupported_gradients = {"torch": ["float16"]} +trunc_divide.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + "handle_backend_invalid", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + + # Array API Standard # # -------------------# @@ -498,6 +512,51 @@ def add( return ivy.current_backend(x1, x2).add(x1, x2, alpha=alpha, out=out) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def angle( + z: Union[ivy.Array, ivy.NativeArray], + /, + *, + deg: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Calculate Element-wise the angle for an array of complex numbers(x+yj). + + Parameters + ---------- + z + Array-like input. + deg + optional bool. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Returns an array of angles for each complex number in the input. + If deg is False(default), angle is calculated in radian and if + deg is True, then angle is calculated in degrees. + + Examples + -------- + >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> z + ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) + >>> ivy.angle(z) + ivy.array([ 2.35619449, 2.35619449, -0.78539816]) + >>> ivy.angle(z,deg=True) + ivy.array([135., 135., -45.]) + """ + return ivy.current_backend(z).angle(z, deg=deg, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -1895,6 +1954,85 @@ def cosh( return ivy.current_backend(x).cosh(x, out=out) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def deg2rad( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Convert the input from degrees to radians. + + Parameters + ---------- + x + input array whose elements are each expressed in degrees. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array with each element in ``x`` converted from degrees to radians. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x=ivy.array([0,90,180,270,360]) + >>> y=ivy.deg2rad(x) + >>> print(y) + ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + + >>> x=ivy.array([0,-1.5,-50,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([ 0., -0.02617994, -0.87266463, nan]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.deg2rad(x, out=x) + >>> print(x) + ivy.array([[ 0.01919862, 0.03839725, 0.05759586], + [-0.07679449, -0.09599311, -0.11519173]]) + + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([0., 0.35081118, nan]) + + With :class:`ivy.Container` input: + + >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), + ... b=ivy.array([0,90,180,270,360])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 0.35081118, -0.88139129, nan]), + b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + } + + >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), + ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), + b: ivy.array([0., -0.02617994, -0.87266463, nan]) + } + """ + return ivy.current_backend(x).deg2rad(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2204,6 +2342,49 @@ def equal( return ivy.current_backend(x1, x2).equal(x1, x2, out=out) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def erf( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Gauss error function of ``x`` element-wise. + + Parameters + ---------- + x + Value to compute exponential for. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The Gauss error function of x. + + Examples + -------- + >>> x = ivy.array([0, 0.3, 0.7, 1.0]) + >>> ivy.erf(x) + ivy.array([0., 0.328, 0.677, 0.842]) + """ + return ivy.current_backend(x).erf(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2351,208 +2532,68 @@ def exp( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def imag( - val: Union[ivy.Array, ivy.NativeArray], +def exp2( + x: Union[ivy.Array, float, list, tuple], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the imaginary part of a complex number for each element ``x_i`` of the input - array ``val``. + Calculate 2**p for all p in the input array. Parameters ---------- - val - input array. Should have a complex floating-point data type. + x + Array-like input. out optional output array, for writing the result to. Returns ------- ret - Returns an array with the imaginary part of complex numbers. - The returned arrau must have a floating-point data type determined by - the precision of ``val`` (e.g., if ``val`` is ``complex64``, - the returned array must be ``float32``). - - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Element-wise 2 to the power x. This is a scalar if x is a scalar. Examples -------- - >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) - >>> b - ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) - >>> ivy.imag(b) - ivy.array([2., 4., 6.]) + >>> x = ivy.array([1, 2, 3]) + >>> ivy.exp2(x) + ivy.array([2., 4., 8.]) + >>> x = [5, 6, 7] + >>> ivy.exp2(x) + ivy.array([32., 64., 128.]) """ - return ivy.current_backend(val).imag(val, out=out) + return ivy.current_backend(x).exp2(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def angle( - z: Union[ivy.Array, ivy.NativeArray], +def expm1( + x: Union[ivy.Array, ivy.NativeArray], /, *, - deg: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate Element-wise the angle for an array of complex numbers(x+yj). + Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain + ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element + ``x_i`` of the input array ``x``. - Parameters - ---------- - z - Array-like input. - deg - optional bool. - out - optional output array, for writing the result to. + .. note:: + The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when + ``x`` is close to zero. Accordingly, conforming implementations should avoid + implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other + IEEE 754-2019 compliant mathematical library, for a potential reference + implementation. - Returns - ------- - ret - Returns an array of angles for each complex number in the input. - If deg is False(default), angle is calculated in radian and if - deg is True, then angle is calculated in degrees. - - Examples - -------- - >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) - >>> z - ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) - >>> ivy.angle(z) - ivy.array([ 2.35619449, 2.35619449, -0.78539816]) - >>> ivy.angle(z,deg=True) - ivy.array([135., 135., -45.]) - """ - return ivy.current_backend(z).angle(z, deg=deg, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def gcd( - x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the greatest common divisor of |x1| and |x2|. - - Parameters - ---------- - x1 - First array-like input. - x2 - Second array-input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Element-wise gcd of |x1| and |x2|. - - Examples - -------- - >>> x1 = ivy.array([1, 2, 3]) - >>> x2 = ivy.array([4, 5, 6]) - >>> ivy.gcd(x1, x2) - ivy.array([1., 1., 3.]) - >>> x1 = ivy.array([1, 2, 3]) - >>> ivy.gcd(x1, 10) - ivy.array([1., 2., 1.]) - """ - return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def exp2( - x: Union[ivy.Array, float, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate 2**p for all p in the input array. - - Parameters - ---------- - x - Array-like input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Element-wise 2 to the power x. This is a scalar if x is a scalar. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.exp2(x) - ivy.array([2., 4., 8.]) - >>> x = [5, 6, 7] - >>> ivy.exp2(x) - ivy.array([32., 64., 128.]) - """ - return ivy.current_backend(x).exp2(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def expm1( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain - ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element - ``x_i`` of the input array ``x``. - - .. note:: - The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when - ``x`` is close to zero. Accordingly, conforming implementations should avoid - implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other - IEEE 754-2019 compliant mathematical library, for a potential reference - implementation. - - .. note:: - For complex floating-point operands, ``expm1(conj(x))`` - must equal ``conj(expm1(x))``. + .. note:: + For complex floating-point operands, ``expm1(conj(x))`` + must equal ``conj(expm1(x))``. .. note:: The exponential function is an entire function @@ -2975,6 +3016,92 @@ def fmin( return ivy.current_backend(x1, x2).fmin(x1, x2, out=out) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fmod( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Compute the element-wise remainder of divisions of two arrays. + + Parameters + ---------- + x1 + First input array. + x2 + Second input array + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array with element-wise remainder of divisions. + + Examples + -------- + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmod(x1, x2) + ivy.array([ 0, 3, 0]) + + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmod(x1, x2) + ivy.array([ nan, nan, nan]) + """ + return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def gcd( + x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the greatest common divisor of |x1| and |x2|. + + Parameters + ---------- + x1 + First array-like input. + x2 + Second array-input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Element-wise gcd of |x1| and |x2|. + + Examples + -------- + >>> x1 = ivy.array([1, 2, 3]) + >>> x2 = ivy.array([4, 5, 6]) + >>> ivy.gcd(x1, x2) + ivy.array([1., 1., 3.]) + >>> x1 = ivy.array([1, 2, 3]) + >>> ivy.gcd(x1, 10) + ivy.array([1., 2., 1.]) + """ + return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -3164,277 +3291,57 @@ def greater_equal( return ivy.current_backend(x1, x2).greater_equal(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def less_equal( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def imag( + val: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 - with the respective element x2_i of the input array x2. + Return the imaginary part of a complex number for each element ``x_i`` of the input + array ``val``. Parameters ---------- - x1 - first input array. May have any data type. - x2 - second input array. Must be compatible with x1 (with Broadcasting). May have any - data type. + val + input array. Should have a complex floating-point data type. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- - ret - an array containing the element-wise results. The returned array must have a - data type of bool. - + ret + Returns an array with the imaginary part of complex numbers. + The returned arrau must have a floating-point data type determined by + the precision of ``val`` (e.g., if ``val`` is ``complex64``, + the returned array must be ``float32``). - This function conforms to the `Array API Standard - `_. This docstring is an extension of the + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.imag.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) - >>> print(x) - ivy.array([True, True, False]) - - >>> x = ivy.array([[10.1, 2.3, -3.6]]) - >>> y = ivy.array([[4.8], [5.2], [6.1]]) - >>> shape = (3,3) - >>> fill_value = False - >>> z = ivy.full(shape, fill_value) - >>> ivy.less_equal(x, y, out=z) - >>> print(z) - ivy.array([[False, True, True], - [False, True, True], - [False, True, True]]) - - >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) - >>> y = ivy.array([[8.4], [2.5], [1.6]]) - >>> ivy.less_equal(x, y, out=x) - >>> print(x) - ivy.array([[[1.], - [0.], - [1.]]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) - >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) - >>> z = ivy.less_equal(x, y) - >>> print(z) - { - a: ivy.array([False, False, False]), - b: ivy.array([True, True, True]) - } - """ - return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def multiply( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Calculate the product for each element x1_i of the input array x1 with the - respective element x2_i of the input array x2. - - .. note:: - Floating-point multiplication is not always associative due to finite precision. - - **Special Cases** - - For real-valued floating-point operands, - - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` and ``x2_i`` have the same mathematical sign, - the result has a positive mathematical sign, unless the result is ``NaN``. - If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` and ``x2_i`` have different mathematical signs, - the result has a negative mathematical sign, - unless the result is ``NaN``. If the result is ``NaN``, - the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, - the result is a signed infinity with the mathematical sign determined by - the rule already stated above. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` - is a nonzero finite number, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - If ``x1_i`` is a nonzero finite number and ``x2_i`` - is either ``+infinity`` or ``-infinity``, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - In the remaining cases, where neither ``infinity`` nor ``NaN`` - is involved, the product must be computed and rounded to the nearest - representable value according to IEEE 754-2019 and a supported - rounding mode. If the magnitude is too large to represent, - the result is an `infinity` of appropriate mathematical sign. - If the magnitude is too small to represent, the result is a zero of - appropriate mathematical sign. - - For complex floating-point operands, multiplication is defined according to the - following table. For real components ``a`` and ``c`` and - imaginary components ``b`` and ``d``, - - +------------+----------------+-----------------+--------------------------+ - | | c | dj | c + dj | - +============+================+=================+==========================+ - | **a** | a * c | (a*d)j | (a*c) + (a*d)j | - +------------+----------------+-----------------+--------------------------+ - | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | - +------------+----------------+-----------------+--------------------------+ - | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | - +------------+----------------+-----------------+--------------------------+ - - In general, for complex floating-point operands, real-valued floating-point - special cases must independently apply to the real and imaginary component - operations involving real numbers as described in the above table. - - When ``a``, ``b``, ``c``, or ``d`` are all finite numbers - (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), - multiplication of complex floating-point operands should be computed - as if calculated according to the textbook formula for complex number multiplication - - .. math:: - (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j - - When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, - ``+infinity``, or ``-infinity``, - - - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, the result is implementation dependent. - - .. note:: - For complex floating-point operands, the results of special cases may be - implementation dependent depending on how an implementation chooses - to model complex numbers and complex infinity - (e.g., complex plane versus Riemann sphere). - For those implementations following C99 and its one-infinity model, - when at least one component is infinite, - even if the other component is ``NaN``, - the complex value is infinite, and the usual arithmetic - rules do not apply to complex-complex multiplication. - In the interest of performance, other implementations - may want to avoid the complex branching logic necessary - to implement the one-infinity model and choose to implement - all complex-complex multiplication according to the textbook formula. - Accordingly, special case behavior is unlikely - to be consistent across implementations. - - Parameters - ---------- - x1 - first input array. Should have a numeric data type. - - x2 - second input array. Must be compatible with ``x1`` - (see :ref'`broadcasting`). Should have a numeric data type - - out - optional output array, for writing the array result to. - It must have a shape that the inputs broadcast to. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Returns - ------- - ret - an array containing the element-wise products. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. - - Examples - -------- - With :code:`ivy.Array` inputs: - - >>> x1 = ivy.array([3., 5., 7.]) - >>> x2 = ivy.array([4., 6., 8.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([12., 30., 56.]) - - With :code:`ivy.NativeArray` inputs: - - >>> x1 = ivy.native_array([1., 3., 9.]) - >>> x2 = ivy.native_array([4., 7.2, 1.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([ 4. , 21.6, 9. ]) - - With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: - - >>> x1 = ivy.array([8., 6., 7.]) - >>> x2 = ivy.native_array([1., 2., 3.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([ 8., 12., 21.]) - - With :code:`ivy.Container` inputs: - - >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) - >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - { - a: ivy.array([12.,12.,24.]), - b: ivy.array([9.,3.,10.]) - } - - With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: - - >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) - >>> x2 = ivy.array([1.,2.,3.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - { - a: ivy.array([3.,8.,15.]), - b: ivy.array([2.,4.,3.]) - } + >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) + >>> b + ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) + >>> ivy.imag(b) + ivy.array([2., 4., 6.]) """ - return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) + return ivy.current_backend(val).imag(val, out=out) @handle_exceptions @@ -3774,28 +3681,27 @@ def isnan( @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def less( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def isreal( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + Test each element ``x_i`` of the input array ``x`` to determine whether the element + is real number. Returns a bool array, where True if input element is real. If + element has complex type with zero complex part, the return value for that element + is True. Parameters ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. + x + input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -3803,16 +3709,126 @@ def less( Returns ------- ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. + an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is + real number and ``False`` otherwise. The returned array should have a data type + of ``bool``. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances in place of + :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints + and also the examples below. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.less(ivy.array([1,2,3]),ivy.array([2,2,2])) - >>> print(x) - ivy.array([ True, False, False]) + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([[[True], [True], [True]]]) + + >>> x = ivy.array([1-0j, 3j, 7+5j]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([ True, False, False]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ + b=ivy.array([5j, 5-6j, 3])) + >>> z = ivy.isreal(x) + >>> print(z) + { + a: ivy.array([False, True, True]), + b: ivy.array([False, False, True]) + } + """ + return ivy.current_backend(x).isreal(x, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lcm( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the element-wise least common multiple (LCM) of x1 and x2. + + Parameters + ---------- + x1 + first input array, must be integers + x2 + second input array, must be integers + out + optional output array, for writing the result to. + + Returns + ------- + ret + an array that includes the element-wise least common multiples of x1 and x2 + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x1=ivy.array([2, 3, 4]) + >>> x2=ivy.array([5, 7, 15]) + >>> x1.lcm(x1, x2) + ivy.array([10, 21, 60]) + """ + return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def less( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + + Parameters + ---------- + x1 + first input array. Should have a numeric data type. + x2 + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of ``bool``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.less(ivy.array([1,2,3]),ivy.array([2,2,2])) + >>> print(x) + ivy.array([ True, False, False]) >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) @@ -3861,6 +3877,93 @@ def less( return ivy.current_backend(x1).less(x1, x2, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def less_equal( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 + with the respective element x2_i of the input array x2. + + Parameters + ---------- + x1 + first input array. May have any data type. + x2 + second input array. Must be compatible with x1 (with Broadcasting). May have any + data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of bool. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) + >>> print(x) + ivy.array([True, True, False]) + + >>> x = ivy.array([[10.1, 2.3, -3.6]]) + >>> y = ivy.array([[4.8], [5.2], [6.1]]) + >>> shape = (3,3) + >>> fill_value = False + >>> z = ivy.full(shape, fill_value) + >>> ivy.less_equal(x, y, out=z) + >>> print(z) + ivy.array([[False, True, True], + [False, True, True], + [False, True, True]]) + + >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) + >>> y = ivy.array([[8.4], [2.5], [1.6]]) + >>> ivy.less_equal(x, y, out=x) + >>> print(x) + ivy.array([[[1.], + [0.], + [1.]]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) + >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) + >>> z = ivy.less_equal(x, y) + >>> print(z) + { + a: ivy.array([False, False, False]), + b: ivy.array([True, True, True]) + } + """ + return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -4779,98 +4882,125 @@ def logical_xor( return ivy.current_backend(x1, x2).logical_xor(x1, x2, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def nan_to_num( - x: Union[ivy.Array, ivy.NativeArray], +def maximum( + x1: Union[ivy.Array, ivy.NativeArray, Number], + x2: Union[ivy.Array, ivy.NativeArray, Number], /, *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Replace NaN with zero and infinity with large finite numbers (default behaviour) or - with the numbers defined by the user using the nan, posinf and/or neginf keywords. + Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. Parameters ---------- - x - Array input. - copy - Whether to create a copy of x (True) or to replace values in-place (False). - The in-place operation only occurs if casting to an array does not require - a copy. Default is True. - nan - Value to be used to fill NaN values. If no value is passed then NaN values - will be replaced with 0.0. - posinf - Value to be used to fill positive infinity values. If no value is passed - then positive infinity values will be replaced with a very large number. - neginf - Value to be used to fill negative infinity values. - If no value is passed then negative infinity values - will be replaced with a very small (or negative) number. + x1 + Input array containing elements to maximum threshold. + x2 + Tensor containing maximum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum + is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Array with the non-finite values replaced. - If copy is False, this may be x itself. + An array with the elements of x1, but clipped to not be lower than the x2 + values. Examples -------- - >>> x = ivy.array([1, 2, 3, nan]) - >>> ivy.nan_to_num(x) - ivy.array([1., 1., 3., 0.0]) - >>> x = ivy.array([1, 2, 3, inf]) - >>> ivy.nan_to_num(x, posinf=5e+100) - ivy.array([1., 2., 3., 5e+100]) - """ - return ivy.current_backend(x).nan_to_num( - x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out - ) - + With :class:`ivy.Array` inputs: -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.maximum(x, y) + >>> print(z) + ivy.array([9, 9, 5]) + + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.maximum(x, y, out=z) + >>> print(z) + ivy.array([[9., 9., 9., 9., 9., 9.], + [3., 5., 9., 8., 3., 7.], + [2., 5., 9., 8., 3., 7.]]) + + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.maximum(x, y, out=x) + >>> print(x) + ivy.array([[7, 7]]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]), + ... b=ivy.array([-5, 9])) + >>> z = ivy.maximum(x, y) + >>> print(z) + { + a: ivy.array([[1, 3], + [2, 4], + [3, 7]]), + b: ivy.array([[1, 9], + [2, 9], + [3, 9]]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) + >>> z = ivy.maximum(x, y) + >>> print(z) + { + a: ivy.array([1, 5, 6]), + b: ivy.array([5, 9, 7]) + } + """ + return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def negative( - x: Union[float, ivy.Array, ivy.NativeArray], +def minimum( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the negative value of each element in ``x``. - - .. note:: - For signed integer data types, the numerical negative of - the minimum representable integer is implementation-dependent. - - .. note:: - If ``x`` has a complex floating-point data type, - both the real and imaginary components for each ``x_i`` - must be negated (a result which follows from the rules of - complex number multiplication). + Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. Parameters ---------- - x - Input array. + x1 + Input array containing elements to minimum threshold. + x2 + Tensor containing minimum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum + is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -4878,53 +5008,62 @@ def negative( Returns ------- ret - A new array with the negative value of each element in ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + An array with the elements of x1, but clipped to not exceed the x2 values. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.array([0,1,1,2]) - >>> y = ivy.negative(x) - >>> print(y) - ivy.array([ 0, -1, -1, -2]) + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.minimum(x, y) + >>> print(z) + ivy.array([7, 3, 2]) - >>> x = ivy.array([0,-1,-0.5,2,3]) - >>> y = ivy.zeros(5) - >>> ivy.negative(x, out=y) - >>> print(y) - ivy.array([-0. , 1. , 0.5, -2. , -3. ]) + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.minimum(x, y, out=z) + >>> print(z) + ivy.array([[1.,5.,9.,8.,3.,7.], + [1.,3.,3.,3.,3.,3.], + [1.,2.,2.,2.,2.,2.]]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.negative(x,out=x) + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.minimum(x, y, out=x) >>> print(x) - ivy.array([[-1.1, -2.2, -3.3], - [4.4, 5.5, 6.6]]) + ivy.array([[0, 3]]) - With :class:`ivy.Container` input: + With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.negative(x) - >>> print(y) + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) + >>> z = ivy.minimum(x, y) + >>> print(z) { - a: ivy.array([-0., -1., -2.]), - b: ivy.array([-3., -4., 5.]) + a: ivy.array([[1, 0], + [1, 0], + [1, 0]]), + b: ivy.array([[-5, 3], + [-5, 4], + [-5, 7]]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([1, 3, 1]), + ... b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]), + ... b=ivy.array([5, 9, 7])) + >>> z = ivy.minimum(x, y) + >>> print(z) + { + a: ivy.array([1, 3, 1]), + b: ivy.array([2, 8, 5]) } """ - return ivy.current_backend(x).negative(x, out=out) + return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) @handle_exceptions @@ -4934,406 +5073,244 @@ def negative( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def not_equal( - x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], - x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], +def multiply( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + r""" + Calculate the product for each element x1_i of the input array x1 with the + respective element x2_i of the input array x2. + + .. note:: + Floating-point multiplication is not always associative due to finite precision. **Special Cases** For real-valued floating-point operands, - - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, - the result is ``True``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, - the result is ``True``. - - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, - and ``x1_i`` does not equal ``x2_i``, the result is ``True``. - - In the remaining cases, the result is ``False``. - - For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, - ``c = real(x2_i)``, ``d = imag(x2_i)``, and + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` and ``x2_i`` have the same mathematical sign, + the result has a positive mathematical sign, unless the result is ``NaN``. + If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` and ``x2_i`` have different mathematical signs, + the result has a negative mathematical sign, + unless the result is ``NaN``. If the result is ``NaN``, + the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, + the result is a signed infinity with the mathematical sign determined by + the rule already stated above. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` + is a nonzero finite number, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - If ``x1_i`` is a nonzero finite number and ``x2_i`` + is either ``+infinity`` or ``-infinity``, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - In the remaining cases, where neither ``infinity`` nor ``NaN`` + is involved, the product must be computed and rounded to the nearest + representable value according to IEEE 754-2019 and a supported + rounding mode. If the magnitude is too large to represent, + the result is an `infinity` of appropriate mathematical sign. + If the magnitude is too small to represent, the result is a zero of + appropriate mathematical sign. - - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. - - In the remaining cases, the result is the logical OR of - the equality comparison between the real values ``a`` and ``c`` - (real components) and between the real values ``b`` and ``d`` - (imaginary components), as described above for real-valued floating-point operands - (i.e., ``a != c OR b != d``). + For complex floating-point operands, multiplication is defined according to the + following table. For real components ``a`` and ``c`` and + imaginary components ``b`` and ``d``, - Parameters - ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + +------------+----------------+-----------------+--------------------------+ + | | c | dj | c + dj | + +============+================+=================+==========================+ + | **a** | a * c | (a*d)j | (a*c) + (a*d)j | + +------------+----------------+-----------------+--------------------------+ + | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | + +------------+----------------+-----------------+--------------------------+ + | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | + +------------+----------------+-----------------+--------------------------+ - Returns - ------- - ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. + In general, for complex floating-point operands, real-valued floating-point + special cases must independently apply to the real and imaginary component + operations involving real numbers as described in the above table. + + When ``a``, ``b``, ``c``, or ``d`` are all finite numbers + (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), + multiplication of complex floating-point operands should be computed + as if calculated according to the textbook formula for complex number multiplication + + .. math:: + (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j + + When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, + ``+infinity``, or ``-infinity``, + + - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, the result is implementation dependent. + + .. note:: + For complex floating-point operands, the results of special cases may be + implementation dependent depending on how an implementation chooses + to model complex numbers and complex infinity + (e.g., complex plane versus Riemann sphere). + For those implementations following C99 and its one-infinity model, + when at least one component is infinite, + even if the other component is ``NaN``, + the complex value is infinite, and the usual arithmetic + rules do not apply to complex-complex multiplication. + In the interest of performance, other implementations + may want to avoid the complex branching logic necessary + to implement the one-infinity model and choose to implement + all complex-complex multiplication according to the textbook formula. + Accordingly, special case behavior is unlikely + to be consistent across implementations. + + Parameters + ---------- + x1 + first input array. Should have a numeric data type. + + x2 + second input array. Must be compatible with ``x1`` + (see :ref'`broadcasting`). Should have a numeric data type + + out + optional output array, for writing the array result to. + It must have a shape that the inputs broadcast to. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.multiply.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x1 = ivy.array([1, 0, 1, 1]) - >>> x2 = ivy.array([1, 0, 0, -1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False, True, True]) - - >>> x1 = ivy.array([1, 0, 1, 0]) - >>> x2 = ivy.array([0, 1, 0, 1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([True, True, True, True]) + Returns + ------- + ret + an array containing the element-wise products. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) + Examples + -------- + With :code:`ivy.Array` inputs: - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.not_equal(x1, x2, out=x1) + >>> x1 = ivy.array([3., 5., 7.]) + >>> x2 = ivy.array([4., 6., 8.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - ivy.array([1, 0, 0, 1]) - - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + ivy.array([12., 30., 56.]) - >>> x1 = ivy.native_array([1, 2]) - >>> x2 = ivy.array([1, 2]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False]) + With :code:`ivy.NativeArray` inputs: - >>> x1 = ivy.native_array([1, -1]) - >>> x2 = ivy.array([0, 1]) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.native_array([1., 3., 9.]) + >>> x2 = ivy.native_array([4., 7.2, 1.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - ivy.array([True, True]) + ivy.array([ 4. , 21.6, 9. ]) - >>> x1 = ivy.native_array([1, -1, 1, -1]) - >>> x2 = ivy.native_array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) + With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: - >>> x1 = ivy.native_array([1, 2, 3, 4]) - >>> x2 = ivy.native_array([0, 2, 3, 4]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) + >>> x1 = ivy.array([8., 6., 7.]) + >>> x2 = ivy.native_array([1., 2., 3.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - ivy.array([1., 0., 0., 0.]) - - With :class:`ivy.Container` input: + ivy.array([ 8., 12., 21.]) - >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - { - a: ivy.array([False, True, False]), - b: ivy.array([False, False, False]), - c: ivy.array([False, False, False]) - } + With :code:`ivy.Container` inputs: - >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1.0, 2.0, 4.0])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.native_array([1.1, 2.1, 3.1]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) + >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) + >>> y = ivy.multiply(x1, x2) >>> print(y) { - a: ivy.array([True, True, True]), - b: ivy.array([True, True, True]), - c: ivy.array([False, False, False]) + a: ivy.array([12.,12.,24.]), + b: ivy.array([9.,3.,10.]) } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 3, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 4, 5])) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - { - a: ivy.array([False, False, False]), - b: ivy.array([False, True, False]) - } + With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: - >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), - ... b=ivy.array([1, 4, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), - ... b=ivy.array([1.0, 4.0, 5.0])) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) + >>> x2 = ivy.array([1.,2.,3.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) { - a: ivy.array([False, False, False]), - b: ivy.array([False, False, False]) + a: ivy.array([3.,8.,15.]), + b: ivy.array([2.,4.,3.]) } """ - return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) + return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def positive( - x: Union[float, ivy.Array, ivy.NativeArray], +def nan_to_num( + x: Union[ivy.Array, ivy.NativeArray], /, *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the positive value of each element in ``x``. + Replace NaN with zero and infinity with large finite numbers (default behaviour) or + with the numbers defined by the user using the nan, posinf and/or neginf keywords. Parameters ---------- x - Input array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - A new array with the positive value of each element in ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3 ,5, 7]) - >>> y = ivy.positive(x) - >>> print(y) - ivy.array([2, 3, 5, 7]) - - >>> x = ivy.array([0, -1, -0.5, 2, 3]) - >>> y = ivy.zeros(5) - >>> ivy.positive(x, out=y) - >>> print(y) - ivy.array([0., -1., -0.5, 2., 3.]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.positive(x,out=x) - >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-4.4, -5.5, -6.6]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.positive(x) - >>> print(y) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., -5.]) - } - """ - return ivy.current_backend(x).positive(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def pow( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate an implementation-dependent approximation of exponentiation by raising - each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` - (the exponent), where ``x2_i`` is the corresponding element of the input array - ``x2``. - - .. note:: - If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when - ``x2_i`` is negative (i.e., less than zero) is unspecified and thus - implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a - floating-point data type, behavior is implementation-dependent (type promotion - between data type "kinds" (integer versus floating-point) is unspecified). - - **Special cases** - - For floating-point operands, - - - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result - is ``+infinity``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result - is ``+0``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. - - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is - ``+0``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is - ``+0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an - odd integer value, the result is ``-infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not - an odd integer value, the result is ``+infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an - odd integer value, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer - value, the result is ``-infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+infinity``. - - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite - number, and ``x2_i`` is not an integer value, the result is ``NaN``. - - For complex floating-point operands, special cases should be handled - as if the operation is implemented as ``exp(x2*log(x1))``. - - .. note:: - Conforming implementations are allowed to treat special cases involving - complex floating-point operands more carefully than as described - in this specification. - - Parameters - ---------- - x1 - first input array whose elements correspond to the exponentiation base. Should - have a numeric data type. - x2 - second input array whose elements correspond to the exponentiation exponent. - Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric - data type. + Array input. + copy + Whether to create a copy of x (True) or to replace values in-place (False). + The in-place operation only occurs if casting to an array does not require + a copy. Default is True. + nan + Value to be used to fill NaN values. If no value is passed then NaN values + will be replaced with 0.0. + posinf + Value to be used to fill positive infinity values. If no value is passed + then positive infinity values will be replaced with a very large number. + neginf + Value to be used to fill negative infinity values. + If no value is passed then negative infinity values + will be replaced with a very small (or negative) number. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array containing the element-wise results. The returned array must have a - data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + Array with the non-finite values replaced. + If copy is False, this may be x itself. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.pow(x, 3) - >>> print(y) - ivy.array([1, 8, 27]) - - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.pow(x, 2, out=y) - >>> print(y) - ivy.array([2.25, 0.64, 0.09]) - - >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) - >>> ivy.pow(x, 2.3, out=x) - >>> print(x) - ivy.array([[ 1.52095687, 4.92457771, 13.49372482], - [ 1. , 8.22738838, 156.5877228 ]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.pow(x, 3) - >>> print(y) - { - a:ivy.array([0,1]), - b:ivy.array([8,27]) - } + >>> x = ivy.array([1, 2, 3, nan]) + >>> ivy.nan_to_num(x) + ivy.array([1., 1., 3., 0.0]) + >>> x = ivy.array([1, 2, 3, inf]) + >>> ivy.nan_to_num(x, posinf=5e+100) + ivy.array([1., 2., 3., 5e+100]) """ - return ivy.current_backend(x1, x2).pow(x1, x2, out=out) - - -pow.unsupported_gradients = {"torch": ["float16"]} + return ivy.current_backend(x).nan_to_num( + x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out + ) @handle_exceptions @@ -5342,23 +5319,31 @@ def pow( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def real( - x: Union[ivy.Array, ivy.NativeArray], +def negative( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Test each element ``x_i`` of the input array ``x`` to take only real part from it. - Returns a float array, where it only contains . If element has complex type with - zero complex part, the return value will be that element, else it only returns real - part. + Return a new array with the negative value of each element in ``x``. + + .. note:: + For signed integer data types, the numerical negative of + the minimum representable integer is implementation-dependent. + + .. note:: + If ``x`` has a complex floating-point data type, + both the real and imaginary components for each ``x_i`` + must be negated (a result which follows from the rules of + complex number multiplication). Parameters ---------- x - input array. + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5366,45 +5351,53 @@ def real( Returns ------- ret - an array containing test results. An element ``out_i`` is - ``real number`` if ``x_i`` contain real number part only - and if it is ``real number with complex part also`` then it - returns the real number part. - The returned array must have a floating-point data type with the - same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, - the returned array must have the floating-point precision of ``float32``). + A new array with the negative value of each element in ``x``. - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input: - >>> x = ivy.array([[[1.1], [2], [-6.3]]]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([[[1.1], [2.], [-6.3]]]) + >>> x = ivy.array([0,1,1,2]) + >>> y = ivy.negative(x) + >>> print(y) + ivy.array([ 0, -1, -1, -2]) - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([4.2, 0., 7.]) + >>> x = ivy.array([0,-1,-0.5,2,3]) + >>> y = ivy.zeros(5) + >>> ivy.negative(x, out=y) + >>> print(y) + ivy.array([-0. , 1. , 0.5, -2. , -3. ]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.negative(x,out=x) + >>> print(x) + ivy.array([[-1.1, -2.2, -3.3], + [4.4, 5.5, 6.6]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ - b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.real(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.negative(x) + >>> print(y) { - a: ivy.array([-6.7, 0.314, 1.23]), - b: ivy.array([0., 5.32, 3.001]) + a: ivy.array([-0., -1., -2.]), + b: ivy.array([-3., -4., 5.]) } """ - return ivy.current_backend(x).real(x, out=out) + return ivy.current_backend(x).negative(x, out=out) @handle_exceptions @@ -5414,127 +5407,181 @@ def real( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def remainder( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def not_equal( + x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], + x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], /, *, - modulus: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the remainder of division for each element ``x1_i`` of the input array ``x1`` - and the respective element ``x2_i`` of the input array ``x2``. - - .. note:: - This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For - input arrays which promote to an integer data type, the result of division by - zero is unspecified and thus implementation-defined. In general, similar to - Python’s ``%`` operator, this function is not recommended for floating-point - operands as semantics do not follow IEEE 754. That this function is specified - to accept floating-point operands is primarily for reasons of backward - compatibility. + Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. **Special Cases** - For floating-point operands, + For real-valued floating-point operands, - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either - ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, - the result is ``NaN``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x2_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x2_i``. (note: this results matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - In the remaining cases, the result must match that of the Python ``%`` operator. + - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, + the result is ``True``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, + the result is ``True``. + - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, + and ``x1_i`` does not equal ``x2_i``, the result is ``True``. + - In the remaining cases, the result is ``False``. + + For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, + ``c = real(x2_i)``, ``d = imag(x2_i)``, and + + - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. + - In the remaining cases, the result is the logical OR of + the equality comparison between the real values ``a`` and ``c`` + (real components) and between the real values ``b`` and ``d`` + (imaginary components), as described above for real-valued floating-point operands + (i.e., ``a != c OR b != d``). Parameters ---------- x1 - dividend input array. Should have a numeric data type. + first input array. Should have a numeric data type. x2 - divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). Should have a numeric data type. - modulus - whether to compute the modulus instead of the remainder. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. - Returns - ------- - ret - an array containing the element-wise results. Each element-wise result must have - the same sign as the respective element ``x2_i``. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of ``bool``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x1 = ivy.array([1, 0, 1, 1]) + >>> x2 = ivy.array([1, 0, 0, -1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False, True, True]) + + >>> x1 = ivy.array([1, 0, 1, 0]) + >>> x2 = ivy.array([0, 1, 0, 1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([True, True, True, True]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.not_equal(x1, x2, out=x1) + >>> print(y) + ivy.array([1, 0, 0, 1]) + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + >>> x1 = ivy.native_array([1, 2]) + >>> x2 = ivy.array([1, 2]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + >>> x1 = ivy.native_array([1, -1]) + >>> x2 = ivy.array([0, 1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([True, True]) - Examples - -------- - With :class:`ivy.Array` inputs: + >>> x1 = ivy.native_array([1, -1, 1, -1]) + >>> x2 = ivy.native_array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) - >>> x1 = ivy.array([2., 5., 15.]) - >>> x2 = ivy.array([3., 2., 4.]) - >>> y = ivy.remainder(x1, x2) + >>> x1 = ivy.native_array([1, 2, 3, 4]) + >>> x2 = ivy.native_array([0, 2, 3, 4]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) >>> print(y) - ivy.array([2., 1., 3.]) + ivy.array([1., 0., 0., 0.]) - With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + With :class:`ivy.Container` input: - >>> x1 = ivy.array([23., 1., 6.]) - >>> x2 = ivy.native_array([11., 2., 4.]) - >>> y = ivy.remainder(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) - ivy.array([1., 1., 2.]) + { + a: ivy.array([False, True, False]), + b: ivy.array([False, False, False]), + c: ivy.array([False, False, False]) + } - With :class:`ivy.Container` inputs: + >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1.0, 2.0, 4.0])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.native_array([1.1, 2.1, 3.1]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + { + a: ivy.array([True, True, True]), + b: ivy.array([True, True, True]), + c: ivy.array([False, False, False]) + } - >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) - >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) - >>> y = ivy.remainder(x1, x2) + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 3, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 4, 5])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) { - a: ivy.array([0., 0., 1.]), - b: ivy.array([0., 2., 1.]) + a: ivy.array([False, False, False]), + b: ivy.array([False, True, False]) + } + + >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), + ... b=ivy.array([1, 4, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), + ... b=ivy.array([1.0, 4.0, 5.0])) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + { + a: ivy.array([False, False, False]), + b: ivy.array([False, False, False]) } """ - return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) + return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) @handle_exceptions @@ -5545,55 +5592,19 @@ def remainder( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def round( - x: Union[ivy.Array, ivy.NativeArray], +def positive( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, - decimals: Optional[int] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued - number. - - .. note:: - For complex floating-point operands, real and imaginary components - must be independently rounded to the nearest integer-valued number. - - Rounded real and imaginary components must be equal - to their equivalent rounded real-valued floating-point - counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` - must equal ``round(real(x)))`` and ``imag(round(x))`` must equal - ``round(imag(x))``). - - **Special cases** - - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - - For floating-point operands, - - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If two integers are equally close to ``x_i``, the result is - the even integer closest to ``x_i``. - - .. note:: - For complex floating-point operands, the following special - cases apply to real and imaginary components independently - (e.g., if ``real(x_i)`` is ``NaN``, the rounded - real component is ``NaN``). - - - If ``x_i`` is already integer-valued, the result is ``x_i``. + Return a new array with the positive value of each element in ``x``. Parameters ---------- x - input array containing elements to round. - decimals - number of decimal places to round to. Default is ``0``. + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5601,121 +5612,145 @@ def round( Returns ------- ret - An array of the same shape and type as x, with the elements rounded to integers. - + A new array with the positive value of each element in ``x``. - Note: PyTorch supports an additional argument :code:`decimals` for the - `round function `_. - It has been deliberately omitted here due to the imprecise - nature of the argument in :code:`torch.round`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.positive.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments - Examples - -------- - With :class:`ivy.Array` input: + Functional Examples + ------------------- - >>> x = ivy.array([1.2, 2.4, 3.6]) - >>> y = ivy.round(x) - >>> print(y) - ivy.array([1.,2.,4.]) + With :class:`ivy.Array` input: - >>> x = ivy.array([-0, 5, 4.5]) - >>> y = ivy.round(x) + >>> x = ivy.array([2, 3 ,5, 7]) + >>> y = ivy.positive(x) >>> print(y) - ivy.array([0.,5.,4.]) + ivy.array([2, 3, 5, 7]) - >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) - >>> y = ivy.zeros(4) - >>> ivy.round(x, out=y) + >>> x = ivy.array([0, -1, -0.5, 2, 3]) + >>> y = ivy.zeros(5) + >>> ivy.positive(x, out=y) >>> print(y) - ivy.array([2.,2.,15.,-5.]) + ivy.array([0., -1., -0.5, 2., 3.]) - >>> x = ivy.array([[0, 5.433, -343.3, 1.5], - ... [-5.5, 44.2, 11.5, 12.01]]) - >>> ivy.round(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.positive(x,out=x) >>> print(x) - ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) + ivy.array([[ 1.1, 2.2, 3.3], + [-4.4, -5.5, -6.6]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), - ... b=ivy.array([-300.9, -527.3, 4.5])) - >>> y = ivy.round(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.positive(x) >>> print(y) { - a:ivy.array([4.,9.,7.,0.]), - b:ivy.array([-301.,-527.,4.]) + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., -5.]) } """ - return ivy.current_backend(x).round(x, decimals=decimals, out=out) + return ivy.current_backend(x).positive(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sign( - x: Union[ivy.Array, ivy.NativeArray], +def pow( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, - np_variant: Optional[bool] = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Return an indication of the sign of a number for each element ``x_i`` of the input - array ``x``. - - The sign function (also known as the **signum function**) - of a number :math:`x_{i}` is defined as - - .. math:: - \operatorname{sign}(x_i) = \begin{cases} - 0 & \textrm{if } x_i = 0 \\ - \frac{x}{|x|} & \textrm{otherwise} - \end{cases} + """ + Calculate an implementation-dependent approximation of exponentiation by raising + each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` + (the exponent), where ``x2_i`` is the corresponding element of the input array + ``x2``. - where :math:`|x_i|` is the absolute value of :math:`x_i`. + .. note:: + If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when + ``x2_i`` is negative (i.e., less than zero) is unspecified and thus + implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a + floating-point data type, behavior is implementation-dependent (type promotion + between data type "kinds" (integer versus floating-point) is unspecified). **Special cases** - - If ``x_i`` is less than ``0``, the result is ``-1``. - - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. - - If ``x_i`` is greater than ``0``, the result is ``+1``. - - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` + For floating-point operands, + + - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result + is ``+infinity``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result + is ``+0``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. + - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is + ``+0``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is + ``+0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an + odd integer value, the result is ``-infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not + an odd integer value, the result is ``+infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an + odd integer value, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer + value, the result is ``-infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+infinity``. + - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite + number, and ``x2_i`` is not an integer value, the result is ``NaN``. - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and + For complex floating-point operands, special cases should be handled + as if the operation is implemented as ``exp(x2*log(x1))``. - - If ``a`` is either ``-0`` or ``+0`` and ``b`` is - either ``-0`` or ``+0``, the result is ``0 + 0j``. - - If ``a`` is ``NaN`` or ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, special cases must be handled - according to the rules of complex number division. + .. note:: + Conforming implementations are allowed to treat special cases involving + complex floating-point operands more carefully than as described + in this specification. Parameters ---------- - x - input array. Should have a numeric data type. - np_variant - Handles complex numbers like numpy does If ``True``, - ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. - otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, - otherwise y = 0.`` - + x1 + first input array whose elements correspond to the exponentiation base. Should + have a numeric data type. + x2 + second input array whose elements correspond to the exponentiation exponent. + Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric + data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5723,53 +5758,54 @@ def sign( Returns ------- ret - an array containing the evaluated result for each element in ``x``. The returned - array must have the same data type as ``x``. + an array containing the element-wise results. The returned array must have a + data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.pow.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([8.3, -0, 6.8, 0.07]) - >>> y = ivy.sign(x) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.pow(x, 3) >>> print(y) - ivy.array([1., 0., 1., 1.]) + ivy.array([1, 8, 27]) - >>> x = ivy.array([[5.78, -4., -6.9, 0], - ... [-.4, 0.5, 8, -0.01]]) - >>> y = ivy.sign(x) + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.pow(x, 2, out=y) >>> print(y) - ivy.array([[ 1., -1., -1., 0.], - [-1., 1., 1., -1.]]) + ivy.array([2.25, 0.64, 0.09]) + + >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) + >>> ivy.pow(x, 2.3, out=x) + >>> print(x) + ivy.array([[ 1.52095687, 4.92457771, 13.49372482], + [ 1. , 8.22738838, 156.5877228 ]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., -0.]), - ... b=ivy.array([1.46, 5.9, -0.0]), - ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) - >>> y = ivy.sign(x) + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.pow(x, 3) >>> print(y) { - a: ivy.array([0., 0.]), - b: ivy.array([1., 1., 0.]), - c: ivy.array([-1., -1., -1., 1.]) + a:ivy.array([0,1]), + b:ivy.array([8,27]) } """ - return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) + return ivy.current_backend(x1, x2).pow(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -5777,37 +5813,19 @@ def sign( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sin( +def rad2deg( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the sine, having domain - ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of - the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. - - .. note:: - The sine is an entire function on the complex plane and has no branch cuts. - - **Special cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - For complex floating-point operands, special cases - must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. + """ + Convert the input from radians to degrees. Parameters ---------- x - input array whose elements are each expressed in radians. Should have a - floating-point data type. + input array whose elements are each expressed in radians. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5815,53 +5833,56 @@ def sin( Returns ------- ret - an array containing the sine of each element in ``x``. The returned array must - have a floating-point data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + an array with each element in ``x`` converted from radians to degrees. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.sin(x) + >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) + >>> y=ivy.rad2deg(x) >>> print(y) - ivy.array([0., 0.841, 0.909]) + ivy.array([ 0., 90., 180., 270., 360.]) - >>> x = ivy.array([0., 1.2, -2.3, 3.6]) - >>> y = ivy.zeros(4) - >>> ivy.sin(x, out=y) + >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.rad2deg(x,out=y) >>> print(y) - ivy.array([0., 0.932, -0.746, -0.443]) + ivy.array([ 0. , -1.5, -50. , nan]) - >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) - >>> ivy.sin(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.rad2deg(x, out=x) >>> print(x) - ivy.array([[0.841, 0.909, 0.141], - [0.757, 0.959, 0.279]]) + ivy.array([[ 63., 126., 189.], + [-252., -315., -378.]]) + + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.rad2deg(x,out=y) + >>> print(y) + ivy.array([ 0., 1150., nan]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), - ... b=ivy.array([-4., -5., -6., -7.])) - >>> y = ivy.sin(x) + >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), + ... b=ivy.array([0., 1., 2., 3., 4.])) + >>> y=ivy.rad2deg(x) >>> print(y) { - a: ivy.array([0., 0.841, 0.909, 0.141]), - b: ivy.array([0.757, 0.959, 0.279, -0.657]) + a: ivy.array([0., 1150., -2890., nan]), + b: ivy.array([0., 57.3, 115., 172., 229.]) + } + + >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), + ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) + >>> y=ivy.rad2deg(x) + >>> print(y) + { + a: ivy.array([0., 573., 10300., 487., 344.]), + b: ivy.array([0., -85.9, 28.6, nan]) } """ - return ivy.current_backend(x).sin(x, out=out) + return ivy.current_backend(x).rad2deg(x, out=out) @handle_exceptions @@ -5870,76 +5891,23 @@ def sin( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def sinh( +def real( x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the hyperbolic sine, having - domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each - element ``x_i`` of the input array ``x``. - - .. math:: - \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} - - .. note:: - The hyperbolic sine is an entire function in the - complex plane and has no branch cuts. - The function is periodic, with period - :math:`2\pi j`, with respect to the imaginary component. - - **Special cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and - - .. note:: - For complex floating-point operands, ``sinh(conj(x))`` - must equal ``conj(sinh(x))``. - - - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. - - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). - - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is ``+0``, - the result is ``+infinity + 0j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, - the result is ``+infinity * cis(b)``. - - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``infinity + NaN j`` (sign of the real component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``infinity + NaN j`` - (sign of the real component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - where ``cis(v)`` is ``cos(v) + sin(v)*1j``. + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Test each element ``x_i`` of the input array ``x`` to take only real part from it. + Returns a float array, where it only contains . If element has complex type with + zero complex part, the return value will be that element, else it only returns real + part. Parameters ---------- x - input array whose elements each represent a hyperbolic angle. Should have a - floating-point data type. + input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5947,45 +5915,45 @@ def sinh( Returns ------- ret - an array containing the hyperbolic sine of each element in ``x``. The returned - array must have a floating-point data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + an array containing test results. An element ``out_i`` is + ``real number`` if ``x_i`` contain real number part only + and if it is ``real number with complex part also`` then it + returns the real number part. + The returned array must have a floating-point data type with the + same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, + the returned array must have the floating-point precision of ``float32``). - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.sinh(x) - >>> print(y) - ivy.array([1.18, 3.63, 10.]) + >>> x = ivy.array([[[1.1], [2], [-6.3]]]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([[[1.1], [2.], [-6.3]]]) - >>> x = ivy.array([0.23, 3., -1.2]) - >>> ivy.sinh(x, out=x) - >>> print(x) - ivy.array([0.232, 10., -1.51]) + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([4.2, 0., 7.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) - >>> y = ivy.sinh(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ + b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.real(x) + >>> print(z) { - a: ivy.array([0.232, -0.253, 1.18]), - b: ivy.array([10., -27.3, 1.62]) + a: ivy.array([-6.7, 0.314, 1.23]), + b: ivy.array([0., 5.32, 3.001]) } """ - return ivy.current_backend(x).sinh(x, out=out) + return ivy.current_backend(x).real(x, out=out) @handle_exceptions @@ -5996,77 +5964,114 @@ def sinh( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sqrt( - x: Union[ivy.Array, ivy.NativeArray], +def reciprocal( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, - +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, - each result must be indistinguishable from the infinitely precise result (as - required by IEEE 754). - - .. note:: - After rounding, each result must be indistinguishable - from the infinitely precise result (as required by IEEE 754). - - .. note:: - For complex floating-point operands, ``sqrt(conj(x))`` - must equal ``conj(sqrt(x))``. + """ + Return a new array with the reciprocal of each element in ``x``. - .. note:: - By convention, the branch cut of the square root is - the negative real axis :math:`(-\infty, 0)`. + Parameters + ---------- + x + Input array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - The square root is a continuous function from above - the branch cut, taking into account the sign of the imaginary component. + Returns + ------- + ret + A new array with the positive value of each element in ``x``. - Accordingly, for complex arguments, the function returns - the square root in the range of the right half-plane, - including the imaginary axis (i.e., the plane defined by - :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` - along the imaginary axis). + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.reciprocal(x) + >>> print(y) + ivy.array([1. , 0.5 , 0.33333333]) + """ + return ivy.current_backend(x).reciprocal(x, out=out) - **Special cases** - For floating-point operands, +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def remainder( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + modulus: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the remainder of division for each element ``x1_i`` of the input array ``x1`` + and the respective element ``x2_i`` of the input array ``x2``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is less than ``0``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + .. note:: + This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For + input arrays which promote to an integer data type, the result of division by + zero is unspecified and thus implementation-defined. In general, similar to + Python’s ``%`` operator, this function is not recommended for floating-point + operands as semantics do not follow IEEE 754. That this function is specified + to accept floating-point operands is primarily for reasons of backward + compatibility. - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and + **Special Cases** - - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, - the result is ``+0 + 0j``. - - If ``a`` is any value (including ``NaN``) and ``b`` is - ``+infinity``, the result is ``+infinity + infinity j``. - - If ``a`` is a finite number and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - If ``a`` ``-infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. - - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, - the result is ``NaN + infinity j`` - (sign of the imaginary component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``+infinity + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is any value, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. + For floating-point operands, + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either + ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, + the result is ``NaN``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x2_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x2_i``. (note: this results matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - In the remaining cases, the result must match that of the Python ``%`` operator. Parameters ---------- - x - input array. Should have a floating-point data type. + x1 + dividend input array. Should have a numeric data type. + x2 + divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). + Should have a numeric data type. + modulus + whether to compute the modulus instead of the remainder. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6074,14 +6079,15 @@ def sqrt( Returns ------- ret - an array containing the square root of each element in ``x``. The returned array - must have a floating-point data type determined by :ref:`type-promotion`. + an array containing the element-wise results. Each element-wise result must have + the same sign as the respective element ``x2_i``. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.remainder.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6090,34 +6096,34 @@ def sqrt( Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.array([0, 4., 8.]) - >>> y = ivy.sqrt(x) + >>> x1 = ivy.array([2., 5., 15.]) + >>> x2 = ivy.array([3., 2., 4.]) + >>> y = ivy.remainder(x1, x2) >>> print(y) - ivy.array([0., 2., 2.83]) + ivy.array([2., 1., 3.]) - >>> x = ivy.array([1, 2., 4.]) - >>> y = ivy.zeros(3) - >>> ivy.sqrt(x, out=y) - ivy.array([1., 1.41, 2.]) + With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - >>> X = ivy.array([40., 24., 100.]) - >>> ivy.sqrt(x, out=x) - >>> ivy.array([6.32455532, 4.89897949, 10.]) + >>> x1 = ivy.array([23., 1., 6.]) + >>> x2 = ivy.native_array([11., 2., 4.]) + >>> y = ivy.remainder(x1, x2) + >>> print(y) + ivy.array([1., 1., 2.]) - With :class:`ivy.Container` input: + With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa - >>> y = ivy.sqrt(x) + >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) + >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) + >>> y = ivy.remainder(x1, x2) >>> print(y) { - a: ivy.array([6.63, 7.48, 13.]), - b: ivy.array([[7., 1.], - [0., 4.47]]) + a: ivy.array([0., 0., 1.]), + b: ivy.array([0., 2., 1.]) } """ - return ivy.current_backend(x).sqrt(x, out=out) + return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) @handle_exceptions @@ -6128,19 +6134,55 @@ def sqrt( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def square( +def round( x: Union[ivy.Array, ivy.NativeArray], /, *, + decimals: Optional[int] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Each element ``x_i`` of the input array ``x``. + Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued + number. + + .. note:: + For complex floating-point operands, real and imaginary components + must be independently rounded to the nearest integer-valued number. + + Rounded real and imaginary components must be equal + to their equivalent rounded real-valued floating-point + counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` + must equal ``round(real(x)))`` and ``imag(round(x))`` must equal + ``round(imag(x))``). + + **Special cases** + + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + + For floating-point operands, + + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If two integers are equally close to ``x_i``, the result is + the even integer closest to ``x_i``. + + .. note:: + For complex floating-point operands, the following special + cases apply to real and imaginary components independently + (e.g., if ``real(x_i)`` is ``NaN``, the rounded + real component is ``NaN``). + + - If ``x_i`` is already integer-valued, the result is ``x_i``. Parameters ---------- x - Input array. Should have a numeric data type. + input array containing elements to round. + decimals + number of decimal places to round to. Default is ``0``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6148,13 +6190,18 @@ def square( Returns ------- ret - an array containing the evaluated result for each element in ``x``. + An array of the same shape and type as x, with the elements rounded to integers. - This method conforms to the `Array API Standard + Note: PyTorch supports an additional argument :code:`decimals` for the + `round function `_. + It has been deliberately omitted here due to the imprecise + nature of the argument in :code:`torch.round`. + + This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.round.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6165,63 +6212,99 @@ def square( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.square(x) + >>> x = ivy.array([1.2, 2.4, 3.6]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([1, 4, 9]) + ivy.array([1.,2.,4.]) - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.square(x, out=y) + >>> x = ivy.array([-0, 5, 4.5]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([2.25, 0.64, 0.09]) + ivy.array([0.,5.,4.]) - >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) - >>> ivy.square(x, out=x) + >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) + >>> y = ivy.zeros(4) + >>> ivy.round(x, out=y) + >>> print(y) + ivy.array([2.,2.,15.,-5.]) + + >>> x = ivy.array([[0, 5.433, -343.3, 1.5], + ... [-5.5, 44.2, 11.5, 12.01]]) + >>> ivy.round(x, out=x) >>> print(x) - ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) + ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.square(x) + >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), + ... b=ivy.array([-300.9, -527.3, 4.5])) + >>> y = ivy.round(x) >>> print(y) { - a:ivy.array([0,1]), - b:ivy.array([4,9]) + a:ivy.array([4.,9.,7.,0.]), + b:ivy.array([-301.,-527.,4.]) } """ - return ivy.current_backend(x).square(x, out=out) + return ivy.current_backend(x).round(x, decimals=decimals, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def subtract( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def sign( + x: Union[ivy.Array, ivy.NativeArray], /, *, - alpha: Optional[Union[int, float]] = None, + np_variant: Optional[bool] = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Calculate the difference for each element ``x1_i`` of the input array ``x1`` with - the respective element ``x2_i`` of the input array ``x2``. + r""" + Return an indication of the sign of a number for each element ``x_i`` of the input + array ``x``. + + The sign function (also known as the **signum function**) + of a number :math:`x_{i}` is defined as + + .. math:: + \operatorname{sign}(x_i) = \begin{cases} + 0 & \textrm{if } x_i = 0 \\ + \frac{x}{|x|} & \textrm{otherwise} + \end{cases} + + where :math:`|x_i|` is the absolute value of :math:`x_i`. + + **Special cases** + + - If ``x_i`` is less than ``0``, the result is ``-1``. + - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. + - If ``x_i`` is greater than ``0``, the result is ``+1``. + - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``-0`` or ``+0`` and ``b`` is + either ``-0`` or ``+0``, the result is ``0 + 0j``. + - If ``a`` is ``NaN`` or ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, special cases must be handled + according to the rules of complex number division. Parameters ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - alpha - optional scalar multiplier for ``x2``. + x + input array. Should have a numeric data type. + np_variant + Handles complex numbers like numpy does If ``True``, + ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. + otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, + otherwise y = 0.`` + out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6229,13 +6312,14 @@ def subtract( Returns ------- ret - an array containing the element-wise differences. + an array containing the evaluated result for each element in ``x``. The returned + array must have the same data type as ``x``. - This method conforms to the `Array API Standard + This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sign.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6244,19 +6328,34 @@ def subtract( Examples -------- - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y) - >>> print(z) - ivy.array([ 1, 5, -3]) + With :class:`ivy.Array` input: - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y, alpha=2) - >>> print(z) - ivy.array([-1, 4, -9]) + >>> x = ivy.array([8.3, -0, 6.8, 0.07]) + >>> y = ivy.sign(x) + >>> print(y) + ivy.array([1., 0., 1., 1.]) + + >>> x = ivy.array([[5.78, -4., -6.9, 0], + ... [-.4, 0.5, 8, -0.01]]) + >>> y = ivy.sign(x) + >>> print(y) + ivy.array([[ 1., -1., -1., 0.], + [-1., 1., 1., -1.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., -0.]), + ... b=ivy.array([1.46, 5.9, -0.0]), + ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) + >>> y = ivy.sign(x) + >>> print(y) + { + a: ivy.array([0., 0.]), + b: ivy.array([1., 1., 0.]), + c: ivy.array([-1., -1., -1., 1.]) + } """ - return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) + return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) @handle_exceptions @@ -6267,28 +6366,19 @@ def subtract( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def tan( +def sin( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r"""Calculate an implementation-dependent approximation to the tangent, having - domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each - element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be - expressed in radians. - .. note:: - Tangent is an analytical function on the complex plane - and has no branch cuts. The function is periodic, - with period :math:`\pi j`, with respect to the real - component and has first order poles along the real - line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. - However, IEEE 754 binary floating-point representation - cannot represent the value :math:`\pi / 2` exactly, and, - thus, no argument value is possible for - which a pole error occurs. + r""" + Calculate an implementation-dependent approximation to the sine, having domain + ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of + the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. - where :math:`{tanh}` is the hyperbolic tangent. + .. note:: + The sine is an entire function on the complex plane and has no branch cuts. **Special cases** @@ -6299,68 +6389,68 @@ def tan( - If ``x_i`` is ``-0``, the result is ``-0``. - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - For complex floating-point operands, special cases must - be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. + For complex floating-point operands, special cases + must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. Parameters ---------- x - input array whose elements are expressed in radians. Should have a + input array whose elements are each expressed in radians. Should have a floating-point data type. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the tangent of each element in ``x``. The return must have a - floating-point data type determined by :ref:`type-promotion`. + an array containing the sine of each element in ``x``. The returned array must + have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sin.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tan(x) + >>> y = ivy.sin(x) >>> print(y) - ivy.array([0., 1.56, -2.19]) + ivy.array([0., 0.841, 0.909]) - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tan(x, out=y) + >>> x = ivy.array([0., 1.2, -2.3, 3.6]) + >>> y = ivy.zeros(4) + >>> ivy.sin(x, out=y) >>> print(y) - ivy.array([0.546, -0.842, -0.916]) + ivy.array([0., 0.932, -0.746, -0.443]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tan(x, out=x) + >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) + >>> ivy.sin(x, out=x) >>> print(x) - ivy.array([[1.96, -1.37, 0.16], - [-3.1, 0.996, -0.328]]) + ivy.array([[0.841, 0.909, 0.141], + [0.757, 0.959, 0.279]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.tan(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), + ... b=ivy.array([-4., -5., -6., -7.])) + >>> y = ivy.sin(x) >>> print(y) { - a: ivy.array([0., 1.56, -2.19]), - b: ivy.array([-0.143, 1.16, -3.38]) + a: ivy.array([0., 0.841, 0.909, 0.141]), + b: ivy.array([0.757, 0.959, 0.279, -0.657]) } """ - return ivy.current_backend(x).tan(x, out=out) + return ivy.current_backend(x).sin(x, out=out) @handle_exceptions @@ -6371,18 +6461,25 @@ def tan( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -@handle_complex_input -def tanh( +def sinh( x: Union[ivy.Array, ivy.NativeArray], /, *, - complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Calculate an implementation-dependent approximation to the hyperbolic tangent, - having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element - ``x_i`` of the input array ``x``. + r""" + Calculate an implementation-dependent approximation to the hyperbolic sine, having + domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each + element ``x_i`` of the input array ``x``. + + .. math:: + \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} + + .. note:: + The hyperbolic sine is an entire function in the + complex plane and has no branch cuts. + The function is periodic, with period + :math:`2\pi j`, with respect to the imaginary component. **Special cases** @@ -6391,80 +6488,62 @@ def tanh( - If ``x_i`` is ``NaN``, the result is ``NaN``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+1``. - - If ``x_i`` is ``-infinity``, the result is ``-1``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. For complex floating-point operands, let ``a = real(x_i)``, ``b = imag(x_i)``, and .. note:: - For complex floating-point operands, ``tanh(conj(x))`` - must equal ``conj(tanh(x))``. + For complex floating-point operands, ``sinh(conj(x))`` + must equal ``conj(sinh(x))``. - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. - - If ``a`` is a nonzero finite number and ``b`` is - ``+infinity``, the result is ``NaN + NaN j``. - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``+0 + NaN j``. - - If ``a`` is a nonzero finite number and ``b`` - is ``NaN``, the result is ``NaN + NaN j``. + the result is ``0 + NaN j`` (sign of the real component is unspecified). - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``+0 + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. + the result is ``0 + NaN j`` (sign of the real component is unspecified). + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is ``+0``, + the result is ``+infinity + 0j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, + the result is ``+infinity * cis(b)``. - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). + the result is ``infinity + NaN j`` (sign of the real component is unspecified). - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, - the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero number, + the result is ``infinity + NaN j`` + (sign of the real component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, the result is ``NaN + NaN j``. - If ``a`` is ``NaN`` and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - .. warning:: - For historical reasons stemming from the C standard, - array libraries may not return the expected - result when ``a`` is ``+0`` and ``b`` is either - ``+infinity`` or ``NaN``. The result should be - ``+0 + NaN j`` in both cases; however, for libraries - compiled against older C versions, the result may be - ``NaN + NaN j``. - - Array libraries are not required to patch these older - C versions, and, thus, users are advised that results - may vary across array library implementations for - these special cases. - + where ``cis(v)`` is ``cos(v) + sin(v)*1j``. Parameters ---------- x input array whose elements each represent a hyperbolic angle. Should have a - real-valued floating-point data - type. - complex_mode - optional specifier for how to handle complex data types. See - ``ivy.func_wrapper.handle_complex_input`` for more detail. + floating-point data type. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the hyperbolic tangent of each element in ``x``. - The returned array must have a real-valued floating-point data type - determined by :ref:`type-promotion`. + an array containing the hyperbolic sine of each element in ``x``. The returned + array must have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sinh.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6475,131 +6554,108 @@ def tanh( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tanh(x) - >>> print(y) - ivy.array([0., 0.762, 0.964]) - - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tanh(x, out=y) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.sinh(x) >>> print(y) - ivy.array([0.462, -0.604, 0.984]) + ivy.array([1.18, 3.63, 10.]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tanh(x, out=x) + >>> x = ivy.array([0.23, 3., -1.2]) + >>> ivy.sinh(x, out=x) >>> print(x) - ivy.array([[0.8, 0.976, 0.997], - [-1., -1., -1.]]) + ivy.array([0.232, 10., -1.51]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.tanh(x) + >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) + >>> y = ivy.sinh(x) >>> print(y) { - a: ivy.array([0., 0.762, 0.964]), - b: ivy.array([0.995, 0.999, 1.]) + a: ivy.array([0.232, -0.253, 1.18]), + b: ivy.array([10., -27.3, 1.62]) } """ - return ivy.current_backend(x).tanh(x, out=out) + return ivy.current_backend(x).sinh(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def trapz( - y: ivy.Array, +def sqrt( + x: Union[ivy.Array, ivy.NativeArray], /, *, - x: Optional[ivy.Array] = None, - dx: float = 1.0, - axis: int = -1, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Integrate along the given axis using the composite trapezoidal rule. - - If x is provided, the integration happens in sequence along its elements - - they are not sorted.. + r""" + Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, + +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, + each result must be indistinguishable from the infinitely precise result (as + required by IEEE 754). - Parameters - ---------- - y - The array that should be integrated. - x - The sample points corresponding to the input array values. - If x is None, the sample points are assumed to be evenly spaced - dx apart. The default is None. - dx - The spacing between sample points when x is None. The default is 1. - axis - The axis along which to integrate. - out - optional output array, for writing the result to. + .. note:: + After rounding, each result must be indistinguishable + from the infinitely precise result (as required by IEEE 754). - Returns - ------- - ret - Definite integral of n-dimensional array as approximated along - a single axis by the trapezoidal rule. If the input array is a - 1-dimensional array, then the result is a float. If n is greater - than 1, then the result is an n-1 dimensional array. + .. note:: + For complex floating-point operands, ``sqrt(conj(x))`` + must equal ``conj(sqrt(x))``. - Examples - -------- - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3]) - 4.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], x=[4, 6, 8]) - 8.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], dx=2) - 8.0 - """ - return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) + .. note:: + By convention, the branch cut of the square root is + the negative real axis :math:`(-\infty, 0)`. + The square root is a continuous function from above + the branch cut, taking into account the sign of the imaginary component. -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def trunc( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Round each element x_i of the input array x to the integer-valued number that is - closest to but no greater than x_i. + Accordingly, for complex arguments, the function returns + the square root in the range of the right half-plane, + including the imaginary axis (i.e., the plane defined by + :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` + along the imaginary axis). **Special cases** - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - For floating-point operands, - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is less than ``0``, the result is ``NaN``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, + the result is ``+0 + 0j``. + - If ``a`` is any value (including ``NaN``) and ``b`` is + ``+infinity``, the result is ``+infinity + infinity j``. + - If ``a`` is a finite number and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - If ``a`` ``-infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. + - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, + the result is ``NaN + infinity j`` + (sign of the imaginary component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``+infinity + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is any value, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + Parameters ---------- x - input array. Should have a numeric data type. + input array. Should have a floating-point data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6607,14 +6663,14 @@ def trunc( Returns ------- ret - an array containing the rounded result for each element in ``x``. - The returned array must have the same data type as ``x``. + an array containing the square root of each element in ``x``. The returned array + must have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sqrt.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6625,38 +6681,32 @@ def trunc( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) - >>> y = ivy.trunc(x) + >>> x = ivy.array([0, 4., 8.]) + >>> y = ivy.sqrt(x) >>> print(y) - ivy.array([-1., 0., 3., -0.]) + ivy.array([0., 2., 2.83]) - >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) - >>> ivy.trunc(x, out=x) - >>> print(x) - ivy.array([ 0., 7., -23., -0.]) + >>> x = ivy.array([1, 2., 4.]) + >>> y = ivy.zeros(3) + >>> ivy.sqrt(x, out=y) + ivy.array([1., 1.41, 2.]) - >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) - >>> y = ivy.zeros([2,3]) - >>> ivy.trunc(x, out=y) - >>> print(y) - ivy.array([[ 0., -8., 0.], - [ 0., 0., 2.]]) + >>> X = ivy.array([40., 24., 100.]) + >>> ivy.sqrt(x, out=x) + >>> ivy.array([6.32455532, 4.89897949, 10.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) - >>> y = ivy.trunc(x) + >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa + >>> y = ivy.sqrt(x) >>> print(y) { - a: ivy.array([-0., 4., 1.]), - b: ivy.array([12., -3., 1.]) + a: ivy.array([6.63, 7.48, 13.]), + b: ivy.array([[7., 1.], + [0., 4.47]]) } """ - return ivy.current_backend(x).trunc(x, out=out) - - -# Extra # -# ------# + return ivy.current_backend(x).sqrt(x, out=out) @handle_exceptions @@ -6667,19 +6717,19 @@ def trunc( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def erf( +def square( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Gauss error function of ``x`` element-wise. + Each element ``x_i`` of the input array ``x``. Parameters ---------- x - Value to compute exponential for. + Input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6687,15 +6737,50 @@ def erf( Returns ------- ret - The Gauss error function of x. + an array containing the evaluated result for each element in ``x``. + + + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([0, 0.3, 0.7, 1.0]) - >>> ivy.erf(x) - ivy.array([0., 0.328, 0.677, 0.842]) + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.square(x) + >>> print(y) + ivy.array([1, 4, 9]) + + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.square(x, out=y) + >>> print(y) + ivy.array([2.25, 0.64, 0.09]) + + >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) + >>> ivy.square(x, out=x) + >>> print(x) + ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.square(x) + >>> print(y) + { + a:ivy.array([0,1]), + b:ivy.array([4,9]) + } """ - return ivy.current_backend(x).erf(x, out=out) + return ivy.current_backend(x).square(x, out=out) @handle_exceptions @@ -6705,26 +6790,27 @@ def erf( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def maximum( - x1: Union[ivy.Array, ivy.NativeArray, Number], - x2: Union[ivy.Array, ivy.NativeArray, Number], +def subtract( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, + alpha: Optional[Union[int, float]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. + Calculate the difference for each element ``x1_i`` of the input array ``x1`` with + the respective element ``x2_i`` of the input array ``x2``. Parameters ---------- x1 - Input array containing elements to maximum threshold. + first input array. Should have a numeric data type. x2 - Tensor containing maximum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum - is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. + alpha + optional scalar multiplier for ``x2``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6732,154 +6818,138 @@ def maximum( Returns ------- ret - An array with the elements of x1, but clipped to not be lower than the x2 - values. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.maximum(x, y) - >>> print(z) - ivy.array([9, 9, 5]) + an array containing the element-wise differences. - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.maximum(x, y, out=z) - >>> print(z) - ivy.array([[9., 9., 9., 9., 9., 9.], - [3., 5., 9., 8., 3., 7.], - [2., 5., 9., 8., 3., 7.]]) - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.maximum(x, y, out=x) - >>> print(x) - ivy.array([[7, 7]]) + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - With one :class:`ivy.Container` input: + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]), - ... b=ivy.array([-5, 9])) - >>> z = ivy.maximum(x, y) + Examples + -------- + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y) >>> print(z) - { - a: ivy.array([[1, 3], - [2, 4], - [3, 7]]), - b: ivy.array([[1, 9], - [2, 9], - [3, 9]]) - } - - With multiple :class:`ivy.Container` inputs: + ivy.array([ 1, 5, -3]) - >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) - >>> z = ivy.maximum(x, y) + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y, alpha=2) >>> print(z) - { - a: ivy.array([1, 5, 6]), - b: ivy.array([5, 9, 7]) - } + ivy.array([-1, 4, -9]) """ - return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) + return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def minimum( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def tan( + x: Union[ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. + r"""Calculate an implementation-dependent approximation to the tangent, having + domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each + element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be + expressed in radians. + .. note:: + Tangent is an analytical function on the complex plane + and has no branch cuts. The function is periodic, + with period :math:`\pi j`, with respect to the real + component and has first order poles along the real + line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. + However, IEEE 754 binary floating-point representation + cannot represent the value :math:`\pi / 2` exactly, and, + thus, no argument value is possible for + which a pole error occurs. + + where :math:`{tanh}` is the hyperbolic tangent. + + **Special cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + + For complex floating-point operands, special cases must + be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. Parameters ---------- - x1 - Input array containing elements to minimum threshold. - x2 - Tensor containing minimum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum - is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. + x + input array whose elements are expressed in radians. Should have a + floating-point data type. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - An array with the elements of x1, but clipped to not exceed the x2 values. + an array containing the tangent of each element in ``x``. The return must have a + floating-point data type determined by :ref:`type-promotion`. - Examples - -------- - With :class:`ivy.Array` inputs: - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.minimum(x, y) - >>> print(z) - ivy.array([7, 3, 2]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.minimum(x, y, out=z) - >>> print(z) - ivy.array([[1.,5.,9.,8.,3.,7.], - [1.,3.,3.,3.,3.,3.], - [1.,2.,2.,2.,2.,2.]]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.minimum(x, y, out=x) - >>> print(x) - ivy.array([[0, 3]]) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.tan(x) + >>> print(y) + ivy.array([0., 1.56, -2.19]) - With one :class:`ivy.Container` input: + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tan(x, out=y) + >>> print(y) + ivy.array([0.546, -0.842, -0.916]) - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) - >>> z = ivy.minimum(x, y) - >>> print(z) - { - a: ivy.array([[1, 0], - [1, 0], - [1, 0]]), - b: ivy.array([[-5, 3], - [-5, 4], - [-5, 7]]) - } + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tan(x, out=x) + >>> print(x) + ivy.array([[1.96, -1.37, 0.16], + [-3.1, 0.996, -0.328]]) - With multiple :class:`ivy.Container` inputs: + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 3, 1]), - ... b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]), - ... b=ivy.array([5, 9, 7])) - >>> z = ivy.minimum(x, y) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.tan(x) + >>> print(y) { - a: ivy.array([1, 3, 1]), - b: ivy.array([2, 8, 5]) + a: ivy.array([0., 1.56, -2.19]), + b: ivy.array([-0.143, 1.16, -3.38]) } """ - return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) + return ivy.current_backend(x).tan(x, out=out) @handle_exceptions @@ -6890,36 +6960,140 @@ def minimum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def reciprocal( - x: Union[float, ivy.Array, ivy.NativeArray], +@handle_complex_input +def tanh( + x: Union[ivy.Array, ivy.NativeArray], /, *, + complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the reciprocal of each element in ``x``. + Calculate an implementation-dependent approximation to the hyperbolic tangent, + having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element + ``x_i`` of the input array ``x``. + + **Special cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``+infinity``, the result is ``+1``. + - If ``x_i`` is ``-infinity``, the result is ``-1``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + .. note:: + For complex floating-point operands, ``tanh(conj(x))`` + must equal ``conj(tanh(x))``. + + - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. + - If ``a`` is a nonzero finite number and ``b`` is + ``+infinity``, the result is ``NaN + NaN j``. + - If ``a`` is ``+0`` and ``b`` is ``+infinity``, + the result is ``+0 + NaN j``. + - If ``a`` is a nonzero finite number and ``b`` + is ``NaN``, the result is ``NaN + NaN j``. + - If ``a`` is ``+0`` and ``b`` is ``NaN``, + the result is ``+0 + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. + - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, + the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero number, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + + .. warning:: + For historical reasons stemming from the C standard, + array libraries may not return the expected + result when ``a`` is ``+0`` and ``b`` is either + ``+infinity`` or ``NaN``. The result should be + ``+0 + NaN j`` in both cases; however, for libraries + compiled against older C versions, the result may be + ``NaN + NaN j``. + + Array libraries are not required to patch these older + C versions, and, thus, users are advised that results + may vary across array library implementations for + these special cases. + Parameters ---------- x - Input array. + input array whose elements each represent a hyperbolic angle. Should have a + real-valued floating-point data + type. + complex_mode + optional specifier for how to handle complex data types. See + ``ivy.func_wrapper.handle_complex_input`` for more detail. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - A new array with the positive value of each element in ``x``. + an array containing the hyperbolic tangent of each element in ``x``. + The returned array must have a real-valued floating-point data type + determined by :ref:`type-promotion`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.reciprocal(x) + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.tanh(x) >>> print(y) - ivy.array([1. , 0.5 , 0.33333333]) + ivy.array([0., 0.762, 0.964]) + + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tanh(x, out=y) + >>> print(y) + ivy.array([0.462, -0.604, 0.984]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tanh(x, out=x) + >>> print(x) + ivy.array([[0.8, 0.976, 0.997], + [-1., -1., -1.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.tanh(x) + >>> print(y) + { + a: ivy.array([0., 0.762, 0.964]), + b: ivy.array([0.995, 0.999, 1.]) + } """ - return ivy.current_backend(x).reciprocal(x, out=out) + return ivy.current_backend(x).tanh(x, out=out) @handle_backend_invalid @@ -6927,80 +7101,61 @@ def reciprocal( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def deg2rad( - x: Union[ivy.Array, ivy.NativeArray], +def trapz( + y: ivy.Array, /, *, + x: Optional[ivy.Array] = None, + dx: float = 1.0, + axis: int = -1, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Convert the input from degrees to radians. + Integrate along the given axis using the composite trapezoidal rule. + + If x is provided, the integration happens in sequence along its elements + - they are not sorted.. Parameters ---------- + y + The array that should be integrated. x - input array whose elements are each expressed in degrees. + The sample points corresponding to the input array values. + If x is None, the sample points are assumed to be evenly spaced + dx apart. The default is None. + dx + The spacing between sample points when x is None. The default is 1. + axis + The axis along which to integrate. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array with each element in ``x`` converted from degrees to radians. + Definite integral of n-dimensional array as approximated along + a single axis by the trapezoidal rule. If the input array is a + 1-dimensional array, then the result is a float. If n is greater + than 1, then the result is an n-1 dimensional array. Examples -------- - With :class:`ivy.Array` input: - - >>> x=ivy.array([0,90,180,270,360]) - >>> y=ivy.deg2rad(x) - >>> print(y) - ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) - - >>> x=ivy.array([0,-1.5,-50,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.deg2rad(x,out=y) - >>> print(y) - ivy.array([ 0., -0.02617994, -0.87266463, nan]) - - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.deg2rad(x, out=x) - >>> print(x) - ivy.array([[ 0.01919862, 0.03839725, 0.05759586], - [-0.07679449, -0.09599311, -0.11519173]]) - - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.deg2rad(x,out=y) - >>> print(y) - ivy.array([0., 0.35081118, nan]) - - With :class:`ivy.Container` input: - - >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), - ... b=ivy.array([0,90,180,270,360])) - >>> y=ivy.deg2rad(x) - >>> print(y) - { - a: ivy.array([0., 0.35081118, -0.88139129, nan]), - b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) - } - - >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), - ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) - >>> y=ivy.deg2rad(x) - >>> print(y) - { - a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), - b: ivy.array([0., -0.02617994, -0.87266463, nan]) - } + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3]) + 4.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], x=[4, 6, 8]) + 8.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], dx=2) + 8.0 """ - return ivy.current_backend(x).deg2rad(x, out=out) + return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -7008,19 +7163,32 @@ def deg2rad( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def rad2deg( +def trunc( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Convert the input from radians to degrees. + Round each element x_i of the input array x to the integer-valued number that is + closest to but no greater than x_i. + + **Special cases** + + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + + For floating-point operands, + + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. Parameters ---------- x - input array whose elements are each expressed in radians. + input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -7028,56 +7196,52 @@ def rad2deg( Returns ------- ret - an array with each element in ``x`` converted from radians to degrees. + an array containing the rounded result for each element in ``x``. + The returned array must have the same data type as ``x``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) - >>> y=ivy.rad2deg(x) - >>> print(y) - ivy.array([ 0., 90., 180., 270., 360.]) - - >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.rad2deg(x,out=y) + >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) + >>> y = ivy.trunc(x) >>> print(y) - ivy.array([ 0. , -1.5, -50. , nan]) + ivy.array([-1., 0., 3., -0.]) - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.rad2deg(x, out=x) + >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) + >>> ivy.trunc(x, out=x) >>> print(x) - ivy.array([[ 63., 126., 189.], - [-252., -315., -378.]]) + ivy.array([ 0., 7., -23., -0.]) - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.rad2deg(x,out=y) + >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) + >>> y = ivy.zeros([2,3]) + >>> ivy.trunc(x, out=y) >>> print(y) - ivy.array([ 0., 1150., nan]) + ivy.array([[ 0., -8., 0.], + [ 0., 0., 2.]]) With :class:`ivy.Container` input: - >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), - ... b=ivy.array([0., 1., 2., 3., 4.])) - >>> y=ivy.rad2deg(x) - >>> print(y) - { - a: ivy.array([0., 1150., -2890., nan]), - b: ivy.array([0., 57.3, 115., 172., 229.]) - } - - >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), - ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) - >>> y=ivy.rad2deg(x) + >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) + >>> y = ivy.trunc(x) >>> print(y) { - a: ivy.array([0., 573., 10300., 487., 344.]), - b: ivy.array([0., -85.9, 28.6, nan]) + a: ivy.array([-0., 4., 1.]), + b: ivy.array([12., -3., 1.]) } """ - return ivy.current_backend(x).rad2deg(x, out=out) + return ivy.current_backend(x).trunc(x, out=out) @handle_exceptions @@ -7124,169 +7288,3 @@ def trunc_divide( ivy.array([ 0., -1., 14.]) """ return ivy.trunc(ivy.divide(x1, x2), out=out) - - -trunc_divide.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - "handle_backend_invalid", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def isreal( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Test each element ``x_i`` of the input array ``x`` to determine whether the element - is real number. Returns a bool array, where True if input element is real. If - element has complex type with zero complex part, the return value for that element - is True. - - Parameters - ---------- - x - input array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is - real number and ``False`` otherwise. The returned array should have a data type - of ``bool``. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances in place of - :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints - and also the examples below. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([[[True], [True], [True]]]) - - >>> x = ivy.array([1-0j, 3j, 7+5j]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([ True, False, False]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ - b=ivy.array([5j, 5-6j, 3])) - >>> z = ivy.isreal(x) - >>> print(z) - { - a: ivy.array([False, True, True]), - b: ivy.array([False, False, True]) - } - """ - return ivy.current_backend(x).isreal(x, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fmod( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Compute the element-wise remainder of divisions of two arrays. - - Parameters - ---------- - x1 - First input array. - x2 - Second input array - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array with element-wise remainder of divisions. - - Examples - -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmod(x1, x2) - ivy.array([ 0, 3, 0]) - - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmod(x1, x2) - ivy.array([ nan, nan, nan]) - """ - return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lcm( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the element-wise least common multiple (LCM) of x1 and x2. - - Parameters - ---------- - x1 - first input array, must be integers - x2 - second input array, must be integers - out - optional output array, for writing the result to. - - Returns - ------- - ret - an array that includes the element-wise least common multiples of x1 and x2 - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x1=ivy.array([2, 3, 4]) - >>> x2=ivy.array([5, 7, 15]) - >>> x1.lcm(x1, x2) - ivy.array([10, 21, 60]) - """ - return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) diff --git a/ivy/functional/ivy/experimental/activations.py b/ivy/functional/ivy/experimental/activations.py index 34203cffbaf6d..30de525bedc40 100644 --- a/ivy/functional/ivy/experimental/activations.py +++ b/ivy/functional/ivy/experimental/activations.py @@ -17,6 +17,68 @@ ) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def elu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + alpha: float = 1.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Apply the elu unit function element-wise. + + Parameters + ---------- + x + Input array. + alpha + scaler for controlling the slope of the function for x <= 0 Default: 1.0 + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The input array with elu applied element-wise. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([0.39, -0.85]) + >>> y = ivy.elu(x) + >>> print(y) + ivy.array([ 0.38999999, -0.57258511]) + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.elu(x, out=y) + >>> print(y) + ivy.array([ 1.5, 0.69999999, -0.90928203]) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.elu(x, out=x) + >>> print(x) + ivy.array([[ 1.10000002, 2.20000005, 3.29999995], + [-0.98772264, -0.99591321, -0.99863964]]) + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.elu(x, out=x) + >>> print(x) + { + a: ivy.array([0., -0.69880581]), + b: ivy.array([0.40000001, -0.18126924]) + } + """ + return current_backend(x).elu(x, alpha=alpha, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -67,6 +129,57 @@ def logit( return current_backend(x).logit(x, eps=eps, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def logsigmoid( + input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply element-wise Log-sigmoid of x. + + logsigmoid(x) = log(1 / (1 + exp(-x)). + + Parameters + ---------- + input + Input array. + + Returns + ------- + Array with same shape as input with Log-sigmoid applied to every element. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-1.31326175, -0.69314718, -0.31326169]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-0.20141329, -0.40318608, -2.48683619]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.logsigmoid(x) + >>> print(x) + { + a: ivy.array([-0.31326169, -1.46328247]), + b: ivy.array([-0.59813893, -0.43748799]) + } + """ + return ivy.current_backend(input).logsigmoid(input, out=out) + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion @@ -125,67 +238,6 @@ def prelu( raise IvyException -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def thresholded_relu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - threshold: Union[int, float] = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Apply the rectified linear unit function with custom threshold. - - Parameters - ---------- - x - input array - threshold - threshold value above which the activation is linear. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the rectified linear unit activation of each element in - ``x``. with custom ``threshold``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.thresholded_relu(x, threshold=0.5) - >>> print(y) - ivy.array([0., 0. , 1.]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.thresholded_relu(x, threshold=1, out = y) - >>> print(y) - ivy.array([ 1.5, 0., 0.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.thresholded_relu(x, threshold=0.5) - >>> print(x) - { - a: ivy.array([1., 0.]), - b: ivy.array([0., 0.6]) - } - """ - return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -232,57 +284,6 @@ def relu6( return current_backend(x).relu6(x, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def logsigmoid( - input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply element-wise Log-sigmoid of x. - - logsigmoid(x) = log(1 / (1 + exp(-x)). - - Parameters - ---------- - input - Input array. - - Returns - ------- - Array with same shape as input with Log-sigmoid applied to every element. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-1.31326175, -0.69314718, -0.31326169]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-0.20141329, -0.40318608, -2.48683619]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.logsigmoid(x) - >>> print(x) - { - a: ivy.array([-0.31326169, -1.46328247]), - b: ivy.array([-0.59813893, -0.43748799]) - } - """ - return ivy.current_backend(input).logsigmoid(input, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -342,6 +343,40 @@ def selu( return current_backend(x).selu(x, out=out) +def sequence_length( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.int64: + """ + Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy + array input. + + Parameters + ---------- + x + Can be a sequence of any tensor type: bool, complex128, + complex64, double, float, float16, int16, int32, int64, + int8, string, uint16, uint32, uint64, uint8 + + Returns + ------- + length + Length of the input sequence, as a scalar (empty shape tensor). + + Examples + -------- + >>> x = ivy.array([True, False, True]) + >>> y = ivy.sequence_length(x) + >>> print(y) + 3 + + >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] + >>> y = ivy.sequence_length(x) + >>> print(y) + 5 + """ + return current_backend(x).sequence_length(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -398,23 +433,23 @@ def silu( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def elu( +@handle_device_shifting +def thresholded_relu( x: Union[ivy.Array, ivy.NativeArray], /, *, - alpha: float = 1.0, + threshold: Union[int, float] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply the elu unit function element-wise. + Apply the rectified linear unit function with custom threshold. Parameters ---------- x - Input array. - alpha - scaler for controlling the slope of the function for x <= 0 Default: 1.0 + input array + threshold + threshold value above which the activation is linear. Default: ``0``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -422,67 +457,32 @@ def elu( Returns ------- ret - The input array with elu applied element-wise. + an array containing the rectified linear unit activation of each element in + ``x``. with custom ``threshold``. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.39, -0.85]) - >>> y = ivy.elu(x) + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.thresholded_relu(x, threshold=0.5) >>> print(y) - ivy.array([ 0.38999999, -0.57258511]) + ivy.array([0., 0. , 1.]) + >>> x = ivy.array([1.5, 0.7, -2.4]) >>> y = ivy.zeros(3) - >>> ivy.elu(x, out=y) + >>> ivy.thresholded_relu(x, threshold=1, out = y) >>> print(y) - ivy.array([ 1.5, 0.69999999, -0.90928203]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.elu(x, out=x) - >>> print(x) - ivy.array([[ 1.10000002, 2.20000005, 3.29999995], - [-0.98772264, -0.99591321, -0.99863964]]) + ivy.array([ 1.5, 0., 0.]) + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.elu(x, out=x) + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.thresholded_relu(x, threshold=0.5) >>> print(x) { - a: ivy.array([0., -0.69880581]), - b: ivy.array([0.40000001, -0.18126924]) + a: ivy.array([1., 0.]), + b: ivy.array([0., 0.6]) } """ - return current_backend(x).elu(x, alpha=alpha, out=out) - - -def sequence_length( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.int64: - """ - Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy - array input. - - Parameters - ---------- - x - Can be a sequence of any tensor type: bool, complex128, - complex64, double, float, float16, int16, int32, int64, - int8, string, uint16, uint32, uint64, uint8 - - Returns - ------- - length - Length of the input sequence, as a scalar (empty shape tensor). - - Examples - -------- - >>> x = ivy.array([True, False, True]) - >>> y = ivy.sequence_length(x) - >>> print(y) - 3 - - >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] - >>> y = ivy.sequence_length(x) - >>> print(y) - 5 - """ - return current_backend(x).sequence_length(x, out=out) + return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) diff --git a/ivy/functional/ivy/experimental/creation.py b/ivy/functional/ivy/experimental/creation.py index bb3fc06d68e66..cf96237fa41af 100644 --- a/ivy/functional/ivy/experimental/creation.py +++ b/ivy/functional/ivy/experimental/creation.py @@ -21,56 +21,31 @@ ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def vorbis_window( - window_length: Union[ivy.Array, ivy.NativeArray], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array that contains a vorbis power complementary window of size - window_length. +# --- Helpers --- # +# --------------- # - Parameters - ---------- - window_length - the length of the vorbis window. - dtype - data type of the returned array. By default float32. - out - optional output array, for writing the result to. - Returns - ------- - ret - Input array with the vorbis window. +def _iter_product(*args, repeat=1): + # itertools.product + pools = [tuple(pool) for pool in args] * repeat + result = [[]] + for pool in pools: + result = [x + [y] for x in result for y in pool] + for prod in result: + yield tuple(prod) - Examples - -------- - >>> ivy.vorbis_window(3) - ivy.array([0.38268346, 1. , 0.38268352]) - >>> ivy.vorbis_window(5) - ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) - """ - return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) +# --- Main --- # +# ------------ # @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @infer_dtype @handle_device_shifting -def hann_window( +def blackman_window( size: int, *, periodic: bool = True, @@ -78,13 +53,14 @@ def hann_window( out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Generate a Hann window. The Hanning window is a taper formed by using a weighted - cosine. + Generate a Blackman window. The Blackman window is a taper formed by using the first + three terms of a summation of cosines. It was designed to have close to the minimal + leakage possible. It is close to optimal, only slightly worse than a Kaiser window. Parameters ---------- - size - the size of the returned window. + window_length + the window_length of the returned window. periodic If True, returns a window to be used as periodic function. If False, return a symmetric window. @@ -97,122 +73,111 @@ def hann_window( ------- ret The array containing the window. - Functional Examples ------------------- - >>> ivy.hann_window(4, periodic = True) - ivy.array([0. , 0.5, 1. , 0.5]) - - >>> ivy.hann_window(7, periodic = False) - ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) + >>> ivy.blackman_window(4, periodic = True) + ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) + >>> ivy.blackman_window(7, periodic = False) + ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, + 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) """ - return ivy.current_backend().hann_window( + return ivy.current_backend().blackman_window( size, periodic=periodic, dtype=dtype, out=out ) @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @infer_dtype -@handle_device_shifting -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +@infer_device +def eye_like( + x: Union[ivy.Array, ivy.NativeArray], *, + k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kaiser window with window length window_length and shape beta. + Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the + same ``shape`` as the first and last dim of input array ``x``. input array ``x`` + should to be 2D. Parameters ---------- - window_length - an int defining the length of the window. - periodic - If True, returns a periodic window suitable for use in spectral analysis. - If False, returns a symmetric window suitable for use in filter design. - beta - a float used as shape parameter for the window. + x + input array from which to derive the output array shape. + k + index of the diagonal. A positive value refers to an upper diagonal, a negative + value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. dtype - data type of the returned array. + output array data type. If dtype is None, the output array data type must be the + default floating-point data type. Default: ``None``. + device + the device on which to place the created array. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The array containing the window. + an array having the same shape as ``x`` and filled with ``ones`` in + diagonal ``k`` and ``zeros`` elsewhere. - Examples - -------- - >>> ivy.kaiser_window(5) - ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) - >>> ivy.kaiser_window(5, True, 5) - ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) - >>> ivy.kaiser_window(5, False, 5) - ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) - """ - return ivy.current_backend().kaiser_window( - window_length, periodic, beta, dtype=dtype, out=out - ) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances as a replacement to any of the arguments. -@handle_exceptions -@handle_nestable -@handle_out_argument -@infer_dtype -def kaiser_bessel_derived_window( - window_length: int, - beta: float = 12.0, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Kaiser bessel derived window with window length window_length and shape - beta. + Functional Examples + ------------------- - Parameters - ---------- - window_length - an int defining the length of the window. - beta - a float used as shape parameter for the window. - dtype - data type of the returned array - out - optional output array, for writing the result to. + With :class:`ivy.Array` input: - Returns - ------- - ret - The array containing the window. + >>> x1 = ivy.array([[0, 1],[2, 3]]) + >>> y1 = ivy.eye_like(x1) + >>> print(y1) + ivy.array([[1., 0.], + [0., 1.]]) - Functional Examples - ------------------- - >>> ivy.kaiser_bessel_derived_window(5) - ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) + >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) + >>> y1 = ivy.eye_like(x1, k=1) + >>> print(y1) + ivy.array([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) - >>> ivy.kaiser_bessel_derived_window(5, 5) - ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) + >>> y = x.eye_like() + >>> print(y) + { + a: ivy.array([[1., 0.], + [0., 1.]]), + b: ivy.array([[1., 0.], + [0., 1.]]) + } """ - if window_length < 2: - result = ivy.array([], dtype=dtype) - if ivy.exists(out): - ivy.inplace_update(out, result) - return result - half_len = window_length // 2 - kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) - kaiser_w_csum = ivy.cumsum(kaiser_w) - half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) - window = ivy.concat((half_w, half_w[::-1]), axis=0) - result = window.astype(dtype) - return result + shape = ivy.shape(x, as_array=True) + dim = len(shape) + if dim <= 1: + cols = dim + else: + cols = int(shape[-1]) + rows = 0 if dim < 1 else int(shape[0]) + return ivy.eye( + rows, + cols, + k=k, + dtype=dtype, + device=device, + out=out, + ) @handle_exceptions @@ -272,214 +237,211 @@ def hamming_window( return result -hamming_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "handle_device_shifting", - ), - "to_skip": (), -} - - @handle_exceptions +@handle_backend_invalid @handle_nestable -@outputs_to_ivy_arrays -@infer_device -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def hann_window( + size: int, *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, -) -> Tuple[ivy.Array, ...]: + periodic: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Return the indices of the lower triangular part of a row by col matrix in a 2-by-N - shape (tuple of two N dimensional arrays), where the first row contains row - coordinates of all indices and the second row contains column coordinates. Indices - are ordered based on rows and then columns. The lower triangular part of the matrix - is defined as the elements on and below the diagonal. The argument k controls which - diagonal to consider. If k = 0, all elements on and below the main diagonal are - retained. A positive value excludes just as many diagonals below the main diagonal, - and similarly a negative value includes just as many diagonals above the main - diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, - n_cols}−1]. - - Notes - ----- - Primary purpose of this function is to slice an array of shape (n,m). See - https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html - for examples - - Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices + Generate a Hann window. The Hanning window is a taper formed by using a weighted + cosine. Parameters ---------- - n_rows - number of rows in the 2-d matrix. - n_cols - number of columns in the 2-d matrix. If None n_cols will be the same as n_rows - k - number of shifts from the main diagonal. k = 0 includes main diagonal, - k > 0 moves downward and k < 0 moves upward - device - device on which to place the created array. Default: ``None``. + size + the size of the returned window. + periodic + If True, returns a window to be used as periodic function. + If False, return a symmetric window. + dtype + The data type to produce. Must be a floating point type. + out + optional output array, for writing the result to. Returns ------- ret - an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) - contains row coordinates of all indices and the second subarray (i.e ret[1]) - contains columns indices. - - Function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - >>> x = ivy.tril_indices(4,4,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + The array containing the window. - >>> x = ivy.tril_indices(4,4,1) - >>> print(x) - (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) + Functional Examples + ------------------- + >>> ivy.hann_window(4, periodic = True) + ivy.array([0. , 0.5, 1. , 0.5]) - >>> x = ivy.tril_indices(4,4,-2) - >>> print(x) - (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) + >>> ivy.hann_window(7, periodic = False) + ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) + """ + return ivy.current_backend().hann_window( + size, periodic=periodic, dtype=dtype, out=out + ) - >>> x = ivy.tril_indices(4,2,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 0, 1])) - >>> x = ivy.tril_indices(2,4,0) - >>> print(x) - (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) +@handle_exceptions +def indices( + dimensions: Sequence[int], + *, + dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, + sparse: bool = False, +) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: + """ + Return an array representing the indices of a grid. - >>> x = ivy.tril_indices(4,-4,0) - >>> print(x) - (ivy.array([]), ivy.array([])) + Parameters + ---------- + dimensions + The shape of the grid. + dtype + The data type of the result. + sparse + Return a sparse representation of the grid instead of a dense representation. - >>> x = ivy.tril_indices(4,4,100) - >>> print(x) - (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + Returns + ------- + ret + If sparse is False, returns one grid indices array of shape + (len(dimensions),) + tuple(dimensions). + If sparse is True, returns a tuple of arrays each of shape + (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. - >>> x = ivy.tril_indices(2,4,-100) - >>> print(x) - (ivy.array([]), ivy.array([])) + Examples + -------- + >>> ivy.indices((3, 2)) + ivy.array([[[0 0] + [1 1] + [2 2]] + [[0 1] + [0 1] + [0 1]]]) + >>> ivy.indices((3, 2), sparse=True) + (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) """ - return current_backend().tril_indices(n_rows, n_cols, k, device=device) + if sparse: + return tuple( + ivy.arange(dim) + .expand_dims( + axis=[j for j in range(len(dimensions)) if i != j], + ) + .astype(dtype) + for i, dim in enumerate(dimensions) + ) + else: + grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") + return ivy.stack(grid, axis=0).astype(dtype) @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays @infer_dtype -@infer_device -def eye_like( - x: Union[ivy.Array, ivy.NativeArray], +def kaiser_bessel_derived_window( + window_length: int, + beta: float = 12.0, *, - k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the - same ``shape`` as the first and last dim of input array ``x``. input array ``x`` - should to be 2D. + Compute the Kaiser bessel derived window with window length window_length and shape + beta. Parameters ---------- - x - input array from which to derive the output array shape. - k - index of the diagonal. A positive value refers to an upper diagonal, a negative - value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. + window_length + an int defining the length of the window. + beta + a float used as shape parameter for the window. dtype - output array data type. If dtype is None, the output array data type must be the - default floating-point data type. Default: ``None``. - device - the device on which to place the created array. + data type of the returned array out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array having the same shape as ``x`` and filled with ``ones`` in - diagonal ``k`` and ``zeros`` elsewhere. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances as a replacement to any of the arguments. + The array containing the window. Functional Examples ------------------- + >>> ivy.kaiser_bessel_derived_window(5) + ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) - With :class:`ivy.Array` input: + >>> ivy.kaiser_bessel_derived_window(5, 5) + ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) + """ + if window_length < 2: + result = ivy.array([], dtype=dtype) + if ivy.exists(out): + ivy.inplace_update(out, result) + return result + half_len = window_length // 2 + kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) + kaiser_w_csum = ivy.cumsum(kaiser_w) + half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) + window = ivy.concat((half_w, half_w[::-1]), axis=0) + result = window.astype(dtype) + return result - >>> x1 = ivy.array([[0, 1],[2, 3]]) - >>> y1 = ivy.eye_like(x1) - >>> print(y1) - ivy.array([[1., 0.], - [0., 1.]]) - >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) - >>> y1 = ivy.eye_like(x1, k=1) - >>> print(y1) - ivy.array([[0., 1., 0.], - [0., 0., 1.], - [0., 0., 0.]]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Kaiser window with window length window_length and shape beta. - With :class:`ivy.Container` input: + Parameters + ---------- + window_length + an int defining the length of the window. + periodic + If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta + a float used as shape parameter for the window. + dtype + data type of the returned array. + out + optional output array, for writing the result to. - >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) - >>> y = x.eye_like() - >>> print(y) - { - a: ivy.array([[1., 0.], - [0., 1.]]), - b: ivy.array([[1., 0.], - [0., 1.]]) - } + Returns + ------- + ret + The array containing the window. + + Examples + -------- + >>> ivy.kaiser_window(5) + ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) + >>> ivy.kaiser_window(5, True, 5) + ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) + >>> ivy.kaiser_window(5, False, 5) + ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) """ - shape = ivy.shape(x, as_array=True) - dim = len(shape) - if dim <= 1: - cols = dim - else: - cols = int(shape[-1]) - rows = 0 if dim < 1 else int(shape[0]) - return ivy.eye( - rows, - cols, - k=k, - dtype=dtype, - device=device, - out=out, + return ivy.current_backend().kaiser_window( + window_length, periodic, beta, dtype=dtype, out=out ) -def _iter_product(*args, repeat=1): - # itertools.product - pools = [tuple(pool) for pool in args] * repeat - result = [[]] - for pool in pools: - result = [x + [y] for x in result for y in pool] - for prod in result: - yield tuple(prod) - - @handle_exceptions @inputs_to_ivy_arrays def ndenumerate( @@ -553,185 +515,60 @@ def ndindex( @handle_exceptions -def indices( - dimensions: Sequence[int], - *, - dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, - sparse: bool = False, -) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: - """ - Return an array representing the indices of a grid. - - Parameters - ---------- - dimensions - The shape of the grid. - dtype - The data type of the result. - sparse - Return a sparse representation of the grid instead of a dense representation. - - Returns - ------- - ret - If sparse is False, returns one grid indices array of shape - (len(dimensions),) + tuple(dimensions). - If sparse is True, returns a tuple of arrays each of shape - (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. - - Examples - -------- - >>> ivy.indices((3, 2)) - ivy.array([[[0 0] - [1 1] - [2 2]] - [[0 1] - [0 1] - [0 1]]]) - >>> ivy.indices((3, 2), sparse=True) - (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) - """ - if sparse: - return tuple( - ivy.arange(dim) - .expand_dims( - axis=[j for j in range(len(dimensions)) if i != j], - ) - .astype(dtype) - for i, dim in enumerate(dimensions) - ) - else: - grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") - return ivy.stack(grid, axis=0).astype(dtype) - - -indices.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} - - -@handle_exceptions -@handle_backend_invalid @handle_nestable -@to_native_arrays_and_back -def unsorted_segment_min( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the minimum along segments of an array. Segments are defined by an integer - array of segment IDs. - - Note - ---- - If the given segment ID `i` is negative, then the corresponding - value is dropped, and will not be included in the result. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented min operation. - For each segment, it computes the min value in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_nestable -@to_native_arrays_and_back -def unsorted_segment_sum( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the sum of elements along segments of an array. Segments are defined by an - integer array of segment IDs. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented sum operation. - For each segment, it computes the sum of values in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back @infer_dtype -@handle_device_shifting -def blackman_window( - size: int, +def random_cp( + shape: Sequence[int], + rank: int, + /, *, - periodic: bool = True, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + full: Optional[bool] = False, + orthogonal: Optional[bool] = False, + seed: Optional[int] = None, + normalise_factors: Optional[bool] = True, +) -> Union[ivy.CPTensor, ivy.Array]: """ - Generate a Blackman window. The Blackman window is a taper formed by using the first - three terms of a summation of cosines. It was designed to have close to the minimal - leakage possible. It is close to optimal, only slightly worse than a Kaiser window. + Generate a random CP tensor. Parameters ---------- - window_length - the window_length of the returned window. - periodic - If True, returns a window to be used as periodic function. - If False, return a symmetric window. - dtype - The data type to produce. Must be a floating point type. - out - optional output array, for writing the result to. + shape + shape of the tensor to generate + rank + rank of the CP decomposition + full + if True, a full tensor is returned + otherwise, the decomposed tensor is returned + orthogonal + if True, creates a tensor with orthogonal components + seed + seed for generating random numbers Returns ------- - ret - The array containing the window. - Functional Examples - ------------------- - >>> ivy.blackman_window(4, periodic = True) - ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) - >>> ivy.blackman_window(7, periodic = False) - ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, - 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) + ivy.CPTensor """ - return ivy.current_backend().blackman_window( - size, periodic=periodic, dtype=dtype, out=out - ) + rank = ivy.CPTensor.validate_cp_rank(shape, rank) + if (rank > min(shape)) and orthogonal: + warnings.warn( + "Can only construct orthogonal tensors when rank <= min(shape) but got " + f"a tensor with min(shape)={min(shape)} < rank={rank}" + ) + + factors = [ + (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape + ] + weights = ivy.ones((rank,), dtype=dtype) + if orthogonal: + factors = [ivy.qr(factor)[0] for factor in factors] + + if full: + return ivy.CPTensor.cp_to_tensor((weights, factors)) + elif normalise_factors: + return ivy.CPTensor.cp_normalize((weights, factors)) + else: + return ivy.CPTensor((weights, factors)) @handle_exceptions @@ -806,59 +643,96 @@ def random_tucker( @handle_exceptions @handle_nestable -@infer_dtype -def random_cp( - shape: Sequence[int], - rank: int, - /, +@outputs_to_ivy_arrays +@infer_device +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - full: Optional[bool] = False, - orthogonal: Optional[bool] = False, - seed: Optional[int] = None, - normalise_factors: Optional[bool] = True, -) -> Union[ivy.CPTensor, ivy.Array]: + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, +) -> Tuple[ivy.Array, ...]: """ - Generate a random CP tensor. + Return the indices of the lower triangular part of a row by col matrix in a 2-by-N + shape (tuple of two N dimensional arrays), where the first row contains row + coordinates of all indices and the second row contains column coordinates. Indices + are ordered based on rows and then columns. The lower triangular part of the matrix + is defined as the elements on and below the diagonal. The argument k controls which + diagonal to consider. If k = 0, all elements on and below the main diagonal are + retained. A positive value excludes just as many diagonals below the main diagonal, + and similarly a negative value includes just as many diagonals above the main + diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, + n_cols}−1]. + + Notes + ----- + Primary purpose of this function is to slice an array of shape (n,m). See + https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html + for examples + + Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices Parameters ---------- - shape - shape of the tensor to generate - rank - rank of the CP decomposition - full - if True, a full tensor is returned - otherwise, the decomposed tensor is returned - orthogonal - if True, creates a tensor with orthogonal components - seed - seed for generating random numbers + n_rows + number of rows in the 2-d matrix. + n_cols + number of columns in the 2-d matrix. If None n_cols will be the same as n_rows + k + number of shifts from the main diagonal. k = 0 includes main diagonal, + k > 0 moves downward and k < 0 moves upward + device + device on which to place the created array. Default: ``None``. Returns ------- - ivy.CPTensor - """ - rank = ivy.CPTensor.validate_cp_rank(shape, rank) - if (rank > min(shape)) and orthogonal: - warnings.warn( - "Can only construct orthogonal tensors when rank <= min(shape) but got " - f"a tensor with min(shape)={min(shape)} < rank={rank}" - ) + ret + an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) + contains row coordinates of all indices and the second subarray (i.e ret[1]) + contains columns indices. - factors = [ - (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape - ] - weights = ivy.ones((rank,), dtype=dtype) - if orthogonal: - factors = [ivy.qr(factor)[0] for factor in factors] + Function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - if full: - return ivy.CPTensor.cp_to_tensor((weights, factors)) - elif normalise_factors: - return ivy.CPTensor.cp_normalize((weights, factors)) - else: - return ivy.CPTensor((weights, factors)) + Examples + -------- + >>> x = ivy.tril_indices(4,4,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,1) + >>> print(x) + (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,-2) + >>> print(x) + (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,2,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 0, 1])) + + >>> x = ivy.tril_indices(2,4,0) + >>> print(x) + (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,-4,0) + >>> print(x) + (ivy.array([]), ivy.array([])) + + >>> x = ivy.tril_indices(4,4,100) + >>> print(x) + (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(2,4,-100) + >>> print(x) + (ivy.array([]), ivy.array([])) + """ + return current_backend().tril_indices(n_rows, n_cols, k, device=device) @handle_nestable @@ -910,3 +784,135 @@ def trilu( instances in place of any of the arguments. """ return current_backend(x).trilu(x, k=k, upper=upper, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_min( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the minimum along segments of an array. Segments are defined by an integer + array of segment IDs. + + Note + ---- + If the given segment ID `i` is negative, then the corresponding + value is dropped, and will not be included in the result. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented min operation. + For each segment, it computes the min value in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_sum( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the sum of elements along segments of an array. Segments are defined by an + integer array of segment IDs. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented sum operation. + For each segment, it computes the sum of values in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def vorbis_window( + window_length: Union[ivy.Array, ivy.NativeArray], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array that contains a vorbis power complementary window of size + window_length. + + Parameters + ---------- + window_length + the length of the vorbis window. + dtype + data type of the returned array. By default float32. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Input array with the vorbis window. + + Examples + -------- + >>> ivy.vorbis_window(3) + ivy.array([0.38268346, 1. , 0.38268352]) + + >>> ivy.vorbis_window(5) + ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) + """ + return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) + + +hamming_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "handle_device_shifting", + ), + "to_skip": (), +} +indices.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} diff --git a/ivy/functional/ivy/experimental/elementwise.py b/ivy/functional/ivy/experimental/elementwise.py index 613fe70acc6b2..e8c672460e7f1 100644 --- a/ivy/functional/ivy/experimental/elementwise.py +++ b/ivy/functional/ivy/experimental/elementwise.py @@ -17,227 +17,214 @@ from ivy.utils.exceptions import handle_exceptions +lerp.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} + + @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument @to_native_arrays_and_back -@handle_array_function -def lgamma( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def allclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> bool: """ - Compute the natural logarithm of the absolute value of the gamma function on x. + Return a True if the two arrays are element-wise equal within given tolerance; + otherwise False. + + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(x2)) and the absolute difference + atol are added together to compare against the absolute difference + between x1 and x2. + The default atol is not appropriate for comparing numbers that are + much smaller than one Parameters ---------- - x - input array. Should have a floating-point data type. + x1 + First input array. + x2 + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in x1 will be + considered equal to NaN's in x2 in the output array. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - an array containing the natural log of Gamma(x) of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. + Returns True if the two arrays are equal within the given tolerance; + False otherwise. Examples -------- - >>> x = ivy.array([1.6, 2.6, 3.5]) - >>> y = x.lgamma() + >>> x1 = ivy.array([1e10, 1e-7]) + >>> x2 = ivy.array([1.00001e10, 1e-8]) + >>> y = ivy.allclose(x1, x2) >>> print(y) - ivy.array([-0.11259177, 0.3574118 , 1.20097363]) + ivy.array(False) - >>> x = ivy.array([1., 2., 3. ]) - >>> y = x.lgamma() + >>> x1 = ivy.array([1.0, ivy.nan]) + >>> x2 = ivy.array([1.0, ivy.nan]) + >>> y = ivy.allclose(x1, x2, equal_nan=True) >>> print(y) - ivy.array([0. ,0. ,0.69314718]) + ivy.array(True) - >>> x = ivy.array([4.5, -4, -5.6]) - >>> x.lgamma(out = x) - >>> print(x) - ivy.array([2.45373654, inf, -4.6477685 ]) + >>> x1 = ivy.array([1e-10, 1e-10]) + >>> x2 = ivy.array([1.00001e-10, 1e-10]) + >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) + >>> print(y) + ivy.array(True) """ - return ivy.current_backend(x).lgamma(x, out=out) + return ivy.current_backend().allclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def sinc( +@inputs_to_ivy_arrays +def binarizer( x: Union[ivy.Array, ivy.NativeArray], /, *, + threshold: float = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate an implementation-dependent approximation of the principal value of the - normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain - ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element - ``x_i`` is assumed to be expressed in radians. - - **Special cases** - - For floating-point operands, - - - If x_i is NaN, the result is NaN. - - If ``x_i`` is ``0``, the result is ``1``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + Map the values of the input tensor to either 0 or 1, element-wise, based on the + outcome of a comparison against a threshold value. Parameters ---------- x - input array. Should have a floating-point data type. + Data to be binarized + threshold + Values greater than this are + mapped to 1, others to 0. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - an array containing the normalized sinc function of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = x.sinc() - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - >>> x = ivy.array([1.5, 0.5, -1.5]) - >>> y = ivy.zeros(3) - >>> ivy.sinc(x, out=y) - >>> print(y) - ivy.array([-0.212,0.637,-0.212]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = ivy.sinc(x) - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), - ... b=ivy.array([3.5, 4.5, 5.5])) - >>> y = x.sinc() - >>> print(y) - { - a: ivy.array([0.637,-0.212,0.127]), - b: ivy.array([-0.0909,0.0707,-0.0579]) - } + Binarized output data """ - return ivy.current_backend(x).sinc(x, out=out) + xc = ivy.copy_array(x, out=out) + if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": + xc = ivy.astype(xc, ivy.default_float_dtype()) + if ivy.is_complex_dtype(xc): + xc = ivy.abs(xc) + ret = ivy.where(xc > threshold, 1, 0) + return ret +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def fmax( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def conj( + x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the - case where one of the elements is NaN. ivy.maximum returns the NaN element while - ivy.fmax returns the non-NaN element. + Return the complex conjugate for each element ``x_i`` of the input array ``x``. - Parameters - ---------- - x1 - First input array. - x2 - Second input array. - out - optional output array, for writing the result to. + For complex number of the form - Returns - ------- - ret - Array with element-wise maximums. + .. math:: + a + bj - Examples - -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmax(x1, x2) - ivy.array([ 2., 5., 4.]) + the complex conjugate is defined as - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmax(x1, x2) - ivy.array([ 0., 0., nan]) - """ - return ivy.current_backend().fmax(x1, x2, out=out) + .. math:: + a - bj + Hence, the returned conjugates must be computed by negating + the imaginary component of each element ``x_i`` -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def float_power( - x1: Union[ivy.Array, float, list, tuple], - x2: Union[ivy.Array, float, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must - be broadcastable to the same shape. This differs from the power function in that - integers, float16, and float32 are promoted to floats with a minimum precision of - float64 so that the result is always inexact. + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Parameters ---------- - x1 - Array-like with elements to raise in power. - x2 - Array-like of exponents. If x1.shape != x2.shape, - they must be broadcastable to a common shape - (which becomes the shape of the output). + x + input array. out optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - The bases in x1 raised to the exponents in x2. - This is a scalar if both x1 and x2 are scalars + an arrray of the same dtype as the input array with + the complex conjugates of the complex values present + in the input array. If x is a scalar then a scalar + will be returned. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. + Examples -------- - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> ivy.float_power(x1, 3) - ivy.array([1., 8., 27., 64., 125.]) - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> x2 = ivy.array([2, 3, 3, 2, 1]) - >>> ivy.float_power(x1, x2) - ivy.array([1., 8., 27., 16., 5.]) + With :class:`ivy.Array` inputs: + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.conj(x) + >>> print(z) + ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) + + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), + ... b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.conj(x) + >>> print(z) + { + a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), + b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) + } """ - return ivy.current_backend().float_power(x1, x2, out=out) + return ivy.current_backend(x).conj(x, out=out) @handle_exceptions @@ -310,239 +297,40 @@ def count_nonzero( ---------- a array for which to count non-zeros. - axis - optional axis or tuple of axes along which to count non-zeros. Default is - None, meaning that non-zeros will be counted along a flattened - version of the input array. - keepdims - optional, if this is set to True, the axes that are counted are left in the - result as dimensions with size one. With this option, the result - will broadcast correctly against the input array. - dtype - optional output dtype. Default is of type integer. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Number of non-zero values in the array along a given axis. Otherwise, - the total number of non-zero values in the array is returned. - - Examples - -------- - >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) - >>> ivy.count_nonzero(a) - ivy.array(7) - >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) - >>> ivy.count_nonzero(a, axis=0) - ivy.array([1, 2, 2, 2]) - >>> a = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) - >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) - ivy.array([[[3, 4]]]) - """ - return ivy.current_backend().count_nonzero( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nansum( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as - zero. - - Parameters - ---------- - x - Input array. - axis - Axis or axes along which the sum is computed. - The default is to compute the sum of the flattened array. - dtype - The type of the returned array and of the accumulator in - which the elements are summed. By default, the dtype of input is used. - keepdims - If this is set to True, the axes which are reduced are left - in the result as dimensions with size one. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - A new array holding the result is returned unless out is specified, - in which it is returned. - - Examples - -------- - >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) - >>> ivy.nansum(a) - 10.0 - >>> ivy.nansum(a, axis=0) - ivy.array([2.1, 5.8, 2.1]) - >>> ivy.nansum(a, axis=1) - ivy.array([5.5, 4.5]) - """ - return ivy.current_backend().nansum( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def isclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a boolean array where two arrays are element-wise equal within a tolerance. - - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(b)) and the absolute difference - atol are added together to compare against the absolute difference - between a and b. - The default atol is not appropriate for comparing numbers that are - much smaller than one - - Parameters - ---------- - a - First input array. - b - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in a will be - considered equal to NaN's in b in the output array. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - Returns a boolean array of where a and b are equal within the given - tolerance. If both a and b are scalars, returns a single boolean value. - - Examples - -------- - >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) - ivy.array([True, False]) - >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) - ivy.array([True, True]) - >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) - ivy.array([False, False]) - >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) - ivy.array([False, True]) - """ - return ivy.current_backend().isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def signbit( - x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return element-wise True where signbit is set (less than zero). - - Parameters - ---------- - x - Array-like input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Output array, or reference to out if that was supplied. - This is a scalar if x is a scalar. - - Examples - -------- - >>> x = ivy.array([1, -2, 3]) - >>> ivy.signbit(x) - ivy.array([False, True, False]) - """ - return ivy.current_backend(x).signbit(x, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def hypot( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Return the hypotenuse given the two sides of a right angle triangle. - - Parameters - ---------- - x1 - The first input array - x2 - The second input array + axis + optional axis or tuple of axes along which to count non-zeros. Default is + None, meaning that non-zeros will be counted along a flattened + version of the input array. + keepdims + optional, if this is set to True, the axes that are counted are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + dtype + optional output dtype. Default is of type integer. + out + optional output array, for writing the result to. Returns ------- ret - An array with the hypotenuse + Number of non-zero values in the array along a given axis. Otherwise, + the total number of non-zero values in the array is returned. Examples -------- - >>> a = ivy.array([3.0, 4.0, 5.0]) - >>> b = ivy.array([4.0, 5.0, 6.0]) - >>> ivy.hypot(a, b) - ivy.array([5.0, 6.4031, 7.8102]) + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a) + ivy.array(7) + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a, axis=0) + ivy.array([1, 2, 2, 2]) + >>> a = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) + ivy.array([[[3, 4]]]) """ - return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) + return ivy.current_backend().count_nonzero( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) @handle_backend_invalid @@ -602,45 +390,26 @@ def diff( @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def allclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def digamma( + x: Union[ivy.Array, ivy.NativeArray], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> bool: +) -> ivy.Array: """ - Return a True if the two arrays are element-wise equal within given tolerance; - otherwise False. + Compute the logarithmic derivative of the gamma function at x. - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(x2)) and the absolute difference - atol are added together to compare against the absolute difference - between x1 and x2. - The default atol is not appropriate for comparing numbers that are - much smaller than one + Note + ---- + The Ivy version only accepts real-valued inputs. Parameters ---------- - x1 - First input array. - x2 - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in x1 will be - considered equal to NaN's in x2 in the output array. + x + Input array. out Alternate output array in which to place the result. The default is None. @@ -648,32 +417,16 @@ def allclose( Returns ------- ret - Returns True if the two arrays are equal within the given tolerance; - False otherwise. + Array with values computed from digamma function from + input arrays' values, element-wise. Examples -------- - >>> x1 = ivy.array([1e10, 1e-7]) - >>> x2 = ivy.array([1.00001e10, 1e-8]) - >>> y = ivy.allclose(x1, x2) - >>> print(y) - ivy.array(False) - - >>> x1 = ivy.array([1.0, ivy.nan]) - >>> x2 = ivy.array([1.0, ivy.nan]) - >>> y = ivy.allclose(x1, x2, equal_nan=True) - >>> print(y) - ivy.array(True) - - >>> x1 = ivy.array([1e-10, 1e-10]) - >>> x2 = ivy.array([1.00001e-10, 1e-10]) - >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) - >>> print(y) - ivy.array(True) + >>> x = ivy.array([.9, 3, 3.2]) + >>> y = ivy.digamma(x) + ivy.array([-0.7549271 0.92278427 0.9988394]) """ - return ivy.current_backend().allclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) + return ivy.current_backend(x).digamma(x, out=out) @handle_backend_invalid @@ -714,22 +467,70 @@ def fix( return ivy.current_backend(x).fix(x, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def nextafter( +def float_power( + x1: Union[ivy.Array, float, list, tuple], + x2: Union[ivy.Array, float, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must + be broadcastable to the same shape. This differs from the power function in that + integers, float16, and float32 are promoted to floats with a minimum precision of + float64 so that the result is always inexact. + + Parameters + ---------- + x1 + Array-like with elements to raise in power. + x2 + Array-like of exponents. If x1.shape != x2.shape, + they must be broadcastable to a common shape + (which becomes the shape of the output). + out + optional output array, for writing the result to. + + Returns + ------- + ret + The bases in x1 raised to the exponents in x2. + This is a scalar if both x1 and x2 are scalars + + Examples + -------- + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> ivy.float_power(x1, 3) + ivy.array([1., 8., 27., 64., 125.]) + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> x2 = ivy.array([2, 3, 3, 2, 1]) + >>> ivy.float_power(x1, x2) + ivy.array([1., 8., 27., 16., 5.]) + """ + return ivy.current_backend().float_power(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fmax( x1: Union[ivy.Array, ivy.NativeArray], x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> bool: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Return the next floating-point value after x1 towards x2, element-wise. + Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the + case where one of the elements is NaN. ivy.maximum returns the NaN element while + ivy.fmax returns the non-NaN element. Parameters ---------- @@ -738,22 +539,26 @@ def nextafter( x2 Second input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. Returns ------- ret - The next representable values of x1 in the direction of x2. + Array with element-wise maximums. Examples -------- - >>> x1 = ivy.array([1.0e-50, 2.0e+50]) - >>> x2 = ivy.array([2.0, 1.0]) - >>> ivy.nextafter(x1, x2) - ivy.array([1.4013e-45., 3.4028e+38]) + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmax(x1, x2) + ivy.array([ 2., 5., 4.]) + + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmax(x1, x2) + ivy.array([ 0., 0., nan]) """ - return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) + return ivy.current_backend().fmax(x1, x2, out=out) @handle_exceptions @@ -763,41 +568,35 @@ def nextafter( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def zeta( +def frexp( x: Union[ivy.Array, ivy.NativeArray], - q: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> bool: + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Compute the Hurwitz zeta function elementwisely with each pair of floats in two - arrays. + Decompose the elements of x into mantissa and twos exponent. Parameters ---------- x - First input array. - q - Second input array, must have the same shape as the first input array + Input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - Array with values computed from zeta function from - input arrays' values. + A tuple of two arrays, the mantissa and the twos exponent. Examples -------- - >>> x = ivy.array([5.0, 3.0]) - >>> q = ivy.array([2.0, 2.0]) - >>> ivy.zeta(x, q) - ivy.array([0.0369, 0.2021]) + >>> x = ivy.array([1, 2, 3]) + >>> ivy.frexp(x) + (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) """ - return ivy.current_backend(x, q).zeta(x, q, out=out) + return ivy.current_backend(x).frexp(x, out=out) @handle_exceptions @@ -880,91 +679,41 @@ def gradient( ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def xlogy( - x: Union[ivy.Array, ivy.NativeArray], - y: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> bool: - """ - Compute x*log(y) element-wise so that the result is 0 if x = 0. - - Parameters - ---------- - x - First input array. - y - Second input array. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - The next representable values of x1 in the direction of x2. - - Examples - -------- - >>> x = ivy.zeros(3) - >>> y = ivy.array([-1.0, 0.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([0.0, 0.0, 0.0]) - - >>> x = ivy.array([1.0, 2.0, 3.0]) - >>> y = ivy.array([3.0, 2.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([1.0986, 1.3863, 0.0000]) - """ - return ivy.current_backend(x, y).xlogy(x, y, out=out) - - -@handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays -def binarizer( - x: Union[ivy.Array, ivy.NativeArray], +@to_native_arrays_and_back +@handle_device_shifting +def hypot( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - threshold: float = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Map the values of the input tensor to either 0 or 1, element-wise, based on the - outcome of a comparison against a threshold value. + Return the hypotenuse given the two sides of a right angle triangle. Parameters ---------- - x - Data to be binarized - threshold - Values greater than this are - mapped to 1, others to 0. - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + x1 + The first input array + x2 + The second input array Returns ------- ret - Binarized output data + An array with the hypotenuse + + Examples + -------- + >>> a = ivy.array([3.0, 4.0, 5.0]) + >>> b = ivy.array([4.0, 5.0, 6.0]) + >>> ivy.hypot(a, b) + ivy.array([5.0, 6.4031, 7.8102]) """ - xc = ivy.copy_array(x, out=out) - if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": - xc = ivy.astype(xc, ivy.default_float_dtype()) - if ivy.is_complex_dtype(xc): - xc = ivy.abs(xc) - ret = ivy.where(xc > threshold, 1, 0) - return ret + return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) @handle_exceptions @@ -974,80 +723,63 @@ def binarizer( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def conj( - x: Union[ivy.Array, ivy.NativeArray], +def isclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the complex conjugate for each element ``x_i`` of the input array ``x``. - - For complex number of the form - - .. math:: - a + bj - - the complex conjugate is defined as - - .. math:: - a - bj - - Hence, the returned conjugates must be computed by negating - the imaginary component of each element ``x_i`` - - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the - `docstring `_ - in the standard. + Return a boolean array where two arrays are element-wise equal within a tolerance. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(b)) and the absolute difference + atol are added together to compare against the absolute difference + between a and b. + The default atol is not appropriate for comparing numbers that are + much smaller than one Parameters ---------- - x - input array. + a + First input array. + b + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in a will be + considered equal to NaN's in b in the output array. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - an arrray of the same dtype as the input array with - the complex conjugates of the complex values present - in the input array. If x is a scalar then a scalar - will be returned. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. - + Returns a boolean array of where a and b are equal within the given + tolerance. If both a and b are scalars, returns a single boolean value. Examples -------- - With :class:`ivy.Array` inputs: - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.conj(x) - >>> print(z) - ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) - - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), - ... b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.conj(x) - >>> print(z) - { - a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), - b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) - } + >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) + ivy.array([True, False]) + >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) + ivy.array([True, True]) + >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) + ivy.array([False, False]) + >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) + ivy.array([False, True]) """ - return ivy.current_backend(x).conj(x, out=out) + return ivy.current_backend().isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) @handle_exceptions @@ -1211,15 +943,151 @@ def lerp( return ivy.add(input, ivy.multiply(weight, ivy.subtract(end, input)), out=out) -lerp.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def lgamma( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the natural logarithm of the absolute value of the gamma function on x. + + Parameters + ---------- + x + input array. Should have a floating-point data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the natural log of Gamma(x) of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. + + Examples + -------- + >>> x = ivy.array([1.6, 2.6, 3.5]) + >>> y = x.lgamma() + >>> print(y) + ivy.array([-0.11259177, 0.3574118 , 1.20097363]) + + >>> x = ivy.array([1., 2., 3. ]) + >>> y = x.lgamma() + >>> print(y) + ivy.array([0. ,0. ,0.69314718]) + + >>> x = ivy.array([4.5, -4, -5.6]) + >>> x.lgamma(out = x) + >>> print(x) + ivy.array([2.45373654, inf, -4.6477685 ]) + """ + return ivy.current_backend(x).lgamma(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def modf( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: + """ + Decompose the elements of x into fractional and integral parts. + + Parameters + ---------- + x + Input array. + out + Optional output array for writing the result to. + It must have a shape that the inputs broadcast to. + + Returns + ------- + ret + A tuple of two arrays, the fractional and integral parts. + + Examples + -------- + >>> x = ivy.array([1.5, 2.7, 3.9]) + >>> ivy.modf(x) + (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) + """ + return ivy.current_backend(x).modf(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def nansum( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as + zero. + + Parameters + ---------- + x + Input array. + axis + Axis or axes along which the sum is computed. + The default is to compute the sum of the flattened array. + dtype + The type of the returned array and of the accumulator in + which the elements are summed. By default, the dtype of input is used. + keepdims + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + A new array holding the result is returned unless out is specified, + in which it is returned. + + Examples + -------- + >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) + >>> ivy.nansum(a) + 10.0 + >>> ivy.nansum(a, axis=0) + ivy.array([2.1, 5.8, 2.1]) + >>> ivy.nansum(a, axis=1) + ivy.array([5.5, 4.5]) + """ + return ivy.current_backend().nansum( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out + ) @handle_exceptions @@ -1229,112 +1097,154 @@ def lerp( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def frexp( - x: Union[ivy.Array, ivy.NativeArray], +def nextafter( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> bool: """ - Decompose the elements of x into mantissa and twos exponent. + Return the next floating-point value after x1 towards x2, element-wise. Parameters ---------- - x - Input array. + x1 + First input array. + x2 + Second input array. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - A tuple of two arrays, the mantissa and the twos exponent. + The next representable values of x1 in the direction of x2. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.frexp(x) - (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) + >>> x1 = ivy.array([1.0e-50, 2.0e+50]) + >>> x2 = ivy.array([2.0, 1.0]) + >>> ivy.nextafter(x1, x2) + ivy.array([1.4013e-45., 3.4028e+38]) """ - return ivy.current_backend(x).frexp(x, out=out) + return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -def modf( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def signbit( + x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], /, *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Decompose the elements of x into fractional and integral parts. + Return element-wise True where signbit is set (less than zero). Parameters ---------- x - Input array. + Array-like input. out - Optional output array for writing the result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - A tuple of two arrays, the fractional and integral parts. + Output array, or reference to out if that was supplied. + This is a scalar if x is a scalar. Examples -------- - >>> x = ivy.array([1.5, 2.7, 3.9]) - >>> ivy.modf(x) - (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) + >>> x = ivy.array([1, -2, 3]) + >>> ivy.signbit(x) + ivy.array([False, True, False]) """ - return ivy.current_backend(x).modf(x, out=out) + return ivy.current_backend(x).signbit(x, out=out) @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -def digamma( +@handle_device_shifting +def sinc( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the logarithmic derivative of the gamma function at x. + Calculate an implementation-dependent approximation of the principal value of the + normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain + ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element + ``x_i`` is assumed to be expressed in radians. - Note - ---- - The Ivy version only accepts real-valued inputs. + **Special cases** + + For floating-point operands, + + - If x_i is NaN, the result is NaN. + - If ``x_i`` is ``0``, the result is ``1``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. Parameters ---------- x - Input array. + input array. Should have a floating-point data type. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Array with values computed from digamma function from - input arrays' values, element-wise. + an array containing the normalized sinc function of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. Examples -------- - >>> x = ivy.array([.9, 3, 3.2]) - >>> y = ivy.digamma(x) - ivy.array([-0.7549271 0.92278427 0.9988394]) + With :class:`ivy.Array` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = x.sinc() + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + >>> x = ivy.array([1.5, 0.5, -1.5]) + >>> y = ivy.zeros(3) + >>> ivy.sinc(x, out=y) + >>> print(y) + ivy.array([-0.212,0.637,-0.212]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = ivy.sinc(x) + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), + ... b=ivy.array([3.5, 4.5, 5.5])) + >>> y = x.sinc() + >>> print(y) + { + a: ivy.array([0.637,-0.212,0.127]), + b: ivy.array([-0.0909,0.0707,-0.0579]) + } """ - return ivy.current_backend(x).digamma(x, out=out) + return ivy.current_backend(x).sinc(x, out=out) @handle_exceptions @@ -1388,3 +1298,93 @@ def sparsify_tensor( tensor = ivy.concat([ivy.zeros(len(x) - card, dtype=x.dtype), x[-card:]], axis=0) return ivy.reshape(tensor, _shape, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def xlogy( + x: Union[ivy.Array, ivy.NativeArray], + y: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> bool: + """ + Compute x*log(y) element-wise so that the result is 0 if x = 0. + + Parameters + ---------- + x + First input array. + y + Second input array. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + The next representable values of x1 in the direction of x2. + + Examples + -------- + >>> x = ivy.zeros(3) + >>> y = ivy.array([-1.0, 0.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([0.0, 0.0, 0.0]) + + >>> x = ivy.array([1.0, 2.0, 3.0]) + >>> y = ivy.array([3.0, 2.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([1.0986, 1.3863, 0.0000]) + """ + return ivy.current_backend(x, y).xlogy(x, y, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def zeta( + x: Union[ivy.Array, ivy.NativeArray], + q: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> bool: + """ + Compute the Hurwitz zeta function elementwisely with each pair of floats in two + arrays. + + Parameters + ---------- + x + First input array. + q + Second input array, must have the same shape as the first input array + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + Array with values computed from zeta function from + input arrays' values. + + Examples + -------- + >>> x = ivy.array([5.0, 3.0]) + >>> q = ivy.array([2.0, 2.0]) + >>> ivy.zeta(x, q) + ivy.array([0.0369, 0.2021]) + """ + return ivy.current_backend(x, q).zeta(x, q, out=out) diff --git a/ivy/functional/ivy/experimental/general.py b/ivy/functional/ivy/experimental/general.py index 0cb749b476ffa..faa8df1e69b2a 100644 --- a/ivy/functional/ivy/experimental/general.py +++ b/ivy/functional/ivy/experimental/general.py @@ -13,6 +13,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _correct_ivy_callable(func): # get the current backend of the given ivy callable if ivy.nested_any( @@ -25,6 +29,10 @@ def _correct_ivy_callable(func): return func +# --- Main --- # +# ------------ # + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py index 120ebe0279784..60cb61ec63724 100644 --- a/ivy/functional/ivy/experimental/layers.py +++ b/ivy/functional/ivy/experimental/layers.py @@ -21,1301 +21,1156 @@ from ivy.functional.ivy.experimental.general import _correct_ivy_callable from ivy.utils.exceptions import handle_exceptions -_min = builtins.min -_slice = builtins.slice -_max = builtins.max +# --- Helpers --- # +# --------------- # -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool1d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D max pool given 3-D input x. - Parameters - ---------- - x - Input image *[batch_size, w, d_in]* if data_format is "NWC". - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) - data_format - "NWC" or "NCW". Defaults to "NWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. +def _cast_init(init, dtype): + if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): + if ivy.is_float_dtype(dtype): + info = ivy.finfo(dtype) + else: + info = ivy.iinfo(dtype) + if "float64" not in str(dtype): + init = info.max if init > 0 else info.min + return ivy.array(init, dtype=dtype) - Returns - ------- - ret - The result of the pooling operation. - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. +def _compute_idx(in_size, out_size, device): + out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) + i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) + maxlength = in_size // out_size + 1 + in_size_mod = in_size % out_size + # adaptive = True iff there are kernels with different lengths + adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) + if adaptive: + maxlength += 1 + elif in_size_mod == 0: + maxlength -= 1 + range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) + idx = ivy.expand_dims(i0, axis=-1) + range_max + if adaptive: + maxval = ivy.full_like(idx, fill_value=in_size - 1) + idx = ivy.minimum(idx, maxval) + i1 = ivy.trunc_divide( + (out_range + 1) * in_size + out_size - 1, out_size + ).astype(ivy.int64) + length = i1 - i0 + else: + length = maxlength + return idx, length, range_max, adaptive - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 4., 5., 6., 7.]], +def _compute_weight_mat( + input_size, + output_size, + scale, + align_corners, + kernel_fn, + antialias: bool, + dim_scale_factor, +): + inv_scale = 1.0 / scale + kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 + if not align_corners: + sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 + x = ( + ivy.abs( + ivy.expand_dims(sample_f) + - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) + / kernel_scale + ) + else: + sample_f = ivy.arange(output_size) * dim_scale_factor + x = ivy.abs( + ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) / (kernel_scale) + weights = kernel_fn(x) + total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) + weights = ivy.where( + ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), + ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), + 0, + ) + input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 + return ivy.where( + ivy.expand_dims( + ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) + ), + weights, + 0, + ) - [[16., 17., 18., 19.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa - ivy.array([[[ 1., 3.], - [ 5., 7.], - [ 9., 11.]], - [[13., 15.], - [17., 19.], - [21., 23.]]]) - """ - return ivy.current_backend(x).max_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): + def _pad(arr, pads, pad_value): + out = ivy.astype( + ivy.pad( + arr, + ivy.maximum(0, pads).to_list(), + mode="constant", + constant_values=ivy.to_scalar(pad_value), + ), + arr.dtype, + ) + slices = tuple( + _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) + for (lo, hi), dim in zip(pads, arr.shape) + ) + return out[slices] + if ( + _min(lhs.ndim, len(rhs_shape)) < 2 + or lhs.ndim != len(rhs_shape) + or lhs.shape[1] != rhs_shape[1] + ): + raise ValueError("Dimension mismatch") + if len(window_strides) != len(rhs_shape) - 2: + raise ValueError("Wrong number of strides for spatial dimensions") + if len(pads) != len(rhs_shape) - 2: + raise ValueError("Wrong number of pads for spatial dimensions") -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_unpool1d( - x: ivy.Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D max unpooling given the 1-D pooled input x and its indices. + lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) + in_shape = lhs.shape[2:] + filter_shape = rhs_shape[2:] + dim = len(filter_shape) - Parameters - ---------- - x - Pooled input image *[batch_size, w, d_in]*. - indices - Indices obtained from the corresponding max pooling operation. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - out - optional output array, for writing the result to. + out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() + view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] - Returns - ------- - ret - The result of the unpooling operation. + out_shape = [ + (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) + ] + view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] - Both the description and the type hints above assume an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + view = ivy.as_strided(lhs, view_shape, view_strides) - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') - >>> print(pool_result) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], + view_axes = list(range(view.ndim)) + sum_axes = view_axes[-dim - 1 :] + rhs_axes = [view.ndim] + sum_axes + out_axes = [0, view.ndim] + list(range(1, dim + 1)) - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') - >>> print(unpool_result) - ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], - [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], + return view, view_axes, rhs_axes, out_axes - [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], - [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) - """ - return ivy.current_backend(x).max_unpool1d( - x, indices, kernel, strides, padding, data_format=data_format, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool2d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 2-D max pool given 4-D input x. - - Parameters - ---------- - x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. - - Returns - ------- - ret - The result of the pooling operation. - - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 2., 3.], - [ 4., 5.], - [ 4., 5.]]], +def _cubic_kernel(x): + out = ((1.5 * x - 2.5) * x) * x + 1.0 + out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) + return ivy.where(x >= 2.0, 0.0, out) - [[[ 8., 9.], - [10., 11.], - [10., 11.]]]]) - - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[16., 17.]], - [[22., 23.]]], +def _dilate(operand, factors, fill_value): + outspace = list(operand.shape[:2]) + [ + shape + (factors[i] - 1) * (shape - 1) + for i, shape in enumerate(operand.shape[2:]) + ] + out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) + lhs_slices = tuple(_slice(None, None, step) for step in factors) + out[(_slice(None),) * 2 + lhs_slices] = operand + return out - [[[40., 41.]], +def _dim_scale_factor(input_size, output_size, align_corners, scales): + if align_corners: + if output_size > 1: + dim_scale_factor = (input_size - 1) / (output_size - 1) + else: + dim_scale_factor = 0.0 + else: + dim_scale_factor = ( + input_size / (input_size * scales) + if scales is not None + else input_size / output_size + ) + return dim_scale_factor - [[46., 47.]]]]) - """ - return ivy.current_backend(x).max_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _expand_to_dim(x, dim): + for _ in range(dim - len(x.shape)): + x = ivy.expand_dims(x, axis=-1) + return x -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool3d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 3-D max pool given 5-D input x. - Parameters - ---------- - x - Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. +def _get_identity(func, dtype, init): + func_name = func.__name__ + if func_name in identities: + identity = identities[func_name] + return _cast_init(identity, dtype) + return init - Returns - ------- - ret - The result of the pooling operation. - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. +def _get_size(scale_factor, size, dims, x_shape): + if scale_factor is not None: + if isinstance(scale_factor, (float, int)): + scale_factor = [scale_factor] * dims + elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: + scale_factor = [scale_factor[0]] * dims - Examples - -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) - ivy.array([[[[[14., 15.]]]], + size = tuple( + [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] + ) + else: + size = (size,) * dims if isinstance(size, int) else tuple(size) + return size +def _interpolate_with_kernel( + x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode +): + spatial_dims = [2 + i for i in range(dims)] + equation = generate_einsum_equation(dims) + kernel_func = get_interpolate_kernel(mode) + output_shape = tuple(input_shape[:2]) + size + operands = [] + for i, d in enumerate(spatial_dims): + m = input_shape[d] + n = output_shape[d] + dim_scale_factor = _dim_scale_factor( + m, + n, + align_corners, + scale_factor[i] if scale_factor is not None else None, + ) + w = _compute_weight_mat( + m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor + ).astype(x.dtype) + operands.append(w) + return ivy.einsum(equation, x, *operands) - [[[[38., 39.]]]]]) - >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) - ivy.array([[[[[14., 15.]]], +def _lanczos_kernel(radius, x): + y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) + out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) + return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) - [[[22., 23.]]]], +def _mask(vals, length, range_max, dim, mask_value=0.0): + if isinstance(length, int): + return vals, length + else: + assert dim < 0 + mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) + if dim == -2: + mask = _expand_to_dim(mask, 4) + vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) + length = _expand_to_dim(length, -dim) + return vals, length - [[[[38., 39.]]], +def _mitchellcubic_kernel(x): + absx = abs(x) + if absx < 1: + return (7 * absx**3 - 12 * absx**2 + 6) / 6 + elif absx < 2: + return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 + else: + return 0 - [[[46., 47.]]]]]) - """ - return ivy.current_backend(x).max_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _output_ceil_shape(w, f, p, s): + return math.ceil((w - f + p) / s) + 1 -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool1d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - division_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D avg pool given 3-D input x. +def _padding_ceil_mode(w, f, p, s, return_added_padding=False): + remaining_pixels = (w - f + sum(p)) % s + added_padding = 0 + if s > 1 and remaining_pixels != 0 and f > 1: + input_size = w + sum(p) + # making sure that the remaining pixels are supposed + # to be covered by the window + # they won't be covered if stride is big enough to skip them + if input_size - remaining_pixels - (f - 1) + s > input_size: + return p + output_shape = _output_ceil_shape( + w, + f, + sum(p), + s, + ) + # calculating new padding with ceil_output_shape + new_pad = (output_shape - 1) * s + f - w + # updating pad_list with new padding by adding it to the end + added_padding = new_pad - sum(p) + p = ( + p[0], + p[1] + added_padding, + ) + if return_added_padding: + return p, added_padding + return p - Parameters - ---------- - x - Input image *[batch_size, w, d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - division_override - If specified, it will be used as the divisor, - otherwise kernel_size will be used. - out - optional output array, for writing the result to. - Returns - ------- - ret - The result of the pooling operation. +def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): + if padding.upper() == "SAME": + out_shape = [ + math.ceil(in_size / stride) + for in_size, stride in zip(in_shape, window_strides) + ] + pad_sizes = [ + _max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size in zip( + out_shape, window_strides, filter_shape, in_shape + ) + ] + return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] + else: + return [(0, 0)] * len(in_shape) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 2., 3., 4., 5.], - [ 8., 9., 10., 11.]], +def _sum_tensors(ts): + return _reduce(ivy.add, ts) - [[14., 15., 16., 17.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 2., 3., 4., 5.]], - [[14., 15., 16., 17.]]]) - """ - return ivy.current_backend(x).avg_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - division_override=division_override, - out=out, +def _tf_area_dim_scale(index, starting_index, scale, ending_index): + if index < starting_index: + dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index + else: + dim_scale = ending_index - index if index + 1 > ending_index else 1.0 + return dim_scale + + +def _tf_area_indices(dim_index, scale): + starting_index = dim_index * scale + ending_index = (dim_index + 1) * scale + rounded_indices = ( + int(starting_index), + math.ceil(ending_index), ) + return starting_index, ending_index, rounded_indices -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool2d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 2-D average pool given 4-D input x. +def _tf_area_interpolate(x, size, dims): + ret = ivy.zeros((x.shape[:2] + size)) + scale = ivy.divide(ivy.shape(x)[2:], size) + area = 1.0 / ivy.prod(scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) + sum_data = ivy.zeros( + ( + d_index[1] - d_index[0], + h_index[1] - h_index[0], + w_index[1] - w_index[0], + ) + ) + for d_ind in range(d_index[0], d_index[1]): + scale_z = _tf_area_dim_scale( + d_ind, d_in, scale[0], d_in1 + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale( + h_ind, h_in, scale[1], h_in1 + ) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[2], w_in1 + ) + sum_data[ + d_ind - d_index[0], + h_ind - h_index[0], + w_ind - w_index[0], + ] = ( + ivy.array(ch[d_ind, h_ind, w_ind]) + * scale_x + * scale_y + * scale_z + * area + ) + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) + sum_data = ivy.zeros( + (h_index[1] - h_index[0], w_index[1] - w_index[0]) + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[1], w_in1 + ) + sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( + ivy.array(ch[h_ind, w_ind]) + * scale_x + * scale_y + * area + ) + ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) + else: + for w_dim in range(size[0]): + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) + sum_data = ivy.zeros((w_index[1] - w_index[0],)) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) + sum_data[w_ind - w_index[0]] = ( + ivy.array(ch[w_ind]) * scale_x * area + ) + ret[i, j, w_dim] = ivy.sum(sum_data) + return ret - Parameters - ---------- - x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimensio paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - out - optional output array, for writing the result to. - Returns - ------- - ret - The result of the pooling operation. +def _triangle_kernel(x): + return ivy.maximum(0, 1 - ivy.abs(x)) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], +def _upsample_bicubic2d_default( + a, + output_size, + align_corners, + scale_h=None, + scale_w=None, +): + N, C, iH, iW = a.shape + oH, oW = output_size - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[ 8., 9.]], + def compute_scale(in_size, out_size, align_corners, scale=None): + if align_corners: + return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 + else: + return 1 / scale if scale is not None and scale > 0 else in_size / out_size - [[14., 15.]]], + def compute_source_index(scale, dst_index, align_corners): + if align_corners: + return scale * dst_index + else: + return scale * (dst_index + 0.5) - 0.5 + height_scale = compute_scale(iH, oH, align_corners, scale_h) + width_scale = compute_scale(iW, oW, align_corners, scale_w) - [[[32., 33.]], + N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) + C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) + out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) + out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) - [[38., 39.]]]]) - """ - return ivy.current_backend(x).avg_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, - ) + real_x = compute_source_index(width_scale, out_x, align_corners) + in_x = ivy.floor(real_x) + t_x = real_x - in_x + ix = ivy.astype(in_x, ivy.int64) + real_y = compute_source_index(height_scale, out_y, align_corners) + in_y = ivy.floor(real_y) + t_y = real_y - in_y + iy = ivy.astype(in_y, ivy.int64) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool3d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 3-D avg pool given 5-D input x. + iys_ofs = (iy - 1, iy, iy + 1, iy + 2) + ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) - Parameters - ---------- - x - Input volume *[batch_size,d,h,w,d_in]*. - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension - paddings. - data_format - NDHWC" or "NCDHW". Defaults to "NDHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - divisor_override - If specified, it will be used as divisor, otherwise kernel_size will be used. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + def load_bounded(ys, xs): + y_idx = ivy.clip(ys, 0, iH - 1) + x_idx = ivy.clip(xs, 0, iW - 1) + return a[N_idx, C_idx, y_idx, x_idx] - Returns - ------- - ret - The result of the pooling operation. + def get_x_interp(y): + coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) + return _upsample_cubic_interp1d(coeffs_x, t_x) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) + result = _upsample_cubic_interp1d(coeffs_y, t_y) - Examples - -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.avg_pool3d(x,2,2,'VALID')) - ivy.array([[[[[ 7., 8.]]]], + return result +def _upsample_cubic_convolution1(x, A): + return ((A + 2) * x - (A + 3)) * x * x + 1 - [[[[31., 32.]]]]]) - >>> print(ivy.avg_pool3d(x,2,2,'SAME')) - ivy.array([[[[[ 7., 8.]]], +def _upsample_cubic_convolution2(x, A): + return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A - [[[19., 20.]]]], +def _upsample_cubic_interp1d(coeffs, ts): + coeffs2 = _upsample_get_cubic_coefficients(ts) + return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) - [[[[31., 32.]]], +def _upsample_get_cubic_coefficients(t): + A = -0.75 + return ( + _upsample_cubic_convolution2(t + 1.0, A), + _upsample_cubic_convolution1(t, A), + _upsample_cubic_convolution1(1.0 - t, A), + _upsample_cubic_convolution2(2.0 - t, A), + ) - [[[43., 44.]]]]]) - """ - return ivy.current_backend(x).avg_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, - ) +# --- Main --- # +# ------------ # -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def pool( - x: Union[ivy.Array, ivy.NativeArray], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, - *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, +@inputs_to_ivy_arrays +def adaptive_avg_pool1d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: int, ) -> ivy.Array: """ - Perform an N-D pooling operation. + Apply a 1D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - Input array to pool over. - window_shape - Shape of the pooling window. - pool_type - Type of pooling operation, either 'MAX' or 'AVG'. - strides - Strides of the pooling operation. - padding - Padding type, either 'VALID' or 'SAME'. - data_format - Data format of the input and output data, either 'NCHW' or 'NHWC'. - dilations - Dilation rate of the pooling operation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + input + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. Returns ------- - ret - The result of the pooling operation. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) - ivy.array([[[[ 8., 9.]], - [[14., 15.]]], - [[[32., 33.]], - [[38., 39.]]]]) + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` """ - return ivy.current_backend(x).pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ceil_mode=ceil_mode, - out=out, - ) + squeeze = False + if input.ndim == 2: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 3: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", + ) + if input.shape[-1] % output_size == 0: + stride = input.shape[-1] // output_size + kernel_size = input.shape[-1] - (output_size - 1) * stride + pooled_output = ivy.avg_pool1d( + input, kernel_size, stride, "VALID", data_format="NCW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output -@handle_exceptions -@handle_backend_invalid + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size, input.device + ) + + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., idxw]) + + if not adaptive_w: + ret = ivy.mean(vals, axis=-1) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret + + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + + ret = None + for i in range(vals.shape[-1]): + if ret is None: + ret = vals[..., i] + else: + ret = ret + vals[..., i] + pooled_output = ret / length_w.astype(ret.dtype) + + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output + + +@handle_exceptions @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def dct( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def adaptive_avg_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], +) -> ivy.Array: """ - Compute the 1D Discrete Cosine Tranformation of a given signal. + Apply a 2D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - The input signal. - type - The type of the dct. Must be 1, 2, 3 or 4. - n - The lenght of the transform. If n is less than the input signal lenght, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the DCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - Array containing the transformed input. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.dct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, - 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], + if isinstance(output_size, int): + output_size = (output_size, output_size) - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.avg_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.dct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.dct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) - With one :class:`ivy.Container` input: + if not adaptive_h and not adaptive_w: + ret = ivy.mean(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.dct(x, type=3, n=None, norm='ortho') - >>> print(y) - { - a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475, 5.19664812, -1.95411837]), - b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, - -1.25980937, 0.64958102, -0.2442648]) - } + vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) - With multiple :class:`ivy.Container` inputs: + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ret + vals[..., i, :, j] + pooled_output = ret / (length_h * length_w).astype(vals.dtype) - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) - >>> print(y) - { - a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, - -26.00041008, 19.75149155, -16.97056389, 10.87819386, - -5.89381361]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } - """ - return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output -@handle_exceptions @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def idct( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +@inputs_to_ivy_arrays +def adaptive_max_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], +): """ - Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. + Apply a 2D adaptive maximum pooling over an input signal composed of several input + planes. Parameters ---------- - x - The input signal. - type - The type of the idct. Must be 1, 2, 3 or 4. - n - The length of the transform. If n is less than the input signal length, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the IDCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - Array containing the transformed input. + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + if isinstance(output_size, int): + output_size = (output_size, output_size) - Examples - -------- - With :class:`ivy.Array` input: + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.max_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.idct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475 , 5.19664812, -1.95411837]) - - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], - - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) - - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.idct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.idct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array( + input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device + ) - With one :class:`ivy.Container` input: + if not adaptive_h and not adaptive_w: + ret = ivy.max(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.idct(x, type=3, n=None, norm='ortho') - >>> print(y) - { - a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, - -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, - -8.80319249e-08, -4.05617893e-01]), - b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, - -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, - -1.10039906e-08, -5.07022366e-02]) - } + vals, length_h = _mask( + vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") + ) + vals, length_w = _mask( + vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") + ) - With multiple :class:`ivy.Container` inputs: + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ivy.maximum(ret, vals[..., i, :, j]) + pooled_output = ret.astype(vals.dtype) - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) - >>> print(y) - { - a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, - -16.18951607, 18.06697273, -17.57439613, 11.68861485, - -4.41308832]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } - """ - return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output -idct.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} +def area_interpolate(x, dims, size, scale): + ret = ivy.zeros((x.shape[:2] + size)) + inv_scale = ivy.divide(1.0, scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_index = ( + int(d_dim * inv_scale[0]), + math.ceil((d_dim + 1) * inv_scale[0]), + ) + h_index = ( + int(h_dim * inv_scale[1]), + math.ceil((h_dim + 1) * inv_scale[1]), + ) + w_index = ( + int(w_dim * scale[2]), + math.ceil((w_dim + 1) * inv_scale[2]), + ) + scale_z = d_index[1] - d_index[0] + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_z * scale_y * scale_x + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( + ch[ + d_index[0] : d_index[1], + h_index[0] : h_index[1], + w_index[0] : w_index[1], + ] + ) * (1 / area) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_index = ( + int(h_dim * inv_scale[0]), + math.ceil((h_dim + 1) * inv_scale[0]), + ) + w_index = ( + int(w_dim * inv_scale[1]), + math.ceil((w_dim + 1) * inv_scale[1]), + ) + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_y * scale_x + ret[i, j, h_dim, w_dim] = ivy.sum( + ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] + ) * (1 / area) + else: + for w_dim in range(size[0]): + w_index = ( + int(w_dim * inv_scale[0]), + math.ceil((w_dim + 1) * inv_scale[0]), + ) + scale_x = w_index[1] - w_index[0] + ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( + 1 / scale_x + ) + return ret -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def fft( +def avg_pool1d( x: Union[ivy.Array, ivy.NativeArray], - dim: int, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + division_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. + """ + Compute a 1-D avg pool given 3-D input x. Parameters ---------- x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT. - dim - The dimension along which to take the one dimensional FFT. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing FFT. - Should be a integer greater than 1. + Input image *[batch_size, w, d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + division_override + If specified, it will be used as the divisor, + otherwise kernel_size will be used. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - The result of the FFT operation. + The result of the pooling operation. + + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, - 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, - 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, - 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, - 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, - 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, - 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, - 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, - 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, - 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 2., 3., 4., 5.], + [ 8., 9., 10., 11.]], + + [[14., 15., 16., 17.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 2., 3., 4., 5.]], + + [[14., 15., 16., 17.]]]) """ - return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) + return ivy.current_backend(x).avg_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + division_override=division_override, + out=out, + ) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout1d( +def avg_pool2d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, /, *, - training: bool = True, - data_format: str = "NWC", + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D - feature map. + Compute a 2-D average pool given 4-D input x. Parameters ---------- x - a 2D or 3D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout1d is performed during training or ignored - during testing. + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimensio paddings. data_format - "NWC" or "NCW". Defaults to "NWC". + NHWC" or "NCHW". Defaults to "NHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. out optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + The result of the pooling operation. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - ivy.array([[[2., 0, 2.]]]) - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[[1, 1, 1]]]) + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[ 8., 9.]], - With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), - ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - { - a: ivy.array([[[200., 400., 0.]]]), - b: ivy.array([[[0., 0., 0.]]]) - } + [[14., 15.]]], + + + [[[32., 33.]], + + [[38., 39.]]]]) """ - return ivy.current_backend(x).dropout1d( - x, prob, training=training, data_format=data_format, out=out + return ivy.current_backend(x).avg_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, ) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout2d( +def avg_pool3d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, /, *, - training: bool = True, - data_format: str = "NHWC", + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D - feature map. + Compute a 3-D avg pool given 5-D input x. Parameters ---------- x - a 3D or 4D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout2d is performed during training or ignored - during testing. + Input volume *[batch_size,d,h,w,d_in]*. + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension + paddings. data_format - "NHWC" or "NCHW". Defaults to "NHWC". + NDHWC" or "NCDHW". Defaults to "NDHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + divisor_override + If specified, it will be used as divisor, otherwise kernel_size will be used. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + The result of the pooling operation. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.avg_pool3d(x,2,2,'VALID')) + ivy.array([[[[[ 7., 8.]]]], - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 0.5) - >>> print(y) - ivy.array([[0., 2., 2.]]) - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[1, 1, 1]]) + + [[[[31., 32.]]]]]) + >>> print(ivy.avg_pool3d(x,2,2,'SAME')) + ivy.array([[[[[ 7., 8.]]], + + + [[[19., 20.]]]], + + + + [[[[31., 32.]]], + + + [[[43., 44.]]]]]) """ - return ivy.current_backend(x).dropout2d( - x, prob, training=training, data_format=data_format, out=out + return ivy.current_backend(x).avg_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout3d( +def dct( x: Union[ivy.Array, ivy.NativeArray], - prob: float, /, *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D - feature map. + Compute the 1D Discrete Cosine Tranformation of a given signal. Parameters ---------- x - a 4D or 5D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout3d is performed during training or ignored - during testing. - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". + The input signal. + type + The type of the dct. Must be 1, 2, 3 or 4. + n + The lenght of the transform. If n is less than the input signal lenght, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the DCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". out optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + Array containing the transformed input. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - """ - return ivy.current_backend(x).dropout3d( - x, prob, training=training, data_format=data_format, out=out - ) + Examples + -------- + With :class:`ivy.Array` input: -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def ifft( - x: Union[ivy.Array, ivy.NativeArray], - dim: int, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs IFFT. - dim - The dimension along which to take the one dimensional IFFT. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing IFFT. - Should be a integer greater than 1. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.dct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, + 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) - Returns - ------- - ret - The result of the IFFT operation. + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], - Examples - -------- - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, - 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, - 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, - 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, - 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, - 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, - 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, - 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, - 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, - 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) - """ - return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) + + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.dct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.dct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def embedding( - weights: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - max_norm: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Embeds a given tensor of indices using a given tensor of weights. + With one :class:`ivy.Container` input: - Parameters - ---------- - weights - The weights tensor. - indices - The indices tensor. - max_norm - The maximum norm of the embeddings. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.dct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475, 5.19664812, -1.95411837]), + b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, + -1.25980937, 0.64958102, -0.2442648]) + } - Returns - ------- - ret - The result of the embedding operation. + With multiple :class:`ivy.Container` inputs: - Examples - -------- - >>> weights = ivy.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) - >>> indices = ivy.array([0, 2]) - >>> print(ivy.embedding(weights, indices, max_norm=5)) - ivy.array([[1. , 2. , 3. ], - [2.51285338, 2.87183261, 3.2308116 ]]) + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, + -26.00041008, 19.75149155, -16.97056389, 10.87819386, + -5.89381361]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } """ - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return ivy.current_backend(indices).embedding( - weights, - indices, - max_norm=max_norm, - out=out, - ) + return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) @handle_exceptions @@ -1397,406 +1252,402 @@ def dft( @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays -def interp(x, xp, fp, left=None, right=None, period=None): - x_arr = ivy.array(x) - fix_later = False - if x_arr.shape == (): - x_arr = ivy.array([x]) - fix_later = True - x = ivy.astype(x_arr, "float64") - xp = ivy.astype(ivy.array(xp), "float64") - fp = ivy.astype(ivy.array(fp), "float64") - ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) - if period is not None: - ivy.utils.assertions.check_equal(period, 0, inverse=True) - period = ivy.abs(period) - x = ivy.remainder(x, period) - xp = ivy.remainder(xp, period) - asort_xp = ivy.argsort(xp) - xp = xp[asort_xp] - fp = fp[asort_xp] - xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) - fp = ivy.concat((fp[-1:], fp, fp[0:1])) - - def interp_inner(value): - value = ivy.array(value) - if value < xp[0]: - return left if left is not None else fp[0] - elif value > xp[-1]: - return right if right is not None else fp[-1] - else: - last = None - if xp.shape[0] < 3: - for i in range(xp.shape[0] - 1, -1, -1): - if xp[i] == value: - return fp[i] - elif xp[i] < value: - last = i - else: - first = 0 - last = xp.shape[0] - while first < last: - midpoint = (first + last) // 2 - if xp[midpoint] == value: - already_exists = ivy.argwhere(xp == value) - if already_exists.shape[0] > 0: - return fp[already_exists[-1][0]] - return fp[midpoint] - else: - if value < xp[midpoint]: - last = midpoint - 1 - else: - first = midpoint + 1 - dist = (value - xp[last]) / (xp[last + 1] - xp[last]) - return (fp[last + 1] - fp[last]) * dist + fp[last] - - ret = ivy.map(interp_inner, unique={"value": x}) - if fix_later: - return ivy.astype(ivy.array(ret[0]), "float64") - else: - return ivy.astype(ivy.array(ret), "float64") +@to_native_arrays_and_back +@handle_device_shifting +def dropout1d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D + feature map. + Parameters + ---------- + x + a 2D or 3D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout1d is performed during training or ignored + during testing. + data_format + "NWC" or "NCW". Defaults to "NWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. -def _tf_area_dim_scale(index, starting_index, scale, ending_index): - if index < starting_index: - dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index - else: - dim_scale = ending_index - index if index + 1 > ending_index else 1.0 - return dim_scale + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. -def _tf_area_indices(dim_index, scale): - starting_index = dim_index * scale - ending_index = (dim_index + 1) * scale - rounded_indices = ( - int(starting_index), - math.ceil(ending_index), - ) - return starting_index, ending_index, rounded_indices + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + ivy.array([[[2., 0, 2.]]]) -def _tf_area_interpolate(x, size, dims): - ret = ivy.zeros((x.shape[:2] + size)) - scale = ivy.divide(ivy.shape(x)[2:], size) - area = 1.0 / ivy.prod(scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) - sum_data = ivy.zeros( - ( - d_index[1] - d_index[0], - h_index[1] - h_index[0], - w_index[1] - w_index[0], - ) - ) - for d_ind in range(d_index[0], d_index[1]): - scale_z = _tf_area_dim_scale( - d_ind, d_in, scale[0], d_in1 - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale( - h_ind, h_in, scale[1], h_in1 - ) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[2], w_in1 - ) - sum_data[ - d_ind - d_index[0], - h_ind - h_index[0], - w_ind - w_index[0], - ] = ( - ivy.array(ch[d_ind, h_ind, w_ind]) - * scale_x - * scale_y - * scale_z - * area - ) - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) - sum_data = ivy.zeros( - (h_index[1] - h_index[0], w_index[1] - w_index[0]) - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[1], w_in1 - ) - sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( - ivy.array(ch[h_ind, w_ind]) - * scale_x - * scale_y - * area - ) - ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) - else: - for w_dim in range(size[0]): - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) - sum_data = ivy.zeros((w_index[1] - w_index[0],)) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) - sum_data[w_ind - w_index[0]] = ( - ivy.array(ch[w_ind]) * scale_x * area - ) - ret[i, j, w_dim] = ivy.sum(sum_data) - return ret + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") + >>> print(y) + ivy.array([[[1, 1, 1]]]) + With one :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), + ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + { + a: ivy.array([[[200., 400., 0.]]]), + b: ivy.array([[[0., 0., 0.]]]) + } + """ + return ivy.current_backend(x).dropout1d( + x, prob, training=training, data_format=data_format, out=out + ) -def nearest_interpolate(x, dims, size, input_shape, exact): - off = 0.5 if exact else 0 - for d in range(dims): - m = input_shape[d + 2] - n = size[d] - offsets = (ivy.arange(n, dtype="float32") + off) * m / n - offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") - x = ivy.gather(x, offsets, axis=d + 2) - return x +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dropout2d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D + feature map. -def _triangle_kernel(x): - return ivy.maximum(0, 1 - ivy.abs(x)) + Parameters + ---------- + x + a 3D or 4D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout2d is performed during training or ignored + during testing. + data_format + "NHWC" or "NCHW". Defaults to "NHWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). -def _cubic_kernel(x): - out = ((1.5 * x - 2.5) * x) * x + 1.0 - out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) - return ivy.where(x >= 2.0, 0.0, out) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + Examples + -------- + With :class:`ivy.Array` input: -def _lanczos_kernel(radius, x): - y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) - out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) - return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 0.5) + >>> print(y) + ivy.array([[0., 2., 2.]]) + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") + >>> print(y) + ivy.array([[1, 1, 1]]) + """ + return ivy.current_backend(x).dropout2d( + x, prob, training=training, data_format=data_format, out=out + ) -def _dim_scale_factor(input_size, output_size, align_corners, scales): - if align_corners: - if output_size > 1: - dim_scale_factor = (input_size - 1) / (output_size - 1) - else: - dim_scale_factor = 0.0 - else: - dim_scale_factor = ( - input_size / (input_size * scales) - if scales is not None - else input_size / output_size - ) - return dim_scale_factor +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dropout3d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D + feature map. -def _mitchellcubic_kernel(x): - absx = abs(x) - if absx < 1: - return (7 * absx**3 - 12 * absx**2 + 6) / 6 - elif absx < 2: - return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 - else: - return 0 + Parameters + ---------- + x + a 4D or 5D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout3d is performed during training or ignored + during testing. + data_format + "NDHWC" or "NCDHW". Defaults to "NDHWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). -def _compute_weight_mat( - input_size, - output_size, - scale, - align_corners, - kernel_fn, - antialias: bool, - dim_scale_factor, -): - inv_scale = 1.0 / scale - kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 - if not align_corners: - sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 - x = ( - ivy.abs( - ivy.expand_dims(sample_f) - - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) - / kernel_scale - ) - else: - sample_f = ivy.arange(output_size) * dim_scale_factor - x = ivy.abs( - ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) / (kernel_scale) - weights = kernel_fn(x) - total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) - weights = ivy.where( - ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), - ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), - 0, - ) - input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 - return ivy.where( - ivy.expand_dims( - ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) - ), - weights, - 0, + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + return ivy.current_backend(x).dropout3d( + x, prob, training=training, data_format=data_format, out=out ) -def _upsample_cubic_convolution1(x, A): - return ((A + 2) * x - (A + 3)) * x * x + 1 - +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def embedding( + weights: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + max_norm: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Embeds a given tensor of indices using a given tensor of weights. -def _upsample_cubic_convolution2(x, A): - return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A + Parameters + ---------- + weights + The weights tensor. + indices + The indices tensor. + max_norm + The maximum norm of the embeddings. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The result of the embedding operation. -def _upsample_get_cubic_coefficients(t): - A = -0.75 - return ( - _upsample_cubic_convolution2(t + 1.0, A), - _upsample_cubic_convolution1(t, A), - _upsample_cubic_convolution1(1.0 - t, A), - _upsample_cubic_convolution2(2.0 - t, A), + Examples + -------- + >>> weights = ivy.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) + >>> indices = ivy.array([0, 2]) + >>> print(ivy.embedding(weights, indices, max_norm=5)) + ivy.array([[1. , 2. , 3. ], + [2.51285338, 2.87183261, 3.2308116 ]]) + """ + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return ivy.current_backend(indices).embedding( + weights, + indices, + max_norm=max_norm, + out=out, ) -def _upsample_cubic_interp1d(coeffs, ts): - coeffs2 = _upsample_get_cubic_coefficients(ts) - return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fft( + x: Union[ivy.Array, ivy.NativeArray], + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT. + dim + The dimension along which to take the one dimensional FFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing FFT. + Should be a integer greater than 1. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. -def _sum_tensors(ts): - return _reduce(ivy.add, ts) + Returns + ------- + ret + The result of the FFT operation. + Examples + -------- + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, + 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, + 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, + 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, + 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, + 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, + 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, + 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, + 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, + 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) + """ + return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) -def _upsample_bicubic2d_default( - a, - output_size, - align_corners, - scale_h=None, - scale_w=None, -): - N, C, iH, iW = a.shape - oH, oW = output_size - def compute_scale(in_size, out_size, align_corners, scale=None): - if align_corners: - return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 - else: - return 1 / scale if scale is not None and scale > 0 else in_size / out_size +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def fft2( + x: Union[ivy.Array, ivy.NativeArray], + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the 2-dimensional discrete Fourier Transform. - def compute_source_index(scale, dst_index, align_corners): - if align_corners: - return scale * dst_index - else: - return scale * (dst_index + 0.5) - 0.5 + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT2. + s + sequence of ints, optional + Shape (length of each transformed axis) of the output (s[0] refers to axis 0, + s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, + if the given shape is smaller than that of the input, the input is cropped. + If it is larger, the input is padded with zeros. if s is not given, the shape + of the input along the axes specified by axes is used. + dim + Axes over which to compute the FFT2. If not given, the last two axes are used. + A repeated index in axes means the transform over that axis is performed + multiple times. A one-element sequence means that a one-dimensional FFT is + performed. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - height_scale = compute_scale(iH, oH, align_corners, scale_h) - width_scale = compute_scale(iW, oW, align_corners, scale_w) + Returns + ------- + ret + The result of the FFT2 operation. - N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) - C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) - out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) - out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) + Examples + -------- + >>> x = ivy.array([[0, 0, 0, 0, 0], + ... [1, 1, 1, 1, 1], + ... [2, 2, 2, 2, 2], + ... [3, 3, 3, 3, 3], + ... [4, 4, 4, 4, 4]]) + >>> y = ivy.fft2(x) + >>> print(y) + ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5+17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5-17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ]]) + """ + return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) - real_x = compute_source_index(width_scale, out_x, align_corners) - in_x = ivy.floor(real_x) - t_x = real_x - in_x - ix = ivy.astype(in_x, ivy.int64) - real_y = compute_source_index(height_scale, out_y, align_corners) - in_y = ivy.floor(real_y) - t_y = real_y - in_y - iy = ivy.astype(in_y, ivy.int64) - - iys_ofs = (iy - 1, iy, iy + 1, iy + 2) - ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) - - def load_bounded(ys, xs): - y_idx = ivy.clip(ys, 0, iH - 1) - x_idx = ivy.clip(xs, 0, iW - 1) - return a[N_idx, C_idx, y_idx, x_idx] - - def get_x_interp(y): - coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) - return _upsample_cubic_interp1d(coeffs_x, t_x) - - coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) - result = _upsample_cubic_interp1d(coeffs_y, t_y) - - return result - - -def area_interpolate(x, dims, size, scale): - ret = ivy.zeros((x.shape[:2] + size)) - inv_scale = ivy.divide(1.0, scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_index = ( - int(d_dim * inv_scale[0]), - math.ceil((d_dim + 1) * inv_scale[0]), - ) - h_index = ( - int(h_dim * inv_scale[1]), - math.ceil((h_dim + 1) * inv_scale[1]), - ) - w_index = ( - int(w_dim * scale[2]), - math.ceil((w_dim + 1) * inv_scale[2]), - ) - scale_z = d_index[1] - d_index[0] - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_z * scale_y * scale_x - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( - ch[ - d_index[0] : d_index[1], - h_index[0] : h_index[1], - w_index[0] : w_index[1], - ] - ) * (1 / area) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_index = ( - int(h_dim * inv_scale[0]), - math.ceil((h_dim + 1) * inv_scale[0]), - ) - w_index = ( - int(w_dim * inv_scale[1]), - math.ceil((w_dim + 1) * inv_scale[1]), - ) - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_y * scale_x - ret[i, j, h_dim, w_dim] = ivy.sum( - ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] - ) * (1 / area) - else: - for w_dim in range(size[0]): - w_index = ( - int(w_dim * inv_scale[0]), - math.ceil((w_dim + 1) * inv_scale[0]), - ) - scale_x = w_index[1] - w_index[0] - ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( - 1 / scale_x - ) - return ret +def generate_einsum_equation(dim): + alphabet = "abcdefghijklmnopqrstuvwxyz" + input_indices = alphabet[: dim + 2] + output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] + contraction_indices = ",".join([input_indices, *output_indices]) + output = input_indices[:2] + "".join([output[-1] for output in output_indices]) + einsum_string = contraction_indices + "->" + output + return einsum_string def get_interpolate_kernel(mode): @@ -1810,38 +1661,336 @@ def get_interpolate_kernel(mode): return kernel_func -def generate_einsum_equation(dim): - alphabet = "abcdefghijklmnopqrstuvwxyz" - input_indices = alphabet[: dim + 2] - output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] - contraction_indices = ",".join([input_indices, *output_indices]) - output = input_indices[:2] + "".join([output[-1] for output in output_indices]) - einsum_string = contraction_indices + "->" + output - return einsum_string +@handle_exceptions +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def idct( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. + Parameters + ---------- + x + The input signal. + type + The type of the idct. Must be 1, 2, 3 or 4. + n + The length of the transform. If n is less than the input signal length, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the IDCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". + out + optional output array, for writing the result to. -def _interpolate_with_kernel( - x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode -): - spatial_dims = [2 + i for i in range(dims)] - equation = generate_einsum_equation(dims) - kernel_func = get_interpolate_kernel(mode) - output_shape = tuple(input_shape[:2]) + size - operands = [] - for i, d in enumerate(spatial_dims): - m = input_shape[d] - n = output_shape[d] - dim_scale_factor = _dim_scale_factor( - m, - n, - align_corners, - scale_factor[i] if scale_factor is not None else None, - ) - w = _compute_weight_mat( - m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor - ).astype(x.dtype) - operands.append(w) - return ivy.einsum(equation, x, *operands) + Returns + ------- + ret + Array containing the transformed input. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.idct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475 , 5.19664812, -1.95411837]) + + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], + + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) + + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.idct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) + + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.idct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.idct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, + -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, + -8.80319249e-08, -4.05617893e-01]), + b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, + -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, + -1.10039906e-08, -5.07022366e-02]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, + -16.18951607, 18.06697273, -17.57439613, 11.68861485, + -4.41308832]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } + """ + return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def ifft( + x: Union[ivy.Array, ivy.NativeArray], + dim: int, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. + + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs IFFT. + dim + The dimension along which to take the one dimensional IFFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing IFFT. + Should be a integer greater than 1. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the IFFT operation. + + Examples + -------- + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, + 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, + 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, + 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, + 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, + 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, + 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, + 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, + 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, + 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) + """ + return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def ifftn( + x: Union[ivy.Array, ivy.NativeArray], + s: Optional[Union[int, Tuple[int, ...]]] = None, + axes: Optional[Union[int, Tuple[int, ...]]] = None, + *, + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the N-dimensional inverse discrete Fourier Transform. + + Parameters + ---------- + x + Input array of complex numbers. + s + Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, + `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, + the input is cropped. If larger, input is padded with zeros. If `s` is not + given, shape of input along axes specified by axes is used. + axes + Axes over which to compute the IFFT. If not given, last `len(s)` axes are + used, or all axes if `s` is also not specified. Repeated indices in axes + means inverse transform over that axis is performed multiple times. + norm + Indicates direction of the forward/backward pair of transforms is scaled + and with what normalization factor. "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" + indicates normalization by $\frac{1}{n}$. + out + Optional output array for writing the result to. It must have a shape that + the inputs broadcast to. + + Returns + ------- + out + The truncated or zero-padded input, transformed along the axes indicated + by axes, or by a combination of s or x, as explained in the parameters + section above. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of axes is larger than the number of axes of x. + + Examples + -------- + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> y = ivy.ifftn(x) + >>> print(y) + ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, + -0.015561 -0.04216015j], + [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, + -0.08371392+0.17252843j], + [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, + 0.05344324+0.07972424j]]) + + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') + >>> print(b) + ivy.array([[ 0.8344667 +0.98222595j], + [-0.48472244+0.30233797j]]) + """ + return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) + + +@handle_exceptions +@handle_nestable +@handle_out_argument +@inputs_to_ivy_arrays +def interp(x, xp, fp, left=None, right=None, period=None): + x_arr = ivy.array(x) + fix_later = False + if x_arr.shape == (): + x_arr = ivy.array([x]) + fix_later = True + x = ivy.astype(x_arr, "float64") + xp = ivy.astype(ivy.array(xp), "float64") + fp = ivy.astype(ivy.array(fp), "float64") + ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) + if period is not None: + ivy.utils.assertions.check_equal(period, 0, inverse=True) + period = ivy.abs(period) + x = ivy.remainder(x, period) + xp = ivy.remainder(xp, period) + asort_xp = ivy.argsort(xp) + xp = xp[asort_xp] + fp = fp[asort_xp] + xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) + fp = ivy.concat((fp[-1:], fp, fp[0:1])) + + def interp_inner(value): + value = ivy.array(value) + if value < xp[0]: + return left if left is not None else fp[0] + elif value > xp[-1]: + return right if right is not None else fp[-1] + else: + last = None + if xp.shape[0] < 3: + for i in range(xp.shape[0] - 1, -1, -1): + if xp[i] == value: + return fp[i] + elif xp[i] < value: + last = i + else: + first = 0 + last = xp.shape[0] + while first < last: + midpoint = (first + last) // 2 + if xp[midpoint] == value: + already_exists = ivy.argwhere(xp == value) + if already_exists.shape[0] > 0: + return fp[already_exists[-1][0]] + return fp[midpoint] + else: + if value < xp[midpoint]: + last = midpoint - 1 + else: + first = midpoint + 1 + dist = (value - xp[last]) / (xp[last + 1] - xp[last]) + return (fp[last + 1] - fp[last]) * dist + fp[last] + + ret = ivy.map(interp_inner, unique={"value": x}) + if fix_later: + return ivy.astype(ivy.array(ret[0]), "float64") + else: + return ivy.astype(ivy.array(ret), "float64") @handle_exceptions @@ -2050,497 +2199,425 @@ def interpolate( return ivy.astype(ret, ivy.dtype(x), out=out) -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} - - -def _get_size(scale_factor, size, dims, x_shape): - if scale_factor is not None: - if isinstance(scale_factor, (float, int)): - scale_factor = [scale_factor] * dims - elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: - scale_factor = [scale_factor[0]] * dims - - size = tuple( - [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] - ) - else: - size = (size,) * dims if isinstance(size, int) else tuple(size) - return size - - -def _output_ceil_shape(w, f, p, s): - return math.ceil((w - f + p) / s) + 1 - - -def _padding_ceil_mode(w, f, p, s, return_added_padding=False): - remaining_pixels = (w - f + sum(p)) % s - added_padding = 0 - if s > 1 and remaining_pixels != 0 and f > 1: - input_size = w + sum(p) - # making sure that the remaining pixels are supposed - # to be covered by the window - # they won't be covered if stride is big enough to skip them - if input_size - remaining_pixels - (f - 1) + s > input_size: - return p - output_shape = _output_ceil_shape( - w, - f, - sum(p), - s, - ) - # calculating new padding with ceil_output_shape - new_pad = (output_shape - 1) * s + f - w - # updating pad_list with new padding by adding it to the end - added_padding = new_pad - sum(p) - p = ( - p[0], - p[1] + added_padding, - ) - if return_added_padding: - return p, added_padding - return p - - -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -def _compute_idx(in_size, out_size, device): - out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) - i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) - maxlength = in_size // out_size + 1 - in_size_mod = in_size % out_size - # adaptive = True iff there are kernels with different lengths - adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) - if adaptive: - maxlength += 1 - elif in_size_mod == 0: - maxlength -= 1 - range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) - idx = ivy.expand_dims(i0, axis=-1) + range_max - if adaptive: - maxval = ivy.full_like(idx, fill_value=in_size - 1) - idx = ivy.minimum(idx, maxval) - i1 = ivy.trunc_divide( - (out_range + 1) * in_size + out_size - 1, out_size - ).astype(ivy.int64) - length = i1 - i0 - else: - length = maxlength - return idx, length, range_max, adaptive - - -def _expand_to_dim(x, dim): - for _ in range(dim - len(x.shape)): - x = ivy.expand_dims(x, axis=-1) - return x - - -def _mask(vals, length, range_max, dim, mask_value=0.0): - if isinstance(length, int): - return vals, length - else: - assert dim < 0 - mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) - if dim == -2: - mask = _expand_to_dim(mask, 4) - vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) - length = _expand_to_dim(length, -dim) - return vals, length - - +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_max_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], -): +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool1d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Apply a 2D adaptive maximum pooling over an input signal composed of several input - planes. + Compute a 1-D max pool given 3-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + Input image *[batch_size, w, d_in]* if data_format is "NWC". + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) + data_format + "NWC" or "NCW". Defaults to "NWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) - - if isinstance(output_size, int): - output_size = (output_size, output_size) - - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.max_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output - - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device - ) - - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array( - input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device - ) - - if not adaptive_h and not adaptive_w: - ret = ivy.max(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + ret + The result of the pooling operation. - vals, length_h = _mask( - vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") - ) - vals, length_w = _mask( - vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") - ) + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ivy.maximum(ret, vals[..., i, :, j]) - pooled_output = ret.astype(vals.dtype) + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 4., 5., 6., 7.]], + [[16., 17., 18., 19.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa + ivy.array([[[ 1., 3.], + [ 5., 7.], + [ 9., 11.]], -adaptive_max_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[13., 15.], + [17., 19.], + [21., 23.]]]) + """ + return ivy.current_backend(x).max_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_avg_pool1d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: int, +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool2d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 1D adaptive average pooling over an input signal composed of several input - planes. + Compute a 2-D max pool given 4-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, L_in) or (C, L_in) where N is - the batch dimension, C is the feature dimension, and L_in is the spatial - dimension. - output_size - Spatial output size. + x + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NHWC" or "NCHW". Defaults to "NHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. Returns ------- - The result of the pooling operation. Will have shape (N, C, L_out) or - (C, L_out), where L_out = `output_size` - """ - squeeze = False - if input.ndim == 2: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 3: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", - ) + ret + The result of the pooling operation. - if input.shape[-1] % output_size == 0: - stride = input.shape[-1] // output_size - kernel_size = input.shape[-1] - (output_size - 1) * stride - pooled_output = ivy.avg_pool1d( - input, kernel_size, stride, "VALID", data_format="NCW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size, input.device - ) + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 2., 3.], + [ 4., 5.], + [ 4., 5.]]], - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., idxw]) - if not adaptive_w: - ret = ivy.mean(vals, axis=-1) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + [[[ 8., 9.], + [10., 11.], + [10., 11.]]]]) - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[16., 17.]], - ret = None - for i in range(vals.shape[-1]): - if ret is None: - ret = vals[..., i] - else: - ret = ret + vals[..., i] - pooled_output = ret / length_w.astype(ret.dtype) + [[22., 23.]]], - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[[40., 41.]], -adaptive_avg_pool1d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[46., 47.]]]]) + """ + return ivy.current_backend(x).max_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) -@handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def adaptive_avg_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool3d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 2D adaptive average pooling over an input signal composed of several input - planes. + Compute a 3-D max pool given 5-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) + data_format + "NDHWC" or "NCDHW". Defaults to "NDHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) - - if isinstance(output_size, int): - output_size = (output_size, output_size) - - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.avg_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output - - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device - ) - - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) + ret + The result of the pooling operation. - if not adaptive_h and not adaptive_w: - ret = ivy.mean(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + Examples + -------- + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) + ivy.array([[[[[14., 15.]]]], - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ret + vals[..., i, :, j] - pooled_output = ret / (length_h * length_w).astype(vals.dtype) - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[[[38., 39.]]]]]) + >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) + ivy.array([[[[[14., 15.]]], -adaptive_avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[[22., 23.]]]], -def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): - def _pad(arr, pads, pad_value): - out = ivy.astype( - ivy.pad( - arr, - ivy.maximum(0, pads).to_list(), - mode="constant", - constant_values=ivy.to_scalar(pad_value), - ), - arr.dtype, - ) - slices = tuple( - _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) - for (lo, hi), dim in zip(pads, arr.shape) - ) - return out[slices] - if ( - _min(lhs.ndim, len(rhs_shape)) < 2 - or lhs.ndim != len(rhs_shape) - or lhs.shape[1] != rhs_shape[1] - ): - raise ValueError("Dimension mismatch") - if len(window_strides) != len(rhs_shape) - 2: - raise ValueError("Wrong number of strides for spatial dimensions") - if len(pads) != len(rhs_shape) - 2: - raise ValueError("Wrong number of pads for spatial dimensions") - lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) - in_shape = lhs.shape[2:] - filter_shape = rhs_shape[2:] - dim = len(filter_shape) + [[[[38., 39.]]], - out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() - view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] - out_shape = [ - (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) - ] - view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] + [[[46., 47.]]]]]) + """ + return ivy.current_backend(x).max_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) - view = ivy.as_strided(lhs, view_shape, view_strides) - view_axes = list(range(view.ndim)) - sum_axes = view_axes[-dim - 1 :] - rhs_axes = [view.ndim] + sum_axes - out_axes = [0, view.ndim] + list(range(1, dim + 1)) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_unpool1d( + x: ivy.Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 1-D max unpooling given the 1-D pooled input x and its indices. - return view, view_axes, rhs_axes, out_axes + Parameters + ---------- + x + Pooled input image *[batch_size, w, d_in]*. + indices + Indices obtained from the corresponding max pooling operation. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + out + optional output array, for writing the result to. + Returns + ------- + ret + The result of the unpooling operation. -def _dilate(operand, factors, fill_value): - outspace = list(operand.shape[:2]) + [ - shape + (factors[i] - 1) * (shape - 1) - for i, shape in enumerate(operand.shape[2:]) - ] - out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) - lhs_slices = tuple(_slice(None, None, step) for step in factors) - out[(_slice(None),) * 2 + lhs_slices] = operand - return out + Both the description and the type hints above assume an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') + >>> print(pool_result) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], -def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): - if padding.upper() == "SAME": - out_shape = [ - math.ceil(in_size / stride) - for in_size, stride in zip(in_shape, window_strides) - ] - pad_sizes = [ - _max((out_size - 1) * stride + filter_size - in_size, 0) - for out_size, stride, filter_size, in_size in zip( - out_shape, window_strides, filter_shape, in_shape - ) - ] - return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] - else: - return [(0, 0)] * len(in_shape) + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') + >>> print(unpool_result) + ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], + [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], + [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], + [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) + """ + return ivy.current_backend(x).max_unpool1d( + x, indices, kernel, strides, padding, data_format=data_format, out=out + ) -identities = { - "max": -float("inf"), - "min": float("inf"), - "add": 0, - "mul": 1, - "multiply": 1, - "logical_and": True, - "logical_or": False, -} +def nearest_interpolate(x, dims, size, input_shape, exact): + off = 0.5 if exact else 0 + for d in range(dims): + m = input_shape[d + 2] + n = size[d] + offsets = (ivy.arange(n, dtype="float32") + off) * m / n + offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") + x = ivy.gather(x, offsets, axis=d + 2) + return x -def _cast_init(init, dtype): - if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): - if ivy.is_float_dtype(dtype): - info = ivy.finfo(dtype) - else: - info = ivy.iinfo(dtype) - if "float64" not in str(dtype): - init = info.max if init > 0 else info.min - return ivy.array(init, dtype=dtype) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def pool( + x: Union[ivy.Array, ivy.NativeArray], + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, + *, + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Perform an N-D pooling operation. -def _get_identity(func, dtype, init): - func_name = func.__name__ - if func_name in identities: - identity = identities[func_name] - return _cast_init(identity, dtype) - return init + Parameters + ---------- + x + Input array to pool over. + window_shape + Shape of the pooling window. + pool_type + Type of pooling operation, either 'MAX' or 'AVG'. + strides + Strides of the pooling operation. + padding + Padding type, either 'VALID' or 'SAME'. + data_format + Data format of the input and output data, either 'NCHW' or 'NHWC'. + dilations + Dilation rate of the pooling operation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The result of the pooling operation. -avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) + ivy.array([[[[ 8., 9.]], + [[14., 15.]]], + [[[32., 33.]], + [[38., 39.]]]]) + """ + return ivy.current_backend(x).pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ceil_mode=ceil_mode, + out=out, + ) @handle_exceptions @@ -2619,178 +2696,6 @@ def reduce_window( return ret.astype(operand.dtype) -reduce_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def fft2( - x: Union[ivy.Array, ivy.NativeArray], - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the 2-dimensional discrete Fourier Transform. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT2. - s - sequence of ints, optional - Shape (length of each transformed axis) of the output (s[0] refers to axis 0, - s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, - if the given shape is smaller than that of the input, the input is cropped. - If it is larger, the input is padded with zeros. if s is not given, the shape - of the input along the axes specified by axes is used. - dim - Axes over which to compute the FFT2. If not given, the last two axes are used. - A repeated index in axes means the transform over that axis is performed - multiple times. A one-element sequence means that a one-dimensional FFT is - performed. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The result of the FFT2 operation. - - Examples - -------- - >>> x = ivy.array([[0, 0, 0, 0, 0], - ... [1, 1, 1, 1, 1], - ... [2, 2, 2, 2, 2], - ... [3, 3, 3, 3, 3], - ... [4, 4, 4, 4, 4]]) - >>> y = ivy.fft2(x) - >>> print(y) - ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5+17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5-17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ]]) - """ - return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) - - -fft2.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def ifftn( - x: Union[ivy.Array, ivy.NativeArray], - s: Optional[Union[int, Tuple[int, ...]]] = None, - axes: Optional[Union[int, Tuple[int, ...]]] = None, - *, - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the N-dimensional inverse discrete Fourier Transform. - - Parameters - ---------- - x - Input array of complex numbers. - s - Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, - `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, - the input is cropped. If larger, input is padded with zeros. If `s` is not - given, shape of input along axes specified by axes is used. - axes - Axes over which to compute the IFFT. If not given, last `len(s)` axes are - used, or all axes if `s` is also not specified. Repeated indices in axes - means inverse transform over that axis is performed multiple times. - norm - Indicates direction of the forward/backward pair of transforms is scaled - and with what normalization factor. "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" - indicates normalization by $\frac{1}{n}$. - out - Optional output array for writing the result to. It must have a shape that - the inputs broadcast to. - - Returns - ------- - out - The truncated or zero-padded input, transformed along the axes indicated - by axes, or by a combination of s or x, as explained in the parameters - section above. - - Raises - ------ - ValueError - If `s` and `axes` have different length. - IndexError - If an element of axes is larger than the number of axes of x. - - Examples - -------- - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> y = ivy.ifftn(x) - >>> print(y) - ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, - -0.015561 -0.04216015j], - [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, - -0.08371392+0.17252843j], - [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, - 0.05344324+0.07972424j]]) - - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') - >>> print(b) - ivy.array([[ 0.8344667 +0.98222595j], - [-0.48472244+0.30233797j]]) - """ - return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2874,3 +2779,87 @@ def rfftn( raise ValueError("s and axes must have the same length.") return ivy.current_backend(x).rfftn(x, s=s, axes=axes, norm=norm, out=out) + + +idct.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +adaptive_max_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +adaptive_avg_pool1d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +adaptive_avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +reduce_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +fft2.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} +_max = builtins.max +_min = builtins.min +_slice = builtins.slice +identities = { + "max": -float("inf"), + "min": float("inf"), + "add": 0, + "mul": 1, + "multiply": 1, + "logical_and": True, + "logical_or": False, +} diff --git a/ivy/functional/ivy/experimental/linear_algebra.py b/ivy/functional/ivy/experimental/linear_algebra.py index 249df35212e17..94a5d48422683 100644 --- a/ivy/functional/ivy/experimental/linear_algebra.py +++ b/ivy/functional/ivy/experimental/linear_algebra.py @@ -17,149 +17,96 @@ ) from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # + +# --- Helpers --- # +# --------------- # def _check_valid_dimension_size(std): ivy.utils.assertions.check_dimensions(std) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@handle_array_function -def eigh_tridiagonal( - alpha: Union[ivy.Array, ivy.NativeArray], - beta: Union[ivy.Array, ivy.NativeArray], - /, - *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] - ] = None, - tol: Optional[float] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: +def _svd_checks(x, n_eigenvecs=None): """ - Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. + Run common checks to all of the SVD methods. Parameters ---------- - alpha - A real or complex array of shape (n), the diagonal elements of the - matrix. If alpha is complex, the imaginary part is ignored - (assumed zero) to satisfy the requirement that the matrix be Hermitian. - beta - A real or complex array of shape (n-1), containing the elements of - the first super-diagonal of the matrix. If beta is complex, the first - sub-diagonal of the matrix is assumed to be the conjugate of beta to - satisfy the requirement that the matrix be Hermitian. - eigvals_only - If False, both eigenvalues and corresponding eigenvectors are - computed. If True, only eigenvalues are computed. Default is True. - select - Optional string with values in {'a', 'v', 'i'} (default is 'a') that - determines which eigenvalues to calculate: 'a': all eigenvalues. - 'v': eigenvalues in the interval (min, max] given by select_range. - 'i': eigenvalues with indices min <= i <= max. - select_range - Size 2 tuple or list or array specifying the range of eigenvalues to - compute together with select. If select is 'a', select_range is ignored. - tol - Optional scalar. Ignored when backend is not Tensorflow. The absolute - tolerance to which each eigenvalue is required. An eigenvalue - (or cluster) is considered to have converged if it lies in an interval - of this width. If tol is None (default), the value eps*|T|_2 is used - where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. + matrix : 2D-array + n_eigenvecs : int, optional, default is None + if specified, number of eigen[vectors-values] to return Returns ------- - eig_vals - The eigenvalues of the matrix in non-decreasing order. - eig_vectors - If eigvals_only is False the eigenvectors are returned in the second output - argument. + n_eigenvecs : int + the number of eigenvectors to solve for + min_dim : int + the minimum dimension of matrix + max_dim : int + the maximum dimension of matrix + """ + # ndims = len(x.shape) + # if ndims != 2: + # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + dim_1, dim_2 = ivy.shape(x)[-2:] + min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) - Examples - -------- - With :class:`ivy.Array` input: + if n_eigenvecs is None: + n_eigenvecs = max_dim - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - ivy.array([0., 0.38196602, 2.61803389]) + if n_eigenvecs > max_dim: + logging.warning( + f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " + f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." + ) + n_eigenvecs = max_dim - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, select='v', - ... select_range=[0.2,3.0]) - >>> print(y) - ivy.array([0.38196602, 2.61803389]) + return n_eigenvecs, min_dim, max_dim - >>> alpha = ivy.array([0., 1., 2., 3.]) - >>> beta = ivy.array([2., 1., 2.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, - ... eigvals_only=False, - ... select='i', - ... select_range=[1,2], - ... tol=1.) - >>> print(y) - (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], - [ 0.06693714, -0.74234426], - [-0.74234426, -0.06693714], - [ 0.56710052, 0.35048741]])) - With :class:`ivy.Container` input: +# TODO uncommment the code below when these svd +# methods have been added +def _svd_interface( + matrix, + method="truncated_svd", + n_eigenvecs=None, + flip_sign=True, + u_based_flip_sign=True, + non_negative=None, + mask=None, + n_iter_mask_imputation=5, + **kwargs, +): + if method == "truncated_svd": + svd_fun = truncated_svd + # elif method == "symeig_svd": + # svd_fun = symeig_svd + # elif method == "randomized_svd": + # svd_fun = randomized_svd + elif callable(method): + svd_fun = method + else: + raise ValueError("Invalid Choice") - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.array([0.,2.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([0., 2., 4.]) - } + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + if mask is not None and n_eigenvecs is not None: + for _ in range(n_iter_mask_imputation): + S = S * ivy.eye(U.shape[-1], V.shape[-2]) + matrix = matrix * mask + (U @ S @ V) * (1 - mask) + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([-0.82842714, 2., 4.82842731]) - } - """ - x = ivy.diag(alpha) - y = ivy.diag(beta, k=1) - z = ivy.diag(beta, k=-1) - w = x + y + z + if flip_sign: + U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) - eigh_out = ivy.linalg.eigh(w) - eigenvalues = eigh_out.eigenvalues - eigenvectors = eigh_out.eigenvectors + if non_negative is not False and non_negative is not None: + U, V = make_svd_non_negative(matrix, U, S, V) - if select == "i": - eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] - eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] - elif select == "v": - condition = ivy.logical_and( - eigenvalues.greater(select_range[0]), - eigenvalues.less_equal(select_range[1]), - ) - eigenvalues = eigenvalues[condition] - eigenvectors = eigenvectors[:, condition] + return U, S, V - if eigvals_only: - return eigenvalues - return eigenvalues, eigenvectors + +# --- Main --- # +# ------------ # @handle_exceptions @@ -169,29 +116,19 @@ def eigh_tridiagonal( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def diagflat( +def adjoint( x: Union[ivy.Array, ivy.NativeArray], /, *, - offset: int = 0, - padding_value: float = 0, - align: str = "RIGHT_LEFT", - num_rows: int = -1, - num_cols: int = -1, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a two-dimensional array with the flattened input as a diagonal. + Compute the complex conjugate transpose of x. Parameters ---------- x - Input data, which is flattened and set as the k-th diagonal of the output. - k - Diagonal to set. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. + An array with more than one dimension. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -199,35 +136,18 @@ def diagflat( Returns ------- ret - The 2-D output array. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.diagflat(x) - ivy.array([[1, 0, 0, 0], - [0, 2, 0, 0], - [0, 0, 3, 0], - [0, 0, 0, 4]]) + the complex conjugate transpose of the input. - >>> x = ivy.array([1,2]) - >>> ivy.diagflat(x, k=1) - ivy.array([[0, 1, 0], - [0, 0, 2], - [0, 0, 0]]) + Examples + -------- + >>> x = np.array([[1.-1.j, 2.+2.j], + [3.+3.j, 4.-4.j]]) + >>> x = ivy.array(x) + >>> ivy.adjoint(x) + ivy.array([[1.+1.j, 3.-3.j], + [2.-2.j, 4.+4.j]]) """ - return current_backend(x).diagflat( - x, - offset=offset, - padding_value=padding_value, - align=align, - num_rows=num_rows, - num_cols=num_cols, - out=out, - ) + return current_backend(x).adjoint(x, out=out) @handle_exceptions @@ -237,23 +157,22 @@ def diagflat( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def kron( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def cond( + x: Union[ivy.Array, ivy.NativeArray], /, *, + p: Optional[Union[int, float, str]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kronecker product, a composite array made of blocks of the second array - scaled by the first. + Compute the condition number of x. Parameters ---------- - a - First input array. - b - Second input array + x + An array with more than one dimension. + p + The order of the norm of the matrix (see :func:`ivy.norm` for details). out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -261,16 +180,21 @@ def kron( Returns ------- ret - Array representing the Kronecker product of the input arrays. + the condition number of the input. Examples -------- - >>> a = ivy.array([1,2]) - >>> b = ivy.array([3,4]) - >>> ivy.kron(a, b) - ivy.array([3, 4, 6, 8]) + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x) + ivy.array(14.933034) + + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x, p=ivy.inf) + ivy.array(21.0) """ - return current_backend(a, b).kron(a, b, out=out) + return current_backend(x).cond(x, p=p, out=out) @handle_exceptions @@ -280,19 +204,29 @@ def kron( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def matrix_exp( +def diagflat( x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, + offset: int = 0, + padding_value: float = 0, + align: str = "RIGHT_LEFT", + num_rows: int = -1, + num_cols: int = -1, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Compute the matrix exponential of a square matrix. + Return a two-dimensional array with the flattened input as a diagonal. Parameters ---------- - a - Square matrix. + x + Input data, which is flattened and set as the k-th diagonal of the output. + k + Diagonal to set. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -300,21 +234,93 @@ def matrix_exp( Returns ------- ret - the matrix exponential of the input. + The 2-D output array. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.diagflat(x) + ivy.array([[1, 0, 0, 0], + [0, 2, 0, 0], + [0, 0, 3, 0], + [0, 0, 0, 4]]) + + >>> x = ivy.array([1,2]) + >>> ivy.diagflat(x, k=1) + ivy.array([[0, 1, 0], + [0, 0, 2], + [0, 0, 0]]) + """ + return current_backend(x).diagflat( + x, + offset=offset, + padding_value=padding_value, + align=align, + num_rows=num_rows, + num_cols=num_cols, + out=out, + ) + + +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_exceptions +def dot( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the dot product between two arrays `a` and `b` using the current backend's + implementation. The dot product is defined as the sum of the element-wise product of + the input arrays. + + Parameters + ---------- + a + First input array. + b + Second input array. + out + Optional output array. If provided, the output array to store the result. + + Returns + ------- + ret + The dot product of the input arrays. Examples -------- - >>> x = ivy.array([[[1., 0.], - [0., 1.]], - [[2., 0.], - [0., 2.]]]) - >>> ivy.matrix_exp(x) - ivy.array([[[2.7183, 1.0000], - [1.0000, 2.7183]], - [[7.3891, 1.0000], - [1.0000, 7.3891]]]) + With :class:`ivy.Array` inputs: + + >>> a = ivy.array([1, 2, 3]) + >>> b = ivy.array([4, 5, 6]) + >>> result = ivy.dot(a, b) + >>> print(result) + ivy.array(32) + + >>> a = ivy.array([[1, 2], [3, 4]]) + >>> b = ivy.array([[5, 6], [7, 8]]) + >>> c = ivy.empty_like(a) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[19, 22], + [43, 50]]) + + >>> a = ivy.array([[1.1, 2.3, -3.6]]) + >>> b = ivy.array([[-4.8], [5.2], [6.1]]) + >>> c = ivy.zeros((1, 1)) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[-15.28]]) """ - return current_backend(x).matrix_exp(x, out=out) + return current_backend(a, b).dot(a, b, out=out) @handle_exceptions @@ -382,244 +388,301 @@ def eig( @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_device_shifting -def eigvals( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> ivy.Array: - """ - Compute eigenvalues of x. Returns a set of eigenvalues. - - Parameters - ---------- - x - An array of shape (..., N, N). - - Returns - ------- - w - Not necessarily ordered array(..., N) of eigenvalues in complex type. - - Functional Examples - ------------------ - With :class:`ivy.Array` inputs: - >>> x = ivy.array([[1,2], [3,4]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array([-0.37228132+0.j, 5.37228132+0.j]) - - >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array( - [ - [-0.37228132+0.j, 5.37228132+0.j], - [ 0. +0.j, 11. +0.j] - ] - ) - """ - return current_backend(x).eigvals(x) - - -@handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def adjoint( - x: Union[ivy.Array, ivy.NativeArray], +@handle_array_function +def eigh_tridiagonal( + alpha: Union[ivy.Array, ivy.NativeArray], + beta: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] + ] = None, + tol: Optional[float] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: """ - Compute the complex conjugate transpose of x. + Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. Parameters ---------- - x - An array with more than one dimension. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + alpha + A real or complex array of shape (n), the diagonal elements of the + matrix. If alpha is complex, the imaginary part is ignored + (assumed zero) to satisfy the requirement that the matrix be Hermitian. + beta + A real or complex array of shape (n-1), containing the elements of + the first super-diagonal of the matrix. If beta is complex, the first + sub-diagonal of the matrix is assumed to be the conjugate of beta to + satisfy the requirement that the matrix be Hermitian. + eigvals_only + If False, both eigenvalues and corresponding eigenvectors are + computed. If True, only eigenvalues are computed. Default is True. + select + Optional string with values in {'a', 'v', 'i'} (default is 'a') that + determines which eigenvalues to calculate: 'a': all eigenvalues. + 'v': eigenvalues in the interval (min, max] given by select_range. + 'i': eigenvalues with indices min <= i <= max. + select_range + Size 2 tuple or list or array specifying the range of eigenvalues to + compute together with select. If select is 'a', select_range is ignored. + tol + Optional scalar. Ignored when backend is not Tensorflow. The absolute + tolerance to which each eigenvalue is required. An eigenvalue + (or cluster) is considered to have converged if it lies in an interval + of this width. If tol is None (default), the value eps*|T|_2 is used + where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. Returns ------- - ret - the complex conjugate transpose of the input. + eig_vals + The eigenvalues of the matrix in non-decreasing order. + eig_vectors + If eigvals_only is False the eigenvectors are returned in the second output + argument. + + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - >>> x = np.array([[1.-1.j, 2.+2.j], - [3.+3.j, 4.-4.j]]) - >>> x = ivy.array(x) - >>> ivy.adjoint(x) - ivy.array([[1.+1.j, 3.-3.j], - [2.-2.j, 4.+4.j]]) - """ - return current_backend(x).adjoint(x, out=out) + With :class:`ivy.Array` input: + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + ivy.array([0., 0.38196602, 2.61803389]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def multi_dot( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the dot product of two or more matrices in a single function call, while - selecting the fastest evaluation order. - - Parameters - ---------- - x - sequence of matrices to multiply. - out - optional output array, for writing the result to. It must have a valid - shape, i.e. the resulting shape after applying regular matrix multiplication - to the inputs. + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, select='v', + ... select_range=[0.2,3.0]) + >>> print(y) + ivy.array([0.38196602, 2.61803389]) - Returns - ------- - ret - dot product of the arrays. + >>> alpha = ivy.array([0., 1., 2., 3.]) + >>> beta = ivy.array([2., 1., 2.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, + ... eigvals_only=False, + ... select='i', + ... select_range=[1,2], + ... tol=1.) + >>> print(y) + (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], + [ 0.06693714, -0.74234426], + [-0.74234426, -0.06693714], + [ 0.56710052, 0.35048741]])) - Examples - -------- - With :class:`ivy.Array` input: + With :class:`ivy.Container` input: - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> ivy.multi_dot((A, B, C)) - ivy.array([[ 26, 49], - [ 80, 148]]) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.array([0.,2.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([0., 2., 4.]) + } - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> D = ivy.zeros((2, 2)) - >>> ivy.multi_dot((A, B, C), out=D) - >>> print(D) - ivy.array([[ 26, 49], - [ 80, 148]]) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([-0.82842714, 2., 4.82842731]) + } """ - return current_backend(x).multi_dot(x, out=out) + x = ivy.diag(alpha) + y = ivy.diag(beta, k=1) + z = ivy.diag(beta, k=-1) + w = x + y + z + eigh_out = ivy.linalg.eigh(w) + eigenvalues = eigh_out.eigenvalues + eigenvectors = eigh_out.eigenvectors -multi_dot.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} + if select == "i": + eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] + eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] + elif select == "v": + condition = ivy.logical_and( + eigenvalues.greater(select_range[0]), + eigenvalues.less_equal(select_range[1]), + ) + eigenvalues = eigenvalues[condition] + eigenvectors = eigenvectors[:, condition] + + if eigvals_only: + return eigenvalues + return eigenvalues, eigenvectors @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def cond( +def eigvals( x: Union[ivy.Array, ivy.NativeArray], /, - *, - p: Optional[Union[int, float, str]] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the condition number of x. + Compute eigenvalues of x. Returns a set of eigenvalues. Parameters ---------- x - An array with more than one dimension. - p - The order of the norm of the matrix (see :func:`ivy.norm` for details). - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + An array of shape (..., N, N). Returns ------- - ret - the condition number of the input. + w + Not necessarily ordered array(..., N) of eigenvalues in complex type. - Examples - -------- - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x) - ivy.array(14.933034) + Functional Examples + ------------------ + With :class:`ivy.Array` inputs: + >>> x = ivy.array([[1,2], [3,4]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array([-0.37228132+0.j, 5.37228132+0.j]) - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x, p=ivy.inf) - ivy.array(21.0) + >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array( + [ + [-0.37228132+0.j, 5.37228132+0.j], + [ 0. +0.j, 11. +0.j] + ] + ) """ - return current_backend(x).cond(x, p=p, out=out) + return current_backend(x).eigvals(x) -# This code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py -@handle_exceptions -@handle_backend_invalid +# This function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 + + +# TODO update svd type hints when other svd methods have been added +# also update the test @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def kronecker( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], - skip_matrix: Optional[int] = None, - reverse: Optional[bool] = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def initialize_tucker( + x: Union[ivy.Array, ivy.NativeArray], + rank: Sequence[int], + modes: Sequence[int], + /, + *, + init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", + seed: Optional[int] = None, + svd: Optional[Literal["truncated_svd"]] = "truncated_svd", + non_negative: Optional[bool] = False, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + svd_mask_repeats: Optional[int] = 5, +) -> Tuple[ivy.Array, Sequence[ivy.Array]]: """ - Kronecker product of a list of matrices. + Initialize core and factors used in `tucker`. The type of initialization is set + using `init`. If `init == 'random'` then initialize factor matrices using + `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the + `rank` left singular vectors of the `m`th unfolding of the input tensor. Parameters ---------- x - Sequence of matrices - - skip_matrix - if not None, index of a matrix to skip - - reverse - if True, the order of the matrices is reversed + input tensor + rank + number of components + modes + modes to consider in the input tensor + seed + Used to create a random seed distribution + when init == 'random' + init + initialization scheme for tucker decomposition. + svd + function to use to compute the SVD + non_negative + if True, non-negative factors are returned + mask + array of booleans with the same shape as ``tensor`` should be 0 where + the values are missing and 1 everywhere else. Note: if tensor is + sparse, then mask should also be sparse with a fill value of 1 (or + True). + svd_mask_repeats + number of iterations for imputing the values in the SVD matrix when + mask is not None Returns ------- - kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` - where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` - and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` + core + initialized core tensor + factors + list of factors """ - if skip_matrix is not None: - x = [x[i] for i in range(len(x)) if i != skip_matrix] + try: + assert len(x.shape) >= 2 + except ValueError: + raise ValueError( + "expected x to have atleast 2 dimensions but it has only" + f" {len(x.shape)} dimension(s)" + ) + + # Initialisation + if init == "svd": + factors = [] + for index, mode in enumerate(modes): + mask_unfold = None if mask is None else ivy.unfold(mask, mode) + U, _, _ = _svd_interface( + ivy.unfold(x, mode), + n_eigenvecs=rank[index], + method=svd, + non_negative=non_negative, + mask=mask_unfold, + n_iter_mask_imputation=svd_mask_repeats, + # random_state=random_state, + ) + factors.append(U) + + # The initial core approximation is needed here for the masking step + core = multi_mode_dot(x, factors, modes=modes, transpose=True) + + elif init == "random": + core = ( + ivy.random_uniform( + shape=[rank[index] for index in range(len(modes))], + dtype=x.dtype, + seed=seed, + ) + + 0.01 + ) + factors = [ + ivy.random_uniform( + shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed + ) + for index, mode in enumerate(modes) + ] - if reverse: - order = -1 else: - order = 1 + (core, factors) = init - for i, matrix in enumerate(x[::order]): - if not i: - res = matrix - else: - res = ivy.kron(res, matrix, out=out) - return res + if non_negative is True: + factors = [ivy.abs(f) for f in factors] + core = ivy.abs(core) + + return (core, factors) # The code has been adapated from tensorly.khatri_rao @@ -670,233 +733,54 @@ def khatri_rao( where ``prod(n_i) = prod([m.shape[0] for m in input])`` i.e. the product of the number of rows of all the input in the product. """ - if skip_matrix is not None: - x = [x[i] for i in range(len(x)) if i != skip_matrix] - - # Khatri-rao of only one matrix: just return that matrix - if len(x) == 1: - if ivy.exists(out): - return ivy.inplace_update(out, x[0]) - return x[0] - - if len(x[0].shape) == 2: - n_columns = x[0].shape[1] - else: - n_columns = 1 - x = [ivy.reshape(m, (-1, 1)) for m in x] - logging.warning( - "Khatri-rao of a series of vectors instead of input. " - "Considering each as a matrix with 1 column." - ) - - # Testing whether the input have the proper size - for i, matrix in enumerate(x): - if len(matrix.shape) != 2: - raise ValueError( - "All the input must have exactly 2 dimensions!" - f"Matrix {i} has dimension {len(matrix.shape)} != 2." - ) - if matrix.shape[1] != n_columns: - raise ValueError( - "All input must have same number of columns!" - f"Matrix {i} has {matrix.shape[1]} columns != {n_columns}." - ) - - for i, e in enumerate(x[1:]): - if not i: - if weights is None: - res = x[0] - else: - res = x[0] * ivy.reshape(weights, (1, -1)) - s1, s2 = ivy.shape(res) - s3, s4 = ivy.shape(e) - - a = ivy.reshape(res, (s1, 1, s2)) - b = ivy.reshape(e, (1, s3, s4)) - res = ivy.reshape(a * b, (-1, n_columns)) - - m = ivy.reshape(mask, (1, -1)) if mask is not None else 1 - - res = res * m - - if ivy.exists(out): - return ivy.inplace_update(out, res) - - return res - - -# The following code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def mode_dot( - x: Union[ivy.Array, ivy.NativeArray], - /, - matrix_or_vector: Union[ivy.Array, ivy.NativeArray], - mode: int, - transpose: Optional[bool] = False, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - N-mode product of a tensor and a matrix or vector at the specified mode. - - Parameters - ---------- - x - tensor of shape ``(i_1, ..., i_k, ..., i_N)`` - matrix_or_vector - 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` - matrix or vectors to which to n-mode multiply the tensor - mode - int in the range(1, N) - transpose - If True, the matrix is transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. - - Returns - ------- - ivy.Array - `mode`-mode product of `tensor` by `matrix_or_vector` - * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` - if matrix_or_vector is a matrix - * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` - if matrix_or_vector is a vector - """ - # the mode along which to fold might decrease if we take product with a vector - fold_mode = mode - new_shape = list(x.shape) - ndims = len(matrix_or_vector.shape) - - if ndims == 2: # Tensor times matrix - # Test for the validity of the operation - dim = 0 if transpose else 1 - if matrix_or_vector.shape[dim] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" - ) - - if transpose: - matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) - - new_shape[mode] = matrix_or_vector.shape[0] - vec = False - - elif ndims == 1: # Tensor times vector - if matrix_or_vector.shape[0] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[0]} (vector size)" - ) - if len(new_shape) > 1: - new_shape.pop(mode) - else: - new_shape = [] - vec = True - - else: - raise ValueError( - "Can only take n_mode_product with a vector or a matrix." - f"Provided array of dimension {ndims} not in [1, 2]." - ) - - res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) - - if vec: # We contracted with a vector, leading to a vector - return ivy.reshape(res, new_shape, out=out) - else: # tensor times vec: refold the unfolding - return ivy.fold(res, fold_mode, new_shape, out=out) - - -# The following code has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def multi_mode_dot( - x: Union[ivy.Array, ivy.NativeArray], - mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], - /, - modes: Optional[Sequence[int]] = None, - skip: Optional[Sequence[int]] = None, - transpose: Optional[bool] = False, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - N-mode product of a tensor and several matrices or vectors over several modes. - - Parameters - ---------- - x - the input tensor - - mat_or_vec_list - sequence of matrices or vectors of length ``tensor.ndim`` - - skip - None or int, optional, default is None - If not None, index of a matrix to skip. - - modes - None or int list, optional, default is None - - transpose - If True, the matrices or vectors in in the list are transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. + if skip_matrix is not None: + x = [x[i] for i in range(len(x)) if i != skip_matrix] - Returns - ------- - ivy.Array - tensor times each matrix or vector in the list at mode `mode` + # Khatri-rao of only one matrix: just return that matrix + if len(x) == 1: + if ivy.exists(out): + return ivy.inplace_update(out, x[0]) + return x[0] - Notes - ----- - If no modes are specified, just assumes there is one matrix or vector per mode and returns: - :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa - """ - if modes is None: - modes = range(len(mat_or_vec_list)) + if len(x[0].shape) == 2: + n_columns = x[0].shape[1] + else: + n_columns = 1 + x = [ivy.reshape(m, (-1, 1)) for m in x] + logging.warning( + "Khatri-rao of a series of vectors instead of input. " + "Considering each as a matrix with 1 column." + ) - decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor + # Testing whether the input have the proper size + for i, matrix in enumerate(x): + if len(matrix.shape) != 2: + raise ValueError( + "All the input must have exactly 2 dimensions!" + f"Matrix {i} has dimension {len(matrix.shape)} != 2." + ) + if matrix.shape[1] != n_columns: + raise ValueError( + "All input must have same number of columns!" + f"Matrix {i} has {matrix.shape[1]} columns != {n_columns}." + ) - res = x + for i, e in enumerate(x[1:]): + if not i: + if weights is None: + res = x[0] + else: + res = x[0] * ivy.reshape(weights, (1, -1)) + s1, s2 = ivy.shape(res) + s3, s4 = ivy.shape(e) - # Order of mode dots doesn't matter for different modes - # Sorting by mode shouldn't change order for equal modes - factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) - for i, (mat_or_vec_list, mode) in enumerate(factors_modes): - ndims = len(mat_or_vec_list.shape) - if (skip is not None) and (i == skip): - continue + a = ivy.reshape(res, (s1, 1, s2)) + b = ivy.reshape(e, (1, s3, s4)) + res = ivy.reshape(a * b, (-1, n_columns)) - if transpose and ndims == 2: - res = mode_dot( - res, - ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), - mode - decrement, - ) - else: - res = mode_dot(res, mat_or_vec_list, mode - decrement) + m = ivy.reshape(mask, (1, -1)) if mask is not None else 1 - if ndims == 1: - decrement += 1 + res = res * m if ivy.exists(out): return ivy.inplace_update(out, res) @@ -904,113 +788,98 @@ def multi_mode_dot( return res -def _svd_checks(x, n_eigenvecs=None): +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def kron( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Run common checks to all of the SVD methods. + Compute the Kronecker product, a composite array made of blocks of the second array + scaled by the first. Parameters ---------- - matrix : 2D-array - n_eigenvecs : int, optional, default is None - if specified, number of eigen[vectors-values] to return + a + First input array. + b + Second input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - n_eigenvecs : int - the number of eigenvectors to solve for - min_dim : int - the minimum dimension of matrix - max_dim : int - the maximum dimension of matrix - """ - # ndims = len(x.shape) - # if ndims != 2: - # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") - - dim_1, dim_2 = ivy.shape(x)[-2:] - min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) - - if n_eigenvecs is None: - n_eigenvecs = max_dim - - if n_eigenvecs > max_dim: - logging.warning( - f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " - f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." - ) - n_eigenvecs = max_dim + ret + Array representing the Kronecker product of the input arrays. - return n_eigenvecs, min_dim, max_dim + Examples + -------- + >>> a = ivy.array([1,2]) + >>> b = ivy.array([3,4]) + >>> ivy.kron(a, b) + ivy.array([3, 4, 6, 8]) + """ + return current_backend(a, b).kron(a, b, out=out) -# This function has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 -@handle_nestable +# This code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def svd_flip( - U: Union[ivy.Array, ivy.NativeArray], - V: Union[ivy.Array, ivy.NativeArray], - /, - u_based_decision: Optional[bool] = True, -) -> Tuple[ivy.Array, ivy.Array]: +def kronecker( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], + skip_matrix: Optional[int] = None, + reverse: Optional[bool] = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Sign correction to ensure deterministic output from SVD. Adjusts the columns of u - and the rows of v such that the loadings in the columns in u that are largest in - absolute value are always positive. This function is borrowed from scikit- - learn/utils/extmath.py. + Kronecker product of a list of matrices. Parameters ---------- - U - left singular matrix output of SVD - V - right singular matrix output of SVD - u_based_decision - If True, use the columns of u as the basis for sign flipping. - Otherwise, use the rows of v. The choice of which variable to base the - decision on is generally algorithm dependent. + x + Sequence of matrices + + skip_matrix + if not None, index of a matrix to skip + + reverse + if True, the order of the matrices is reversed Returns ------- - u_adjusted, v_adjusted : arrays with the same dimensions as the input. + kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` + where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` + and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` """ - if u_based_decision: - # columns of U, rows of V - max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) - signs = ivy.sign( - ivy.array( - [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], - ) - ) - U = U * signs - if ivy.shape(V)[0] > ivy.shape(U)[1]: - signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) - V = V * signs[: ivy.shape(V)[0]][:, None] + if skip_matrix is not None: + x = [x[i] for i in range(len(x)) if i != skip_matrix] + + if reverse: + order = -1 else: - # rows of V, columns of U - max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) - signs = ivy.sign( - ivy.array( - [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], - ) - ) - V = V * signs[:, None] - if ivy.shape(U)[1] > ivy.shape(V)[0]: - signs = ivy.concat( - ( - signs, - ivy.ones( - ivy.shape(U)[1] - ivy.shape(V)[0], - ), - ) - ) - U = U * signs[: ivy.shape(U)[1]] + order = 1 - return U, V + for i, matrix in enumerate(x[::order]): + if not i: + res = matrix + else: + res = ivy.kron(res, matrix, out=out) + return res # This function has been adapted from TensorLy @@ -1105,210 +974,282 @@ def make_svd_non_negative( ' "nndsvda")' ) - return W, H + return W, H + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def matrix_exp( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the matrix exponential of a square matrix. + + Parameters + ---------- + a + Square matrix. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + the matrix exponential of the input. + + Examples + -------- + >>> x = ivy.array([[[1., 0.], + [0., 1.]], + [[2., 0.], + [0., 2.]]]) + >>> ivy.matrix_exp(x) + ivy.array([[[2.7183, 1.0000], + [1.0000, 2.7183]], + [[7.3891, 1.0000], + [1.0000, 7.3891]]]) + """ + return current_backend(x).matrix_exp(x, out=out) -# The following function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 +# The following code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def truncated_svd( +def mode_dot( x: Union[ivy.Array, ivy.NativeArray], /, - compute_uv: bool = True, - n_eigenvecs: Optional[int] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: + matrix_or_vector: Union[ivy.Array, ivy.NativeArray], + mode: int, + transpose: Optional[bool] = False, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Compute a truncated SVD on `x` using the standard SVD. + N-mode product of a tensor and a matrix or vector at the specified mode. Parameters ---------- x - 2D-array - compute_uv - If ``True`` then left and right singular vectors will be computed and returned - in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be - computed, which can be significantly faster. - n_eigenvecs - if specified, number of eigen[vectors-values] to return - else full matrices will be returned + tensor of shape ``(i_1, ..., i_k, ..., i_N)`` + matrix_or_vector + 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` + matrix or vectors to which to n-mode multiply the tensor + mode + int in the range(1, N) + transpose + If True, the matrix is transposed. + For complex tensors, the conjugate transpose is used. + out + optional output array, for writing the result to. It must have a shape that the + result can broadcast to. Returns ------- - ret - a namedtuple ``(U, S, Vh)`` - Each returned array must have the same floating-point data type as ``x``. + ivy.Array + `mode`-mode product of `tensor` by `matrix_or_vector` + * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` + if matrix_or_vector is a matrix + * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` + if matrix_or_vector is a vector """ - n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) - full_matrices = True if n_eigenvecs > min_dim else False + # the mode along which to fold might decrease if we take product with a vector + fold_mode = mode + new_shape = list(x.shape) + ndims = len(matrix_or_vector.shape) - if compute_uv: - U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) - return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] - else: - S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) - return S[:n_eigenvecs] + if ndims == 2: # Tensor times matrix + # Test for the validity of the operation + dim = 0 if transpose else 1 + if matrix_or_vector.shape[dim] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" + ) + if transpose: + matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) + + new_shape[mode] = matrix_or_vector.shape[0] + vec = False + + elif ndims == 1: # Tensor times vector + if matrix_or_vector.shape[0] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[0]} (vector size)" + ) + if len(new_shape) > 1: + new_shape.pop(mode) + else: + new_shape = [] + vec = True -# TODO uncommment the code below when these svd -# methods have been added -def _svd_interface( - matrix, - method="truncated_svd", - n_eigenvecs=None, - flip_sign=True, - u_based_flip_sign=True, - non_negative=None, - mask=None, - n_iter_mask_imputation=5, - **kwargs, -): - if method == "truncated_svd": - svd_fun = truncated_svd - # elif method == "symeig_svd": - # svd_fun = symeig_svd - # elif method == "randomized_svd": - # svd_fun = randomized_svd - elif callable(method): - svd_fun = method else: - raise ValueError("Invalid Choice") + raise ValueError( + "Can only take n_mode_product with a vector or a matrix." + f"Provided array of dimension {ndims} not in [1, 2]." + ) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) - if mask is not None and n_eigenvecs is not None: - for _ in range(n_iter_mask_imputation): - S = S * ivy.eye(U.shape[-1], V.shape[-2]) - matrix = matrix * mask + (U @ S @ V) * (1 - mask) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) - if flip_sign: - U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) + if vec: # We contracted with a vector, leading to a vector + return ivy.reshape(res, new_shape, out=out) + else: # tensor times vec: refold the unfolding + return ivy.fold(res, fold_mode, new_shape, out=out) - if non_negative is not False and non_negative is not None: - U, V = make_svd_non_negative(matrix, U, S, V) - return U, S, V +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def multi_dot( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the dot product of two or more matrices in a single function call, while + selecting the fastest evaluation order. + + Parameters + ---------- + x + sequence of matrices to multiply. + out + optional output array, for writing the result to. It must have a valid + shape, i.e. the resulting shape after applying regular matrix multiplication + to the inputs. + + Returns + ------- + ret + dot product of the arrays. + Examples + -------- + With :class:`ivy.Array` input: -# This function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> ivy.multi_dot((A, B, C)) + ivy.array([[ 26, 49], + [ 80, 148]]) + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> D = ivy.zeros((2, 2)) + >>> ivy.multi_dot((A, B, C), out=D) + >>> print(D) + ivy.array([[ 26, 49], + [ 80, 148]]) + """ + return current_backend(x).multi_dot(x, out=out) -# TODO update svd type hints when other svd methods have been added -# also update the test + +# The following code has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def initialize_tucker( +def multi_mode_dot( x: Union[ivy.Array, ivy.NativeArray], - rank: Sequence[int], - modes: Sequence[int], + mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], /, + modes: Optional[Sequence[int]] = None, + skip: Optional[Sequence[int]] = None, + transpose: Optional[bool] = False, *, - init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", - seed: Optional[int] = None, - svd: Optional[Literal["truncated_svd"]] = "truncated_svd", - non_negative: Optional[bool] = False, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - svd_mask_repeats: Optional[int] = 5, -) -> Tuple[ivy.Array, Sequence[ivy.Array]]: - """ - Initialize core and factors used in `tucker`. The type of initialization is set - using `init`. If `init == 'random'` then initialize factor matrices using - `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the - `rank` left singular vectors of the `m`th unfolding of the input tensor. + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + N-mode product of a tensor and several matrices or vectors over several modes. Parameters ---------- x - input tensor - rank - number of components + the input tensor + + mat_or_vec_list + sequence of matrices or vectors of length ``tensor.ndim`` + + skip + None or int, optional, default is None + If not None, index of a matrix to skip. + modes - modes to consider in the input tensor - seed - Used to create a random seed distribution - when init == 'random' - init - initialization scheme for tucker decomposition. - svd - function to use to compute the SVD - non_negative - if True, non-negative factors are returned - mask - array of booleans with the same shape as ``tensor`` should be 0 where - the values are missing and 1 everywhere else. Note: if tensor is - sparse, then mask should also be sparse with a fill value of 1 (or - True). - svd_mask_repeats - number of iterations for imputing the values in the SVD matrix when - mask is not None + None or int list, optional, default is None + + transpose + If True, the matrices or vectors in in the list are transposed. + For complex tensors, the conjugate transpose is used. + out + optional output array, for writing the result to. It must have a shape that the + result can broadcast to. Returns ------- - core - initialized core tensor - factors - list of factors + ivy.Array + tensor times each matrix or vector in the list at mode `mode` + + Notes + ----- + If no modes are specified, just assumes there is one matrix or vector per mode and returns: + :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa """ - try: - assert len(x.shape) >= 2 - except ValueError: - raise ValueError( - "expected x to have atleast 2 dimensions but it has only" - f" {len(x.shape)} dimension(s)" - ) + if modes is None: + modes = range(len(mat_or_vec_list)) - # Initialisation - if init == "svd": - factors = [] - for index, mode in enumerate(modes): - mask_unfold = None if mask is None else ivy.unfold(mask, mode) - U, _, _ = _svd_interface( - ivy.unfold(x, mode), - n_eigenvecs=rank[index], - method=svd, - non_negative=non_negative, - mask=mask_unfold, - n_iter_mask_imputation=svd_mask_repeats, - # random_state=random_state, - ) - factors.append(U) + decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor - # The initial core approximation is needed here for the masking step - core = multi_mode_dot(x, factors, modes=modes, transpose=True) + res = x - elif init == "random": - core = ( - ivy.random_uniform( - shape=[rank[index] for index in range(len(modes))], - dtype=x.dtype, - seed=seed, - ) - + 0.01 - ) - factors = [ - ivy.random_uniform( - shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed + # Order of mode dots doesn't matter for different modes + # Sorting by mode shouldn't change order for equal modes + factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) + for i, (mat_or_vec_list, mode) in enumerate(factors_modes): + ndims = len(mat_or_vec_list.shape) + if (skip is not None) and (i == skip): + continue + + if transpose and ndims == 2: + res = mode_dot( + res, + ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), + mode - decrement, ) - for index, mode in enumerate(modes) - ] + else: + res = mode_dot(res, mat_or_vec_list, mode - decrement) - else: - (core, factors) = init + if ndims == 1: + decrement += 1 - if non_negative is True: - factors = [ivy.abs(f) for f in factors] - core = ivy.abs(core) + if ivy.exists(out): + return ivy.inplace_update(out, res) - return (core, factors) + return res # This function has been adpated from TensorLy @@ -1462,6 +1403,122 @@ def partial_tucker( return (core, factors) +# This function has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def svd_flip( + U: Union[ivy.Array, ivy.NativeArray], + V: Union[ivy.Array, ivy.NativeArray], + /, + u_based_decision: Optional[bool] = True, +) -> Tuple[ivy.Array, ivy.Array]: + """ + Sign correction to ensure deterministic output from SVD. Adjusts the columns of u + and the rows of v such that the loadings in the columns in u that are largest in + absolute value are always positive. This function is borrowed from scikit- + learn/utils/extmath.py. + + Parameters + ---------- + U + left singular matrix output of SVD + V + right singular matrix output of SVD + u_based_decision + If True, use the columns of u as the basis for sign flipping. + Otherwise, use the rows of v. The choice of which variable to base the + decision on is generally algorithm dependent. + + Returns + ------- + u_adjusted, v_adjusted : arrays with the same dimensions as the input. + """ + if u_based_decision: + # columns of U, rows of V + max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) + signs = ivy.sign( + ivy.array( + [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], + ) + ) + U = U * signs + if ivy.shape(V)[0] > ivy.shape(U)[1]: + signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) + V = V * signs[: ivy.shape(V)[0]][:, None] + else: + # rows of V, columns of U + max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) + signs = ivy.sign( + ivy.array( + [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], + ) + ) + V = V * signs[:, None] + if ivy.shape(U)[1] > ivy.shape(V)[0]: + signs = ivy.concat( + ( + signs, + ivy.ones( + ivy.shape(U)[1] - ivy.shape(V)[0], + ), + ) + ) + U = U * signs[: ivy.shape(U)[1]] + + return U, V + + +# The following function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def truncated_svd( + x: Union[ivy.Array, ivy.NativeArray], + /, + compute_uv: bool = True, + n_eigenvecs: Optional[int] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: + """ + Compute a truncated SVD on `x` using the standard SVD. + + Parameters + ---------- + x + 2D-array + compute_uv + If ``True`` then left and right singular vectors will be computed and returned + in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be + computed, which can be significantly faster. + n_eigenvecs + if specified, number of eigen[vectors-values] to return + else full matrices will be returned + + Returns + ------- + ret + a namedtuple ``(U, S, Vh)`` + Each returned array must have the same floating-point data type as ``x``. + """ + n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) + full_matrices = True if n_eigenvecs > min_dim else False + + if compute_uv: + U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) + return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] + else: + S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) + return S[:n_eigenvecs] + + @handle_nestable @handle_exceptions @handle_array_like_without_promotion @@ -1610,59 +1667,7 @@ def tucker( return ivy.TuckerTensor((core, factors)) -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_exceptions -def dot( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the dot product between two arrays `a` and `b` using the current backend's - implementation. The dot product is defined as the sum of the element-wise product of - the input arrays. - - Parameters - ---------- - a - First input array. - b - Second input array. - out - Optional output array. If provided, the output array to store the result. - - Returns - ------- - ret - The dot product of the input arrays. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> a = ivy.array([1, 2, 3]) - >>> b = ivy.array([4, 5, 6]) - >>> result = ivy.dot(a, b) - >>> print(result) - ivy.array(32) - - >>> a = ivy.array([[1, 2], [3, 4]]) - >>> b = ivy.array([[5, 6], [7, 8]]) - >>> c = ivy.empty_like(a) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[19, 22], - [43, 50]]) - - >>> a = ivy.array([[1.1, 2.3, -3.6]]) - >>> b = ivy.array([[-4.8], [5.2], [6.1]]) - >>> c = ivy.zeros((1, 1)) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[-15.28]]) - """ - return current_backend(a, b).dot(a, b, out=out) +multi_dot.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} diff --git a/ivy/functional/ivy/experimental/losses.py b/ivy/functional/ivy/experimental/losses.py index 7b9f55d010c5a..f873fb6e6d1dd 100644 --- a/ivy/functional/ivy/experimental/losses.py +++ b/ivy/functional/ivy/experimental/losses.py @@ -12,86 +12,67 @@ from ivy.utils.exceptions import handle_exceptions -# log_poisson_loss @handle_exceptions @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def log_poisson_loss( +def huber_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - compute_full_loss: bool = False, - axis: int = -1, - reduction: str = "none", + delta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the log-likelihood loss between the prediction and the target under the - assumption that the target has a Poisson distribution. Caveat: By default, this is - not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect - for optimization, but does not play well with relative loss comparisons. To compute - an approximation of the log factorial term, specify ``compute_full_loss=True`` to - enable Stirling's Approximation. + Compute the Huber loss (smooth L1 loss) between true and predicted values. Parameters ---------- - true - input array containing true labels. - pred - input array containing Predicted labels. - compute_full_loss - whether to compute the full loss. If false, a constant term is dropped - in favor of more efficient optimization. Default: ``False``. - axis - the axis along which to compute the log-likelihood loss. If axis is ``-1``, - the log-likelihood loss will be computed along the last dimension. - Default: ``-1``. - reduction - ``'none'``: No reduction will be applied to the output. - ``'mean'``: The output will be averaged. - ``'sum'``: The output will be summed. Default: ``'none'``. - out - optional output array, for writing the result to. It must have a shape + true: array_like + The true (ground truth) values. + pred : array_like + The predicted values by the model. + delta : float, optional + The threshold parameter that determines the point where the loss transitions fro + -m + squared error to absolute error. Default is 1.0. + reduction : str, optional + The type of reduction to apply to the loss. Possible values are "mean" (default) + and "sum". + out : array_like, optional + Optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret - The binary log-likelihood loss between the given distributions. - + ret : array_like + The Huber loss between the true and predicted values. Examples -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.log_poisson_loss(x, z)) - ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) + >>> true = ivy.array([2, 4, 7, 1]) + >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) + >>> huber_loss(true, pred, delta=1.0) + ivy.array([0.125, 0.125, 0.5 , 0.125]) - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) - ivy.array(1.1573164) + >>> huber_loss(true, pred, delta=2.0) + ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + + >>> huber_loss(true, pred, delta=0.5) + ivy.array([0.25 , 0.25 , 0. , 0.125]) """ - try: - assert true.shape == pred.shape - except ValueError: - raise ValueError( - "`pred` and `true` must have the same shape, received " - f"({pred.shape} vs {true.shape})." - ) + abs_diff = ivy.abs(true - pred) + quadratic_loss = 0.5 * (abs_diff**2) + linear_loss = delta * (abs_diff - 0.5 * delta) + loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) - loss = ivy.exp(pred) - pred * true - if compute_full_loss: - stirling_approx = ( - (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) - ) - cond = ivy.logical_and(true >= 0.0, true <= 1.0) - loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, axis=axis, out=out) + return ivy.sum(loss, out=out) elif reduction == "mean": - return ivy.mean(loss, axis=axis, out=out) + return ivy.mean(loss, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss @@ -154,67 +135,86 @@ def l1_loss( return ivy.inplace_update(out, loss) if out is not None else loss +# log_poisson_loss @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def huber_loss( +def log_poisson_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - delta: Optional[float] = 1.0, - reduction: Optional[str] = "mean", + compute_full_loss: bool = False, + axis: int = -1, + reduction: str = "none", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Huber loss (smooth L1 loss) between true and predicted values. + Compute the log-likelihood loss between the prediction and the target under the + assumption that the target has a Poisson distribution. Caveat: By default, this is + not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect + for optimization, but does not play well with relative loss comparisons. To compute + an approximation of the log factorial term, specify ``compute_full_loss=True`` to + enable Stirling's Approximation. Parameters ---------- - true: array_like - The true (ground truth) values. - pred : array_like - The predicted values by the model. - delta : float, optional - The threshold parameter that determines the point where the loss transitions fro - -m - squared error to absolute error. Default is 1.0. - reduction : str, optional - The type of reduction to apply to the loss. Possible values are "mean" (default) - and "sum". - out : array_like, optional - Optional output array, for writing the result to. It must have a shape + true + input array containing true labels. + pred + input array containing Predicted labels. + compute_full_loss + whether to compute the full loss. If false, a constant term is dropped + in favor of more efficient optimization. Default: ``False``. + axis + the axis along which to compute the log-likelihood loss. If axis is ``-1``, + the log-likelihood loss will be computed along the last dimension. + Default: ``-1``. + reduction + ``'none'``: No reduction will be applied to the output. + ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'``. + out + optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret : array_like - The Huber loss between the true and predicted values. + ret + The binary log-likelihood loss between the given distributions. + Examples -------- - >>> true = ivy.array([2, 4, 7, 1]) - >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) - >>> huber_loss(true, pred, delta=1.0) - ivy.array([0.125, 0.125, 0.5 , 0.125]) - - >>> huber_loss(true, pred, delta=2.0) - ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.log_poisson_loss(x, z)) + ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) - >>> huber_loss(true, pred, delta=0.5) - ivy.array([0.25 , 0.25 , 0. , 0.125]) + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) + ivy.array(1.1573164) """ - abs_diff = ivy.abs(true - pred) - quadratic_loss = 0.5 * (abs_diff**2) - linear_loss = delta * (abs_diff - 0.5 * delta) - loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) + try: + assert true.shape == pred.shape + except ValueError: + raise ValueError( + "`pred` and `true` must have the same shape, received " + f"({pred.shape} vs {true.shape})." + ) + loss = ivy.exp(pred) - pred * true + if compute_full_loss: + stirling_approx = ( + (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) + ) + cond = ivy.logical_and(true >= 0.0, true <= 1.0) + loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, out=out) + return ivy.sum(loss, axis=axis, out=out) elif reduction == "mean": - return ivy.mean(loss, out=out) + return ivy.mean(loss, axis=axis, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index 46b3e98ff1e70..c3a21d4a63e19 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -33,667 +33,91 @@ from ivy.utils.exceptions import handle_exceptions -@handle_exceptions -@handle_nestable -@handle_partial_mixed_function -@handle_array_like_without_promotion -@handle_view -@inputs_to_ivy_arrays -@handle_array_function -def flatten( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - start_dim: Optional[int] = 0, - end_dim: Optional[int] = -1, - order: Optional[str] = "C", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Flattens input by reshaping it into a one-dimensional tensor. If start_dim or - end_dim are passed, only dimensions starting with start_dim and ending with end_dim - are flattened. The order of elements in input is unchanged. - - Parameters - ---------- - x - input array to flatten. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - start_dim - first dim to flatten. If not set, defaults to 0. - end_dim - last dim to flatten. If not set, defaults to -1. - order - Read the elements of the input container using this index order, - and place the elements into the reshaped array using this index order. - ‘C’ means to read / write the elements using C-like index order, - with the last axis index changing fastest, back to the first axis index - changing slowest. - ‘F’ means to read / write the elements using Fortran-like index order, with - the first index changing fastest, and the last index changing slowest. - Note that the ‘C’ and ‘F’ options take no account of the memory layout - of the underlying array, and only refer to the order of indexing. - Default order is 'C' - out - optional output array, for writing the result to. - - Returns - ------- - ret - the flattened array over the specified dimensions. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x) - ivy.array([1, 2, 3, 4]) - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x, order='F') - ivy.array([1, 3, 2, 4]) - - >>> x = ivy.array( - [[[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12]], - - [[ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]]], - - - [[[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17]], - - [[ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]]], - - - [[[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1]], - - [[19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]]] - ) - >>> ivy.flatten(x, start_dim = 1, end_dim = 2) - ivy.array( - [[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12], - [ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]], - - [[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17], - [ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]], - - [[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1], - [19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]])) - """ - if x.shape == (): - x = ivy.reshape(x, (1, -1))[0, :] - if start_dim == end_dim: - return ivy.inplace_update(out, x) if ivy.exists(out) else x - if start_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" - ) - if end_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" - ) - if start_dim < 0: - start_dim = len(x.shape) + start_dim - if end_dim < 0: - end_dim = len(x.shape) + end_dim - c = 1 - for i in range(start_dim, end_dim + 1): - c *= x.shape[i] - lst = [c] - if start_dim != 0: - for i in range(0, start_dim): - lst.insert(i, x.shape[i]) - for i in range(end_dim + 1, len(x.shape)): - lst.insert(i, x.shape[i]) - return ivy.reshape(x, tuple(lst), order=order, out=out) - - -flatten.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def moveaxis( - a: Union[ivy.Array, ivy.NativeArray], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Move axes of an array to new positions.. - - Parameters - ---------- - a - The array whose axes should be reordered. - source - Original positions of the axes to move. These must be unique. - destination - Destination positions for each of the original axes. - These must also be unique. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array with moved axes. This array is a view of the input array. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.zeros((3, 4, 5)) - >>> ivy.moveaxis(x, 0, -1).shape - (4, 5, 3) - >>> ivy.moveaxis(x, -1, 0).shape - (5, 3, 4) - """ - return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def heaviside( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Heaviside step function for each element in x1. - - Parameters - ---------- - x1 - input array. - x2 - values to use where x1 is zero. - out - optional output array, for writing the result to. - - Returns - ------- - ret - output array with element-wise Heaviside step function of x1. - This is a scalar if both x1 and x2 are scalars. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([0.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0.0000, 0.5000, 1.0000]) - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([1.2, -2.0, 3.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0., -2., 1.]) - """ - return ivy.current_backend().heaviside(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def flipud( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Flip array in the up/down direction. Flip the entries in each column in the up/down - direction. Rows are preserved, but appear in a different order than before. - - Parameters - ---------- - m - The array to be flipped. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array corresponding to input array with elements - order reversed along axis 0. - - Examples - -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.flipud(m) - ivy.array([[ 0., 0., 3.], - [ 0., 2., 0.], - [ 1., 0., 0.]]) - """ - return ivy.current_backend().flipud(m, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def vstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: - """ - Stack arrays in sequence vertically (row wise). - - Parameters - ---------- - arrays - Sequence of arrays to be stacked. - - Returns - ------- - ret - The array formed by stacking the given arrays. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.vstack((x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4]]) - >>> ivy.vstack((x, y, x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [1, 2, 3], - [2, 3, 4]]) - - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.vstack(y)) - ivy.array([[5, 6], - [7, 8]]) - """ - return ivy.current_backend().vstack(arrays, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def hstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: - """ - Stack arrays in sequence horizotally (column wise). - - Parameters - ---------- - arrays - Sequence of arrays to be stacked. - - Returns - ------- - ret - The array formed by stacking the given arrays. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.hstack((x, y)) - ivy.array([1, 2, 3, 2, 3, 4]) - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0, 0, 0]) - >>> ivy.hstack((x, y, x)) - ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.hstack(y)) - ivy.array([[5, 6, 7, 8]]) - """ - return ivy.current_backend().hstack(arrays, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def rot90( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is - from the first towards the second axis. - - Parameters - ---------- - m - Input array of two or more dimensions. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - k - Number of times the array is rotated by 90 degrees. - axes - The array is rotated in the plane defined by the axes. Axes must be - different. - out - optional output container, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - A rotated view of m. - - Examples - -------- - With :code:`ivy.Array` input: - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) - With :code:`ivy.NativeArray` input: - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.native_array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) - """ - return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def top_k( - x: Union[ivy.Array, ivy.NativeArray], - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[tuple] = None, -) -> Tuple[ivy.Array, ivy.NativeArray]: - """ - Return the `k` largest elements of the given input array along a given axis. - - Parameters - ---------- - x - The array to compute top_k for. - k - Number of top elements to retun must not exceed the array size. - axis - The axis along which we must return the top elements default value is 1. - largest - If largest is set to False we return k smallest elements of the array. - sorted - If sorted is set to True we return the elements in sorted order. - out: - Optional output tuple, for writing the result to. Must have two arrays inside, - with a shape that the returned tuple broadcast to. - - Returns - ------- - ret - A named tuple with values and indices of top k elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 2) - >>> print(y) - top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) - - >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) - >>> y = ivy.top_k(x, 2, axis=1, largest=False) - >>> print(y) - top_k(values=ivy.array([[-2., 0.], - [-8., -1.]]), indices=ivy.array([[0, 3], - [0, 2]])) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 3) - >>> print(y) - top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) - >>> y = x.top_k(2) - >>> print(y) - [{ - a: ivy.array([2, -1]), - b: ivy.array([5., 4.]) - }, { - a: ivy.array([1, 0]), - b: ivy.array([1, 0]) - }] - """ - return current_backend(x).top_k( - x, k, axis=axis, largest=largest, sorted=sorted, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fliplr( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Flip array in the left/right direction. Flip the entries in each column in the - left/right direction. Columns are preserved, but appear in a different order than - before. - - Parameters - ---------- - m - The array to be flipped. Must be at least 2-D. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array corresponding to input array with elements - order reversed along axis 1. - - Examples - -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.fliplr(m) - ivy.array([[0, 0, 1], - [0, 2, 0], - [3, 0, 0]]) - """ - return ivy.current_backend().fliplr(m, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def i0( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Bessel i0 function of x element-wise. - - Parameters - ---------- - x - Array input. - out - optional output array, for writing the result to. +# --- Helpers --- # +# --------------- # - Returns - ------- - ret - Array with the modified Bessel function - evaluated at each of the elements of x. - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.i0(x) - ivy.array([1.26606588, 2.2795853 , 4.88079259]) - """ - return ivy.current_backend(x).i0(x, out=out) +def _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, +): + ivy.utils.assertions.check_true( + callable(mode) + or mode + in [ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + message="the provided mode is not supported", + ) + _check_tuple_arg(pad_width, "pad_width") + if mode not in ["dilated"]: + ivy.utils.assertions.check_true( + all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), + message="the pad_widths must be greater or equal to zero", + ) + if mode in ["maximum", "mean", "median", "minimum"]: + _check_tuple_arg(stat_length, "stat_length") + ivy.utils.assertions.check_true( + all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), + message="the stat lengths must be greater than zero", + ) + elif mode == "constant": + _check_tuple_arg(constant_values, "constant_values", force_integer=False) + elif mode == "linear_ramp": + _check_tuple_arg(end_values, "end_values", force_integer=False) + ivy.utils.assertions.check_true( + reflect_type in ["even", "odd"], + message="the provided reflect_type is not supported", + ) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) +def _check_bounds(shape0, shape1, strides1, itemsize): + numel0 = math.prod(shape0) + ndim1 = len(shape1) + return ( + sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize + <= numel0 * itemsize + ) -def _set_pad_area(padded, axis, width_pair, value_pair): - if width_pair[0] > 0: - left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) - padded[left_slice] = value_pair[0] - if width_pair[1] > 0: - right_slice = _slice_at_axis( - slice(padded.shape[axis] - width_pair[1], None), axis - ) - padded[right_slice] = value_pair[1] - return padded +def _check_tuple_arg(arg, name, force_integer=True): + is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype + flag_assert = False + if isinstance(arg, (tuple, list)): + for nested in arg: + if isinstance(nested, (tuple, list)): + for sub_nested in nested: + if not is_scalar(sub_nested): + flag_assert = True + break + elif not is_scalar(nested): + flag_assert = True + elif not is_scalar(arg): + flag_assert = True + if flag_assert: + if not force_integer: + raise ivy.utils.exceptions.IvyException( + name + " should be scalar, tuple of scalars or tuple of scalar tuples" + ) + else: + raise ivy.utils.exceptions.IvyException( + name + " should be int, tuple of ints or tuple of int tuples" + ) def _get_edges(padded, axis, width_pair): @@ -758,6 +182,82 @@ def _get_stats(padded, axis, width_pair, length_pair, stat_func): return left_stat, right_stat +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = ivy.full(new_shape, padding_value, dtype=operand.dtype) + src_indices = ivy.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) + + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = ivy.constant_pad(padded, pad_width, value=padding_value) + return padded + + +def _interleave(a, b, axis): + assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 + a_pad = [(0, 0, 0)] * a.ndim + b_pad = [(0, 0, 0)] * b.ndim + a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) + b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) + a = _interior_pad(a, 0.0, a_pad) + b = _interior_pad(b, 0.0, b_pad) + return ivy.add(a, b) + + +def _pad_simple(array, pad_width, fill_value=None): + new_shape = tuple( + [left + size + right for size, (left, right) in zip(array.shape, pad_width)] + ) + padded = ivy.zeros(new_shape, dtype=array.dtype) + if fill_value is not None: + padded = ivy.ones_like(padded) * fill_value + original_area_slice = tuple( + [ + slice(left, left + size) + for size, (left, right) in zip(array.shape, pad_width) + ] + ) + padded[original_area_slice] = array + return padded, original_area_slice + + +def _set_pad_area(padded, axis, width_pair, value_pair): + if width_pair[0] > 0: + left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) + padded[left_slice] = value_pair[0] + if width_pair[1] > 0: + right_slice = _slice_at_axis( + slice(padded.shape[axis] - width_pair[1], None), axis + ) + padded[right_slice] = value_pair[1] + return padded + + def _set_reflect_both(padded, axis, width_pair, method, include_edge=False): left_pad, right_pad = width_pair old_length = padded.shape[axis] - right_pad - left_pad @@ -835,35 +335,42 @@ def _set_wrap_both(padded, axis, width_pair): return new_left_pad, new_right_pad, padded -def _pad_simple(array, pad_width, fill_value=None): - new_shape = tuple( - [left + size + right for size, (left, right) in zip(array.shape, pad_width)] - ) - padded = ivy.zeros(new_shape, dtype=array.dtype) - if fill_value is not None: - padded = ivy.ones_like(padded) * fill_value - original_area_slice = tuple( - [ - slice(left, left + size) - for size, (left, right) in zip(array.shape, pad_width) - ] - ) - padded[original_area_slice] = array - return padded, original_area_slice +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides + + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] -def _to_pairs(x, n): +def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): + if axis >= 0: + slices = [slice(None)] * axis + [slice(start, stop, stride)] + else: + slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) + return x[tuple(slices)] + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +def _to_dilated(x, n): if ivy.isscalar(x): - return ((x, x),) * n - elif len(x) == 2 and ivy.isscalar(x[0]): - return ((x[0], x[1]),) * n + return ((x, x, x),) * n + elif len(x) == 3 and ivy.isscalar(x[0]): + return ((x[0], x[1], x[2]),) * n elif len(x) != n: ivy.utils.assertions.check_equal( ivy.asarray(list(x)).shape, - (n, 2), + (n, 3), message=( "tuple argument should contain " - "ndim pairs where ndim is the number of " + "ndim groups where ndim is the number of " "the input's dimensions" ), as_array=False, @@ -871,18 +378,18 @@ def _to_pairs(x, n): return x -def _to_dilated(x, n): +def _to_pairs(x, n): if ivy.isscalar(x): - return ((x, x, x),) * n - elif len(x) == 3 and ivy.isscalar(x[0]): - return ((x[0], x[1], x[2]),) * n + return ((x, x),) * n + elif len(x) == 2 and ivy.isscalar(x[0]): + return ((x[0], x[1]),) * n elif len(x) != n: ivy.utils.assertions.check_equal( ivy.asarray(list(x)).shape, - (n, 3), + (n, 2), message=( "tuple argument should contain " - "ndim groups where ndim is the number of " + "ndim pairs where ndim is the number of " "the input's dimensions" ), as_array=False, @@ -890,395 +397,422 @@ def _to_dilated(x, n): return x -def _check_tuple_arg(arg, name, force_integer=True): - is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype - flag_assert = False - if isinstance(arg, (tuple, list)): - for nested in arg: - if isinstance(nested, (tuple, list)): - for sub_nested in nested: - if not is_scalar(sub_nested): - flag_assert = True - break - elif not is_scalar(nested): - flag_assert = True - elif not is_scalar(arg): - flag_assert = True - if flag_assert: - if not force_integer: - raise ivy.utils.exceptions.IvyException( - name + " should be scalar, tuple of scalars or tuple of scalar tuples" - ) - else: - raise ivy.utils.exceptions.IvyException( - name + " should be int, tuple of ints or tuple of int tuples" - ) +# --- Main --- # +# ------------ # + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@inputs_to_native_shapes +def as_strided( + x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + strides: Sequence[int], + /, +) -> ivy.Array: + """ + Create a copy of the input array with the given shape and strides. + + Parameters + ---------- + x + Input Array. + shape + The shape of the new array. + strides + The strides of the new array (specified in bytes). + + Returns + ------- + ret + Output Array + + Examples + -------- + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> ivy.as_strided(x, (4, 3), (8, 8)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [3, 4, 5], + [4, 5, 6]]) + """ + itemsize = x.itemsize + if not _check_bounds(x.shape, shape, strides, itemsize): + raise ivy.exceptions.IvyException("attempted unsafe memory access") + if any(strides[i] % itemsize != 0 for i in range(len(strides))): + raise ivy.exceptions.IvyException("strides must be multiple of itemsize") + + src = memoryview(ivy.to_numpy(x)).cast("b") + + src_ind = ivy.inner( + ivy.indices(shape).reshape((len(shape), -1)).T, + ivy.array(strides), + ) + src_ind = ivy.expand_dims(src_ind, axis=-1) + src_ind = src_ind + ivy.arange(itemsize) + src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() + temp_list = [src[i] for i in src_ind] + temp_array = ivy.asarray(temp_list, dtype=ivy.int8) + result = bytearray(temp_array.to_numpy()) -def _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, -): - ivy.utils.assertions.check_true( - callable(mode) - or mode - in [ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - message="the provided mode is not supported", - ) - _check_tuple_arg(pad_width, "pad_width") - if mode not in ["dilated"]: - ivy.utils.assertions.check_true( - all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), - message="the pad_widths must be greater or equal to zero", - ) - if mode in ["maximum", "mean", "median", "minimum"]: - _check_tuple_arg(stat_length, "stat_length") - ivy.utils.assertions.check_true( - all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), - message="the stat lengths must be greater than zero", - ) - elif mode == "constant": - _check_tuple_arg(constant_values, "constant_values", force_integer=False) - elif mode == "linear_ramp": - _check_tuple_arg(end_values, "end_values", force_integer=False) - ivy.utils.assertions.check_true( - reflect_type in ["even", "odd"], - message="the provided reflect_type is not supported", + return ivy.reshape( + ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), + shape, ) @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def pad( - input: Union[ivy.Array, ivy.NativeArray], - pad_width: Union[Iterable[Tuple[int]], int], +def associative_scan( + x: Union[ivy.Array, ivy.NativeArray], + fn: Callable, /, *, - mode: Union[ - Literal[ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - Callable, - ] = "constant", - stat_length: Union[Iterable[Tuple[int]], int] = 1, - constant_values: Union[Iterable[Tuple[Number]], Number] = 0, - end_values: Union[Iterable[Tuple[Number]], Number] = 0, - reflect_type: Literal["even", "odd"] = "even", - **kwargs: Optional[Any], + reverse: bool = False, + axis: int = 0, ) -> ivy.Array: """ - Pad an array. + Perform an associative scan over the given array. Parameters ---------- - input - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths - for each axis. - - ((before, after),) yields same before and after pad for each axis. - - pad (integer) is shortcut for before = after = pad width for all axes. - mode - One of the following string values or a user-supplied function. - - "constant": Pads with a constant value. - - "edge": Pads with the input's edge values. - - "linear_ramp": Pads with the linear ramp between end_value - and the input's edge value. - - "maximum": Pads with the maximum value of all or part of the vector - along each axis. - - "mean": Pads with the mean value of all or part of the vector along - each axis. - - "median": Pads with the median value of all or part of the vector - along each axis. - - "minimum": Pads with the minimum value of all or part of the vector - along each axis. - - "reflect": Pads with the reflection mirrored on the first and last - values of the vector along each axis. - - "symmetric": Pads with the reflection of the vector mirrored along - the edge of the input. - - "wrap": Pads with the wrap of the vector along the axis. - The first values are used to pad the end and the end values are used - to pad the beginning. - - "empty": Pads with undefined values. - - : Pads with a user-defined padding function. The padding - function should modify a rank 1 array following the signature - `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: - - `vector` is a rank 1 array already padded with zeros. Padded - values are `vector[:iaxis_pad_width[0]]` and - `vector[-iaxis_pad_width[1]:]`. - - `iaxis_pad_width` is a 2-tuple of ints, where - `iaxis_pad_width[0]` represents the number of values padded at - the beginning of `vector` and `iaxis_pad_width[1]` represents - the number of values padded at the end of `vector`. - - `iaxis` is the axis currently being calculated. - - `kwargs` is a dict of keyword arguments the function requires. - stat_length - Used in "maximum", "mean", "median", and "minimum". Number of values at edge - of each axis used to calculate the statistic value. - - ((before_1, after_1), … (before_N, after_N)) yields unique statistic - lengths for each axis. - - ((before, after),) yields same before and after statistic lengths for - each axis. - - stat_length (integer) is a shortcut for before = after = stat_length - length for all axes. - - None uses the entire axis. - constant_values - Used in "constant". The values to set the padded values for each axis. - - ((before_1, after_1), ... (before_N, after_N)) yields unique pad - constants for each axis. - - ((before, after),) yields same before and after constants for each axis. - - constant (integer) is a shortcut for before = after = constant for - all axes. - end_values - Used in "linear_ramp". The values used for the ending value of the linear_ramp - and that will form the edge of the padded array. - - ((before_1, after_1), ... (before_N, after_N)) yields unique end values - for each axis. - - ((before, after),) yields same before and after end values for each axis - - end (integer) is a shortcut for before = after = end for all axes. - reflect_type - Used in "reflect", and "symmetric". The "even" style is the default with an - unaltered reflection around the edge value. For the "odd" style, the extended - part of the array is created by subtracting the reflected values from two - times the edge value. + x + The array to scan over. + fn + The associative function to apply. + reverse + Whether to scan in reverse with respect to the given axis. + axis + The axis to scan over. + + Returns + ------- + ret + The result of the scan. + """ + elems = [x] + + if reverse: + elems = [ivy.flip(elem, axis=[axis]) for elem in elems] + + def _combine(a, b): + a = a[0] + b = b[0] + if a.shape[axis] == 0: + return [a] + c = fn(a, b) + return [c] + + def _scan(elems): + num_elems = elems[0].shape[axis] + + if num_elems < 2: + return elems + + reduced_elems = _combine( + [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], + [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], + ) + + odd_elems = _scan(reduced_elems) + + if num_elems % 2 == 0: + even_elems = _combine( + [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + else: + even_elems = _combine( + odd_elems, + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + even_elems = [ + ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) + for (elem, result) in zip(elems, even_elems) + ] + return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) + + scans = _scan(elems) + + if reverse: + scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] + + return ivy.reshape(ivy.asarray(scans), elems[0].shape) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_1d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least one dimension. Scalar inputs are converted to + 1-dimensional arrays, whilst higher-dimensional inputs are preserved. + + Parameters + ---------- + arys + One or more input arrays. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - Padded array of the same rank as the input but with shape increased according - to pad_width. + An array, or list of arrays, each with atleast 1D. + Copies are made only if necessary. + + Examples + -------- + >>> ary1 = ivy.array(5) + >>> ivy.atleast_1d(ary1) + ivy.array([5]) + >>> ary2 = ivy.array([[3,4]]) + >>> ivy.atleast_1d(ary2) + ivy.array([[3, 4]]) + >>> ivy.atleast_1d(6,7,8) + [ivy.array([6]), ivy.array([7]), ivy.array([8])] + """ + return ivy.current_backend().atleast_1d(*arys, copy=copy) - Both the description and the type hints above assume an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_2d( + *arys: Union[ivy.Array, ivy.NativeArray], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least two dimension. Scalar inputs are converted to + 2-dimensional arrays, whilst higher-dimensional inputs are preserved. + + Parameters + ---------- + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have two or more + dimensions are preserved. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + + Returns + ------- + ret + An array, or list of arrays, each with atleast 2D. + Copies are made only if necessary. Examples -------- - With :class:`ivy.Array` input: + >>> ary1 = ivy.array(5) + >>> ivy.atleast_2d(ary1) + ivy.array([[5]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_2d(ary2) + ivy.array([[[3, 4]]]) + >>> ivy.atleast_2d(6,7,8) + [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] + """ + return ivy.current_backend().atleast_2d(*arys, copy=copy) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 3, 0, 0], - [0, 0, 4, 5, 6, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="reflect") - >>> print(y) - ivy.array([[6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1], - [6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1]]) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_3d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least three dimension. Scalar inputs are converted + to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="symmetric") - >>> print(y) - ivy.array([[2, 1, 1, 2, 3, 3, 2], - [2, 1, 1, 2, 3, 3, 2], - [5, 4, 4, 5, 6, 6, 5], - [5, 4, 4, 5, 6, 6, 5]]) + Parameters + ---------- + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have three or more + dimensions are preserved. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. - With :class:`ivy.NativeArray` input: + Returns + ------- + ret + An array, or list of arrays, each with a.ndim >= 3. Copies + are avoided where possible, and views with three or more + dimensions are returned. For example, a 1-D array of shape + (N,) becomes a view of shape (1, N, 1), and a 2-D array of + shape (M, N) becomes a view of shape (M, N, 1). - >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) - >>> print(y) - ivy.array([[7, 7, 7, 7, 7, 7, 7], - [7, 7, 1, 2, 3, 7, 7], - [7, 7, 4, 5, 6, 7, 7], - [7, 7, 7, 7, 7, 7, 7]]) + Examples + -------- + >>> ary1 = ivy.array([5,6]) + >>> ivy.atleast_3d(ary1) + ivy.array([[[5], + [6]]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_3d(ary2) + ivy.array([[[3, 4]]]) + >>> ary3 = ivy.array([[3,4],[9,10]]) + >>> ivy.atleast_3d(6,7,ary3) + [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], + [ 4]], - With :class:`ivy.Container` input: + [[ 9], + [10]]])] + """ + return ivy.current_backend().atleast_3d(*arys, copy=copy) - >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) - >>> padding = (1, 1) - >>> y = ivy.pad(x, padding, mode="constant") - >>> print(y) - { - a: ivy.array([0, 0, 1, 2, 0]), - b: ivy.array([0, 4, 5, 6, 0]) - } + +@handle_exceptions +@inputs_to_native_shapes +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: """ - _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, - ) - if mode == "dilated": - pad_width = _to_dilated(pad_width, input.ndim) - if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: - constant_values = ivy.asarray(constant_values, dtype=input.dtype) - return _interior_pad(input, constant_values, pad_width) - pad_width = _to_pairs(pad_width, len(input.shape)) - if callable(mode): - func = mode - padded, _ = _pad_simple(input, pad_width, fill_value=0) - for axis in range(padded.ndim): - padded = ivy.moveaxis(padded, axis, -1) - inds = ivy.ndindex(padded.shape[:-1]) - for ind in inds: - padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) - return padded - padded, original_area_slice = _pad_simple(input, pad_width) - axes = range(padded.ndim) - stat_functions = { - "maximum": ivy.max, - "minimum": ivy.min, - "mean": ivy.mean, - "median": ivy.median, - } - if mode == "constant": - constant_values = _to_pairs(constant_values, padded.ndim) - constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) - for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): - padded = _set_pad_area(padded, axis, width_pair, value_pair) - elif mode == "empty": - pass - elif mode == "edge": - for axis, width_pair in zip(axes, pad_width): - edge_pair = _get_edges(padded, axis, width_pair) - padded = _set_pad_area(padded, axis, width_pair, edge_pair) - elif mode == "linear_ramp": - end_values = _to_pairs(end_values, padded.ndim) - for axis, width_pair, value_pair in zip(axes, pad_width, end_values): - ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) - padded = _set_pad_area(padded, axis, width_pair, ramp_pair) - elif mode in stat_functions: - func = stat_functions[mode] - stat_length = _to_pairs(stat_length, padded.ndim) - if mode == "median": - ivy.utils.assertions.check_true( - ivy.is_float_dtype(input), - message="median interpolation is only supported for floats", - ) - for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): - stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) - padded = _set_pad_area(padded, axis, width_pair, stat_pair) - elif mode in {"reflect", "symmetric"}: - include_edge = True if mode == "symmetric" else False - for axis, (left_index, right_index) in zip(axes, pad_width): - if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): - edge_pair = _get_edges(padded, axis, (left_index, right_index)) - padded = _set_pad_area( - padded, axis, (left_index, right_index), edge_pair - ) - continue - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_reflect_both( - padded, axis, (left_index, right_index), reflect_type, include_edge - ) - elif mode == "wrap": - for axis, (left_index, right_index) in zip(axes, pad_width): - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_wrap_both( - padded, axis, (left_index, right_index) - ) - return padded + Broadcasts shapes. + Parameters + ---------- + shapes + The shapes to broadcast. -pad.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Returns + ------- + ret + The broadcasted shape. + + Examples + -------- + >>> x = [(3, 3), (3, 1)] + >>> print(ivy.broadcast_shapes(*x)) + (3, 3) + + >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) + (3, 3) + """ + return ivy.current_backend().broadcast_shapes(*shapes) @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view +@handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def vsplit( - ary: Union[ivy.Array, ivy.NativeArray], - indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], +def choose( + arr: Union[ivy.Array, ivy.NativeArray], + choices: Union[ivy.Array, ivy.NativeArray], /, *, - copy: Optional[bool] = None, -) -> List[ivy.Array]: + out: None = None, + mode: Union[str, None] = None, +) -> ivy.Array: """ - Split an array vertically into multiple sub-arrays. + Take values from the input array by matching 1d index and data slices. Parameters ---------- - ary - Array input. - indices_or_sections - If indices_or_sections is an integer n, the array is split into n - equal sections, provided that n must be a divisor of the split axis. - If indices_or_sections is a sequence of ints or 1-D array, - then input is split at each of the indices. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + arr + The source array. + choices + The indices of the values to extract. + out + The output array. + mode + One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices + will be handled. Returns ------- ret - input array split vertically. + The returned array has the same shape as `indices`. Examples -------- - >>> ary = ivy.array( - [[[0., 1.], - [2., 3.]], - [[4., 5.], - [6., 7.]]] - ) - >>> ivy.vsplit(ary, 2) - [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) + >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], + [20, 21, 22, 23], [30, 31, 32, 33]]) + >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) + ivy.array([20, 1, 12, 3]) """ - return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) + return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def concat_from_sequence( + input_sequence: Union[ + Tuple[Union[ivy.Array, ivy.NativeArray]], + List[Union[ivy.Array, ivy.NativeArray]], + ], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Concatenate a sequence of arrays along a new or an existing axis. + + Parameters + ---------- + input_sequence + A sequence of arrays. + new_axis + Insert and concatenate on a new axis or not, + default 0 means do not insert new axis. + new_axis = 0: concatenate + new_axis = 1: stack + axis + axis along which the arrays will be concatenated. + + out + optional output array, for writing the result to. + + Returns + ------- + ret + Output Array + """ + return current_backend(input_sequence).concat_from_sequence( + input_sequence, new_axis=new_axis, axis=axis, out=out + ) @handle_exceptions @@ -1336,243 +870,451 @@ def dsplit( return ivy.current_backend(ary).dsplit(ary, indices_or_sections, copy=copy) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Stack arrays in sequence depth wise (along third axis). + + Parameters + ---------- + arrays + Sequence of arrays to be stacked. + + Returns + ------- + ret + The array formed by stacking the given arrays. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2], + [2, 3], + [3, 4]]]) + >>> x = ivy.array([[1], [2], [3]]) + >>> y = ivy.array([[2], [3], [4]]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2]], + [[2, 3]], + [[3, 4]]]) + """ + return ivy.current_backend().dstack(arrays, out=out) + + +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_view +@handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_device_shifting -def atleast_1d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], +def expand( + x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape], + /, + *, + copy: Optional[bool] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Broadcast the input Array following the given shape and the broadcast rule. + + Parameters + ---------- + x + Array input. + shape + A 1-D Array indicates the shape you want to expand to, + following the broadcast rule. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Output Array + """ + return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back +@handle_array_function +def fill_diagonal( + a: Union[ivy.Array, ivy.NativeArray], + v: Union[int, float], + /, + *, + wrap: bool = False, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Fill the main diagonal of the given array of any dimensionality.. + + Parameters + ---------- + a + Array at least 2D. + v + The value to write on the diagonal. + wrap + The diagonal ‘wrapped’ after N columns for tall matrices. + + Returns + ------- + ret + Array with the diagonal filled. + """ + return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) + + +@handle_exceptions +@handle_nestable +@handle_partial_mixed_function +@handle_array_like_without_promotion +@handle_view +@inputs_to_ivy_arrays +@handle_array_function +def flatten( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, copy: Optional[bool] = None, -) -> List[ivy.Array]: + start_dim: Optional[int] = 0, + end_dim: Optional[int] = -1, + order: Optional[str] = "C", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert inputs to arrays with at least one dimension. Scalar inputs are converted to - 1-dimensional arrays, whilst higher-dimensional inputs are preserved. + Flattens input by reshaping it into a one-dimensional tensor. If start_dim or + end_dim are passed, only dimensions starting with start_dim and ending with end_dim + are flattened. The order of elements in input is unchanged. Parameters ---------- - arys - One or more input arrays. + x + input array to flatten. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. + start_dim + first dim to flatten. If not set, defaults to 0. + end_dim + last dim to flatten. If not set, defaults to -1. + order + Read the elements of the input container using this index order, + and place the elements into the reshaped array using this index order. + ‘C’ means to read / write the elements using C-like index order, + with the last axis index changing fastest, back to the first axis index + changing slowest. + ‘F’ means to read / write the elements using Fortran-like index order, with + the first index changing fastest, and the last index changing slowest. + Note that the ‘C’ and ‘F’ options take no account of the memory layout + of the underlying array, and only refer to the order of indexing. + Default order is 'C' + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with atleast 1D. - Copies are made only if necessary. + the flattened array over the specified dimensions. Examples -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_1d(ary1) - ivy.array([5]) - >>> ary2 = ivy.array([[3,4]]) - >>> ivy.atleast_1d(ary2) - ivy.array([[3, 4]]) - >>> ivy.atleast_1d(6,7,8) - [ivy.array([6]), ivy.array([7]), ivy.array([8])] + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x) + ivy.array([1, 2, 3, 4]) + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x, order='F') + ivy.array([1, 3, 2, 4]) + + >>> x = ivy.array( + [[[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12]], + + [[ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]]], + + + [[[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17]], + + [[ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]]], + + + [[[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1]], + + [[19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]]] + ) + >>> ivy.flatten(x, start_dim = 1, end_dim = 2) + ivy.array( + [[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12], + [ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]], + + [[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17], + [ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]], + + [[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1], + [19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]])) """ - return ivy.current_backend().atleast_1d(*arys, copy=copy) + if x.shape == (): + x = ivy.reshape(x, (1, -1))[0, :] + if start_dim == end_dim: + return ivy.inplace_update(out, x) if ivy.exists(out) else x + if start_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" + ) + if end_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" + ) + if start_dim < 0: + start_dim = len(x.shape) + start_dim + if end_dim < 0: + end_dim = len(x.shape) + end_dim + c = 1 + for i in range(start_dim, end_dim + 1): + c *= x.shape[i] + lst = [c] + if start_dim != 0: + for i in range(0, start_dim): + lst.insert(i, x.shape[i]) + for i in range(end_dim + 1, len(x.shape)): + lst.insert(i, x.shape[i]) + return ivy.reshape(x, tuple(lst), order=order, out=out) @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion +@handle_view @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dstack( - arrays: Sequence[ivy.Array], +def fliplr( + m: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Stack arrays in sequence depth wise (along third axis). + Flip array in the left/right direction. Flip the entries in each column in the + left/right direction. Columns are preserved, but appear in a different order than + before. Parameters ---------- - arrays - Sequence of arrays to be stacked. + m + The array to be flipped. Must be at least 2-D. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - The array formed by stacking the given arrays. + Array corresponding to input array with elements + order reversed along axis 1. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2], - [2, 3], - [3, 4]]]) - >>> x = ivy.array([[1], [2], [3]]) - >>> y = ivy.array([[2], [3], [4]]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2]], - [[2, 3]], - [[3, 4]]]) + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.fliplr(m) + ivy.array([[0, 0, 1], + [0, 2, 0], + [3, 0, 0]]) """ - return ivy.current_backend().dstack(arrays, out=out) + return ivy.current_backend().fliplr(m, copy=copy, out=out) @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_view +@handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def atleast_2d( - *arys: Union[ivy.Array, ivy.NativeArray], +def flipud( + m: Union[ivy.Array, ivy.NativeArray], + /, + *, copy: Optional[bool] = None, -) -> List[ivy.Array]: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Convert inputs to arrays with at least two dimension. Scalar inputs are converted to - 2-dimensional arrays, whilst higher-dimensional inputs are preserved. + Flip array in the up/down direction. Flip the entries in each column in the up/down + direction. Rows are preserved, but appear in a different order than before. Parameters ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have two or more - dimensions are preserved. + m + The array to be flipped. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with atleast 2D. - Copies are made only if necessary. + Array corresponding to input array with elements + order reversed along axis 0. Examples -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_2d(ary1) - ivy.array([[5]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_2d(ary2) - ivy.array([[[3, 4]]]) - >>> ivy.atleast_2d(6,7,8) - [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.flipud(m) + ivy.array([[ 0., 0., 3.], + [ 0., 2., 0.], + [ 1., 0., 0.]]) """ - return ivy.current_backend().atleast_2d(*arys, copy=copy) + return ivy.current_backend().flipud(m, copy=copy, out=out) -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def atleast_3d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], - copy: Optional[bool] = None, -) -> List[ivy.Array]: +def fold( + x: Union[ivy.Array, ivy.NativeArray], + /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert inputs to arrays with at least three dimension. Scalar inputs are converted - to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. + Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, + refolds the n-mode unfolded tensor into the original tensor of the specified shape. - Parameters - ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have three or more - dimensions are preserved. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + Parameters + ---------- + input + unfolded tensor of shape ``(shape[mode], -1)`` + mode + the mode of the unfolding + shape + shape of the original tensor before unfolding + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with a.ndim >= 3. Copies - are avoided where possible, and views with three or more - dimensions are returned. For example, a 1-D array of shape - (N,) becomes a view of shape (1, N, 1), and a 2-D array of - shape (M, N) becomes a view of shape (M, N, 1). - - Examples - -------- - >>> ary1 = ivy.array([5,6]) - >>> ivy.atleast_3d(ary1) - ivy.array([[[5], - [6]]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_3d(ary2) - ivy.array([[[3, 4]]]) - >>> ary3 = ivy.array([[3,4],[9,10]]) - >>> ivy.atleast_3d(6,7,ary3) - [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], - [ 4]], - - [[ 9], - [10]]])] + folded_tensor of shape `shape` """ - return ivy.current_backend().atleast_3d(*arys, copy=copy) + full_shape = list(shape) + mode_dim = full_shape.pop(mode) + full_shape.insert(0, mode_dim) + return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def take_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - axis: int, +def heaviside( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Take values from the input array by matching 1d index and data slices. + Compute the Heaviside step function for each element in x1. Parameters ---------- - arr - The source array. - indices - The indices of the values to extract. - axis - The axis over which to select values. - If axis is None, arr is treated as a flattened 1D array. - mode - One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices - will be handled. + x1 + input array. + x2 + values to use where x1 is zero. out - The output array. + optional output array, for writing the result to. Returns ------- ret - The returned array has the same shape as `indices`. + output array with element-wise Heaviside step function of x1. + This is a scalar if both x1 and x2 are scalars. Examples -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> y = ivy.take_along_axis(arr, indices, 1) - >>> print(y) - ivy.array([[4, 3, 3], [1, 1, 1]]) + With :class:`ivy.Array` input: + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([0.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0.0000, 0.5000, 1.0000]) + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([1.2, -2.0, 3.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0., -2., 1.]) """ - return ivy.current_backend(arr).take_along_axis( - arr, indices, axis, mode=mode, out=out - ) + return ivy.current_backend().heaviside(x1, x2, out=out) @handle_exceptions @@ -1634,524 +1376,595 @@ def hsplit( return ivy.current_backend(ary).hsplit(ary, indices_or_sections, copy=copy) -@handle_exceptions -@inputs_to_native_shapes -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def hstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> ivy.Array: """ - Broadcasts shapes. + Stack arrays in sequence horizotally (column wise). Parameters ---------- - shapes - The shapes to broadcast. + arrays + Sequence of arrays to be stacked. Returns ------- ret - The broadcasted shape. + The array formed by stacking the given arrays. Examples -------- - >>> x = [(3, 3), (3, 1)] - >>> print(ivy.broadcast_shapes(*x)) - (3, 3) - - >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) - (3, 3) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.hstack((x, y)) + ivy.array([1, 2, 3, 2, 3, 4]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0, 0, 0]) + >>> ivy.hstack((x, y, x)) + ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.hstack(y)) + ivy.array([[5, 6, 7, 8]]) """ - return ivy.current_backend().broadcast_shapes(*shapes) + return ivy.current_backend().hstack(arrays, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_device_shifting -def expand( +def i0( x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape], /, *, - copy: Optional[bool] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Broadcast the input Array following the given shape and the broadcast rule. + Compute the Bessel i0 function of x element-wise. Parameters ---------- x Array input. - shape - A 1-D Array indicates the shape you want to expand to, - following the broadcast rule. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. out optional output array, for writing the result to. Returns ------- ret - Output Array + Array with the modified Bessel function + evaluated at each of the elements of x. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> ivy.i0(x) + ivy.array([1.26606588, 2.2795853 , 4.88079259]) """ - return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) + return ivy.current_backend(x).i0(x, out=out) -@handle_exceptions @handle_nestable +@handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays -def put_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - values: Union[ivy.Array, ivy.NativeArray], - axis: int, +@handle_array_function +@handle_device_shifting +def matricize( + x: Union[ivy.Array, ivy.NativeArray], /, + row_modes: Sequence[int], + column_modes: Optional[Sequence[int]] = None, *, - mode: str = "raise", out: Optional[ivy.Array] = None, -) -> None: +) -> ivy.Array: """ - Put values into the input array by matching 1d index and data slices along a - specified axis. + Matricizes the given tensor. Parameters ---------- - arr : array_like - The input array to modify. - indices : array_like - The indices of the values to put into `arr`. - values : array_like - The values to put into `arr`. - axis : int - The axis over which to put the `values`. - mode : {'raise', 'wrap', 'clip'}, optional - Specifies how out-of-bounds indices will be handled. - The following modes are available: - - - 'raise': a `ValueError` is raised when an index is out of bounds. - - 'wrap': the index is wrapped around to the corresponding index - at the other end of the axis. - - 'clip': the index is clipped to the closest in-bounds index. - out : ndarray, optional - Output array in which to place the result. - If not specified, a new array is created. + x + the input tensor + row_modes + modes to use as row of the matrix (in the desired order) + column_modes + modes to use as column of the matrix, in the desired order + if None, the modes not in `row_modes` will be used in ascending order + out + optional output array, for writing the result to. - Returns + ret ------- - None - - Examples - -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) - >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') - >>> print(arr) - ivy.array([[3, 7, 5], - [6, 4, 1]]) + ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) """ - if out is None: - out = ivy.zeros_like(arr) - - indices = ivy.expand_dims(indices, axis=axis) - values = ivy.expand_dims(values, axis=axis) - - stacked = ivy.concat((arr, values), axis=axis) - - sorted_indices = ivy.argsort(indices, axis=axis) - sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), - sorted_stacked, - arr, - ) - - if mode == "clip": - indices = ivy.clip(indices, 0, arr.shape[axis] - 1) - elif mode == "wrap": - indices = ivy.mod(indices, arr.shape[axis]) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values - ) + ndims = len(x.shape) + row_indices = list(row_modes) - ivy.assign(out, arr) + if column_modes: + column_indices = list(column_modes) + else: + column_indices = [i for i in range(ndims) if i not in row_indices] + if sorted(column_indices + row_indices) != list(range(ndims)): + msg = ( + "If you provide both column and row modes for the matricization then" + " column_modes + row_modes must contain all the modes of the tensor." + f" Yet, got row_modes={row_modes} and column_modes={column_modes}." + ) + raise ValueError(msg) + row_size, column_size = 1, 1 + row_size = int(ivy.prod([x.shape[i] for i in row_indices])) + column_size = int(ivy.prod([x.shape[i] for i in column_indices])) -def _check_bounds(shape0, shape1, strides1, itemsize): - numel0 = math.prod(shape0) - ndim1 = len(shape1) - return ( - sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize - <= numel0 * itemsize + return ivy.reshape( + ivy.permute_dims(x, row_indices + column_indices), + (row_size, column_size), + out=out, ) -@handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@inputs_to_native_shapes -def as_strided( - x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - strides: Sequence[int], +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def moveaxis( + a: Union[ivy.Array, ivy.NativeArray], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], /, -) -> ivy.Array: + *, + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Create a copy of the input array with the given shape and strides. + Move axes of an array to new positions.. Parameters ---------- - x - Input Array. - shape - The shape of the new array. - strides - The strides of the new array (specified in bytes). + a + The array whose axes should be reordered. + source + Original positions of the axes to move. These must be unique. + destination + Destination positions for each of the original axes. + These must also be unique. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - Output Array + Array with moved axes. This array is a view of the input array. Examples -------- - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> ivy.as_strided(x, (4, 3), (8, 8)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [3, 4, 5], - [4, 5, 6]]) + With :class:`ivy.Array` input: + >>> x = ivy.zeros((3, 4, 5)) + >>> ivy.moveaxis(x, 0, -1).shape + (4, 5, 3) + >>> ivy.moveaxis(x, -1, 0).shape + (5, 3, 4) """ - itemsize = x.itemsize - if not _check_bounds(x.shape, shape, strides, itemsize): - raise ivy.exceptions.IvyException("attempted unsafe memory access") - if any(strides[i] % itemsize != 0 for i in range(len(strides))): - raise ivy.exceptions.IvyException("strides must be multiple of itemsize") - - src = memoryview(ivy.to_numpy(x)).cast("b") - - src_ind = ivy.inner( - ivy.indices(shape).reshape((len(shape), -1)).T, - ivy.array(strides), - ) - src_ind = ivy.expand_dims(src_ind, axis=-1) - src_ind = src_ind + ivy.arange(itemsize) - src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() - - temp_list = [src[i] for i in src_ind] - temp_array = ivy.asarray(temp_list, dtype=ivy.int8) - result = bytearray(temp_array.to_numpy()) - - return ivy.reshape( - ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), - shape, - ) - - -as_strided.unsupported_dtypes = ("bfloat16",) -as_strided.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back +@handle_array_like_without_promotion +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def concat_from_sequence( - input_sequence: Union[ - Tuple[Union[ivy.Array, ivy.NativeArray]], - List[Union[ivy.Array, ivy.NativeArray]], - ], +def pad( + input: Union[ivy.Array, ivy.NativeArray], + pad_width: Union[Iterable[Tuple[int]], int], /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[ivy.Array] = None, + mode: Union[ + Literal[ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + Callable, + ] = "constant", + stat_length: Union[Iterable[Tuple[int]], int] = 1, + constant_values: Union[Iterable[Tuple[Number]], Number] = 0, + end_values: Union[Iterable[Tuple[Number]], Number] = 0, + reflect_type: Literal["even", "odd"] = "even", + **kwargs: Optional[Any], ) -> ivy.Array: """ - Concatenate a sequence of arrays along a new or an existing axis. + Pad an array. Parameters ---------- - input_sequence - A sequence of arrays. - new_axis - Insert and concatenate on a new axis or not, - default 0 means do not insert new axis. - new_axis = 0: concatenate - new_axis = 1: stack - axis - axis along which the arrays will be concatenated. - - out - optional output array, for writing the result to. + input + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths + for each axis. + - ((before, after),) yields same before and after pad for each axis. + - pad (integer) is shortcut for before = after = pad width for all axes. + mode + One of the following string values or a user-supplied function. + - "constant": Pads with a constant value. + - "edge": Pads with the input's edge values. + - "linear_ramp": Pads with the linear ramp between end_value + and the input's edge value. + - "maximum": Pads with the maximum value of all or part of the vector + along each axis. + - "mean": Pads with the mean value of all or part of the vector along + each axis. + - "median": Pads with the median value of all or part of the vector + along each axis. + - "minimum": Pads with the minimum value of all or part of the vector + along each axis. + - "reflect": Pads with the reflection mirrored on the first and last + values of the vector along each axis. + - "symmetric": Pads with the reflection of the vector mirrored along + the edge of the input. + - "wrap": Pads with the wrap of the vector along the axis. + The first values are used to pad the end and the end values are used + to pad the beginning. + - "empty": Pads with undefined values. + - : Pads with a user-defined padding function. The padding + function should modify a rank 1 array following the signature + `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: + - `vector` is a rank 1 array already padded with zeros. Padded + values are `vector[:iaxis_pad_width[0]]` and + `vector[-iaxis_pad_width[1]:]`. + - `iaxis_pad_width` is a 2-tuple of ints, where + `iaxis_pad_width[0]` represents the number of values padded at + the beginning of `vector` and `iaxis_pad_width[1]` represents + the number of values padded at the end of `vector`. + - `iaxis` is the axis currently being calculated. + - `kwargs` is a dict of keyword arguments the function requires. + stat_length + Used in "maximum", "mean", "median", and "minimum". Number of values at edge + of each axis used to calculate the statistic value. + - ((before_1, after_1), … (before_N, after_N)) yields unique statistic + lengths for each axis. + - ((before, after),) yields same before and after statistic lengths for + each axis. + - stat_length (integer) is a shortcut for before = after = stat_length + length for all axes. + - None uses the entire axis. + constant_values + Used in "constant". The values to set the padded values for each axis. + - ((before_1, after_1), ... (before_N, after_N)) yields unique pad + constants for each axis. + - ((before, after),) yields same before and after constants for each axis. + - constant (integer) is a shortcut for before = after = constant for + all axes. + end_values + Used in "linear_ramp". The values used for the ending value of the linear_ramp + and that will form the edge of the padded array. + - ((before_1, after_1), ... (before_N, after_N)) yields unique end values + for each axis. + - ((before, after),) yields same before and after end values for each axis + - end (integer) is a shortcut for before = after = end for all axes. + reflect_type + Used in "reflect", and "symmetric". The "even" style is the default with an + unaltered reflection around the edge value. For the "odd" style, the extended + part of the array is created by subtracting the reflected values from two + times the edge value. Returns ------- ret - Output Array - """ - return current_backend(input_sequence).concat_from_sequence( - input_sequence, new_axis=new_axis, axis=axis, out=out - ) + Padded array of the same rank as the input but with shape increased according + to pad_width. -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides + Both the description and the type hints above assume an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 3, 0, 0], + [0, 0, 4, 5, 6, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="reflect") + >>> print(y) + ivy.array([[6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1], + [6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1]]) -def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): - if axis >= 0: - slices = [slice(None)] * axis + [slice(start, stop, stride)] - else: - slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) - return x[tuple(slices)] + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="symmetric") + >>> print(y) + ivy.array([[2, 1, 1, 2, 3, 3, 2], + [2, 1, 1, 2, 3, 3, 2], + [5, 4, 4, 5, 6, 6, 5], + [5, 4, 4, 5, 6, 6, 5]]) + With :class:`ivy.NativeArray` input: -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = ivy.full( - new_shape, padding_value, dtype=operand.dtype - ) - src_indices = ivy.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array + >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) + >>> print(y) + ivy.array([[7, 7, 7, 7, 7, 7, 7], + [7, 7, 1, 2, 3, 7, 7], + [7, 7, 4, 5, 6, 7, 7], + [7, 7, 7, 7, 7, 7, 7]]) - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) + With :class:`ivy.Container` input: - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = ivy.constant_pad(padded, pad_width, value=padding_value) + >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) + >>> padding = (1, 1) + >>> y = ivy.pad(x, padding, mode="constant") + >>> print(y) + { + a: ivy.array([0, 0, 1, 2, 0]), + b: ivy.array([0, 4, 5, 6, 0]) + } + """ + _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, + ) + if mode == "dilated": + pad_width = _to_dilated(pad_width, input.ndim) + if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: + constant_values = ivy.asarray(constant_values, dtype=input.dtype) + return _interior_pad(input, constant_values, pad_width) + pad_width = _to_pairs(pad_width, len(input.shape)) + if callable(mode): + func = mode + padded, _ = _pad_simple(input, pad_width, fill_value=0) + for axis in range(padded.ndim): + padded = ivy.moveaxis(padded, axis, -1) + inds = ivy.ndindex(padded.shape[:-1]) + for ind in inds: + padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) + return padded + padded, original_area_slice = _pad_simple(input, pad_width) + axes = range(padded.ndim) + stat_functions = { + "maximum": ivy.max, + "minimum": ivy.min, + "mean": ivy.mean, + "median": ivy.median, + } + if mode == "constant": + constant_values = _to_pairs(constant_values, padded.ndim) + constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) + for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): + padded = _set_pad_area(padded, axis, width_pair, value_pair) + elif mode == "empty": + pass + elif mode == "edge": + for axis, width_pair in zip(axes, pad_width): + edge_pair = _get_edges(padded, axis, width_pair) + padded = _set_pad_area(padded, axis, width_pair, edge_pair) + elif mode == "linear_ramp": + end_values = _to_pairs(end_values, padded.ndim) + for axis, width_pair, value_pair in zip(axes, pad_width, end_values): + ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) + padded = _set_pad_area(padded, axis, width_pair, ramp_pair) + elif mode in stat_functions: + func = stat_functions[mode] + stat_length = _to_pairs(stat_length, padded.ndim) + if mode == "median": + ivy.utils.assertions.check_true( + ivy.is_float_dtype(input), + message="median interpolation is only supported for floats", + ) + for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): + stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) + padded = _set_pad_area(padded, axis, width_pair, stat_pair) + elif mode in {"reflect", "symmetric"}: + include_edge = True if mode == "symmetric" else False + for axis, (left_index, right_index) in zip(axes, pad_width): + if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): + edge_pair = _get_edges(padded, axis, (left_index, right_index)) + padded = _set_pad_area( + padded, axis, (left_index, right_index), edge_pair + ) + continue + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_reflect_both( + padded, axis, (left_index, right_index), reflect_type, include_edge + ) + elif mode == "wrap": + for axis, (left_index, right_index) in zip(axes, pad_width): + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_wrap_both( + padded, axis, (left_index, right_index) + ) return padded -def _interleave(a, b, axis): - assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 - a_pad = [(0, 0, 0)] * a.ndim - b_pad = [(0, 0, 0)] * b.ndim - a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) - b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) - a = _interior_pad(a, 0.0, a_pad) - b = _interior_pad(b, 0.0, b_pad) - return ivy.add(a, b) - - -@handle_exceptions @handle_nestable +@handle_exceptions +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def associative_scan( +@handle_device_shifting +def partial_fold( x: Union[ivy.Array, ivy.NativeArray], - fn: Callable, /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, - reverse: bool = False, - axis: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Perform an associative scan over the given array. + Re-folds a partially unfolded tensor. Parameters ---------- x - The array to scan over. - fn - The associative function to apply. - reverse - Whether to scan in reverse with respect to the given axis. - axis - The axis to scan over. + a partially unfolded tensor + mode + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions left untouched at the beginning + out + optional output array, for writing the result to. Returns ------- ret - The result of the scan. + partially re-folded tensor """ - elems = [x] - - if reverse: - elems = [ivy.flip(elem, axis=[axis]) for elem in elems] - - def _combine(a, b): - a = a[0] - b = b[0] - if a.shape[axis] == 0: - return [a] - c = fn(a, b) - return [c] - - def _scan(elems): - num_elems = elems[0].shape[axis] - - if num_elems < 2: - return elems - - reduced_elems = _combine( - [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], - [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], - ) - - odd_elems = _scan(reduced_elems) - - if num_elems % 2 == 0: - even_elems = _combine( - [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], - ) - else: - even_elems = _combine( - odd_elems, - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], - ) - even_elems = [ - ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) - for (elem, result) in zip(elems, even_elems) - ] - return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) - - scans = _scan(elems) - - if reverse: - scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] - - return ivy.reshape(ivy.asarray(scans), elems[0].shape) + transposed_shape = list(shape) + mode_dim = transposed_shape.pop(skip_begin + mode) + transposed_shape.insert(skip_begin, mode_dim) + return ivy.moveaxis( + ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unique_consecutive( +def partial_tensor_to_vec( x: Union[ivy.Array, ivy.NativeArray], /, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, *, - axis: Optional[int] = None, -) -> Tuple[ - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], -]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Eliminates all but the first element from every consecutive group of equivalent - elements in ``x``. + Partial vectorization of a tensor while ignoring the specified dimension at the + beginning and the end. Parameters ---------- x - input array. - - axis - the axis to apply unique on. If None, unique is applied on flattened ``x``. + tensor to partially vectorise + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + out + optional output array, for writing the result to. Returns ------- ret - a namedtuple ``(output, inverse_indices, counts)`` whose - - first element has the field name ``output`` and is an array - containing ``x`` with its equivalent consecutive elements eliminated. - - second element has the field name ``inverse_indices`` and is an - array containing the indices of ``output`` that reconstruct ``x``. - - third element has the field name ``counts`` and is an array - containing the number of occurrences for each unique value or array in ``x``. - - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) - >>> ivy..unique_consecutive(x) - Results(values=ivy.array([1, 2, 3, 1, 2]), - inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), - counts=ivy.array([2, 2, 1, 2, 1])) + partially vectorised tensor with the + `skip_begin` first and `skip_end` last dimensions untouched """ - return ivy.current_backend(x).unique_consecutive(x, axis=axis) + return partial_unfold( + x, + mode=0, + skip_begin=skip_begin, + skip_end=skip_end, + ravel_tensors=True, + out=out, + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -def fill_diagonal( - a: Union[ivy.Array, ivy.NativeArray], - v: Union[int, float], +@handle_device_shifting +def partial_unfold( + x: Union[ivy.Array, ivy.NativeArray], /, + mode: Optional[int] = 0, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, + ravel_tensors: Optional[bool] = False, *, - wrap: bool = False, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Fill the main diagonal of the given array of any dimensionality.. + Partial unfolding of a tensor while ignoring the specified number of dimensions at + the beginning and the end. For instance, if the first dimension of the tensor is the + number of samples, to unfold each sample, set skip_begin=1. This would, for each i + in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. Parameters ---------- - a - Array at least 2D. - v - The value to write on the diagonal. - wrap - The diagonal ‘wrapped’ after N columns for tall matrices. + x + tensor of shape n_samples x n_1 x n_2 x ... x n_i + mode + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + ravel_tensors + if True, the unfolded tensors are also flattened + out + optional output array, for writing the result to. Returns ------- ret - Array with the diagonal filled. + partially unfolded tensor """ - return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) + if ravel_tensors: + new_shape = [-1] + else: + new_shape = [x.shape[mode + skip_begin], -1] + + if skip_begin: + new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape + + if skip_end: + new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] + + return ivy.reshape( + ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out + ) @handle_nestable @@ -2160,71 +1973,205 @@ def fill_diagonal( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unfold( +def partial_vec_to_tensor( x: Union[ivy.Array, ivy.NativeArray], /, - mode: Optional[int] = 0, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. + Refolds a partially vectorised tensor into a full one. Parameters ---------- x - input tensor to be unfolded - mode - indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` + a partially vectorised tensor + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions to leave untouched at the beginning out optional output array, for writing the result to. Returns ------- ret - unfolded_tensor of shape ``(tensor.shape[mode], -1)`` + full tensor """ - return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) + return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) -@handle_nestable @handle_exceptions +@handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays -@handle_array_function +def put_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + values: Union[ivy.Array, ivy.NativeArray], + axis: int, + /, + *, + mode: str = "raise", + out: Optional[ivy.Array] = None, +) -> None: + """ + Put values into the input array by matching 1d index and data slices along a + specified axis. + + Parameters + ---------- + arr : array_like + The input array to modify. + indices : array_like + The indices of the values to put into `arr`. + values : array_like + The values to put into `arr`. + axis : int + The axis over which to put the `values`. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will be handled. + The following modes are available: + + - 'raise': a `ValueError` is raised when an index is out of bounds. + - 'wrap': the index is wrapped around to the corresponding index + at the other end of the axis. + - 'clip': the index is clipped to the closest in-bounds index. + out : ndarray, optional + Output array in which to place the result. + If not specified, a new array is created. + + Returns + ------- + None + + Examples + -------- + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) + >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') + >>> print(arr) + ivy.array([[3, 7, 5], + [6, 4, 1]]) + """ + if out is None: + out = ivy.zeros_like(arr) + + indices = ivy.expand_dims(indices, axis=axis) + values = ivy.expand_dims(values, axis=axis) + + stacked = ivy.concat((arr, values), axis=axis) + + sorted_indices = ivy.argsort(indices, axis=axis) + sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), + sorted_stacked, + arr, + ) + + if mode == "clip": + indices = ivy.clip(indices, 0, arr.shape[axis] - 1) + elif mode == "wrap": + indices = ivy.mod(indices, arr.shape[axis]) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values + ) + + ivy.assign(out, arr) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def fold( - x: Union[ivy.Array, ivy.NativeArray], +def rot90( + m: Union[ivy.Array, ivy.NativeArray], /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, - refolds the n-mode unfolded tensor into the original tensor of the specified shape. + Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is + from the first towards the second axis. Parameters ---------- - input - unfolded tensor of shape ``(shape[mode], -1)`` - mode - the mode of the unfolding - shape - shape of the original tensor before unfolding + m + Input array of two or more dimensions. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + k + Number of times the array is rotated by 90 degrees. + axes + The array is rotated in the plane defined by the axes. Axes must be + different. out - optional output array, for writing the result to. + optional output container, for writing the result to. It must have a shape + that the inputs broadcast to. Returns ------- ret - folded_tensor of shape `shape` + A rotated view of m. + + Examples + -------- + With :code:`ivy.Array` input: + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) + With :code:`ivy.NativeArray` input: + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.native_array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) """ - full_shape = list(shape) - mode_dim = full_shape.pop(mode) - full_shape.insert(0, mode_dim) - return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) + return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) @handle_nestable @@ -2233,144 +2180,190 @@ def fold( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def partial_unfold( +def soft_thresholding( x: Union[ivy.Array, ivy.NativeArray], /, - mode: Optional[int] = 0, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, - ravel_tensors: Optional[bool] = False, + threshold: Union[float, ivy.Array, ivy.NativeArray], *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Partial unfolding of a tensor while ignoring the specified number of dimensions at - the beginning and the end. For instance, if the first dimension of the tensor is the - number of samples, to unfold each sample, set skip_begin=1. This would, for each i - in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. + Soft-thresholding operator. + + sign(tensor) * max[abs(tensor) - threshold, 0] Parameters ---------- x - tensor of shape n_samples x n_1 x n_2 x ... x n_i - mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - ravel_tensors - if True, the unfolded tensors are also flattened + input array + threshold + float or array with shape tensor.shape + * If float the threshold is applied to the whole tensor + * If array, one threshold is applied per elements, 0 values are ignored out optional output array, for writing the result to. Returns ------- - ret - partially unfolded tensor - """ - if ravel_tensors: - new_shape = [-1] - else: - new_shape = [x.shape[mode + skip_begin], -1] + ivy.Array + thresholded tensor on which the operator has been applied - if skip_begin: - new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape + Examples + -------- + Basic shrinkage - if skip_end: - new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] + >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) + >>> soft_thresholding(x, 1.1) + array([[ 0. , -0.9, 0.4], + [-2.9, 1.9, 0. ]]) - return ivy.reshape( - ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out - ) + + Example with missing values + + >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) + >>> soft_thresholding(x, mask*1.1) + array([[ 1. , -2. , 0.4], + [-2.9, 3. , 0. ]]) + """ + res = ivy.abs(x) - threshold + res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) + + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def partial_fold( - x: Union[ivy.Array, ivy.NativeArray], +def take_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + axis: int, /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, *, + mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Re-folds a partially unfolded tensor. + Take values from the input array by matching 1d index and data slices. Parameters ---------- - x - a partially unfolded tensor + arr + The source array. + indices + The indices of the values to extract. + axis + The axis over which to select values. + If axis is None, arr is treated as a flattened 1D array. mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions left untouched at the beginning + One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices + will be handled. out - optional output array, for writing the result to. + The output array. Returns ------- ret - partially re-folded tensor + The returned array has the same shape as `indices`. + + Examples + -------- + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> y = ivy.take_along_axis(arr, indices, 1) + >>> print(y) + ivy.array([[4, 3, 3], [1, 1, 1]]) """ - transposed_shape = list(shape) - mode_dim = transposed_shape.pop(skip_begin + mode) - transposed_shape.insert(skip_begin, mode_dim) - return ivy.moveaxis( - ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out + return ivy.current_backend(arr).take_along_axis( + arr, indices, axis, mode=mode, out=out ) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def partial_tensor_to_vec( +def top_k( x: Union[ivy.Array, ivy.NativeArray], + k: int, /, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[tuple] = None, +) -> Tuple[ivy.Array, ivy.NativeArray]: """ - Partial vectorization of a tensor while ignoring the specified dimension at the - beginning and the end. + Return the `k` largest elements of the given input array along a given axis. Parameters ---------- x - tensor to partially vectorise - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - out - optional output array, for writing the result to. + The array to compute top_k for. + k + Number of top elements to retun must not exceed the array size. + axis + The axis along which we must return the top elements default value is 1. + largest + If largest is set to False we return k smallest elements of the array. + sorted + If sorted is set to True we return the elements in sorted order. + out: + Optional output tuple, for writing the result to. Must have two arrays inside, + with a shape that the returned tuple broadcast to. Returns ------- ret - partially vectorised tensor with the - `skip_begin` first and `skip_end` last dimensions untouched - """ - return partial_unfold( - x, - mode=0, - skip_begin=skip_begin, - skip_end=skip_end, - ravel_tensors=True, - out=out, + A named tuple with values and indices of top k elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 2) + >>> print(y) + top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) + + >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) + >>> y = ivy.top_k(x, 2, axis=1, largest=False) + >>> print(y) + top_k(values=ivy.array([[-2., 0.], + [-8., -1.]]), indices=ivy.array([[0, 3], + [0, 2]])) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 3) + >>> print(y) + top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) + >>> y = x.top_k(2) + >>> print(y) + [{ + a: ivy.array([2, -1]), + b: ivy.array([5., 4.]) + }, { + a: ivy.array([1, 0]), + b: ivy.array([1, 0]) + }] + """ + return current_backend(x).top_k( + x, k, axis=axis, largest=largest, sorted=sorted, out=out ) @@ -2380,200 +2373,209 @@ def partial_tensor_to_vec( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def partial_vec_to_tensor( +def unfold( x: Union[ivy.Array, ivy.NativeArray], /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, + mode: Optional[int] = 0, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Refolds a partially vectorised tensor into a full one. + Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. Parameters ---------- x - a partially vectorised tensor - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions to leave untouched at the beginning + input tensor to be unfolded + mode + indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` out optional output array, for writing the result to. Returns ------- ret - full tensor + unfolded_tensor of shape ``(tensor.shape[mode], -1)`` """ - return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) + return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def matricize( +def unique_consecutive( x: Union[ivy.Array, ivy.NativeArray], /, - row_modes: Sequence[int], - column_modes: Optional[Sequence[int]] = None, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + axis: Optional[int] = None, +) -> Tuple[ + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], +]: """ - Matricizes the given tensor. + Eliminates all but the first element from every consecutive group of equivalent + elements in ``x``. Parameters ---------- x - the input tensor - row_modes - modes to use as row of the matrix (in the desired order) - column_modes - modes to use as column of the matrix, in the desired order - if None, the modes not in `row_modes` will be used in ascending order - out - optional output array, for writing the result to. + input array. - ret - ------- - ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) - """ - ndims = len(x.shape) - row_indices = list(row_modes) + axis + the axis to apply unique on. If None, unique is applied on flattened ``x``. - if column_modes: - column_indices = list(column_modes) - else: - column_indices = [i for i in range(ndims) if i not in row_indices] - if sorted(column_indices + row_indices) != list(range(ndims)): - msg = ( - "If you provide both column and row modes for the matricization then" - " column_modes + row_modes must contain all the modes of the tensor." - f" Yet, got row_modes={row_modes} and column_modes={column_modes}." - ) - raise ValueError(msg) + Returns + ------- + ret + a namedtuple ``(output, inverse_indices, counts)`` whose + - first element has the field name ``output`` and is an array + containing ``x`` with its equivalent consecutive elements eliminated. + - second element has the field name ``inverse_indices`` and is an + array containing the indices of ``output`` that reconstruct ``x``. + - third element has the field name ``counts`` and is an array + containing the number of occurrences for each unique value or array in ``x``. - row_size, column_size = 1, 1 - row_size = int(ivy.prod([x.shape[i] for i in row_indices])) - column_size = int(ivy.prod([x.shape[i] for i in column_indices])) - return ivy.reshape( - ivy.permute_dims(x, row_indices + column_indices), - (row_size, column_size), - out=out, - ) + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) + >>> ivy..unique_consecutive(x) + Results(values=ivy.array([1, 2, 3, 1, 2]), + inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), + counts=ivy.array([2, 2, 1, 2, 1])) + """ + return ivy.current_backend(x).unique_consecutive(x, axis=axis) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_view +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def soft_thresholding( - x: Union[ivy.Array, ivy.NativeArray], +def vsplit( + ary: Union[ivy.Array, ivy.NativeArray], + indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], /, - threshold: Union[float, ivy.Array, ivy.NativeArray], *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + copy: Optional[bool] = None, +) -> List[ivy.Array]: """ - Soft-thresholding operator. - - sign(tensor) * max[abs(tensor) - threshold, 0] + Split an array vertically into multiple sub-arrays. Parameters ---------- - x - input array - threshold - float or array with shape tensor.shape - * If float the threshold is applied to the whole tensor - * If array, one threshold is applied per elements, 0 values are ignored - out - optional output array, for writing the result to. + ary + Array input. + indices_or_sections + If indices_or_sections is an integer n, the array is split into n + equal sections, provided that n must be a divisor of the split axis. + If indices_or_sections is a sequence of ints or 1-D array, + then input is split at each of the indices. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- - ivy.Array - thresholded tensor on which the operator has been applied + ret + input array split vertically. Examples -------- - Basic shrinkage - - >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) - >>> soft_thresholding(x, 1.1) - array([[ 0. , -0.9, 0.4], - [-2.9, 1.9, 0. ]]) - - - Example with missing values - - >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) - >>> soft_thresholding(x, mask*1.1) - array([[ 1. , -2. , 0.4], - [-2.9, 3. , 0. ]]) + >>> ary = ivy.array( + [[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]] + ) + >>> ivy.vsplit(ary, 2) + [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) """ - res = ivy.abs(x) - threshold - res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) - - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def choose( - arr: Union[ivy.Array, ivy.NativeArray], - choices: Union[ivy.Array, ivy.NativeArray], +def vstack( + arrays: Sequence[ivy.Array], /, *, - out: None = None, - mode: Union[str, None] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Take values from the input array by matching 1d index and data slices. + Stack arrays in sequence vertically (row wise). Parameters ---------- - arr - The source array. - choices - The indices of the values to extract. - out - The output array. - mode - One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices - will be handled. + arrays + Sequence of arrays to be stacked. Returns ------- ret - The returned array has the same shape as `indices`. + The array formed by stacking the given arrays. Examples -------- - >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], - [20, 21, 22, 23], [30, 31, 32, 33]]) - >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) - ivy.array([20, 1, 12, 3]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.vstack((x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4]]) + >>> ivy.vstack((x, y, x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [1, 2, 3], + [2, 3, 4]]) + + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.vstack(y)) + ivy.array([[5, 6], + [7, 8]]) """ - return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) + return ivy.current_backend().vstack(arrays, out=out) + + +flatten.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +pad.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +as_strided.unsupported_dtypes = ("bfloat16",) +as_strided.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} diff --git a/ivy/functional/ivy/experimental/norms.py b/ivy/functional/ivy/experimental/norms.py index 4d59a2aaa6907..d37b4d6c0e954 100644 --- a/ivy/functional/ivy/experimental/norms.py +++ b/ivy/functional/ivy/experimental/norms.py @@ -17,89 +17,35 @@ from ivy.utils.exceptions import handle_exceptions -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l1_normalize( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[int, Tuple[int, ...]]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """Normalize the input array along the given axis to have L1 norm equal to - 1. - - Parameters - ---------- - x - Input array. - axis - Axis or axes along which to normalize. If ``None``, - the whole array is normalized. - out - Optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - The normalized array. - - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l1_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.33333334, 1.33333337], - [1.28571439, 2.28571439]]) - """ - return current_backend(x).l1_normalize(x, axis=axis, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l2_normalize( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """Normalize the input array along the given axis to have L2 norm equal to - 1. - - Parameters - ---------- - x - Input array. - axis - Axis along which to normalize. If ``None``, the whole array is normalized. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The normalized array. - - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l2_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.44721359, 0.89442718], - [0.60000002, 0.80000001]]) - """ - return current_backend(x).l2_normalize(x, axis=axis, out=out) +batch_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +instance_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +group_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_exceptions @@ -212,16 +158,80 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def group_norm( + x: Union[ivy.NativeArray, ivy.Array], + num_groups: int = 1, + /, + *, + offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Apply group normalization to the input array and returns the normalized input. + + Parameters + ---------- + x + Input array of default shape (N, *S, C), where N is the batch dimension, + *S corresponds to any number of spatial dimensions and + C corresponds to the channel dimension. + num_groups + number of groups to separate the channels into + offset + An offset array of size C. If present, will be added + to the normalized input. + scale + A scale array of size C. If present, the scale is + applied to the normalized input. + eps + A small float number to avoid dividing by 0. + data_format + The ordering of the dimensions in the input, one of "NSC" or "NCS", + where N is the batch dimension, S represents any number of spatial + dimensions and C is the channel dimension. Default is "NSC". + out + optional output arrays, for writing the result to. + + Returns + ------- + ret + The normalized array. + """ + xdims = ivy.get_num_dims(x) + if data_format == "NSC": + x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) + N = x.shape[0] + C = x.shape[1] + S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 + assert C % num_groups == 0 + x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) + mean = ivy.mean(x_, axis=(2, 3), keepdims=True) + var = ivy.var(x_, axis=(2, 3), keepdims=True) + x_normalized = (x_ - mean) / ivy.sqrt(var + eps) + x_normalized = ivy.reshape(x_normalized, x.shape) + + if ivy.exists(scale): + scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized * scale + + if ivy.exists(offset): + offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized + offset + + if data_format == "NSC": + x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) + + if ivy.exists(out): + x_normalized = ivy.inplace_update(out, x_normalized) + return x_normalized @handle_exceptions @@ -344,103 +354,89 @@ def instance_norm( return (xnormalized, runningmean, runningvariance) -instance_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def group_norm( - x: Union[ivy.NativeArray, ivy.Array], - num_groups: int = 1, +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l1_normalize( + x: Union[ivy.Array, ivy.NativeArray], /, *, - offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", + axis: Optional[Union[int, Tuple[int, ...]]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Apply group normalization to the input array and returns the normalized input. + """Normalize the input array along the given axis to have L1 norm equal to + 1. Parameters ---------- x - Input array of default shape (N, *S, C), where N is the batch dimension, - *S corresponds to any number of spatial dimensions and - C corresponds to the channel dimension. - num_groups - number of groups to separate the channels into - offset - An offset array of size C. If present, will be added - to the normalized input. - scale - A scale array of size C. If present, the scale is - applied to the normalized input. - eps - A small float number to avoid dividing by 0. - data_format - The ordering of the dimensions in the input, one of "NSC" or "NCS", - where N is the batch dimension, S represents any number of spatial - dimensions and C is the channel dimension. Default is "NSC". + Input array. + axis + Axis or axes along which to normalize. If ``None``, + the whole array is normalized. out - optional output arrays, for writing the result to. + Optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret The normalized array. - """ - xdims = ivy.get_num_dims(x) - if data_format == "NSC": - x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) - N = x.shape[0] - C = x.shape[1] - S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 - assert C % num_groups == 0 - x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) - mean = ivy.mean(x_, axis=(2, 3), keepdims=True) - var = ivy.var(x_, axis=(2, 3), keepdims=True) - x_normalized = (x_ - mean) / ivy.sqrt(var + eps) - x_normalized = ivy.reshape(x_normalized, x.shape) - if ivy.exists(scale): - scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized * scale + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l1_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.33333334, 1.33333337], + [1.28571439, 2.28571439]]) + """ + return current_backend(x).l1_normalize(x, axis=axis, out=out) - if ivy.exists(offset): - offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized + offset - if data_format == "NSC": - x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l2_normalize( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """Normalize the input array along the given axis to have L2 norm equal to + 1. - if ivy.exists(out): - x_normalized = ivy.inplace_update(out, x_normalized) - return x_normalized + Parameters + ---------- + x + Input array. + axis + Axis along which to normalize. If ``None``, the whole array is normalized. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The normalized array. -group_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l2_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.44721359, 0.89442718], + [0.60000002, 0.80000001]]) + """ + return current_backend(x).l2_normalize(x, axis=axis, out=out) @handle_exceptions diff --git a/ivy/functional/ivy/experimental/random.py b/ivy/functional/ivy/experimental/random.py index 70586db0f3ccf..ff5a5d32f56bb 100644 --- a/ivy/functional/ivy/experimental/random.py +++ b/ivy/functional/ivy/experimental/random.py @@ -14,69 +14,65 @@ from ivy.utils.exceptions import handle_exceptions -# dirichlet @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back +@infer_dtype @handle_device_shifting -def dirichlet( - alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], - /, +@infer_device +def bernoulli( + probs: Union[float, ivy.Array, ivy.NativeArray], *, - size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- - distributed random variable can be seen as a multivariate generalization of a Beta - distribution. The Dirichlet distribution is a conjugate prior of a multinomial - distribution in Bayesian inference. + Draws samples from Bernoulli distrubution paramterized by probs or logits (but not + both) Parameters ---------- - alpha - Sequence of floats of length k - size - optional int or tuple of ints, Output shape. If the given shape is, - e.g., (m, n), then m * n * k samples are drawn. Default is None, - in which case a vector of length k is returned. + logits + An N-D Array representing the log-odds of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution where the probability of an event is sigmoid + (logits). Only one of logits or probs should be passed in. + probs + An N-D Array representing the probability of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution. Only one of logits or probs should be passed in + shape + If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. + (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + (Default value = None). dtype output array data type. If ``dtype`` is ``None``, the output array data type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The drawn samples, of shape (size, k). - - Functional Examples - ------------------- - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha) - ivy.array([0.10598304, 0.21537054, 0.67864642]) - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha, size = (2,3)) - ivy.array([[[0.48006698, 0.07472073, 0.44521229], - [0.55479872, 0.05426367, 0.39093761], - [0.19531053, 0.51675832, 0.28793114]], - - [[0.12315625, 0.29823365, 0.5786101 ], - [0.15564976, 0.50542368, 0.33892656], - [0.1325352 , 0.44439589, 0.42306891]]]) + Drawn samples from the Bernoulli distribution """ - return ivy.current_backend().dirichlet( - alpha, - size=size, + return ivy.current_backend(probs).bernoulli( + probs, + logits=logits, + shape=shape, + device=device, dtype=dtype, seed=seed, out=out, @@ -138,6 +134,75 @@ def beta( ) +# dirichlet +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dirichlet( + alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- + distributed random variable can be seen as a multivariate generalization of a Beta + distribution. The Dirichlet distribution is a conjugate prior of a multinomial + distribution in Bayesian inference. + + Parameters + ---------- + alpha + Sequence of floats of length k + size + optional int or tuple of ints, Output shape. If the given shape is, + e.g., (m, n), then m * n * k samples are drawn. Default is None, + in which case a vector of length k is returned. + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. + + Returns + ------- + ret + The drawn samples, of shape (size, k). + + Functional Examples + ------------------- + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha) + ivy.array([0.10598304, 0.21537054, 0.67864642]) + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha, size = (2,3)) + ivy.array([[[0.48006698, 0.07472073, 0.44521229], + [0.55479872, 0.05426367, 0.39093761], + [0.19531053, 0.51675832, 0.28793114]], + + [[0.12315625, 0.29823365, 0.5786101 ], + [0.15564976, 0.50542368, 0.33892656], + [0.1325352 , 0.44439589, 0.42306891]]]) + """ + return ivy.current_backend().dirichlet( + alpha, + size=size, + dtype=dtype, + seed=seed, + out=out, + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -260,68 +325,3 @@ def poisson( fill_value=fill_value, out=out, ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -@infer_device -def bernoulli( - probs: Union[float, ivy.Array, ivy.NativeArray], - *, - logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Draws samples from Bernoulli distrubution paramterized by probs or logits (but not - both) - - Parameters - ---------- - logits - An N-D Array representing the log-odds of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution where the probability of an event is sigmoid - (logits). Only one of logits or probs should be passed in. - probs - An N-D Array representing the probability of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution. Only one of logits or probs should be passed in - shape - If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. - (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Drawn samples from the Bernoulli distribution - """ - return ivy.current_backend(probs).bernoulli( - probs, - logits=logits, - shape=shape, - device=device, - dtype=dtype, - seed=seed, - out=out, - ) diff --git a/ivy/functional/ivy/experimental/sparse_array.py b/ivy/functional/ivy/experimental/sparse_array.py index 46e6eb64fa428..785eccaa6b3bd 100644 --- a/ivy/functional/ivy/experimental/sparse_array.py +++ b/ivy/functional/ivy/experimental/sparse_array.py @@ -4,356 +4,6 @@ from ivy.utils.exceptions import handle_exceptions -# helpers -def _verify_coo_components(indices=None, values=None, dense_shape=None): - ivy.utils.assertions.check_all_or_any_fn( - indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message="indices, values and dense_shape must all be specified", - ) - # coordinates style (COO), must be shaped (x, y) - ivy.utils.assertions.check_equal( - len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.to_ivy_shape(dense_shape)), - ivy.shape(indices)[0], - message="shape and indices shape do not match", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(values)[0], - ivy.shape(indices)[1], - message="values and indices do not match", - as_array=False, - ) - for i in range(ivy.shape(indices)[0]): - ivy.utils.assertions.check_less( - indices[i], - ivy.to_ivy_shape(dense_shape)[i], - message="indices is larger than shape", - ) - - -def _verify_common_row_format_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" -): - ivy.utils.assertions.check_all_or_any_fn( - crow_indices, - col_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "crow_indices, col_indices, values and dense_shape must all be specified." - ), - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(crow_indices)), - 1, - message="crow_indices must be 1D.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(col_indices)), - 1, - message="col_indices must be 1D.", - as_array=False, - ) - - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", - as_array=False, - ) - - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - crow_indices[-1], - message="size of col_indices does not match with last element of crow_indices", - ) - - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - ivy.shape(values)[0], - message="values and col_indices do not match", - as_array=False, - ) - - # index in crow_indices must not exceed length of col_indices - ivy.utils.assertions.check_less( - crow_indices, - ivy.shape(col_indices)[0], - allow_equal=True, - message="index in crow_indices does not match the number of col_indices", - ) - - -def _verify_csr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="csr", - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False - ) - # number of intervals must be equal to x in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False - ) - - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1], - message="index in col_indices does not match shape", - ) - - -def _verify_bsr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="bsr", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="The number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="The number of cols of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False - ) - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1] // ncolblocks, - message="index in col_indices does not match shape", - ) - - -def _verify_common_column_format_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" -): - ivy.utils.assertions.check_all_or_any_fn( - ccol_indices, - row_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "ccol_indices, row_indices, values and dense_shape must all be specified" - ), - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(ccol_indices)), - 1, - message="ccol_indices must be 1D", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False - ) - - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(row_indices)[0], - ivy.shape(values)[0], - message="values and row_indices do not match", - as_array=False, - ) - # index in ccol_indices must not exceed length of row_indices - ivy.utils.assertions.check_less( - ccol_indices, - ivy.shape(row_indices)[0], - allow_equal=True, - message="index in ccol_indices does not match the number of row_indices", - ) - - -def _verify_csc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="csc", - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0], - message="index in row_indices does not match shape", - ) - - -def _verify_bsc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="bsc", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="number of cols of array must be divisible by that of block.", - as_array=False, - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0] // nrowblocks, - message="index in row_indices does not match shape", - ) - - -def _is_data_not_indices_values_and_shape( - data=None, - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format=None, -): - if data is not None: - ivy.utils.assertions.check_all_or_any_fn( - coo_indices, - crow_indices, - col_indices, - ccol_indices, - row_indices, - values, - dense_shape, - format, - fn=ivy.exists, - type="any", - limit=[0], - message=( - "Only specify data, coo_indices for COO format, crow_indices and" - " col_indices for CSR and BSR, ccol_indices and row_indicesfor CSC and" - " BSC." - ), - ) - return True - return False - - -def _is_valid_format( - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format="coo", -): - valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] - - if not isinstance(format, str) or not format.lower() in valid_formats: - return False - - if format.endswith("o"): - # format is coo - return ( - ivy.exists(coo_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and crow_indices is None - and col_indices is None - and ccol_indices is None - and row_indices is None - ) - - if format.endswith("r"): - # format is either csr or bsr - return ( - ivy.exists(crow_indices) - and ivy.exists(col_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and ccol_indices is None - and row_indices is None - ) - # format is either csc or bsc - return ( - ivy.exists(ccol_indices) - and ivy.exists(row_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and crow_indices is None - and col_indices is None - ) - - class SparseArray: def __init__( self, @@ -550,240 +200,598 @@ def coo_indices(self): def crow_indices(self): return self._crow_indices - @property - def col_indices(self): - return self._col_indices + @property + def col_indices(self): + return self._col_indices + + @property + def ccol_indices(self): + return self._ccol_indices + + @property + def row_indices(self): + return self._row_indices + + @property + def values(self): + return self._values + + @property + def dense_shape(self): + return self._dense_shape + + @property + def format(self): + return self._format + + # Setters # + # --------# + + @data.setter + def data(self, data): + self._init_data(data) + + @coo_indices.setter + def coo_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + _verify_coo_components( + indices=indices, values=self._values, dense_shape=self._dense_shape + ) + self._coo_indices = indices + + @crow_indices.setter + def crow_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._crow_indices = indices + + @col_indices.setter + def col_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._col_indices = indices + + @ccol_indices.setter + def ccol_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._ccol_indices = indices + + @row_indices.setter + def row_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._row_indices = indices + + @values.setter + def values(self, values): + values = ivy.array(values) + _verify_coo_components( + indices=self._coo_indices, values=values, dense_shape=self._dense_shape + ) + self._values = values + + @dense_shape.setter + def dense_shape(self, dense_shape): + dense_shape = ivy.Shape(dense_shape) + _verify_coo_components( + indices=self._coo_indices, values=self._values, dense_shape=dense_shape + ) + self._dense_shape = dense_shape + + @format.setter + def format(self, format): + self._format = format + + # Instance Methods # + # ---------------- # + + def _coo_to_dense_coordinates(self): + all_coordinates = [] + for i in range(self._values.shape[0]): + coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) + coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) + all_coordinates.append(coordinate.to_list()) + return all_coordinates + + def _csr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._col_indices.to_list() + all_cols = self._crow_indices.to_list() + for row in range(total_rows): + cols = all_rows[all_cols[row] : all_cols[row + 1]] + for col in cols: + all_coordinates.append([row, col]) + return all_coordinates + + def _csc_to_dense_coordinates(self): + # CSC sparse array + all_coordinates = [] + total_rows = self._dense_shape[1] + all_cols = self._row_indices.to_list() + all_rows = self._ccol_indices.to_list() + for col in range(total_rows): + rows = all_cols[all_rows[col] : all_rows[col + 1]] + for row in rows: + all_coordinates.append([row, col]) + return all_coordinates + + def _bsr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._crow_indices.to_list() + all_cols = self._col_indices.to_list() + + nblockrows, nblockcols = self._values.shape[-2:] + + for row in range(total_rows // nblockrows): + cols = all_cols[all_rows[row] : all_rows[row + 1]] + for col in cols: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates + + def _bsc_to_dense_coordinates(self): + all_coordinates = [] + total_cols = self._dense_shape[1] + all_rows = self._row_indices.to_list() + all_cols = self._ccol_indices.to_list() + + nblockrows, nblockcols = self._values.shape[-2:] + + for col in range(total_cols // nblockcols): + rows = all_rows[all_cols[col] : all_cols[col + 1]] + for row in rows: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates + + def to_dense_array(self, *, native=False): + if self._format == "coo": + all_coordinates = self._coo_to_dense_coordinates() + elif self._format == "csr": + all_coordinates = self._csr_to_dense_coordinates() + elif self._format == "csc": + all_coordinates = self._csc_to_dense_coordinates() + elif self._format == "bsc": + all_coordinates = self._bsc_to_dense_coordinates() + else: + all_coordinates = self._bsr_to_dense_coordinates() + + # make dense array + ret = ivy.scatter_nd( + ivy.array(all_coordinates), + ivy.flatten(self._values), + ivy.array(self._dense_shape), + ) + return ret.to_native() if native else ret + + +class NativeSparseArray: + pass + + +# --- Helpers --- # +# --------------- # + + +def _is_data_not_indices_values_and_shape( + data=None, + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format=None, +): + if data is not None: + ivy.utils.assertions.check_all_or_any_fn( + coo_indices, + crow_indices, + col_indices, + ccol_indices, + row_indices, + values, + dense_shape, + format, + fn=ivy.exists, + type="any", + limit=[0], + message=( + "Only specify data, coo_indices for COO format, crow_indices and" + " col_indices for CSR and BSR, ccol_indices and row_indicesfor CSC and" + " BSC." + ), + ) + return True + return False + - @property - def ccol_indices(self): - return self._ccol_indices +def _is_valid_format( + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format="coo", +): + valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] - @property - def row_indices(self): - return self._row_indices + if not isinstance(format, str) or not format.lower() in valid_formats: + return False - @property - def values(self): - return self._values + if format.endswith("o"): + # format is coo + return ( + ivy.exists(coo_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and crow_indices is None + and col_indices is None + and ccol_indices is None + and row_indices is None + ) - @property - def dense_shape(self): - return self._dense_shape + if format.endswith("r"): + # format is either csr or bsr + return ( + ivy.exists(crow_indices) + and ivy.exists(col_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and ccol_indices is None + and row_indices is None + ) + # format is either csc or bsc + return ( + ivy.exists(ccol_indices) + and ivy.exists(row_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and crow_indices is None + and col_indices is None + ) - @property - def format(self): - return self._format - # Setters # - # --------# +def _verify_bsc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="bsc", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="number of cols of array must be divisible by that of block.", + as_array=False, + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0] // nrowblocks, + message="index in row_indices does not match shape", + ) - @data.setter - def data(self, data): - self._init_data(data) - @coo_indices.setter - def coo_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - _verify_coo_components( - indices=indices, values=self._values, dense_shape=self._dense_shape - ) - self._coo_indices = indices +def _verify_bsr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="bsr", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="The number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="The number of cols of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False + ) + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1] // ncolblocks, + message="index in col_indices does not match shape", + ) - @crow_indices.setter - def crow_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._crow_indices = indices - @col_indices.setter - def col_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._col_indices = indices +def _verify_common_column_format_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" +): + ivy.utils.assertions.check_all_or_any_fn( + ccol_indices, + row_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "ccol_indices, row_indices, values and dense_shape must all be specified" + ), + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(ccol_indices)), + 1, + message="ccol_indices must be 1D", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False + ) - @ccol_indices.setter - def ccol_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._ccol_indices = indices + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(row_indices)[0], + ivy.shape(values)[0], + message="values and row_indices do not match", + as_array=False, + ) + # index in ccol_indices must not exceed length of row_indices + ivy.utils.assertions.check_less( + ccol_indices, + ivy.shape(row_indices)[0], + allow_equal=True, + message="index in ccol_indices does not match the number of row_indices", + ) - @row_indices.setter - def row_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._row_indices = indices - @values.setter - def values(self, values): - values = ivy.array(values) - _verify_coo_components( - indices=self._coo_indices, values=values, dense_shape=self._dense_shape - ) - self._values = values +def _verify_common_row_format_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" +): + ivy.utils.assertions.check_all_or_any_fn( + crow_indices, + col_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "crow_indices, col_indices, values and dense_shape must all be specified." + ), + ) - @dense_shape.setter - def dense_shape(self, dense_shape): - dense_shape = ivy.Shape(dense_shape) - _verify_coo_components( - indices=self._coo_indices, values=self._values, dense_shape=dense_shape - ) - self._dense_shape = dense_shape + ivy.utils.assertions.check_equal( + len(ivy.shape(crow_indices)), + 1, + message="crow_indices must be 1D.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(col_indices)), + 1, + message="col_indices must be 1D.", + as_array=False, + ) - @format.setter - def format(self, format): - self._format = format + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", + as_array=False, + ) - # Instance Methods # - # ---------------- # + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + crow_indices[-1], + message="size of col_indices does not match with last element of crow_indices", + ) - def _coo_to_dense_coordinates(self): - all_coordinates = [] - for i in range(self._values.shape[0]): - coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) - coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) - all_coordinates.append(coordinate.to_list()) - return all_coordinates + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + ivy.shape(values)[0], + message="values and col_indices do not match", + as_array=False, + ) - def _csr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._col_indices.to_list() - all_cols = self._crow_indices.to_list() - for row in range(total_rows): - cols = all_rows[all_cols[row] : all_cols[row + 1]] - for col in cols: - all_coordinates.append([row, col]) - return all_coordinates + # index in crow_indices must not exceed length of col_indices + ivy.utils.assertions.check_less( + crow_indices, + ivy.shape(col_indices)[0], + allow_equal=True, + message="index in crow_indices does not match the number of col_indices", + ) - def _csc_to_dense_coordinates(self): - # CSC sparse array - all_coordinates = [] - total_rows = self._dense_shape[1] - all_cols = self._row_indices.to_list() - all_rows = self._ccol_indices.to_list() - for col in range(total_rows): - rows = all_cols[all_rows[col] : all_rows[col + 1]] - for row in rows: - all_coordinates.append([row, col]) - return all_coordinates - def _bsr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._crow_indices.to_list() - all_cols = self._col_indices.to_list() +# helpers +def _verify_coo_components(indices=None, values=None, dense_shape=None): + ivy.utils.assertions.check_all_or_any_fn( + indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message="indices, values and dense_shape must all be specified", + ) + # coordinates style (COO), must be shaped (x, y) + ivy.utils.assertions.check_equal( + len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.to_ivy_shape(dense_shape)), + ivy.shape(indices)[0], + message="shape and indices shape do not match", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(values)[0], + ivy.shape(indices)[1], + message="values and indices do not match", + as_array=False, + ) + for i in range(ivy.shape(indices)[0]): + ivy.utils.assertions.check_less( + indices[i], + ivy.to_ivy_shape(dense_shape)[i], + message="indices is larger than shape", + ) - nblockrows, nblockcols = self._values.shape[-2:] - for row in range(total_rows // nblockrows): - cols = all_cols[all_rows[row] : all_rows[row + 1]] - for col in cols: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates +def _verify_csc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="csc", + ) - def _bsc_to_dense_coordinates(self): - all_coordinates = [] - total_cols = self._dense_shape[1] - all_rows = self._row_indices.to_list() - all_cols = self._ccol_indices.to_list() + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0], + message="index in row_indices does not match shape", + ) - nblockrows, nblockcols = self._values.shape[-2:] - for col in range(total_cols // nblockcols): - rows = all_rows[all_cols[col] : all_cols[col + 1]] - for row in rows: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates +def _verify_csr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="csr", + ) - def to_dense_array(self, *, native=False): - if self._format == "coo": - all_coordinates = self._coo_to_dense_coordinates() - elif self._format == "csr": - all_coordinates = self._csr_to_dense_coordinates() - elif self._format == "csc": - all_coordinates = self._csc_to_dense_coordinates() - elif self._format == "bsc": - all_coordinates = self._bsc_to_dense_coordinates() - else: - all_coordinates = self._bsr_to_dense_coordinates() + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False + ) + # number of intervals must be equal to x in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False + ) - # make dense array - ret = ivy.scatter_nd( - ivy.array(all_coordinates), - ivy.flatten(self._values), - ivy.array(self._dense_shape), - ) - return ret.to_native() if native else ret + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1], + message="index in col_indices does not match shape", + ) -class NativeSparseArray: - pass +# --- Main --- # +# ------------ # def is_ivy_sparse_array(x): diff --git a/ivy/functional/ivy/experimental/statistical.py b/ivy/functional/ivy/experimental/statistical.py index 1349cc5752326..e9392dcff411a 100644 --- a/ivy/functional/ivy/experimental/statistical.py +++ b/ivy/functional/ivy/experimental/statistical.py @@ -13,141 +13,47 @@ from ivy.utils.exceptions import handle_exceptions -# TODO: Make bins optional by offering an automatic bins creation like numpy. -# Make density argument work in tensorflow -# Bins as str is not defined (check Numpy implementation). -# Permit multiple axis. -# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def histogram( - a: Union[ivy.Array, ivy.NativeArray], +def bincount( + x: ivy.Array, /, *, - bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - density: Optional[bool] = False, + weights: Optional[ivy.Array] = None, + minlength: int = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the histogram of the array ``a``. - - .. note:: - Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), - ..., I_{K-1} = [c_{K-1}, cK]. + Count the number of occurrences of each value in an integer array. Parameters ---------- - a - input array. - bins - if ``bins`` is an int, it defines the number of equal-width bins in the given - range. - if ``bins`` is an array, it defines a monotonically increasing array of bin - edges, including the rightmost edge, allowing for non-uniform bin widths. - axis - dimension along which maximum values must be computed. By default, the maximum - value must be computed over the entire array. Default: ``None``. - extend_lower_interval - if True, extend the lowest interval I0 to (-inf, c1]. - extend_upper_interval - ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). - dtype - the output type. - range - the lower and upper range of the bins. The first element of the range must be - less than or equal to the second. + self + Input array. weights - each value in ``a`` only contributes its associated weight towards the bin count - (instead of 1). Must be of the same shape as a. - density - if True, the result is the value of the probability density function at the - bin, normalized such that the integral over the range of bins is 1. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + An optional input array. + minlength + A minimum number of bins for the output array. Returns ------- ret - a tuple containing the values of the histogram and the bin edges. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The bincount of the array elements. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) - >>> z = ivy.histogram(x, bins=y) - >>> print(z) - ivy.array([1., 0., 1., 1.]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [4.4, 5.5, .6]]) - >>> bins = 4 - >>> range = (0., 5.) - >>> dtype = ivy.int32 - >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) - >>> print(y) - ivy.array([2, 1, 1, 1]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> axis = 1 - >>> extend_lower_interval = True - >>> extend_upper_interval = True - >>> dtype = ivy.float32 - >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) - >>> z = ivy.histogram( - ... x, - ... bins=y, - ... axis=axis, - ... extend_lower_interval=extend_lower_interval, - ... extend_upper_interval=extend_upper_interval, - ... dtype=dtype, - ... weights=weights) - >>> print(z) - ivy.array([[0., 3.], - [1., 0.], - [1., 0.], - [1., 0.], - [0., 0.]]) - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> dtype = ivy.int32 - >>> z = ivy.histogram(x, bins=y, dtype=dtype) - >>> print(z) - { - a: ivy.array([1, 1, 1, 0, 0]), - b: ivy.array([0, 0, 0, 1, 2]) - } + >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) + >>> a.bincount(a) + 3.0 + >>> a.bincount(a, axis=0) + array([6.5, 2. , 2.5]) """ - return ivy.current_backend(a).histogram( - a, - bins=bins, - axis=axis, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - range=range, - weights=weights, - density=density, - out=out, + return ivy.current_backend(x).bincount( + x, weights=weights, minlength=minlength, out=out ) @@ -157,276 +63,533 @@ def histogram( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def median( - input: ivy.Array, +def corrcoef( + x: ivy.Array, /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, + y: Optional[ivy.Array] = None, + rowvar: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute the median along the specified axis. - - Parameters - ---------- - input - Input array. - axis - Axis or axes along which the medians are computed. The default is to compute - the median along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. - out - optional output array, for writing the result to. - - Returns - ------- - ret - The median of the array elements. - - Functional Examples - ------------------- - >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> ivy.median(a) - 3.5 - >>> ivy.median(a, axis=0) - ivy.array([6.5, 4.5, 2.5]) - """ - return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) + return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_out_argument +@handle_array_like_without_promotion @to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nanmean( - a: ivy.Array, +def cov( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray] = None, /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[ivy.Array] = None, + aweights: Optional[ivy.Array] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the mean of all non-NaN elements along the specified dimensions. + Compute the covariance of matrix x1, or variables x1 and x2. Parameters ---------- - a - Input array. - axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of sub-classes - of ndarray. If the sub-classes methods does not implement keepdims any - exceptions will be raised. + x1 + a 1D or 2D input array, with a numeric data type. + x2 + optional second 1D or 2D input array, with a numeric data type. + Must have the same shape as ``self``. + rowVar + optional variable where each row of input is interpreted as a variable + (default = True). If set to False, each column is instead interpreted + as a variable. + bias + optional variable for normalizing input (default = False) by (N - 1) where + N is the number of given observations. If set to True, then normalization + is instead by N. Can be overridden by keyword ``ddof``. + ddof + optional variable to override ``bias`` (default = None). ddof=1 will return + the unbiased estimate, even with fweights and aweights given. ddof=0 will + return the simple average. + fweights + optional 1D array of integer frequency weights; the number of times each + observation vector should be repeated. + aweights + optional 1D array of observation vector weights. These relative weights are + typically large for observations considered "important" and smaller for + observations considered less "important". If ddof=0 is specified, the array + of weights can be used to assign probabilities to observation vectors. dtype - The desired data type of returned tensor. Default is None. + optional variable to set data-type of the result. By default, data-type + will have at least ``numpy.float64`` precision. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that + the inputs broadcast to. Returns ------- ret - The nanmean of the array elements. + an array containing the covariance matrix of an input matrix, or the + covariance matrix of two variables. The returned array must have a + floating-point data type determined by Type Promotion Rules and must be + a square matrix of shape (N, N), where N is the number of variables in the + input(s). - Functional Examples - ------------------- - >>> a = ivy.array([[1, ivy.nan], [3, 4]]) - >>> ivy.nanmean(a) - 2.6666666666666665 - >>> ivy.nanmean(a, axis=0) - ivy.array([2., 4.]) - """ - return ivy.current_backend(a).nanmean( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out - ) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([[1, 2, 3], + ... [4, 5, 6]]) + >>> y = x[0].cov(x[1]) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Container` inputs: + >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) + >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) + >>> z = ivy.Container.static_cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) + >>> z = ivy.cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With :class:`ivy.Array` input and rowVar flag set to False (True by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], rowVar=False) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Array` input and bias flag set to True (False by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], bias=True) + >>> print(y) + ivy.array([[0.66666667, 0.66666667], + [0.66666667, 0.66666667]]) + + With :class:`ivy.Array` input with both fweights and aweights given: + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> fw = ivy.array([1,2,3]) + >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) + >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) + >>> print(y) + ivy.array([[0.48447205, 0.48447205], + [0.48447205, 0.48447205]]) + """ + return ivy.current_backend(x1).cov( + x1, + x2, + rowVar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, + dtype=dtype, + ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def quantile( - a: ivy.Array, - q: Union[ivy.Array, float], +@handle_array_function +def cummax( + x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: bool = False, - interpolation: str = "linear", + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the q-th quantile of the data along the specified axis. + Return a tuple containing the cumulative maximum of elements of input along the + given axis and index location of each maximum value found along the given axis. Parameters ---------- - a + x Input array. - q - Quantile or sequence of quantiles to compute, which must be - between 0 and 1 inclusive. axis - Axis or axes along which the quantiles are computed. The default - is to compute the quantile(s) along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original array a. - interpolation - {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. - Default value: 'linear'. - This specifies the interpolation method to use when the desired quantile lies - between two data points i < j: - - linear: i + (j - i) * fraction, where fraction is the fractional part of the - index surrounded by i and j. - - lower: i. - - higher: j. - - nearest: i or j, whichever is nearest. - - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with - integer dtypes. - - nearest_jax: provides jax-like computation for interpolation='nearest'. + Axis along which the cumulative maximum is computed. Default is ``0``. + exclusive + Whether to perform cummax exclusively. Default is ``False``. + reverse + Whether to perform the cummax from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis - is None, a rank(q) array. The first rank(q) dimensions index quantiles for - different values of q. + Array which holds the result of applying cummax at each + original array elements along the specified axis. Examples -------- - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = ivy.array(0.5) - >>> ivy.quantile(a, q) - ivy.array(3.5) + With :class:`ivy.Array` input: - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = 0.5 - >>> ivy.quantile(a, q) - ivy.array(3.5) + >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) + >>> y = ivy.cummax(x, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), + ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) - >>> ivy.quantile(a, q, axis=0) - ivy.array([6.5, 4.5, 2.5]) + >>> x = ivy.array([ 14, 15, 49, -24, -39]) + >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) - >>> ivy.quantile(a, q, axis=1) - ivy.array([7., 2.]) + >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) + >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) + >>> print(x) + ivy.array([[0, 0, 0, 0], + [0, 1, 1, 1]]) - >>> ivy.quantile(a, q, axis=1, keepdims=True) - ivy.array([[7.],[2.]]) + >>> x = ivy.array([[-36, 83, -81], + ... [ 23, 29, 63], + ... [-83, 85, 2], + ... [ 31, 25, -86], + ... [-10, -52, 0], + ... [ 22, 38, 55], + ... [ 33, 54, -16]]) + >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) + >>> print(y) + (ivy.array([[ 0, 0, 83], + [ 0, 23, 29], + [ 0, 0, 85], + [ 0, 31, 31], + [ 0, 0, 0], + [ 0, 22, 38], + [ 0, 33, 54]]), ivy.array([[0, 0, 2], + [0, 1, 2], + [0, 0, 2], + [0, 1, 1], + [0, 0, 0], + [0, 1, 2], + [0, 1, 2]])) - >>> a = ivy.array([1., 2., 3., 4.]) - >>> q = ivy.array([0.3, 0.7]) - >>> ivy.quantile(a, q, interpolation='lower') - ivy.array([1., 3.]) + >>> x = ivy.array([73, 15, 47]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) + >>> print(y) + (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) + + >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) + >>> print(y) + (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) """ - return ivy.current_backend(a).quantile( - a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out + return ivy.current_backend(x).cummax( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def corrcoef( - x: ivy.Array, +@handle_array_function +def cummin( + x: Union[ivy.Array, ivy.NativeArray], /, *, - y: Optional[ivy.Array] = None, - rowvar: bool = True, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) + """ + Return the cumulative minimum of the elements along a given axis. + + Parameters + ---------- + x + Input array. + axis + Axis along which the cumulative minimum is computed. Default is ``0``. + reverse + Whether to perform the cummin from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Array which holds the result of applying cummin at each + original array elements along the specified axis. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cummin(x) + >>> print(y) + ivy.array([1, 1, 1, 0]) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cummin(x, axis=0, reverse=True, out=y) + >>> print(y) + ivy.array([[1., 3., 0.], + [1., 3., 0.]]) + + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[ 2, 4, 5], + [ 3, 5, 5], + [ 1, 3, 10]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cummin(x, axis= 0) + >>> print(y) + { + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) + } + + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cummin(x,axis=1,reverse=True, out=y) + >>> print(y) + { + a: ivy.array([[1., 3., 4.]]), + b: ivy.array([[3., 5., 8.], + [5., 5., 5.]]), + c: ivy.array([[1., 1., 1.], + [3., 6., 9.], + [0., 2., 3.]]) + } + + >>> x = ivy.Container(a=ivy.array([[0],[5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cummin(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [0]]), + b: ivy.array([[6, 8, 7], + [4, 2, 3]]), + c: ivy.array([[1, 2], + [1, 2], + [1, 2]]) + } + """ + return ivy.current_backend(x).cummin( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + ) +# TODO: Make bins optional by offering an automatic bins creation like numpy. +# Make density argument work in tensorflow +# Bins as str is not defined (check Numpy implementation). +# Permit multiple axis. +# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def nanmedian( - input: ivy.Array, +def histogram( + a: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, + bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + density: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the - function, and so the docstring for ivy.nanmedian also applies to this method with - minimal changes. + Compute the histogram of the array ``a``. + + .. note:: + Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), + ..., I_{K-1} = [c_{K-1}, cK]. Parameters ---------- - self - Input array. + a + input array. + bins + if ``bins`` is an int, it defines the number of equal-width bins in the given + range. + if ``bins`` is an array, it defines a monotonically increasing array of bin + edges, including the rightmost edge, allowing for non-uniform bin widths. axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of - sub-classes of ndarray. If the sub-classes methods does not implement - keepdims any exceptions will be raised. - overwrite_input - If True, then allow use of memory of input array a for calculations. - The input array will be modified by the call to median. This will - save memory when you do not need to preserve the contents of the input array. - Treat the input as undefined, but it will probably be fully or partially sorted. - Default is False. If overwrite_input is True and a is not already an ndarray, - an error will be raised. + dimension along which maximum values must be computed. By default, the maximum + value must be computed over the entire array. Default: ``None``. + extend_lower_interval + if True, extend the lowest interval I0 to (-inf, c1]. + extend_upper_interval + ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). + dtype + the output type. + range + the lower and upper range of the bins. The first element of the range must be + less than or equal to the second. + weights + each value in ``a`` only contributes its associated weight towards the bin count + (instead of 1). Must be of the same shape as a. + density + if True, the result is the value of the probability density function at the + bin, normalized such that the integral over the range of bins is 1. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A new array holding the result. If the input contains integers + a tuple containing the values of the histogram and the bin edges. - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) - >>> ivy.nanmedian(x) - ivy.array(23.) - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), - b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) - >>> ivy.nanmedian(x) + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) + >>> z = ivy.histogram(x, bins=y) + >>> print(z) + ivy.array([1., 0., 1., 1.]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [4.4, 5.5, .6]]) + >>> bins = 4 + >>> range = (0., 5.) + >>> dtype = ivy.int32 + >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) + >>> print(y) + ivy.array([2, 1, 1, 1]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> axis = 1 + >>> extend_lower_interval = True + >>> extend_upper_interval = True + >>> dtype = ivy.float32 + >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) + >>> z = ivy.histogram( + ... x, + ... bins=y, + ... axis=axis, + ... extend_lower_interval=extend_lower_interval, + ... extend_upper_interval=extend_upper_interval, + ... dtype=dtype, + ... weights=weights) + >>> print(z) + ivy.array([[0., 3.], + [1., 0.], + [1., 0.], + [1., 0.], + [0., 0.]]) + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> dtype = ivy.int32 + >>> z = ivy.histogram(x, bins=y, dtype=dtype) + >>> print(z) { - a: ivy.array(3.), - b: ivy.array(23.) + a: ivy.array([1, 1, 1, 0, 0]), + b: ivy.array([0, 0, 0, 1, 2]) } """ - return ivy.current_backend().nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + return ivy.current_backend(a).histogram( + a, + bins=bins, + axis=axis, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + range=range, + weights=weights, + density=density, + out=out, ) @@ -436,42 +599,39 @@ def nanmedian( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def bincount( - x: ivy.Array, +def igamma( + a: Union[ivy.Array, ivy.NativeArray], /, *, - weights: Optional[ivy.Array] = None, - minlength: int = 0, - out: Optional[ivy.Array] = None, + x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Count the number of occurrences of each value in an integer array. + Compute the regularized lower gamma function of ``a`` and ``x``. Parameters ---------- self Input array. - weights - An optional input array. - minlength - A minimum number of bins for the output array. + x + An additional input array. + `x` has the same type as `a`. + out + optional output array, for writing the result to. Returns ------- ret - The bincount of the array elements. + The lower incomplete gamma function of the array elements. Examples -------- - >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) - >>> a.bincount(a) - 3.0 - >>> a.bincount(a, axis=0) - array([6.5, 2. , 2.5]) + >>> a = ivy.array([2.5]) + >>> x = ivy.array([1.7, 1.2]) + >>> a.igamma(x) + ivy.array([0.3614, 0.2085]) """ - return ivy.current_backend(x).bincount( - x, weights=weights, minlength=minlength, out=out - ) + return ivy.current_backend().igamma(a, x=x, out=out) @handle_exceptions @@ -480,417 +640,257 @@ def bincount( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def igamma( - a: Union[ivy.Array, ivy.NativeArray], +def median( + input: ivy.Array, /, *, - x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the regularized lower gamma function of ``a`` and ``x``. + Compute the median along the specified axis. Parameters ---------- - self + input Input array. - x - An additional input array. - `x` has the same type as `a`. + axis + Axis or axes along which the medians are computed. The default is to compute + the median along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. out optional output array, for writing the result to. Returns ------- ret - The lower incomplete gamma function of the array elements. + The median of the array elements. - Examples - -------- - >>> a = ivy.array([2.5]) - >>> x = ivy.array([1.7, 1.2]) - >>> a.igamma(x) - ivy.array([0.3614, 0.2085]) + Functional Examples + ------------------- + >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> ivy.median(a) + 3.5 + >>> ivy.median(a, axis=0) + ivy.array([6.5, 4.5, 2.5]) """ - return ivy.current_backend().igamma(a, x=x, out=out) + return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -def cov( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray] = None, +@infer_dtype +@handle_device_shifting +def nanmean( + a: ivy.Array, /, *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[ivy.Array] = None, - aweights: Optional[ivy.Array] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the covariance of matrix x1, or variables x1 and x2. + Compute the mean of all non-NaN elements along the specified dimensions. Parameters ---------- - x1 - a 1D or 2D input array, with a numeric data type. - x2 - optional second 1D or 2D input array, with a numeric data type. - Must have the same shape as ``self``. - rowVar - optional variable where each row of input is interpreted as a variable - (default = True). If set to False, each column is instead interpreted - as a variable. - bias - optional variable for normalizing input (default = False) by (N - 1) where - N is the number of given observations. If set to True, then normalization - is instead by N. Can be overridden by keyword ``ddof``. - ddof - optional variable to override ``bias`` (default = None). ddof=1 will return - the unbiased estimate, even with fweights and aweights given. ddof=0 will - return the simple average. - fweights - optional 1D array of integer frequency weights; the number of times each - observation vector should be repeated. - aweights - optional 1D array of observation vector weights. These relative weights are - typically large for observations considered "important" and smaller for - observations considered less "important". If ddof=0 is specified, the array - of weights can be used to assign probabilities to observation vectors. + a + Input array. + axis + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of sub-classes + of ndarray. If the sub-classes methods does not implement keepdims any + exceptions will be raised. dtype - optional variable to set data-type of the result. By default, data-type - will have at least ``numpy.float64`` precision. + The desired data type of returned tensor. Default is None. out - optional output array, for writing the result to. It must have a shape that - the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array containing the covariance matrix of an input matrix, or the - covariance matrix of two variables. The returned array must have a - floating-point data type determined by Type Promotion Rules and must be - a square matrix of shape (N, N), where N is the number of variables in the - input(s). - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([[1, 2, 3], - ... [4, 5, 6]]) - >>> y = x[0].cov(x[1]) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) - >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) - >>> z = ivy.Container.static_cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) - >>> z = ivy.cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With :class:`ivy.Array` input and rowVar flag set to False (True by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], rowVar=False) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input and bias flag set to True (False by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], bias=True) - >>> print(y) - ivy.array([[0.66666667, 0.66666667], - [0.66666667, 0.66666667]]) + The nanmean of the array elements. - With :class:`ivy.Array` input with both fweights and aweights given: - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> fw = ivy.array([1,2,3]) - >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) - >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) - >>> print(y) - ivy.array([[0.48447205, 0.48447205], - [0.48447205, 0.48447205]]) + Functional Examples + ------------------- + >>> a = ivy.array([[1, ivy.nan], [3, 4]]) + >>> ivy.nanmean(a) + 2.6666666666666665 + >>> ivy.nanmean(a, axis=0) + ivy.array([2., 4.]) """ - return ivy.current_backend(x1).cov( - x1, - x2, - rowVar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - dtype=dtype, + return ivy.current_backend(a).nanmean( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummax( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def nanmedian( + input: ivy.Array, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a tuple containing the cumulative maximum of elements of input along the - given axis and index location of each maximum value found along the given axis. + ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the + function, and so the docstring for ivy.nanmedian also applies to this method with + minimal changes. Parameters ---------- - x + self Input array. axis - Axis along which the cumulative maximum is computed. Default is ``0``. - exclusive - Whether to perform cummax exclusively. Default is ``False``. - reverse - Whether to perform the cummax from last to first element in the selected - axis. Default is ``False`` (from first to last element) + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of + sub-classes of ndarray. If the sub-classes methods does not implement + keepdims any exceptions will be raised. + overwrite_input + If True, then allow use of memory of input array a for calculations. + The input array will be modified by the call to median. This will + save memory when you do not need to preserve the contents of the input array. + Treat the input as undefined, but it will probably be fully or partially sorted. + Default is False. If overwrite_input is True and a is not already an ndarray, + an error will be raised. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cummax at each - original array elements along the specified axis. + A new array holding the result. If the input contains integers + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) - >>> y = ivy.cummax(x, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), - ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) - - >>> x = ivy.array([ 14, 15, 49, -24, -39]) - >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) - - >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) - >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) - >>> print(x) - ivy.array([[0, 0, 0, 0], - [0, 1, 1, 1]]) - - >>> x = ivy.array([[-36, 83, -81], - ... [ 23, 29, 63], - ... [-83, 85, 2], - ... [ 31, 25, -86], - ... [-10, -52, 0], - ... [ 22, 38, 55], - ... [ 33, 54, -16]]) - >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) - >>> print(y) - (ivy.array([[ 0, 0, 83], - [ 0, 23, 29], - [ 0, 0, 85], - [ 0, 31, 31], - [ 0, 0, 0], - [ 0, 22, 38], - [ 0, 33, 54]]), ivy.array([[0, 0, 2], - [0, 1, 2], - [0, 0, 2], - [0, 1, 1], - [0, 0, 0], - [0, 1, 2], - [0, 1, 2]])) - - >>> x = ivy.array([73, 15, 47]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) - >>> print(y) - (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) - - >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) - >>> print(y) - (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) + >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) + >>> ivy.nanmedian(x) + ivy.array(23.) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), + b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) + >>> ivy.nanmedian(x) + { + a: ivy.array(3.), + b: ivy.array(23.) + } """ - return ivy.current_backend(x).cummax( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + return ivy.current_backend().nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummin( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def quantile( + a: ivy.Array, + q: Union[ivy.Array, float], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: bool = False, + interpolation: str = "linear", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative minimum of the elements along a given axis. + Compute the q-th quantile of the data along the specified axis. Parameters ---------- - x + a Input array. + q + Quantile or sequence of quantiles to compute, which must be + between 0 and 1 inclusive. axis - Axis along which the cumulative minimum is computed. Default is ``0``. - reverse - Whether to perform the cummin from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. + Axis or axes along which the quantiles are computed. The default + is to compute the quantile(s) along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original array a. + interpolation + {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. + Default value: 'linear'. + This specifies the interpolation method to use when the desired quantile lies + between two data points i < j: + - linear: i + (j - i) * fraction, where fraction is the fractional part of the + index surrounded by i and j. + - lower: i. + - higher: j. + - nearest: i or j, whichever is nearest. + - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with + integer dtypes. + - nearest_jax: provides jax-like computation for interpolation='nearest'. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cummin at each - original array elements along the specified axis. + A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis + is None, a rank(q) array. The first rank(q) dimensions index quantiles for + different values of q. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cummin(x) - >>> print(y) - ivy.array([1, 1, 1, 0]) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cummin(x, axis=0, reverse=True, out=y) - >>> print(y) - ivy.array([[1., 3., 0.], - [1., 3., 0.]]) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = ivy.array(0.5) + >>> ivy.quantile(a, q) + ivy.array(3.5) - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[ 2, 4, 5], - [ 3, 5, 5], - [ 1, 3, 10]]) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = 0.5 + >>> ivy.quantile(a, q) + ivy.array(3.5) - With :class:`ivy.Container` input: + >>> ivy.quantile(a, q, axis=0) + ivy.array([6.5, 4.5, 2.5]) - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cummin(x, axis= 0) - >>> print(y) - { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) - } + >>> ivy.quantile(a, q, axis=1) + ivy.array([7., 2.]) - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cummin(x,axis=1,reverse=True, out=y) - >>> print(y) - { - a: ivy.array([[1., 3., 4.]]), - b: ivy.array([[3., 5., 8.], - [5., 5., 5.]]), - c: ivy.array([[1., 1., 1.], - [3., 6., 9.], - [0., 2., 3.]]) - } + >>> ivy.quantile(a, q, axis=1, keepdims=True) + ivy.array([[7.],[2.]]) - >>> x = ivy.Container(a=ivy.array([[0],[5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cummin(x,axis=0,out=x) - >>> print(x) - { - a: ivy.array([[0], - [0]]), - b: ivy.array([[6, 8, 7], - [4, 2, 3]]), - c: ivy.array([[1, 2], - [1, 2], - [1, 2]]) - } + >>> a = ivy.array([1., 2., 3., 4.]) + >>> q = ivy.array([0.3, 0.7]) + >>> ivy.quantile(a, q, interpolation='lower') + ivy.array([1., 3.]) """ - return ivy.current_backend(x).cummin( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + return ivy.current_backend(a).quantile( + a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out ) diff --git a/ivy/functional/ivy/general.py b/ivy/functional/ivy/general.py index 1267e86d08269..4d0aa36bfef8a 100644 --- a/ivy/functional/ivy/general.py +++ b/ivy/functional/ivy/general.py @@ -46,23 +46,18 @@ FN_CACHE = dict() INF = float("inf") - -precise_mode_stack = list() -queue_timeout_stack = list() array_mode_stack = list() -shape_array_mode_stack = list() -nestable_mode_stack = list() exception_trace_mode_stack = list() inplace_mode_stack = list() -trace_mode_dict = dict() -trace_mode_dict["frontend"] = "ivy/functional/frontends" -trace_mode_dict["ivy"] = "ivy/" -trace_mode_dict["full"] = "" -trace_mode_dict["none"] = "" -show_func_wrapper_trace_mode_stack = list() -min_denominator_stack = list() min_base_stack = list() +min_denominator_stack = list() +nestable_mode_stack = list() +precise_mode_stack = list() +queue_timeout_stack = list() +shape_array_mode_stack = list() +show_func_wrapper_trace_mode_stack = list() tmp_dir_stack = list() +trace_mode_dict = dict() # Extra # @@ -88,78 +83,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self -ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True - - -@handle_exceptions -def set_precise_mode(mode: bool) -> None: - """ - Set the mode of whether to use a promotion table that avoids any precision loss or a - compute effecient table that avoids most wider-than-necessary promotions. - - Parameter - --------- - mode - boolean whether to use high precision promtion table - - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False - - >>> ivy.set_precise_mode(True) - >>> ivy.precise_mode - True - """ - global precise_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - precise_mode_stack.append(mode) - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) - - -@handle_exceptions -def unset_precise_mode() -> None: - """ - Reset the mode of whether to use a promotion table that avoids any precision loss or - a compute effecient table that avoids most wider-than-necessary promotions. - - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False - - >>> ivy.unset_precise_mode() - >>> ivy.precise_mode - True - """ - global precise_mode_stack - if precise_mode_stack: - precise_mode_stack.pop(-1) - mode = precise_mode_stack[-1] if precise_mode_stack else True - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) - - -def _update_promotion_table(precise): - """Update the current datatype promotion table.""" - if precise: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.precise_extra_promotion_table, - } - - else: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.extra_promotion_table, - } - - class ArrayMode: """Array Mode Context Manager.""" @@ -179,426 +102,371 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self -def get_referrers_recursive( - item, depth=0, max_depth=None, seen_set=None, local_set=None -): - """ - Summary. +# --- Helpers --- # +# --------------- # - Parameters - ---------- - item - depth - (Default value = 0) - max_depth - (Default value = None) - seen_set - (Default value = None) - local_set - (Default value = None`) - """ - seen_set = ivy.default(seen_set, set()) - local_set = ivy.default(local_set, set()) - ret_cont = ivy.Container( - repr=str(item).replace(" ", ""), - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - referrers = [ - ref - for ref in gc.get_referrers(item) - if not ( - isinstance(ref, dict) - and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) - ) - ] - local_set.add(str(id(referrers))) - for ref in referrers: - ref_id = str(id(ref)) - if ref_id in local_set or hasattr(ref, "cell_contents"): - continue - seen = ref_id in seen_set - seen_set.add(ref_id) - refs_rec = lambda: get_referrers_recursive( - ref, depth + 1, max_depth, seen_set, local_set - ) - this_repr = "tracked" if seen else str(ref).replace(" ", "") - if not seen and (not max_depth or depth < max_depth): - val = ivy.Container( - repr=this_repr, - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - refs = refs_rec() - for k, v in refs.items(): - val[k] = v - else: - val = this_repr - ret_cont[str(ref_id)] = val - return ret_cont +def _all_dnd_combinations(): + all_comb = {} + for device in ivy.all_devices: + all_comb[device] = ivy.all_dtypes + return all_comb -@handle_exceptions -@handle_backend_invalid -def is_native_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False -) -> bool: - """ - Determine whether the input x is an :class:`ivy.NativeArray` instance. +def _broadcast_to(input, target_shape): + input = ivy.squeeze(input) + if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): + return ivy.reshape(input, target_shape) + else: + input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input + new_dims = () + i_i = len(input.shape) - 1 + for i_t in range(len(target_shape) - 1, -1, -1): + if len(input.shape) + len(new_dims) >= len(target_shape): + break + if i_i < 0 or target_shape[i_t] != input.shape[i_i]: + new_dims += (i_t,) + else: + i_i -= 1 + input = ivy.expand_dims(input, axis=new_dims) + return ivy.broadcast_to(input, target_shape) - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. - Returns - ------- - ret - Boolean, whether or not x is an :class:`ivy.NativeArray`. +def _dnd_dict_difference(a, b): + res = a + for device in list(a): + if device in b: + difference = set.difference(set(a[device]), set(b[device])) + if difference: + res[device] = tuple(difference) + else: + del res[device] + return res - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_native_array(x) - False - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_native_array(x, exclusive=True) - True - """ - try: - return current_backend(x).is_native_array(x, exclusive=exclusive) - except ValueError: - return False +def _dnd_dict_intersection(a, b): + res = {} + for device in a: + if device in b: + intersection = set.intersection(set(a[device]), set(b[device])) + if intersection: + res[device] = tuple(intersection) + return res -@handle_exceptions -@handle_backend_invalid -def is_ivy_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False -) -> bool: - """ - Determine whether the input x is a valid Ivy Array. +def _dnd_dict_union(a, b): + res = {} + for device in set(list(a) + list(b)): + u1 = set(a.get(device, ())) + u2 = set(b.get(device, ())) + res[device] = tuple(set.union(u1, u2)) - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + return res - Returns - ------- - ret - Boolean, whether or not x is a valid Ivy Array. - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_ivy_array(x) - True +def _get_devices_and_dtypes(fn, recurse=False, complement=True): + supported_devices = ivy.function_supported_devices(fn, recurse=recurse) + supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_ivy_array(x, exclusive=True) - False - """ - return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) + if hasattr(fn, "partial_mixed_handler"): + supported_devices = supported_devices["primary"] + supported_dtypes = supported_dtypes["primary"] + supported = {} + # Generate a base supported set from other attributes + for device in supported_devices: + supported[device] = supported_dtypes -@handle_exceptions -@handle_backend_invalid -def is_array(x: Any, /, *, exclusive: bool = False) -> bool: - """ - Determine whether the input x is either an Ivy Array or a Native Array. + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + backend = ivy.current_backend_str() - Returns - ------- - ret - Boolean, whether or not x is an array. + # Their values are formatted like either + # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype.__get__() - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> print(ivy.is_array(x)) - True + if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): + fn_supported_dnd = fn_supported_dnd.get(backend, supported) - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> print(ivy.is_array(x, exclusive=True)) - True + ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) + # dict intersection + supported = _dnd_dict_intersection(supported, fn_supported_dnd) - >>> x = [2, 3] - >>> print(ivy.is_array(x)) - False - """ - return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( - x, exclusive=exclusive - ) + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() + if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) -@handle_exceptions -def is_ivy_container(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Container. + ivy.utils.assertions.check_isinstance( + list(fn_unsupported_dnd.values())[0], tuple + ) + # dict difference + supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - Parameters - ---------- - x - The input to check + if complement: + # dict difference + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported - Returns - ------- - ret - Boolean, whether or not x is an ivy container. - Examples - -------- - >>> x = ivy.Container() - >>> print(ivy.is_ivy_container(x)) - True +def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: + fn_unsupported_dnd = {} + fn_supported_dnd = {} + backend = ivy.current_backend_str() + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): + fn_supported_dnd = fn_supported_dnd.get(backend, {}) - >>> x = [2, 3] - >>> print(ivy.is_ivy_container(x)) - False - """ - return isinstance(x, ivy.Container) + ivy.utils.assertions.check_false( + fn_unsupported_dnd and fn_supported_dnd, + "unsupported_device_and_dtype and supported_device_and_dtype cannot" + " both be defined for the same function", + ) + us = "unsupported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") -ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True + ss = "supported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") + return True -@handle_exceptions -def set_array_mode(mode: bool) -> None: - """ - Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs - back to ivy.Array. - It Stops the conversion of ivy.NativeArray to ivy.Array in the - case when it is set to False. +def _numel(shape): + shape = tuple(shape) + return ivy.prod(shape).to_scalar() if shape != () else 1 - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions - Examples - -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False +def _parse_query(query, x_shape): + query = (query,) if not isinstance(query, tuple) else query + query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) - >>> ivy.set_array_mode(True) - >>> ivy.array_mode - True - """ - global array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - array_mode_stack.append(mode) - ivy.__setattr__("array_mode", mode, True) + # array containing all of x's flat indices + x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) + # use numpy's __getitem__ to get the queried indices + x_idxs = ivy.array(x_.to_numpy()[query_]) + target_shape = x_idxs.shape -@handle_exceptions -def unset_array_mode() -> None: - """ - Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back - to ivy.Array to the previous state. + if 0 in x_idxs.shape or 0 in x_shape: + return None, target_shape - Examples - -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False + # convert the flat indices to multi-D indices + x_idxs = ivy.unravel_index(x_idxs, x_shape) - >>> ivy.unset_shape_array_mode() - >>> ivy.array_mode - True - """ - global array_mode_stack - if array_mode_stack: - array_mode_stack.pop(-1) - mode = array_mode_stack[-1] if array_mode_stack else True - ivy.__setattr__("array_mode", mode, True) + # stack the multi-D indices to bring them to gather_nd/scatter_nd format + x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) + return x_idxs, target_shape -ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True +def _update_promotion_table(precise): + """Update the current datatype promotion table.""" + if precise: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.precise_extra_promotion_table, + } -@handle_exceptions -def set_nestable_mode(mode: bool) -> None: - """ - Set the mode of whether to check if function inputs are ivy.Container. + else: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.extra_promotion_table, + } - Parameter - --------- - mode - boolean whether to check if function inputs are ivy.Container - Examples - -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False +def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): + attr_list = () + if hasattr(fn, other_attr_name): + attr_list = getattr(fn, other_attr_name) + if isinstance(attr_list, dict): + attr_list = attr_list.get(backend, ()) + ivy.utils.assertions.check_false( + dnd_dict and attr_list, + f"Cannot specify both {first_attr_name} and {other_attr_name} " + "cannot both be defined for the same function", + ) - >>> ivy.set_nestable_mode(True) - >>> ivy.nestable_mode - True - """ - global nestable_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - nestable_mode_stack.append(mode) - ivy.__setattr__("nestable_mode", mode, True) + +# --- Main --- # +# ------------ # @handle_exceptions -def unset_nestable_mode() -> None: +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +def all_equal( + *xs: Iterable[Any], equality_matrix: bool = False +) -> Union[bool, ivy.Array, ivy.NativeArray]: """ - Reset the mode of whether to check if function inputs are ivy.Container to the - previous state. + Determine whether the inputs are all equal. + + Parameters + ---------- + xs + inputs to compare. + equality_matrix + Whether to return a matrix of equalities comparing each input with every other. + Default is ``False``. + + Returns + ------- + ret + Boolean, whether or not the inputs are equal, or matrix array of booleans if + equality_matrix=True is set. Examples -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False + With :class:`ivy.Array` inputs: - >>> ivy.unset_nestable_mode() - >>> ivy.nestable_mode + >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> y = ivy.all_equal(x1, x2) + >>> print(y) True - """ - global nestable_mode_stack - if nestable_mode_stack: - nestable_mode_stack.pop(-1) - mode = nestable_mode_stack[-1] if nestable_mode_stack else True - ivy.__setattr__("nestable_mode", mode, True) - - -ivy.exception_trace_mode = ( - exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" -) + >>> x1 = ivy.array([0, 0]) + >>> x2 = ivy.array([0, 0]) + >>> x3 = ivy.array([1, 0]) + >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) + >>> print(y) + ivy.array([[ True, True, False], + [ True, True, False], + [False, False, True]]) -@handle_exceptions -def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: - """ - Set the mode of whether to show frontend-truncated exception stack traces, ivy- - truncated exception stack traces or full exception stack traces. + With one :class:`ivy.Container` inputs: - Parameter - --------- - mode - str exeption trace mode, one of `ivy`, `full` or `frontend` + >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), + ... b=ivy.array([0, 0, -1, 1, 0])) + >>> x2 = ivy.array([0, 0, -1, 1, 0]) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) + { + a: True, + b: True + } - Examples - -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' + With multiple :class:`ivy.Container` inputs: - >>> ivy.set_exception_trace_mode("full") - >>> ivy.exception_trace_mode - 'full' + >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, 0, 1])) + >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, -1, -1])) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) + { + a: True, + b: False + } """ - global exception_trace_mode_stack - trace_modes = list(trace_mode_dict.keys()) - ivy.utils.assertions.check_elem_in_list( - mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) - ) - exception_trace_mode_stack.append(mode) - ivy.__setattr__("exception_trace_mode", mode, True) + equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b + if equality_matrix: + num_arrays = len(xs) + mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] + for i, xa in enumerate(xs): + for j_, xb in enumerate(xs[i:]): + j = j_ + i + res = equality_fn(xa, xb) + if ivy.is_native_array(res): + # noinspection PyTypeChecker + res = ivy.to_scalar(res) + # noinspection PyTypeChecker + mat[i][j] = res + # noinspection PyTypeChecker + mat[j][i] = res + return ivy.array(mat) + x0 = xs[0] + for x in xs[1:]: + if not equality_fn(x0, x): + return False + return True @handle_exceptions -def unset_exception_trace_mode() -> None: +def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): """ - Reset the trace mode to the previously set mode. + Return the index and `inspect.Parameter` representation of the specified argument. + In the form of a dict with keys "idx" and "param". - Examples - -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' + Parameters + ---------- + fn + The function to retrieve the argument information for + name + The name of the argument + idx + the index of the argument in the inputs - >>> ivy.unset_exception_trace_mode() - >>> ivy.exception_trace_mode - 'full' + Returns + ------- + ret + a `dict` containing the idx, and the `inspect.Parameter` for the argument, + which itself contains the parameter name, type, and other helpful information. """ - global exception_trace_mode_stack - if exception_trace_mode_stack: - exception_trace_mode_stack.pop(-1) - mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" - ivy.__setattr__("exception_trace_mode", mode, True) - - -ivy.show_func_wrapper_trace_mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True -) + ivy.utils.assertions.check_all_or_any_fn( + name, + idx, + fn=ivy.exists, + type="any", + limit=[1], + message="exactly one of the keyword arguments name or idx must be provided", + as_array=False, + ) + params = inspect.signature(fn).parameters + if ivy.exists(name): + return {"idx": list(params).index(name), "param": params[name]} + return {"idx": idx, "param": list(params.values())[idx]} @handle_exceptions -def set_show_func_wrapper_trace_mode(mode: bool) -> None: - """ - Set the mode of whether to show the full stack trace with function wrapping traces. - - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions - - Examples - -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False - - >>> ivy.set_show_func_wrapper_trace_mode(True) - >>> ivy.show_func_wrapper_trace_mode - True +def arg_names(receiver): """ - global show_func_wrapper_trace_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - show_func_wrapper_trace_mode_stack.append(mode) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + Get the expected keyword arguments for a function or class constructor. + Parameters + ---------- + receiver + Function or class constructor -@handle_exceptions -def unset_show_func_wrapper_trace_mode() -> None: - """ - Reset the mode of whether to show the full stack trace with function wrapping - traces. + Returns + ------- + ret + List containing the keyword arguments' names for a function or class constructor Examples -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False + >>> x = ivy.arg_names(ivy.tan) + >>> print(x) + ['x', 'out'] - >>> ivy.unset_show_func_wrapper_trace_mode() - >>> ivy.show_func_wrapper_trace_mode - True + >>> x = ivy.arg_names(ivy.optimizers.Adam) + >>> print(x) + ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', + 'stop_gradients', 'compile_on_next_step', 'device'] """ - global show_func_wrapper_trace_mode_stack - if show_func_wrapper_trace_mode_stack: - show_func_wrapper_trace_mode_stack.pop(-1) - mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True - ) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + return list(inspect.signature(receiver).parameters.keys()) @handle_exceptions @@ -653,347 +521,254 @@ def array_equal( @handle_exceptions @handle_nestable -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_array_function -def all_equal( - *xs: Iterable[Any], equality_matrix: bool = False -) -> Union[bool, ivy.Array, ivy.NativeArray]: +def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: """ - Determine whether the inputs are all equal. + Assert that inplace operations are supported for x. Parameters ---------- - xs - inputs to compare. - equality_matrix - Whether to return a matrix of equalities comparing each input with every other. - Default is ``False``. + x + Input variable or array to check for inplace support for. Returns ------- ret - Boolean, whether or not the inputs are equal, or matrix array of booleans if - equality_matrix=True is set. + True if supports, raises IvyBackendException otherwise + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> y = ivy.all_equal(x1, x2) - >>> print(y) + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) True - >>> x1 = ivy.array([0, 0]) - >>> x2 = ivy.array([0, 0]) - >>> x3 = ivy.array([1, 0]) - >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) - >>> print(y) - ivy.array([[ True, True, False], - [ True, True, False], - [False, False, True]]) + With :class:`ivy.Array` input and default backend set as `torch`: - With one :class:`ivy.Container` inputs: + >>> ivy.set_backend("torch") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) + True - >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), - ... b=ivy.array([0, 0, -1, 1, 0])) - >>> x2 = ivy.array([0, 0, -1, 1, 0]) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) + With :class:`ivy.Container` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) { a: True, b: True } - With multiple :class:`ivy.Container` inputs: + With :class:`ivy.Container` input and default backend set as `torch`: - >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, 0, 1])) - >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, -1, -1])) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) { a: True, - b: False + b: True } """ - equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b - if equality_matrix: - num_arrays = len(xs) - mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] - for i, xa in enumerate(xs): - for j_, xb in enumerate(xs[i:]): - j = j_ + i - res = equality_fn(xa, xb) - if ivy.is_native_array(res): - # noinspection PyTypeChecker - res = ivy.to_scalar(res) - # noinspection PyTypeChecker - mat[i][j] = res - # noinspection PyTypeChecker - mat[j][i] = res - return ivy.array(mat) - x0 = xs[0] - for x in xs[1:]: - if not equality_fn(x0, x): - return False + ivy.utils.assertions.check_true( + ivy.supports_inplace_updates(x), + "Inplace operations are not supported {} types with {} backend".format( + type(x), ivy.current_backend_str() + ), + ) return True @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays -@handle_array_function -@handle_device_shifting -def to_numpy( - x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True -) -> np.ndarray: +def cache_fn(func: Callable) -> Callable: """ - Convert an array into a numpy array. + Cache function outputs. + + A decorator to wrap a function, such that computed outputs are cached to avoid + recalculating them later. Parameters ---------- - x - input array - copy - whether to copy the array to a new address or not. - Default is ``True``. + func + The function to wrap, whose output should be cached for later. Returns ------- ret - a numpy array copying all the element of the array ``x``. + The newly cache wrapped function. Examples -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [-1 0 1] - - >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [[-1 0 1] - [-1 0 1] - [ 1 0 -1]] + With positional arguments only: - With :class:`ivy.Container` input: + >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 + >>> cached_sum = ivy.cache_fn(my_sum) + >>> print(cached_sum(3, 5)) + 8 - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([-1, 0, 1], dtype=int32) - } + With keyword arguments: - >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), - ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([[-1., 0., 1.], - [-1., 0., 1.], - [1., 0., -1.]], dtype=float32), - b: array([[-1, 0, 0], - [1, 0, 1], - [1, 1, 1]], dtype=int32) - } + >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc + >>> cached_line_eq = ivy.cache_fn(line_eq) + >>> print(cached_line_eq(3, itc=5, slp=2)) + 11 """ - return current_backend(x).to_numpy(x, copy=copy) + global FN_CACHE + if func not in FN_CACHE: + FN_CACHE[func] = dict() + @wraps(func) + def cached_fn(*args, **kwargs): + key = "".join( + [str(i) + ", " for i in args] + + [" kw, "] + + [str(i) + ", " for i in sorted(kwargs.items())] + ) + cache = FN_CACHE[func] + if key in cache: + return cache[key] + ret = func(*args, **kwargs) + cache[key] = ret + return ret -@handle_exceptions -@handle_nestable -def isscalar(x: Any, /) -> bool: - return np.isscalar(x) + return cached_fn @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: +def clip_matrix_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, + /, + *, + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert an array with a single element into a scalar. + Clips (limits) the matrix norm of an array. Parameters ---------- x - Input array with a single element. + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - a scalar copying the element of the array ``x``. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + An array with the matrix norm downscaled to the max norm if needed. Functional Examples ------------------- With :class:`ivy.Array` input: - >>> x = ivy.array([3]) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[0., 1., 2.]]) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) - 3 + ivy.array([[0. , 0.894, 1.79 ]]) - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) + >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) + >>> print(y) + ivy.array([[ 0.0353, -0.424 , 1.31 ], + [ 0. , 2.58 , -0.176 ]]) - >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[[5., 4.], [-2., 6.]], + ... [[3., 7.], [0., -5.]]]) + >>> y = ivy.empty((2, 2, 2)) + >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) >>> print(y) - { - a: -1, - b: 3 - } + ivy.array([[[ 0.339, 0.271], + [-0.135, 0.406]], + [[ 0.168, 0.391], + [ 0. , -0.279]]]) - >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), - ... c=ivy.array([-1])) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[0., 1.], + ... [2., 3.]]) + >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) + >>> print(x) + ivy.array([[0., 1.], + [2., 3.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), + ... b=ivy.array([[3., 4., 5.]])) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) { - a: 1, - b: 0, - c: -1 + a: ivy.array([[0., 0.894, 1.79]]), + b: ivy.array([[0.849, 1.13, 1.41]]) } """ - return current_backend(x).to_scalar(x) + norms = ivy.matrix_norm(x, ord=p, keepdims=True) + ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) + return ivy.multiply(ratios, x, out=out) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: +def clip_vector_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, + /, + *, + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Create a (possibly nested) list from input array. + Clips (limits) the vector p-norm of an array. Parameters ---------- x - Input array. + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - A list representation of the input array ``x``. + An array with the vector norm downscaled to the max norm if needed. + + Functional Examples + ------------------ - Examples - -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_list(x) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.clip_vector_norm(x, 2.0) >>> print(y) - [-1, 0, 1] + ivy.array([0. , 0.894, 1.79 ]) - >>> x = ivy.array([[ 1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.to_list(x) + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) >>> print(y) - [[1.100000023841858,2.200000047683716,3.299999952316284], - [-4.400000095367432,-5.5,-6.599999904632568]] - - >>> x = ivy.array([[[-1, 0, 1], - ... [ 1, 0, -1]], - ... [[ 1, -1, 0], - ... [ 1, 0, -1]]]) - >>> y = ivy.to_list(x) - >>> print(y) - [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [-1, 0, 1] - } - - >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], - ... [-1, 0, 1], - ... [1, 0, -1]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] - } - - >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], - ... [[1, -1, 0],[1, 0, -1]]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - } - """ - return current_backend(x).to_list(x) - - -@handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def clip_vector_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, - /, - *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Clips (limits) the vector p-norm of an array. - - Parameters - ---------- - x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - An array with the vector norm downscaled to the max norm if needed. - - Functional Examples - ------------------ - - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.clip_vector_norm(x, 2.0) - >>> print(y) - ivy.array([0. , 0.894, 1.79 ]) - - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) - >>> print(y) - ivy.array([ 0.417, -0.583, 2. ]) + ivy.array([ 0.417, -0.583, 2. ]) >>> x = ivy.array([[[0., 0.], [1., 3.], [2., 6.]], ... [[3., 9.], [4., 12.], [5., 15.]]]) @@ -1037,85 +812,139 @@ def clip_vector_norm( @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def clip_matrix_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, +def container_types(): + """ + Summary. + + Returns + ------- + ret + a key-value structure, and exposes public methods .keys(), .values() and + items(). + """ + # noinspection PyBroadException + try: + return current_backend().container_types() + except ValueError: + return [] + + +@handle_exceptions +def current_backend_str() -> Union[str, None]: + """ + Return framework string. + + Returns + ------- + ret + The framework string. + """ + fw = current_backend() + if not backend_stack: + return "" + return fw.current_backend_str() + + +@handle_exceptions +def default( + x: Any, /, + default_val: Any, *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + catch_exceptions: bool = False, + rev: bool = False, + with_callable: bool = False, +) -> Any: """ - Clips (limits) the matrix norm of an array. + Return x provided it exists (is not None), else returns default value. Parameters ---------- x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input which may or may not exist (be None). + default_val + The default value. + catch_exceptions + Whether to catch exceptions from callable x. + Default is ``False``. + rev + Whether to reverse the input x and default_val. + Default is ``False``. + with_callable + Whether either of the arguments might be callable functions. + Default is ``False``. Returns ------- ret - An array with the matrix norm downscaled to the max norm if needed. + x if x exists (is not None), else default. Functional Examples - ------------------- + ------------------ + With :code:`Any` input: - With :class:`ivy.Array` input: + >>> x = None + >>> y = ivy.default(x, "default_string") + >>> print(y) + default_string - >>> x = ivy.array([[0., 1., 2.]]) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = "" + >>> y = ivy.default(x, "default_string") >>> print(y) - ivy.array([[0. , 0.894, 1.79 ]]) - >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) - >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) + + >>> x = ivy.array([4, 5, 6]) + >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) >>> print(y) - ivy.array([[ 0.0353, -0.424 , 1.31 ], - [ 0. , 2.58 , -0.176 ]]) + ivy.array([1, 2, 3]) - >>> x = ivy.array([[[5., 4.], [-2., 6.]], - ... [[3., 7.], [0., -5.]]]) - >>> y = ivy.empty((2, 2, 2)) - >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) + >>> x = lambda: ivy.array([1, 2, 3]) + >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) >>> print(y) - ivy.array([[[ 0.339, 0.271], - [-0.135, 0.406]], - [[ 0.168, 0.391], - [ 0. , -0.279]]]) + ivy.array([1, 2, 3]) - >>> x = ivy.array([[0., 1.], - ... [2., 3.]]) - >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) - >>> print(x) - ivy.array([[0., 1.], - [2., 3.]]) + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) + >>> print(y) + ivy.array([1, 2, 3]) - With :class:`ivy.Container` input: + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) + >>> print(y) + ivy.array([1, 2, 3]) - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), - ... b=ivy.array([[3., 4., 5.]])) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True) >>> print(y) - { - a: ivy.array([[0., 0.894, 1.79]]), - b: ivy.array([[0.849, 1.13, 1.41]]) - } + ivy.array([1, 2, 3]) + + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True, rev=True) + >>> print(y) + ivy.array([1, 2, 3]) """ - norms = ivy.matrix_norm(x, ord=p, keepdims=True) - ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) - return ivy.multiply(ratios, x, out=out) + with_callable = catch_exceptions or with_callable + if rev: + tmp = x + x = default_val + default_val = tmp + if with_callable: + x_callable = callable(x) + default_callable = callable(default_val) + else: + x_callable = False + default_callable = False + if catch_exceptions: + # noinspection PyBroadException + try: + x = x() if x_callable else x + except Exception: + return default_val() if default_callable else default_val + else: + x = x() if x_callable else x + return x if exists(x) else default_val() if default_callable else default_val @handle_exceptions @@ -1123,234 +952,267 @@ def clip_matrix_norm( @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def fourier_encode( +def einops_rearrange( x: Union[ivy.Array, ivy.NativeArray], - max_freq: Union[float, ivy.Array, ivy.NativeArray], + pattern: str, /, *, - num_bands: int = 4, - linear: bool = False, - concat: bool = True, - flatten: bool = False, -) -> Union[ivy.Array, ivy.NativeArray, Tuple]: + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Pad an array with fourier encodings. + Perform einops rearrange operation on input array x. Parameters ---------- x - Input array to encode. - max_freq - The maximum frequency of the encoding. - num_bands - The number of frequency bands for the encoding. - Default is 4. - linear - Whether to space the frequency bands linearly as opposed to geometrically. - Default is ``False``. - concat - Whether to concatenate the position, sin and cos values, or return seperately. - Default is ``True``. - flatten - Whether to flatten the position dimension into the batch dimension. - Default is False. + Input array to be re-arranged. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - New array with the final dimension expanded, and the encodings stored in this - channel. + New array with einops.rearrange having been applied. Examples -------- - >>> x = ivy.array([1,2,3]) - >>> y = 1.5 - >>> z = ivy.fourier_encode(x,y) - >>> print(z) - ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00]]) + With :class:`ivy.Array` instance method: + >>> x = ivy.array([[1, 2, 3], + ... [-4, -5, -6]]) + >>> y = x.einops_rearrange("height width -> width height") + >>> print(y) + ivy.array([[ 1, -4], + [ 2, -5], + [ 3, -6]]) - >>> x = ivy.array([3,10]) - >>> y = 2.5 - >>> z = ivy.fourier_encode(x, y, num_bands=3) - >>> print(z) - ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, - -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], - [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, - 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) - """ - x_in = x - dim = x.shape[-1] - x = ivy.expand_dims(x, axis=-1) - orig_x = x - if linear: - scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) - else: - if ivy.backend == "torch" and isinstance(max_freq, float): - scales = ivy.logspace( - 0.0, - ivy.log(ivy.array(max_freq / 2)) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - else: - scales = ivy.logspace( - 0.0, - ivy.log(max_freq / 2) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - scales = ivy.astype(scales, ivy.dtype(x)) - scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] - x = x * scales * math.pi - sin_x = ivy.sin(x) - cos_x = ivy.cos(x) - if flatten: - orig_x = x_in - sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) - cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) - if concat: - return ivy.concat([orig_x, sin_x, cos_x], axis=-1) - return sin_x, cos_x + >>> x = ivy.array([[[ 1, 2, 3], + ... [ 4, 5, 6]], + ... [[ 7, 8, 9], + ... [10, 11, 12]]]) + >>> y = x.einops_rearrange("c h w -> c (h w)") + >>> print(y) + ivy.array([[ 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12]]) + >>> x = ivy.array([[1, 2, 3, 4, 5, 6], + ... [7, 8, 9, 10, 11, 12]]) + >>> y = ivy.zeros((4,3)) + >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) + >>> print(y) + ivy.array([[ 1, 2, 3], + [ 4, 5, 6], + [ 7, 8, 9], + [10, 11, 12]]) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def value_is_nan( - x: Union[ivy.Array, ivy.NativeArray, Number], - /, - *, - include_infs: bool = True, -) -> bool: - """ - Determine whether the single valued array or scalar is of nan type. + With :class:`ivy.Container` input: - Parameters - ---------- - x - The input to check Input array. - include_infs - Whether to include infs and -infs in the check. - Default is ``True``. + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> y = ivy.einops_rearrange(x, 'a b -> b a') + >>> print(y) + { + a: ivy.array([[-4.46999979, 3.66000009], + [0.93000001, 24.29000092], + [-3.33999991, 3.6400001]]), + b: ivy.array([[4.96000004, 4.36000013], + [1.51999998, 13.96000004], + [-10.67000008, 0.30000001]]) + } - Returns - ------- - ret - Boolean as to whether the input value is a nan or not. + With varying pattern: - Examples - -------- - >>> x = ivy.array([451]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False + Suppose we have a set of 32 images in "h w c" format (height-width-channel) + and concatenate images along height (vertical axis), 960 = 32 * 30 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') + >>> print(x.shape) + (960, 40, 3) - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - True + # Concatenate images along horizontal axis, 1280 = 32 * 40 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') + >>> print(x.shape) + (30, 1280, 3) - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - False + # Reorder axes to "b c h w" format for deep learning + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') + >>> print(x.shape) + (32, 3, 30, 40) - >>> x = ivy.array([float('nan')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - True + # Flatten each image into a vector, 3600 = 30 * 40 * 3 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') + >>> print(x.shape) + (32, 3600) - >>> x = ivy.array([0]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False + # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), + # 128 = 32 * 2 * 2 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', + ... h1=2, w1=2) + >>> print(x.shape) + (128, 15, 20, 3) + + # Space-to-depth operation + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, + ... w1=2) + >>> print(x.shape) + (32, 15, 20, 12) """ - x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x - if not x_scalar == x: - return True - if include_infs and (x_scalar == INF or x_scalar == -INF): - return True - return False + ret = einops.rearrange(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret @handle_exceptions @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_array_function -def has_nans( - x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True -) -> bool: +def einops_reduce( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + reduction: Union[str, Callable], + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Determine whether the array contains any nans, as well as infs or -infs if - specified. + Perform einops reduce operation on input array x. Parameters ---------- x - Input array. - include_infs - Whether to include ``+infinity`` and ``-infinity`` in the check. - Default is ``True``. + Input array to be reduced. + pattern + Reduction pattern. + reduction + One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Boolean as to whether the array contains nans. + New array with einops.reduce having been applied. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - False + >>> x = ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]) + >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') + >>> print(reduced) + ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) - >>> x = ivy.array([float('nan'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True + With :class:`ivy.Container` input: - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') + >>> print(reduced) + { + a: ivy.array([-2.29333329, 10.53000069]), + b: ivy.array([-1.39666676, 6.20666695]) + } + """ + ret = einops.reduce(x, pattern, reduction, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x, include_infs=False) - >>> print(y) - False + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def einops_repeat( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: + """ + Perform einops repeat operation on input array x. + + Parameters + ---------- + x + Input array to be repeated. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + New array with einops.repeat having been applied. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4]) + >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) + >>> print(repeated) + ivy.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.has_nans(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[4,5], + ... [1, 3]]), + ... b=ivy.array([[9, 10], + ... [4, 2]])) + >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) + >>> print(repeated) { - a: False, - b: False + a: ivy.array([[4, 5, 4, 5], + [1, 3, 1, 3]]), + b: ivy.array([[9, 10, 9, 10], + [4, 2, 4, 2]]) } """ - return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) + ret = einops.repeat(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret @handle_exceptions @@ -1428,986 +1290,1254 @@ def exists(x: Any) -> bool: @handle_exceptions -def default( - x: Any, +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def fourier_encode( + x: Union[ivy.Array, ivy.NativeArray], + max_freq: Union[float, ivy.Array, ivy.NativeArray], /, - default_val: Any, *, - catch_exceptions: bool = False, - rev: bool = False, - with_callable: bool = False, -) -> Any: + num_bands: int = 4, + linear: bool = False, + concat: bool = True, + flatten: bool = False, +) -> Union[ivy.Array, ivy.NativeArray, Tuple]: """ - Return x provided it exists (is not None), else returns default value. + Pad an array with fourier encodings. Parameters ---------- x - Input which may or may not exist (be None). - default_val - The default value. - catch_exceptions - Whether to catch exceptions from callable x. - Default is ``False``. - rev - Whether to reverse the input x and default_val. - Default is ``False``. - with_callable - Whether either of the arguments might be callable functions. + Input array to encode. + max_freq + The maximum frequency of the encoding. + num_bands + The number of frequency bands for the encoding. + Default is 4. + linear + Whether to space the frequency bands linearly as opposed to geometrically. Default is ``False``. + concat + Whether to concatenate the position, sin and cos values, or return seperately. + Default is ``True``. + flatten + Whether to flatten the position dimension into the batch dimension. + Default is False. Returns ------- ret - x if x exists (is not None), else default. - - Functional Examples - ------------------ - With :code:`Any` input: + New array with the final dimension expanded, and the encodings stored in this + channel. - >>> x = None - >>> y = ivy.default(x, "default_string") - >>> print(y) - default_string + Examples + -------- + >>> x = ivy.array([1,2,3]) + >>> y = 1.5 + >>> z = ivy.fourier_encode(x,y) + >>> print(z) + ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00]]) - >>> x = "" - >>> y = ivy.default(x, "default_string") - >>> print(y) + >>> x = ivy.array([3,10]) + >>> y = 2.5 + >>> z = ivy.fourier_encode(x, y, num_bands=3) + >>> print(z) + ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, + -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], + [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, + 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) + """ + x_in = x + dim = x.shape[-1] + x = ivy.expand_dims(x, axis=-1) + orig_x = x + if linear: + scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) + else: + if ivy.backend == "torch" and isinstance(max_freq, float): + scales = ivy.logspace( + 0.0, + ivy.log(ivy.array(max_freq / 2)) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + else: + scales = ivy.logspace( + 0.0, + ivy.log(max_freq / 2) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + scales = ivy.astype(scales, ivy.dtype(x)) + scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] + x = x * scales * math.pi + sin_x = ivy.sin(x) + cos_x = ivy.cos(x) + if flatten: + orig_x = x_in + sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) + cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) + if concat: + return ivy.concat([orig_x, sin_x, cos_x], axis=-1) + return sin_x, cos_x - >>> x = ivy.array([4, 5, 6]) - >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) - >>> print(y) - ivy.array([1, 2, 3]) - >>> x = lambda: ivy.array([1, 2, 3]) - >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) - >>> print(y) - ivy.array([1, 2, 3]) +@handle_exceptions +@handle_nestable +def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: + """ + Return the supported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the supported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True, rev=True) - >>> print(y) - ivy.array([1, 2, 3]) - """ - with_callable = catch_exceptions or with_callable - if rev: - tmp = x - x = default_val - default_val = tmp - if with_callable: - x_callable = callable(x) - default_callable = callable(default_val) - else: - x_callable = False - default_callable = False - if catch_exceptions: - # noinspection PyBroadException - try: - x = x() if x_callable else x - except Exception: - return default_val() if default_callable else default_val - else: - x = x() if x_callable else x - return x if exists(x) else default_val() if default_callable else default_val - - -@handle_exceptions -def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: - """ - Return the input shape in ivy.Shape form. - - Parameters - ---------- - shape - The input to be converted + Parameters + ---------- + fn + The function to check for the supported device and dtype attribute + recurse + Whether to recurse into used ivy functions. + Default is ``True``. Returns ------- - ret - the input in ivy.Shape form + ret + Tuple or dict containing the supported devices and dtypes of the function """ - if isinstance(shape, ivy.Shape): - return shape - return ivy.Shape(shape) + ivy.utils.assertions.check_true( + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", + ) + + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=False), + } + else: + supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) + if recurse: + supported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + supported_devices_dtypes, + _dnd_dict_intersection, + function_supported_devices_and_dtypes, + wrapper=lambda x: x, + ) + + return supported_devices_dtypes @handle_exceptions -def to_native_shape( - shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] -) -> ivy.NativeShape: +@handle_nestable +def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: """ - Return the input shape in its native backend framework form. + Return the unsupported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the unsupported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. Parameters ---------- - shape - The input to be converted + fn + The function to check for the unsupported device and dtype attribute + recurse + Whether to recurse into used ivy functions. + Default is ``True``. Returns ------- - ret - the input in its native framework form + ret + Tuple or dict containing the unsupported devices and dtypes of the function """ - native_shape_type = (ivy.NativeShape,) - if ivy.current_backend_str() == "torch": - native_shape_type += (tuple,) - if len(backend_stack) != 0 and isinstance(shape, native_shape_type): - return shape - ivy.utils.assertions.check_isinstance( - shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) - ) - if isinstance(shape, int): - shape = (shape,) - elif isinstance(shape, list): - shape = tuple(shape) - elif is_array(shape): - shape = ivy.to_numpy(shape).tolist() - elif isinstance(shape, ivy.Shape): - shape = shape.shape - ivy.utils.assertions.check_all( - [isinstance(v, int) for v in shape if not is_array(v)], - "shape must take integers only", - as_array=False, - ) ivy.utils.assertions.check_true( - not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", ) - return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=True), + } + else: + unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) + if recurse: + unsupported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + unsupported_devices_dtypes, + _dnd_dict_union, + function_unsupported_devices_and_dtypes, + wrapper=lambda x: x, + ) + return unsupported_devices_dtypes @handle_exceptions +@handle_backend_invalid @handle_nestable -def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def gather( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Try and return the function, otherwise return None if an exception was raised during - function execution. + Gather slices from params at axis according to indices. Parameters ---------- - fn - Function to try and call and return. - args - list of arguments. - kwargs - dictionay of keyword arguments + params + The array from which to gather values. + indices + The array which indicates the indices that will be gathered along + the specified axis. + axis + optional int, the axis from which to gather from. + Default is ``-1``. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional array, for writing the result to. It must have a shape + that the inputs broadcast to. Returns ------- - Either the function itself or None if an exception was raised - during function execution. + ret + New array with the values gathered at the specified indices along the + specified axis. + + + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - with a function that is executed without any exception: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> z = ivy.try_else_none(ivy.add, x, y) - >>> print(z.__name__) - add + With :class:`ivy.Array` input: - with a function that is executed with an exception: + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.array([1, 2]) + >>> print(ivy.gather(x, y)) + ivy.array([1., 2.]) - >>> x = ivy.array([1, 2, 3]) - >>> y = 'hemant' - >>> z = ivy.try_else_none(ivy.add,x, y) + >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) + >>> y = ivy.array([[0, 1],[1, 2]]) + >>> z = ivy.zeros((2, 2, 2)) + >>> ivy.gather(x, y, out=z) >>> print(z) - None - """ - try: - _ = fn(*args, **kwargs) - return fn - except Exception: - return None - + ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) -@handle_exceptions -def arg_names(receiver): - """ - Get the expected keyword arguments for a function or class constructor. + >>> x = ivy.array([[[0., 1.], [2., 3.]], + ... [[8., 9.], [10., 11.]]]) + >>> y = ivy.array([[0, 1]]) + >>> z = ivy.zeros((1, 2, 2, 2)) + >>> ivy.gather(x, y, axis=0, out=z) + >>> print(z) + ivy.array( + [[[[ 0., 1.], + [ 2., 3.]], + [[ 8., 9.], + [10., 11.]]]]) - Parameters - ---------- - receiver - Function or class constructor + >>> x = ivy.array([[0, 10, 20, 0, 0], + ... [0, 0, 0, 30, 40], + ... [0, 10, 0, 0, 40]]) + >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) + >>> z = ivy.gather(x, y, batch_dims=1) + >>> print(z) + ivy.array([[10, 20], [30, 40],[10, 40]]) - Returns - ------- - ret - List containing the keyword arguments' names for a function or class constructor + With :class:`ivy.Container` input: - Examples - -------- - >>> x = ivy.arg_names(ivy.tan) - >>> print(x) - ['x', 'out'] + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.Container(a = ivy.array([0, 1]), + ... b = ivy.array([1, 2])) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([5., 6.]) + } - >>> x = ivy.arg_names(ivy.optimizers.Adam) - >>> print(x) - ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', - 'stop_gradients', 'compile_on_next_step', 'device'] + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.array([0, 1]) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([4., 5.]) + } """ - return list(inspect.signature(receiver).parameters.keys()) + return current_backend(params, indices).gather( + params, indices, axis=axis, batch_dims=batch_dims, out=out + ) @handle_exceptions -def match_kwargs( - kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False -) -> Union[List[Dict], Dict]: +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def gather_nd( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + batch_dims: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Match keyword arguments to either class or function receivers. + Gather slices from params into a array with shape specified by indices. Parameters ---------- - kwargs - Keyword arguments to match. - receivers - Functions and/or classes to match the keyword arguments to. - allow_duplicates - Whether to allow one keyword argument to be used for multiple receivers. - Default is ``False``. + params + The array from which to gather values. + indices + Index array. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Sequence of keyword arguments split as best as possible. + New array of given shape, with the values gathered at the indices. Examples -------- - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) - >>> print(x) - [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] - - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) - >>> print(x) - [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] - """ - split_kwargs = list() - for receiver in receivers: - expected_kwargs = arg_names(receiver) - found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} - if not allow_duplicates: - for k in found_kwargs.keys(): - del kwargs[k] - split_kwargs.append(found_kwargs) - if len(split_kwargs) == 1: - return split_kwargs[0] - return split_kwargs - - -@handle_exceptions -def cache_fn(func: Callable) -> Callable: - """ - Cache function outputs. - - A decorator to wrap a function, such that computed outputs are cached to avoid - recalculating them later. + With :class:`ivy.Array` input: - Parameters - ---------- - func - The function to wrap, whose output should be cached for later. + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) + ivy.array(1.) - Returns - ------- - ret - The newly cache wrapped function. + >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) + >>> y = ivy.array([[0],[1],[1]], dtype='int32') + >>> z = ivy.gather_nd(x,y,batch_dims=1) + ivy.array([0., 3., 5.]) - Examples - -------- - With positional arguments only: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 - >>> cached_sum = ivy.cache_fn(my_sum) - >>> print(cached_sum(3, 5)) - 8 + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) + { + a: ivy.array(1.), + b: ivy.array(5.) + } - With keyword arguments: + With :class:`ivy.Container` input: - >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc - >>> cached_line_eq = ivy.cache_fn(line_eq) - >>> print(cached_line_eq(3, itc=5, slp=2)) - 11 + >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), + ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) + >>> y = ivy.Container(a=ivy.array([1,0]), + ... b=ivy.array([0])) + >>> print(ivy.gather_nd(x, y)) + { + a: ivy.array(30.), + b: ivy.array([0., 100., 200.]) + } """ - global FN_CACHE - if func not in FN_CACHE: - FN_CACHE[func] = dict() - - @wraps(func) - def cached_fn(*args, **kwargs): - key = "".join( - [str(i) + ", " for i in args] - + [" kw, "] - + [str(i) + ", " for i in sorted(kwargs.items())] - ) - cache = FN_CACHE[func] - if key in cache: - return cache[key] - ret = func(*args, **kwargs) - cache[key] = ret - return ret - - return cached_fn + res = current_backend(params, indices).gather_nd( + params, indices, batch_dims=batch_dims + ) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res @handle_exceptions -def current_backend_str() -> Union[str, None]: +def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: """ - Return framework string. + Get all arrays which are currently alive. Returns ------- ret - The framework string. + All arrays which are alive. + + Examples + -------- + >>> ivy.get_all_arrays_in_memory() + [] + >>> x = ivy.get_all_arrays_in_memory() + >>> x + [] + >>> y = ivy.array([0, 1, 2]) + >>> x + [ivy.array([0, 1, 2])] """ - fw = current_backend() - if not backend_stack: - return "" - return fw.current_backend_str() + all_arrays = list() + for obj in gc.get_objects(): + try: + if ivy.current_backend_str() in ["", "numpy"]: + if ivy.is_ivy_array(obj): + all_arrays.append(obj) + else: + if ivy.is_native_array(obj): + all_arrays.append(obj) + + except Exception: + pass + return all_arrays -@handle_exceptions @handle_nestable -@handle_array_like_without_promotion +@handle_partial_mixed_function +@handle_view_indexing @inputs_to_ivy_arrays @handle_array_function -def einops_rearrange( +@handle_device_shifting +def get_item( x: Union[ivy.Array, ivy.NativeArray], - pattern: str, /, + query: Union[ivy.Array, ivy.NativeArray, Tuple], *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], + copy: Optional[bool] = None, ) -> ivy.Array: """ - Perform einops rearrange operation on input array x. + Gather slices from x according to query array, identical to x[query]. Parameters ---------- x - Input array to be re-arranged. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + array, the array from which to gather values. + query + array, index array, integer indices or boolean mask. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - New array with einops.rearrange having been applied. + New array with the values gathered at the specified indices. - Examples - -------- - With :class:`ivy.Array` instance method: + Functional Examples + ------------------- - >>> x = ivy.array([[1, 2, 3], - ... [-4, -5, -6]]) - >>> y = x.einops_rearrange("height width -> width height") - >>> print(y) - ivy.array([[ 1, -4], - [ 2, -5], - [ 3, -6]]) + >>> x = ivy.array([0, -1, 20]) + >>> query = ivy.array([0, 1]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 0, -1]) - >>> x = ivy.array([[[ 1, 2, 3], - ... [ 4, 5, 6]], - ... [[ 7, 8, 9], - ... [10, 11, 12]]]) - >>> y = x.einops_rearrange("c h w -> c (h w)") - >>> print(y) - ivy.array([[ 1, 2, 3, 4, 5, 6], - [ 7, 8, 9, 10, 11, 12]]) - - >>> x = ivy.array([[1, 2, 3, 4, 5, 6], - ... [7, 8, 9, 10, 11, 12]]) - >>> y = ivy.zeros((4,3)) - >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) - >>> print(y) - ivy.array([[ 1, 2, 3], - [ 4, 5, 6], - [ 7, 8, 9], - [10, 11, 12]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> y = ivy.einops_rearrange(x, 'a b -> b a') - >>> print(y) - { - a: ivy.array([[-4.46999979, 3.66000009], - [0.93000001, 24.29000092], - [-3.33999991, 3.6400001]]), - b: ivy.array([[4.96000004, 4.36000013], - [1.51999998, 13.96000004], - [-10.67000008, 0.30000001]]) - } - - With varying pattern: - - Suppose we have a set of 32 images in "h w c" format (height-width-channel) - and concatenate images along height (vertical axis), 960 = 32 * 30 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') - >>> print(x.shape) - (960, 40, 3) - - # Concatenate images along horizontal axis, 1280 = 32 * 40 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') - >>> print(x.shape) - (30, 1280, 3) - - # Reorder axes to "b c h w" format for deep learning - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') - >>> print(x.shape) - (32, 3, 30, 40) - - # Flatten each image into a vector, 3600 = 30 * 40 * 3 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') - >>> print(x.shape) - (32, 3600) - - # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), - # 128 = 32 * 2 * 2 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', - ... h1=2, w1=2) - >>> print(x.shape) - (128, 15, 20, 3) - - # Space-to-depth operation - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, - ... w1=2) - >>> print(x.shape) - (32, 15, 20, 12) + >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) + >>> query = ivy.array([[True, False], [False, False], [True, True]]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 4, -2, -10]) """ - ret = einops.rearrange(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return ivy.array([], shape=(0,), dtype=x.dtype) + return ivy.expand_dims(x, axis=0) + query = ivy.nonzero(query, as_tuple=False) + ret = ivy.gather_nd(x, query) + else: + indices, target_shape = _parse_query(query, x.shape) + if indices is None: + return ivy.empty(target_shape, dtype=x.dtype) + ret = ivy.gather_nd(x, indices) + ret = ivy.reshape(ret, target_shape) return ret -@handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function -def einops_reduce( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - reduction: Union[str, Callable], - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: +@handle_device_shifting +def get_num_dims( + x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False +) -> int: """ - Perform einops reduce operation on input array x. + Return the number of dimensions of the array x. Parameters ---------- x - Input array to be reduced. - pattern - Reduction pattern. - reduction - One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array to infer the number of dimensions for. + as_array + Whether to return the shape as a array, default False. Returns ------- ret - New array with einops.reduce having been applied. + Shape of the array - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]) - >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') - >>> print(reduced) - ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) + >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) + >>> b = ivy.get_num_dims(a, as_array=False) + >>> print(b) + 3 With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') - >>> print(reduced) + >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) + >>> print(ivy.get_num_dims(a)) { - a: ivy.array([-2.29333329, 10.53000069]), - b: ivy.array([-1.39666676, 6.20666695]) + b: 2 + } + + >>> b = ivy.get_num_dims(a, as_array=True) + >>> print(b) + { + b: ivy.array(2) } """ - ret = einops.reduce(x, pattern, reduction, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + return current_backend(x).get_num_dims(x, as_array=as_array) -# IMPORTANT: assign attribute directly to function instead of wrapper here -einops_reduce.unsupported_dtypes = { - "torch": ("float16",), - "tensorflow": ("complex",), - "paddle": ("complex", "uint8", "int8", "int16", "float16"), -} +def get_referrers_recursive( + item, depth=0, max_depth=None, seen_set=None, local_set=None +): + """ + Summary. + + Parameters + ---------- + item + + depth + (Default value = 0) + max_depth + (Default value = None) + seen_set + (Default value = None) + local_set + (Default value = None`) + """ + seen_set = ivy.default(seen_set, set()) + local_set = ivy.default(local_set, set()) + ret_cont = ivy.Container( + repr=str(item).replace(" ", ""), + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + referrers = [ + ref + for ref in gc.get_referrers(item) + if not ( + isinstance(ref, dict) + and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) + ) + ] + local_set.add(str(id(referrers))) + for ref in referrers: + ref_id = str(id(ref)) + if ref_id in local_set or hasattr(ref, "cell_contents"): + continue + seen = ref_id in seen_set + seen_set.add(ref_id) + refs_rec = lambda: get_referrers_recursive( + ref, depth + 1, max_depth, seen_set, local_set + ) + this_repr = "tracked" if seen else str(ref).replace(" ", "") + if not seen and (not max_depth or depth < max_depth): + val = ivy.Container( + repr=this_repr, + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + refs = refs_rec() + for k, v in refs.items(): + val[k] = v + else: + val = this_repr + ret_cont[str(ref_id)] = val + return ret_cont @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def einops_repeat( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: +def has_nans( + x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True +) -> bool: """ - Perform einops repeat operation on input array x. + Determine whether the array contains any nans, as well as infs or -infs if + specified. Parameters ---------- x - Input array to be repeated. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array. + include_infs + Whether to include ``+infinity`` and ``-infinity`` in the check. + Default is ``True``. Returns ------- ret - New array with einops.repeat having been applied. + Boolean as to whether the array contains nans. - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3, 4]) - >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) - >>> print(repeated) - ivy.array([[1, 2, 3, 4], - [1, 2, 3, 4]]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + False + + >>> x = ivy.array([float('nan'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x, include_infs=False) + >>> print(y) + False With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[4,5], - ... [1, 3]]), - ... b=ivy.array([[9, 10], - ... [4, 2]])) - >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) - >>> print(repeated) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.has_nans(x) + >>> print(y) { - a: ivy.array([[4, 5, 4, 5], - [1, 3, 1, 3]]), - b: ivy.array([[9, 10, 9, 10], - [4, 2, 4, 2]]) + a: False, + b: False } """ - ret = einops.repeat(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) -ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 +@handle_exceptions +def inplace_arrays_supported() -> bool: + """ + Determine whether inplace arrays are supported for the current backend framework. + + Returns + ------- + ret + Boolean, whether or not inplace arrays are supported. + """ + return current_backend().inplace_arrays_supported() @handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays @handle_array_function -def set_min_denominator(val: float) -> None: +@handle_device_shifting +def inplace_decrement( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Set the global minimum denominator used by ivy for numerically stable division. + Perform in-place decrement for the input array. Parameters ---------- + x + The input array to be decremented by the defined value. val - The value to set the global minimum denominator to. - - Examples - -------- - >>> x = ivy.min_denominator - >>> print(x) - 1e-12 + The value of decrement. - >>> ivy.set_min_denominator(1e-13) - >>> y = ivy.min_denominator - >>> print(y) - 1e-13 - """ - global min_denominator_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_denominator_stack.append(val) - ivy.__setattr__("min_denominator", val, True) + Returns + ------- + ret + The array following the in-place decrement. -@handle_exceptions -def unset_min_denominator() -> None: - """ - Reset the global minimum denominator used by ivy for numerically stable division to - the previous value. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> ivy.set_min_denominator(1e-10) - >>> y = ivy.min_denominator + With :class:`ivy.Array` input: + + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_decrement(x, 1.25) >>> print(y) - 1e-10 + ivy.array([[ 4.05, 5.75, -1.25], + [ 5.55, 6.75, 2.65], + [-1.25, 8.75, 5.05]]) - >>> ivy.unset_min_denominator() - >>> ivy.min_denominator - 1e-12 - """ - global min_denominator_stack - if min_denominator_stack: - min_denominator_stack.pop(-1) - val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 - ivy.__setattr__("min_denominator", val, True) + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) + >>> y = ivy.inplace_decrement(x, 1.5) + >>> print(y) + { + a: ivy.array([-1., -6.5, 28.5]), + b: ivy.array([-1.5, -26.5, 48.5]) + } -ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([0., 0., 0.]), + b: ivy.array([0., 0., 0.]) + } + + >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) + >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([1., 1.5, 3.]), + b: ivy.array([0., 50., 3.5]) + } + """ + return current_backend(x).inplace_decrement(x, val) @handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays @handle_array_function -def set_min_base(val: float) -> None: +@handle_device_shifting +def inplace_increment( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Set the global minimum base used by ivy for numerically stable power raising. + Perform in-place increment for the input array. Parameters ---------- + x + The input array to be incremented by the defined value. val - The new value to set the minimum base to. + The value of increment. + + Returns + ------- + ret + The array following the in-place increment. Examples -------- - >>> x = ivy.min_base - >>> print(x) - 1e-05 + With :class:`ivy.Array` input: - >>> ivy.set_min_base(1e-04) - >>> y = ivy.min_base + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_increment(x, 3.) >>> print(y) - 1e-04 - """ - global min_base_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_base_stack.append(val) - ivy.__setattr__("min_base", val, True) - + ivy.array([[ 8.3, 10., 3.], + [ 9.8, 11., 6.9], + [ 3., 13., 9.3]]) -@handle_exceptions -def unset_min_base() -> None: - """ - Reset the global minimum base used by ivy for numerically stable power raising to - the previous value. + With :class:`ivy.Container` input: - Examples - -------- - >>> ivy.set_min_base(1e-07) - >>> y = ivy.min_base + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.inplace_increment(x, 2.5) >>> print(y) - 1e-07 + { + a: ivy.array([2.5, 17.5, 32.5]), + b: ivy.array([2.5, 27.5, 52.5]) + } - >>> ivy.unset_min_base() - >>> ivy.min_base - 1e-05 + + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_increment(x, y) + >>> print(z) + { + a: ivy.array([0., 30., 60.]), + b: ivy.array([0., 50., 100.]) + } """ - global min_base_stack - if min_base_stack: - min_base_stack.pop(-1) - val = min_base_stack[-1] if min_base_stack else 1e-05 - ivy.__setattr__("min_base", val, True) + return current_backend(x).inplace_increment(x, val) @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def stable_divide( - numerator: Union[Number, ivy.Array, ivy.NativeArray], - denominator: Union[Number, ivy.Array, ivy.NativeArray], +# @handle_device_shifting +def inplace_update( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], /, *, - min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, -) -> Union[Number, ivy.Array]: + ensure_in_backend: bool = False, + keep_input_dtype: bool = False, +) -> ivy.Array: """ - Divide the numerator by the denominator, with min denominator added to the - denominator for numerical stability. + Perform in-place update for the input array. + + This will always be performed on ivy.Array instances pass in the input, and will + also be performed on the native array classes in the backend when the backend + supports this. If the backend does not natively support inplace updates, and x is an + ivy.NativeArray instance, then an + exception will be thrown. Parameters ---------- - numerator - The numerator of the division. - denominator - The denominator of the division. - min_denominator - The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) - by default. + x + The variable to update. + val + The array to update the variable with. + ensure_in_backend + Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. + In cases where it should be, backends which do not natively support inplace + updates will raise an exception. + keep_input_dtype + Whether or not to preserve `x` data type after the update, otherwise `val` + data type will be applied. Defaults to False. Returns ------- ret - The new item following the numerically stable division. + The array following the in-place update. + + Raises + ------ + IvyException + If backend set doesn't natively support inplace updates and ensure_in_backend is + True, above exception will be raised. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the arguments. Examples -------- - With :code:`int` input: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x = ivy.stable_divide(1, 2) + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0]) + >>> ivy.inplace_update(x, y) >>> print(x) - 0.49999999999975 + ivy.array([0]) - >>> x = ivy.stable_divide(1, 4, min_denominator=1) + With :class:`ivy.Array` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) + >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) + >>> ivy.inplace_update(x, y, keep_input_dtype=True) >>> print(x) - 0.2 + ivy.array([0., 0., 0.]) - With float input: + With :class:`ivy.Container` instances:, and backend set as `torch`: - >>> x = ivy.stable_divide(5.0, 3.33) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) + >>> ivy.inplace_update(x, y) >>> print(x) - 1.5015015015010504 + { + a: ivy.array([1, 1]), + b: ivy.array([2, 2]) + } - With :code:`complex` input: + With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend + set as `torch`: - >>> x = ivy.stable_divide(1+1j, 1-1j) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.array([1, 2]) + >>> ivy.inplace_update(x, y) >>> print(x) - (5.000444502911705e-13+0.9999999999995j) + { + a: ivy.array([1, 2]), + b: ivy.array([1, 2]) + } + """ + return current_backend(x).inplace_update( + x, + val, + ensure_in_backend=ensure_in_backend, + keep_input_dtype=keep_input_dtype, + ) - With :class:`ivy.Array` input: - >>> x = ivy.asarray([[10., 20., 30.], - ... [40., 50., 60.]]) - >>> y = ivy.stable_divide(x, 10.) - >>> print(y) - ivy.array([[1., 2., 3.], - [4., 5., 6.]]) +@handle_exceptions +def inplace_variables_supported() -> bool: + """ + Determine whether inplace variables are supported for the current backend framework. + Returns + ------- + ret + Boolean, whether or not inplace variables are supported. + """ + return current_backend().inplace_variables_supported() - >>> x = ivy.asarray([1,2,3]) - >>> y = np.array((1., 3., 5.)) - >>> z = ivy.stable_divide(x, y) - >>> print(z) - ivy.array([1. , 0.667, 0.6 ]) - >>> x = ivy.asarray([1., 2., 4.]) - >>> y = ivy.asarray([1., 0.5, 0.25]) - >>> z = ivy.asarray([0.01, 0.02, 0.03]) - >>> w = ivy.stable_divide(x, y, min_denominator=z) - >>> print(w) - ivy.array([ 0.99, 3.85, 14.3 ]) +@handle_exceptions +@handle_backend_invalid +def is_array(x: Any, /, *, exclusive: bool = False) -> bool: + """ + Determine whether the input x is either an Ivy Array or a Native Array. - With :class:`ivy.Container` input: + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. - >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) - >>> y = ivy.stable_divide(x, 0.5) - >>> print(y) - { - a: ivy.array([20., 30.]), - b: ivy.array([40., 50.]) - } + Returns + ------- + ret + Boolean, whether or not x is an array. + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> print(ivy.is_array(x)) + True - >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) - >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) - >>> z = ivy.stable_divide(x, y) - >>> print(z) - { - a: ivy.array([2., 0.8]), - b: ivy.array([0.857, 10.]) - } + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> print(ivy.is_array(x, exclusive=True)) + True + + >>> x = [2, 3] + >>> print(ivy.is_array(x)) + False """ - return numerator / (denominator + default(min_denominator, ivy.min_denominator)) + return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( + x, exclusive=exclusive + ) @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def stable_pow( - base: Union[Number, ivy.Array, ivy.NativeArray], - exponent: Union[Number, ivy.Array, ivy.NativeArray], - /, - *, - min_base: float = None, -) -> Any: +@handle_backend_invalid +def is_ivy_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False +) -> bool: """ - Raise the base by the power, with ivy.min_base added to the base when exponent > 1 - for numerical stability. + Determine whether the input x is a valid Ivy Array. Parameters ---------- - base - The base number. - exponent - The exponent number. - min_base - The minimum base to use, use global ivy.min_base by default. + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. Returns ------- ret - The new item following the numerically stable power. + Boolean, whether or not x is a valid Ivy Array. Examples -------- - With :code:`int` input: + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_ivy_array(x) + True - >>> x = ivy.stable_pow(2, 2) - >>> print(x) - ivy.array(4.00004) + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_ivy_array(x, exclusive=True) + False + """ + return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) - >>> x = ivy.stable_pow(2, 2, min_base=2) - >>> print(x) - ivy.array(16) - With float input: +@handle_exceptions +def is_ivy_container(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Container. - >>> x = ivy.stable_pow(4.0, .5) - >>> print(x) - ivy.array(2.00000262) + Parameters + ---------- + x + The input to check - With :code:`complex` input: + Returns + ------- + ret + Boolean, whether or not x is an ivy container. - >>> x = ivy.stable_pow(3+4j, 2j) - >>> print(x) - ivy.array(-0.15605032-0.01208451j) + Examples + -------- + >>> x = ivy.Container() + >>> print(ivy.is_ivy_container(x)) + True - With :class:`ivy.Array` input: + >>> x = [2, 3] + >>> print(ivy.is_ivy_container(x)) + False + """ + return isinstance(x, ivy.Container) - >>> x = ivy.asarray([[2, 4], - ... [6, 8]]) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) - ivy.array([[ 4.00004, 16.00008], - [36.00012, 64.00016]]) - >>> x = ivy.asarray([2, 4, 6]) - >>> y = ivy.asarray([2, 3, 4]) - >>> z = ivy.stable_pow(x, y) - >>> print(z) - ivy.array([ 4.00004, 64.00048, 1296.00864]) +def is_ivy_nested_array(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Nested Array. - With :class:`ivy.Container` input: + Parameters + ---------- + x + The input to check + Returns + ------- + ret + Boolean, whether or not x is an ivy nested array. + """ + return isinstance(x, ivy.NestedArray) - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) - { - a: ivy.array([4.00004, 16.00008]), - b: ivy.array([36.00012, 64.00016]) - } - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) - >>> z = ivy.stable_pow(x, y) - >>> print(z) - { - a: ivy.array([2.00001, 64.00048]), - b: ivy.array([1296.00864, 32768.2048]) - } +@handle_exceptions +@handle_backend_invalid +def is_native_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False +) -> bool: """ - return_dtype = ivy.promote_types( - ivy.default_dtype(item=base), - ivy.default_dtype(item=default(min_base, ivy.min_base)), + Determine whether the input x is an :class:`ivy.NativeArray` instance. + + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. + + Returns + ------- + ret + Boolean, whether or not x is an :class:`ivy.NativeArray`. + + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_native_array(x) + False + + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_native_array(x, exclusive=True) + True + """ + try: + return current_backend(x).is_native_array(x, exclusive=exclusive) + except ValueError: + return False + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +@handle_device_shifting +def isin( + elements: Union[ivy.Array, ivy.NativeArray], + test_elements: Union[ivy.Array, ivy.NativeArray], + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> ivy.Array: + """ + Test if each element of elements is in test_elements. + + Parameters + ---------- + elements + input array + test_elements + values against which to test for each input element + assume_unique + If True, assumes both elements and test_elements contain unique elements, + which can speed up the calculation. Default value is False. + invert + If True, inverts the boolean return array, resulting in True values for + elements not in test_elements. Default value is False. + + Returns + ------- + ret + output a boolean array of the same shape as elements that is True for elements + in test_elements and False otherwise. + + Examples + -------- + >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y) + ivy.array([[False, False, False], [ True, True, True]]) + + >>> x = ivy.array([3, 2, 1, 0]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y, invert=True) + ivy.array([False, False, False, True]) + """ + return ivy.current_backend(elements, test_elements).isin( + elements, test_elements, assume_unique=assume_unique, invert=invert ) - return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) - ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) - return ret.astype(return_dtype) -stable_pow.unsupported_dtypes = ("bfloat16",) +@handle_exceptions +@handle_nestable +def isscalar(x: Any, /) -> bool: + return np.isscalar(x) @handle_exceptions -def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def itemsize( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> int: """ - Get all arrays which are currently alive. + Return the size of the input array's elements. + + Parameters + ---------- + x + The input array. Returns ------- ret - All arrays which are alive. + An integer specifying the element size in bytes. Examples -------- - >>> ivy.get_all_arrays_in_memory() - [] - >>> x = ivy.get_all_arrays_in_memory() - >>> x - [] - >>> y = ivy.array([0, 1, 2]) - >>> x - [ivy.array([0, 1, 2])] + >>> x = ivy.array([1,2,3], dtype=ivy.float64) + >>> ivy.itemsize(x) + 8 + + >>> x = ivy.array([1,2,3], dtype=ivy.complex128) + >>> ivy.itemsize(x) + 16 """ - all_arrays = list() - for obj in gc.get_objects(): - try: - if ivy.current_backend_str() in ["", "numpy"]: - if ivy.is_ivy_array(obj): - all_arrays.append(obj) - else: - if ivy.is_native_array(obj): - all_arrays.append(obj) + return ivy.current_backend(x).itemsize(x) - except Exception: - pass - return all_arrays + +@handle_exceptions +def match_kwargs( + kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False +) -> Union[List[Dict], Dict]: + """ + Match keyword arguments to either class or function receivers. + + Parameters + ---------- + kwargs + Keyword arguments to match. + receivers + Functions and/or classes to match the keyword arguments to. + allow_duplicates + Whether to allow one keyword argument to be used for multiple receivers. + Default is ``False``. + + Returns + ------- + ret + Sequence of keyword arguments split as best as possible. + + Examples + -------- + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) + >>> print(x) + [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] + + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) + >>> print(x) + [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] + """ + split_kwargs = list() + for receiver in receivers: + expected_kwargs = arg_names(receiver) + found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} + if not allow_duplicates: + for k in found_kwargs.keys(): + del kwargs[k] + split_kwargs.append(found_kwargs) + if len(split_kwargs) == 1: + return split_kwargs[0] + return split_kwargs + + +@handle_exceptions +@handle_nestable +@handle_array_function +def multiprocessing(context: Optional[str] = None): + """ + Return backend-specific multiprocessing module. + + Parameters + ---------- + context + The context of the multiprocessing, either fork, forkserver or spawn. + Default is ``None``. + + Returns + ------- + ret + Multiprocessing module + """ + return current_backend().multiprocessing(context) @handle_exceptions @@ -2447,365 +2577,286 @@ def print_all_arrays_in_memory(): print(type(arr), arr.shape) -ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 - - @handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back @handle_array_function -def set_queue_timeout(timeout: float): +@handle_device_shifting +def scatter_flat( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], + /, + *, + size: Optional[int] = None, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Set a timeout value (in seconds) for the global queue. - - Set the global queue timeout value (in seconds) Default value without this function - being called is 15 seconds. + Scatter flat updates into a new flat array according to flat indices. Parameters ---------- - timeout - The timeout when waiting for containers to arrive from the queues. - To be set in seconds. - - Examples - -------- - >>> x = ivy.set_queue_timeout(10) - >>> x = ivy.queue_timeout - >>> print(x) - 10.0 - - >>> ivy.set_queue_timeout(30) - >>> y = ivy.queue_timeout - >>> print(y) - 30 - """ - global queue_timeout_stack - ivy.utils.assertions.check_isinstance(timeout, (int, float)) - queue_timeout_stack.append(timeout) - ivy.__setattr__("queue_timeout", timeout, True) + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + size + The size of the result. Default is `None`, in which case tensor + argument out must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + New array of given shape, with the values scattered at the indices. -@handle_exceptions -def unset_queue_timeout() -> None: - """ - Reset the global queue timeout value (in seconds) to the previous state. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - >>> ivy.set_queue_timeout(10.0) - >>> y = ivy.queue_timeout - >>> print(y) - 10.0 + With :class:`ivy.Array` input: + >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) + >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) + >>> ivy.scatter_flat(indices, updates, out=out) + >>> print(out) + ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) - >>> ivy.unset_queue_timeout() - >>> ivy.queue_timeout - 15.0 - """ - global queue_timeout_stack - if queue_timeout_stack: - queue_timeout_stack.pop(-1) - timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 - ivy.__setattr__("queue_timeout", timeout, True) + + With :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) -ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" + With :class:`ivy.Container` and :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + { + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + } + + + With :class:`ivy.Container` input: + >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), + ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + { + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + } + """ + return current_backend(indices).scatter_flat( + indices, updates, size=size, reduction=reduction, out=out + ) @handle_exceptions -def set_tmp_dir(tmp_dr: str) -> None: +@handle_backend_invalid +@handle_nestable +@inputs_to_native_shapes +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def scatter_nd( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], + /, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + *, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Set the directory for saving temporary files. + Scatter updates into a new array according to indices. Parameters ---------- - tmp_dr - The new directory for saving temporary files + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + shape + The shape of the result. Default is ``None``, in which case tensor + argument must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + New array of given shape, with the values scattered at the indices. Examples -------- - >>> x = ivy.tmp_dir - >>> print(x) - /tmp + scatter values into an empty array, With :class:`ivy.Array` input: - >>> ivy.set_tmp_dir("/my_tmp") - >>> y = ivy.tmp_dir - >>> print(y) - /my_tmp - """ - global tmp_dir_stack - ivy.utils.assertions.check_isinstance(tmp_dr, str) - tmp_dir_stack.append(tmp_dr) - ivy.__setattr__("tmp_dir", tmp_dr, True) + >>> indices = ivy.array([[4], [3], [1], [7]]) + >>> updates = ivy.array([9, 10, 11, 12]) + >>> shape = ivy.array([8]) + >>> scatter = ivy.scatter_nd(indices, updates, shape) + >>> print(scatter) + ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) + With scatter into an empty array, With :class:`ivy.Container` input: -@handle_exceptions -def unset_tmp_dir() -> None: - """ - Reset the directory for saving temporary files to the previous value. + >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), + ... b=ivy.array([[5],[1],[2]])) + >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), + ... b=ivy.array([20, 30, 40])) + >>> shape = ivy.Container(a=ivy.array([10]), + ... b = ivy.array([10])) + >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') + >>> print(z) + { + a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), + b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) + } - Examples - -------- - >>> ivy.set_tmp_dir("/my_dir") - >>> y = ivy.tmp_dir - >>> print(y) - /my_dir + With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> ivy.unset_tmp_dir() - >>> ivy.tmp_dir - /tmp + >>> indices = ivy.array([[4],[3],[1]]) + >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), + ... b=ivy.array([200, 300, 400])) + >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), + ... b=ivy.array([10, 20, 30, 40, 50])) + >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) + >>> print(z) + { + a: ivy.array([1, 30, 3, 20, 10]), + b: ivy.array([10, 400, 30, 300, 200]) + } """ - global tmp_dir_stack - if tmp_dir_stack: - tmp_dir_stack.pop(-1) - tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" - ivy.__setattr__("tmp_dir", tmp_dr, True) + return current_backend(indices).scatter_nd( + indices, updates, shape=shape, reduction=reduction, out=out + ) @handle_exceptions -def container_types(): +def set_array_mode(mode: bool) -> None: """ - Summary. + Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs + back to ivy.Array. - Returns - ------- - ret - a key-value structure, and exposes public methods .keys(), .values() and - items(). - """ - # noinspection PyBroadException - try: - return current_backend().container_types() - except ValueError: - return [] + It Stops the conversion of ivy.NativeArray to ivy.Array in the + case when it is set to False. + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions -@handle_exceptions -def inplace_arrays_supported() -> bool: - """ - Determine whether inplace arrays are supported for the current backend framework. + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - Returns - ------- - ret - Boolean, whether or not inplace arrays are supported. + >>> ivy.set_array_mode(True) + >>> ivy.array_mode + True """ - return current_backend().inplace_arrays_supported() + global array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + array_mode_stack.append(mode) + ivy.__setattr__("array_mode", mode, True) @handle_exceptions -def inplace_variables_supported() -> bool: +def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: """ - Determine whether inplace variables are supported for the current backend framework. + Set the mode of whether to show frontend-truncated exception stack traces, ivy- + truncated exception stack traces or full exception stack traces. - Returns - ------- - ret - Boolean, whether or not inplace variables are supported. - """ - return current_backend().inplace_variables_supported() + Parameter + --------- + mode + str exeption trace mode, one of `ivy`, `full` or `frontend` + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -@handle_array_function -def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: - """ - Return if in-place operations are supported for x's data type. - - Determine whether in-place operations are supported for x's data type, by the - current backend framework setting. - - Parameters - ---------- - x - Input variable for whose data type we check whether the current backend - framework supports in-place operations. - - Returns - ------- - ret - Value depends on whether in-place operations are supported for - data type of x. - - Raises - ------ - IvyException - If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will - be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. - - Examples - -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) - True - - With :class:`ivy.Container` input and backend set as `torch`: - - >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) - { - a: True, - b: True - } - - With `ivy.Array` input and backend set as "tensorflow": - - >>> x = ivy.array([1., 4.2, 2.2]) - >>> ret = x.supports_inplace_updates() - >>> print(ret) - False + >>> ivy.set_exception_trace_mode("full") + >>> ivy.exception_trace_mode + 'full' """ - if _is_variable(x): - return ivy.inplace_variables_supported() - elif ivy.is_native_array(x): - return ivy.inplace_arrays_supported() - raise ivy.utils.exceptions.IvyException( - "Input x must be either a variable or an array." + global exception_trace_mode_stack + trace_modes = list(trace_mode_dict.keys()) + ivy.utils.assertions.check_elem_in_list( + mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) ) + exception_trace_mode_stack.append(mode) + ivy.__setattr__("exception_trace_mode", mode, True) @handle_exceptions -@handle_nestable -@inputs_to_native_arrays -@handle_array_function -def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: +def set_inplace_mode(mode: str = "lenient") -> None: """ - Assert that inplace operations are supported for x. + Set the memory management behavior for in-place updates in Ivy. + + By default, Ivy creates new arrays in the backend for in-place updates. + However, this behavior can be controlled by the user + using the 'inplace_mode' parameter. Parameters ---------- - x - Input variable or array to check for inplace support for. + mode : str + The mode for memory management during in-place updates. + - 'lenient': (Default) In this mode, new arrays will be created during + in-place updates to avoid breaking existing code. + This is the default behavior. + - 'strict': In this mode, an error will be raised if the + 'inplace_update' function is called + in a backend that doesn't support inplace updates natively. Returns ------- - ret - True if supports, raises IvyBackendException otherwise - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + None Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) - True - - With :class:`ivy.Array` input and default backend set as `torch`: - - >>> ivy.set_backend("torch") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) - True - - With :class:`ivy.Container` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) - { - a: True, - b: True - } + >>> set_inplace_mode('lenient') + >>> ivy.inplace_mode + 'lenient' - With :class:`ivy.Container` input and default backend set as `torch`: + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) - { - a: True, - b: True - } + Note + ---- + Enabling strict mode can help users have more control over memory management + but may lead to errors if the backend doesn't support inplace updates natively. """ - ivy.utils.assertions.check_true( - ivy.supports_inplace_updates(x), - "Inplace operations are not supported {} types with {} backend".format( - type(x), ivy.current_backend_str() - ), + global inplace_mode_stack + inplace_modes = ["lenient", "strict"] + ivy.utils.assertions.check_elem_in_list( + mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" ) - return True - - -@handle_nestable -@handle_partial_mixed_function -@handle_view_indexing -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def get_item( - x: Union[ivy.Array, ivy.NativeArray], - /, - query: Union[ivy.Array, ivy.NativeArray, Tuple], - *, - copy: Optional[bool] = None, -) -> ivy.Array: - """ - Gather slices from x according to query array, identical to x[query]. - - Parameters - ---------- - x - array, the array from which to gather values. - query - array, index array, integer indices or boolean mask. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - - Returns - ------- - ret - New array with the values gathered at the specified indices. - - Functional Examples - ------------------- - - >>> x = ivy.array([0, -1, 20]) - >>> query = ivy.array([0, 1]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 0, -1]) - - >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) - >>> query = ivy.array([[True, False], [False, False], [True, True]]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 4, -2, -10]) - """ - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return ivy.array([], shape=(0,), dtype=x.dtype) - return ivy.expand_dims(x, axis=0) - query = ivy.nonzero(query, as_tuple=False) - ret = ivy.gather_nd(x, query) - else: - indices, target_shape = _parse_query(query, x.shape) - if indices is None: - return ivy.empty(target_shape, dtype=x.dtype) - ret = ivy.gather_nd(x, indices) - ret = ivy.reshape(ret, target_shape) - return ret - - -get_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + inplace_mode_stack.append(mode) + ivy.__setattr__("inplace_mode", mode, True) @handle_nestable @@ -2879,758 +2930,711 @@ def set_item( return ret -set_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - -def _parse_query(query, x_shape): - query = (query,) if not isinstance(query, tuple) else query - query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) - - # array containing all of x's flat indices - x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) +@handle_exceptions +@handle_array_function +def set_min_base(val: float) -> None: + """ + Set the global minimum base used by ivy for numerically stable power raising. - # use numpy's __getitem__ to get the queried indices - x_idxs = ivy.array(x_.to_numpy()[query_]) - target_shape = x_idxs.shape + Parameters + ---------- + val + The new value to set the minimum base to. - if 0 in x_idxs.shape or 0 in x_shape: - return None, target_shape + Examples + -------- + >>> x = ivy.min_base + >>> print(x) + 1e-05 - # convert the flat indices to multi-D indices - x_idxs = ivy.unravel_index(x_idxs, x_shape) - - # stack the multi-D indices to bring them to gather_nd/scatter_nd format - x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) - - return x_idxs, target_shape - - -def _numel(shape): - shape = tuple(shape) - return ivy.prod(shape).to_scalar() if shape != () else 1 - - -def _broadcast_to(input, target_shape): - input = ivy.squeeze(input) - if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): - return ivy.reshape(input, target_shape) - else: - input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input - new_dims = () - i_i = len(input.shape) - 1 - for i_t in range(len(target_shape) - 1, -1, -1): - if len(input.shape) + len(new_dims) >= len(target_shape): - break - if i_i < 0 or target_shape[i_t] != input.shape[i_i]: - new_dims += (i_t,) - else: - i_i -= 1 - input = ivy.expand_dims(input, axis=new_dims) - return ivy.broadcast_to(input, target_shape) + >>> ivy.set_min_base(1e-04) + >>> y = ivy.min_base + >>> print(y) + 1e-04 + """ + global min_base_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_base_stack.append(val) + ivy.__setattr__("min_base", val, True) @handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_ivy_arrays @handle_array_function -# @handle_device_shifting -def inplace_update( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], - /, - *, - ensure_in_backend: bool = False, - keep_input_dtype: bool = False, -) -> ivy.Array: +def set_min_denominator(val: float) -> None: """ - Perform in-place update for the input array. - - This will always be performed on ivy.Array instances pass in the input, and will - also be performed on the native array classes in the backend when the backend - supports this. If the backend does not natively support inplace updates, and x is an - ivy.NativeArray instance, then an - exception will be thrown. + Set the global minimum denominator used by ivy for numerically stable division. Parameters ---------- - x - The variable to update. val - The array to update the variable with. - ensure_in_backend - Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. - In cases where it should be, backends which do not natively support inplace - updates will raise an exception. - keep_input_dtype - Whether or not to preserve `x` data type after the update, otherwise `val` - data type will be applied. Defaults to False. - - Returns - ------- - ret - The array following the in-place update. - - Raises - ------ - IvyException - If backend set doesn't natively support inplace updates and ensure_in_backend is - True, above exception will be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the arguments. + The value to set the global minimum denominator to. Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0]) - >>> ivy.inplace_update(x, y) + >>> x = ivy.min_denominator >>> print(x) - ivy.array([0]) + 1e-12 - With :class:`ivy.Array` input and default backend set as `numpy`: + >>> ivy.set_min_denominator(1e-13) + >>> y = ivy.min_denominator + >>> print(y) + 1e-13 + """ + global min_denominator_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_denominator_stack.append(val) + ivy.__setattr__("min_denominator", val, True) - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) - >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) - >>> ivy.inplace_update(x, y, keep_input_dtype=True) - >>> print(x) - ivy.array([0., 0., 0.]) - With :class:`ivy.Container` instances:, and backend set as `torch`: +@handle_exceptions +def set_nestable_mode(mode: bool) -> None: + """ + Set the mode of whether to check if function inputs are ivy.Container. - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) - >>> ivy.inplace_update(x, y) - >>> print(x) - { - a: ivy.array([1, 1]), - b: ivy.array([2, 2]) - } + Parameter + --------- + mode + boolean whether to check if function inputs are ivy.Container - With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend - set as `torch`: + Examples + -------- + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.array([1, 2]) - >>> ivy.inplace_update(x, y) - >>> print(x) - { - a: ivy.array([1, 2]), - b: ivy.array([1, 2]) - } + >>> ivy.set_nestable_mode(True) + >>> ivy.nestable_mode + True """ - return current_backend(x).inplace_update( - x, - val, - ensure_in_backend=ensure_in_backend, - keep_input_dtype=keep_input_dtype, - ) + global nestable_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + nestable_mode_stack.append(mode) + ivy.__setattr__("nestable_mode", mode, True) -inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} +@handle_exceptions +def set_precise_mode(mode: bool) -> None: + """ + Set the mode of whether to use a promotion table that avoids any precision loss or a + compute effecient table that avoids most wider-than-necessary promotions. -ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" + Parameter + --------- + mode + boolean whether to use high precision promtion table + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.set_precise_mode(True) + >>> ivy.precise_mode + True + """ + global precise_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + precise_mode_stack.append(mode) + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) @handle_exceptions -def set_inplace_mode(mode: str = "lenient") -> None: +@handle_array_function +def set_queue_timeout(timeout: float): """ - Set the memory management behavior for in-place updates in Ivy. + Set a timeout value (in seconds) for the global queue. - By default, Ivy creates new arrays in the backend for in-place updates. - However, this behavior can be controlled by the user - using the 'inplace_mode' parameter. + Set the global queue timeout value (in seconds) Default value without this function + being called is 15 seconds. Parameters ---------- - mode : str - The mode for memory management during in-place updates. - - 'lenient': (Default) In this mode, new arrays will be created during - in-place updates to avoid breaking existing code. - This is the default behavior. - - 'strict': In this mode, an error will be raised if the - 'inplace_update' function is called - in a backend that doesn't support inplace updates natively. - - Returns - ------- - None + timeout + The timeout when waiting for containers to arrive from the queues. + To be set in seconds. Examples -------- - >>> set_inplace_mode('lenient') - >>> ivy.inplace_mode - 'lenient' - - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' + >>> x = ivy.set_queue_timeout(10) + >>> x = ivy.queue_timeout + >>> print(x) + 10.0 - Note - ---- - Enabling strict mode can help users have more control over memory management - but may lead to errors if the backend doesn't support inplace updates natively. + >>> ivy.set_queue_timeout(30) + >>> y = ivy.queue_timeout + >>> print(y) + 30 """ - global inplace_mode_stack - inplace_modes = ["lenient", "strict"] - ivy.utils.assertions.check_elem_in_list( - mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" - ) - inplace_mode_stack.append(mode) - ivy.__setattr__("inplace_mode", mode, True) + global queue_timeout_stack + ivy.utils.assertions.check_isinstance(timeout, (int, float)) + queue_timeout_stack.append(timeout) + ivy.__setattr__("queue_timeout", timeout, True) @handle_exceptions -def unset_inplace_mode() -> None: +def set_shape_array_mode(mode: bool) -> None: """ - Reset the memory management behavior for in-place updates in Ivy to the previous - state. + Set the mode of returning shape as ivy.Array to the given mode instance. + + Parameter + --------- + mode + boolean whether to return shape as ivy.Array Examples -------- - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' + >>> ivy.set_shape_array_mode(False) + >>> ivy.shape_array_mode + False - >>> unset_inplace_mode() - >>> ivy.inplace_mode - 'lenient' + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True """ - global inplace_mode_stack - if inplace_mode_stack: - inplace_mode_stack.pop(-1) - mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" - ivy.__setattr__("inplace_mode", mode, True) + global shape_array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + shape_array_mode_stack.append(mode) + ivy.__setattr__("shape_array_mode", mode, True) + + +@handle_exceptions +def set_show_func_wrapper_trace_mode(mode: bool) -> None: + """ + Set the mode of whether to show the full stack trace with function wrapping traces. + + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions + + Examples + -------- + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False + + >>> ivy.set_show_func_wrapper_trace_mode(True) + >>> ivy.show_func_wrapper_trace_mode + True + """ + global show_func_wrapper_trace_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + show_func_wrapper_trace_mode_stack.append(mode) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + + +@handle_exceptions +def set_tmp_dir(tmp_dr: str) -> None: + """ + Set the directory for saving temporary files. + + Parameters + ---------- + tmp_dr + The new directory for saving temporary files + + Examples + -------- + >>> x = ivy.tmp_dir + >>> print(x) + /tmp + + >>> ivy.set_tmp_dir("/my_tmp") + >>> y = ivy.tmp_dir + >>> print(y) + /my_tmp + """ + global tmp_dir_stack + ivy.utils.assertions.check_isinstance(tmp_dr, str) + tmp_dir_stack.append(tmp_dr) + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions @handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@handle_array_like_without_promotion +@inputs_to_native_arrays +@outputs_to_ivy_shapes +@outputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def inplace_decrement( +def shape( x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: + /, + *, + as_array: bool = False, +) -> Union[ivy.Shape, ivy.NativeShape]: """ - Perform in-place decrement for the input array. + Return the shape of the array ``x``. Parameters ---------- x - The input array to be decremented by the defined value. - val - The value of decrement. + Input array to infer the shape of. + as_array + Whether to return the shape as an array. + Default is False. Returns ------- ret - The array following the in-place decrement. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Shape of the array ``x``. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_decrement(x, 1.25) - >>> print(y) - ivy.array([[ 4.05, 5.75, -1.25], - [ 5.55, 6.75, 2.65], - [-1.25, 8.75, 5.05]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) - >>> y = ivy.inplace_decrement(x, 1.5) + >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) + >>> y = ivy.shape(x) + >>> z = ivy.shape(x, as_array = True) >>> print(y) - { - a: ivy.array([-1., -6.5, 28.5]), - b: ivy.array([-1.5, -26.5, 48.5]) - } - - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_decrement(x, y) - >>> print(z) - { - a: ivy.array([0., 0., 0.]), - b: ivy.array([0., 0., 0.]) - } + (2, 3) - >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) - >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) - >>> z = ivy.inplace_decrement(x, y) >>> print(z) - { - a: ivy.array([1., 1.5, 3.]), - b: ivy.array([0., 50., 3.5]) - } + ivy.array([2, 3]) """ - return current_backend(x).inplace_decrement(x, val) + return current_backend(x).shape(x, as_array=as_array) @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def inplace_increment( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: +def stable_divide( + numerator: Union[Number, ivy.Array, ivy.NativeArray], + denominator: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, +) -> Union[Number, ivy.Array]: """ - Perform in-place increment for the input array. + Divide the numerator by the denominator, with min denominator added to the + denominator for numerical stability. Parameters ---------- - x - The input array to be incremented by the defined value. - val - The value of increment. + numerator + The numerator of the division. + denominator + The denominator of the division. + min_denominator + The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) + by default. Returns ------- ret - The array following the in-place increment. + The new item following the numerically stable division. Examples -------- + With :code:`int` input: + + >>> x = ivy.stable_divide(1, 2) + >>> print(x) + 0.49999999999975 + + >>> x = ivy.stable_divide(1, 4, min_denominator=1) + >>> print(x) + 0.2 + + With float input: + + >>> x = ivy.stable_divide(5.0, 3.33) + >>> print(x) + 1.5015015015010504 + + With :code:`complex` input: + + >>> x = ivy.stable_divide(1+1j, 1-1j) + >>> print(x) + (5.000444502911705e-13+0.9999999999995j) + With :class:`ivy.Array` input: - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_increment(x, 3.) + >>> x = ivy.asarray([[10., 20., 30.], + ... [40., 50., 60.]]) + >>> y = ivy.stable_divide(x, 10.) >>> print(y) - ivy.array([[ 8.3, 10., 3.], - [ 9.8, 11., 6.9], - [ 3., 13., 9.3]]) + ivy.array([[1., 2., 3.], + [4., 5., 6.]]) + + + >>> x = ivy.asarray([1,2,3]) + >>> y = np.array((1., 3., 5.)) + >>> z = ivy.stable_divide(x, y) + >>> print(z) + ivy.array([1. , 0.667, 0.6 ]) + + >>> x = ivy.asarray([1., 2., 4.]) + >>> y = ivy.asarray([1., 0.5, 0.25]) + >>> z = ivy.asarray([0.01, 0.02, 0.03]) + >>> w = ivy.stable_divide(x, y, min_denominator=z) + >>> print(w) + ivy.array([ 0.99, 3.85, 14.3 ]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.inplace_increment(x, 2.5) + >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) + >>> y = ivy.stable_divide(x, 0.5) >>> print(y) { - a: ivy.array([2.5, 17.5, 32.5]), - b: ivy.array([2.5, 27.5, 52.5]) + a: ivy.array([20., 30.]), + b: ivy.array([40., 50.]) } - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_increment(x, y) + >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) + >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) + >>> z = ivy.stable_divide(x, y) >>> print(z) { - a: ivy.array([0., 30., 60.]), - b: ivy.array([0., 50., 100.]) + a: ivy.array([2., 0.8]), + b: ivy.array([0.857, 10.]) } """ - return current_backend(x).inplace_increment(x, val) + return numerator / (denominator + default(min_denominator, ivy.min_denominator)) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def scatter_flat( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], +def stable_pow( + base: Union[Number, ivy.Array, ivy.NativeArray], + exponent: Union[Number, ivy.Array, ivy.NativeArray], /, *, - size: Optional[int] = None, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: + min_base: float = None, +) -> Any: """ - Scatter flat updates into a new flat array according to flat indices. + Raise the base by the power, with ivy.min_base added to the base when exponent > 1 + for numerical stability. Parameters ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - size - The size of the result. Default is `None`, in which case tensor - argument out must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + base + The base number. + exponent + The exponent number. + min_base + The minimum base to use, use global ivy.min_base by default. Returns ------- ret - New array of given shape, with the values scattered at the indices. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + The new item following the numerically stable power. Examples -------- - With :class:`ivy.Array` input: - >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) - >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) - >>> ivy.scatter_flat(indices, updates, out=out) - >>> print(out) - ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) + With :code:`int` input: + + >>> x = ivy.stable_pow(2, 2) + >>> print(x) + ivy.array(4.00004) + >>> x = ivy.stable_pow(2, 2, min_base=2) + >>> print(x) + ivy.array(16) + + With float input: + + >>> x = ivy.stable_pow(4.0, .5) + >>> print(x) + ivy.array(2.00000262) + + With :code:`complex` input: + + >>> x = ivy.stable_pow(3+4j, 2j) + >>> print(x) + ivy.array(-0.15605032-0.01208451j) With :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) - ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) + >>> x = ivy.asarray([[2, 4], + ... [6, 8]]) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) + ivy.array([[ 4.00004, 16.00008], + [36.00012, 64.00016]]) - With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) + >>> x = ivy.asarray([2, 4, 6]) + >>> y = ivy.asarray([2, 3, 4]) + >>> z = ivy.stable_pow(x, y) + >>> print(z) + ivy.array([ 4.00004, 64.00048, 1296.00864]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + a: ivy.array([4.00004, 16.00008]), + b: ivy.array([36.00012, 64.00016]) } - - With :class:`ivy.Container` input: - >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), - ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) + >>> z = ivy.stable_pow(x, y) + >>> print(z) { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + a: ivy.array([2.00001, 64.00048]), + b: ivy.array([1296.00864, 32768.2048]) } """ - return current_backend(indices).scatter_flat( - indices, updates, size=size, reduction=reduction, out=out + return_dtype = ivy.promote_types( + ivy.default_dtype(item=base), + ivy.default_dtype(item=default(min_base, ivy.min_base)), ) + return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) + ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) + return ret.astype(return_dtype) @handle_exceptions -@handle_backend_invalid @handle_nestable -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def scatter_nd( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], +def strides( + x: Union[ivy.Array, ivy.NativeArray], /, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - *, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[int]: """ - Scatter updates into a new array according to indices. + Return the input array's strides across each dimension. Parameters ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - shape - The shape of the result. Default is ``None``, in which case tensor - argument must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + x + The input array. Returns ------- ret - New array of given shape, with the values scattered at the indices. + A tuple containing the strides. Examples -------- - scatter values into an empty array, With :class:`ivy.Array` input: - - >>> indices = ivy.array([[4], [3], [1], [7]]) - >>> updates = ivy.array([9, 10, 11, 12]) - >>> shape = ivy.array([8]) - >>> scatter = ivy.scatter_nd(indices, updates, shape) - >>> print(scatter) - ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) - - With scatter into an empty array, With :class:`ivy.Container` input: - - >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), - ... b=ivy.array([[5],[1],[2]])) - >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), - ... b=ivy.array([20, 30, 40])) - >>> shape = ivy.Container(a=ivy.array([10]), - ... b = ivy.array([10])) - >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') - >>> print(z) - { - a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), - b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) - } - - With :class:`ivy.Container` and :class:`ivy.Array` input: - - >>> indices = ivy.array([[4],[3],[1]]) - >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), - ... b=ivy.array([200, 300, 400])) - >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), - ... b=ivy.array([10, 20, 30, 40, 50])) - >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) - >>> print(z) - { - a: ivy.array([1, 30, 3, 20, 10]), - b: ivy.array([10, 400, 30, 300, 200]) - } + >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) + >>> ivy.strides(x) + (4, 8) """ - return current_backend(indices).scatter_nd( - indices, updates, shape=shape, reduction=reduction, out=out - ) + if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): + return ivy.to_numpy(x).strides + # if x is an ivy array with a base, + # convert it to a numpy array with the same base: + ret = ivy.to_numpy(x.base) + ivy_numpy = ivy.with_backend("numpy") + for fn, args, kwargs, index in x._manipulation_stack: + ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) + ret = ret[index] if ivy.exists(index) else ret + return ret.to_native().strides @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function -@handle_device_shifting -def gather( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: """ - Gather slices from params at axis according to indices. + Return if in-place operations are supported for x's data type. + + Determine whether in-place operations are supported for x's data type, by the + current backend framework setting. Parameters ---------- - params - The array from which to gather values. - indices - The array which indicates the indices that will be gathered along - the specified axis. - axis - optional int, the axis from which to gather from. - Default is ``-1``. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional array, for writing the result to. It must have a shape - that the inputs broadcast to. + x + Input variable for whose data type we check whether the current backend + framework supports in-place operations. Returns ------- ret - New array with the values gathered at the specified indices along the - specified axis. + Value depends on whether in-place operations are supported for + data type of x. + Raises + ------ + IvyException + If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will + be raised. - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.array([1, 2]) - >>> print(ivy.gather(x, y)) - ivy.array([1., 2.]) + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) + True - >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) - >>> y = ivy.array([[0, 1],[1, 2]]) - >>> z = ivy.zeros((2, 2, 2)) - >>> ivy.gather(x, y, out=z) - >>> print(z) - ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) - - >>> x = ivy.array([[[0., 1.], [2., 3.]], - ... [[8., 9.], [10., 11.]]]) - >>> y = ivy.array([[0, 1]]) - >>> z = ivy.zeros((1, 2, 2, 2)) - >>> ivy.gather(x, y, axis=0, out=z) - >>> print(z) - ivy.array( - [[[[ 0., 1.], - [ 2., 3.]], - [[ 8., 9.], - [10., 11.]]]]) - - >>> x = ivy.array([[0, 10, 20, 0, 0], - ... [0, 0, 0, 30, 40], - ... [0, 10, 0, 0, 40]]) - >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) - >>> z = ivy.gather(x, y, batch_dims=1) - >>> print(z) - ivy.array([[10, 20], [30, 40],[10, 40]]) - - With :class:`ivy.Container` input: + With :class:`ivy.Container` input and backend set as `torch`: - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.Container(a = ivy.array([0, 1]), - ... b = ivy.array([1, 2])) - >>> print(ivy.gather(x, y)) + >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) { - a: ivy.array([0., 1.]), - b: ivy.array([5., 6.]) + a: True, + b: True } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With `ivy.Array` input and backend set as "tensorflow": - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.array([0, 1]) - >>> print(ivy.gather(x, y)) - { - a: ivy.array([0., 1.]), - b: ivy.array([4., 5.]) - } + >>> x = ivy.array([1., 4.2, 2.2]) + >>> ret = x.supports_inplace_updates() + >>> print(ret) + False """ - return current_backend(params, indices).gather( - params, indices, axis=axis, batch_dims=batch_dims, out=out + if _is_variable(x): + return ivy.inplace_variables_supported() + elif ivy.is_native_array(x): + return ivy.inplace_arrays_supported() + raise ivy.utils.exceptions.IvyException( + "Input x must be either a variable or an array." ) +@handle_exceptions +def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: + """ + Return the input shape in ivy.Shape form. + + Parameters + ---------- + shape + The input to be converted + + Returns + ------- + ret + the input in ivy.Shape form + """ + if isinstance(shape, ivy.Shape): + return shape + return ivy.Shape(shape) + + @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function @handle_device_shifting -def gather_nd( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - batch_dims: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: """ - Gather slices from params into a array with shape specified by indices. + Create a (possibly nested) list from input array. Parameters ---------- - params - The array from which to gather values. - indices - Index array. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + x + Input array. Returns ------- ret - New array of given shape, with the values gathered at the indices. + A list representation of the input array ``x``. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) - ivy.array(1.) + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_list(x) + >>> print(y) + [-1, 0, 1] - >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) - >>> y = ivy.array([[0],[1],[1]], dtype='int32') - >>> z = ivy.gather_nd(x,y,batch_dims=1) - ivy.array([0., 3., 5.]) + >>> x = ivy.array([[ 1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.to_list(x) + >>> print(y) + [[1.100000023841858,2.200000047683716,3.299999952316284], + [-4.400000095367432,-5.5,-6.599999904632568]] - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([[[-1, 0, 1], + ... [ 1, 0, -1]], + ... [[ 1, -1, 0], + ... [ 1, 0, -1]]]) + >>> y = ivy.to_list(x) + >>> print(y) + [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_list(x) + >>> print(y) { - a: ivy.array(1.), - b: ivy.array(5.) + a: [-1, 0, 1] } - With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], + ... [-1, 0, 1], + ... [1, 0, -1]])) + >>> y = ivy.to_list(x) + >>> print(y) + { + a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] + } - >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), - ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) - >>> y = ivy.Container(a=ivy.array([1,0]), - ... b=ivy.array([0])) - >>> print(ivy.gather_nd(x, y)) + >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], + ... [[1, -1, 0],[1, 0, -1]]])) + >>> y = ivy.to_list(x) + >>> print(y) { - a: ivy.array(30.), - b: ivy.array([0., 100., 200.]) + a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] } """ - res = current_backend(params, indices).gather_nd( - params, indices, batch_dims=batch_dims - ) - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + return current_backend(x).to_list(x) @handle_exceptions -@handle_nestable -@handle_array_function -def multiprocessing(context: Optional[str] = None): +def to_native_shape( + shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] +) -> ivy.NativeShape: """ - Return backend-specific multiprocessing module. + Return the input shape in its native backend framework form. Parameters ---------- - context - The context of the multiprocessing, either fork, forkserver or spawn. - Default is ``None``. + shape + The input to be converted Returns ------- - ret - Multiprocessing module + ret + the input in its native framework form """ - return current_backend().multiprocessing(context) + native_shape_type = (ivy.NativeShape,) + if ivy.current_backend_str() == "torch": + native_shape_type += (tuple,) + if len(backend_stack) != 0 and isinstance(shape, native_shape_type): + return shape + ivy.utils.assertions.check_isinstance( + shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) + ) + if isinstance(shape, int): + shape = (shape,) + elif isinstance(shape, list): + shape = tuple(shape) + elif is_array(shape): + shape = ivy.to_numpy(shape).tolist() + elif isinstance(shape, ivy.Shape): + shape = shape.shape + ivy.utils.assertions.check_all( + [isinstance(v, int) for v in shape if not is_array(v)], + "shape must take integers only", + as_array=False, + ) + ivy.utils.assertions.check_true( + not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" + ) + return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) @handle_exceptions @@ -3638,422 +3642,490 @@ def multiprocessing(context: Optional[str] = None): @handle_nestable @handle_array_like_without_promotion @inputs_to_native_arrays -@outputs_to_ivy_shapes -@outputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def shape( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - as_array: bool = False, -) -> Union[ivy.Shape, ivy.NativeShape]: +def to_numpy( + x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True +) -> np.ndarray: """ - Return the shape of the array ``x``. + Convert an array into a numpy array. Parameters ---------- x - Input array to infer the shape of. - as_array - Whether to return the shape as an array. - Default is False. + input array + copy + whether to copy the array to a new address or not. + Default is ``True``. Returns ------- ret - Shape of the array ``x``. + a numpy array copying all the element of the array ``x``. Examples -------- - >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) - >>> y = ivy.shape(x) - >>> z = ivy.shape(x, as_array = True) + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_numpy(x, copy=True) >>> print(y) - (2, 3) + [-1 0 1] - >>> print(z) - ivy.array([2, 3]) - """ - return current_backend(x).shape(x, as_array=as_array) + >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) + >>> y = ivy.to_numpy(x, copy=True) + >>> print(y) + [[-1 0 1] + [-1 0 1] + [ 1 0 -1]] + With :class:`ivy.Container` input: -ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False - - -@handle_exceptions -def set_shape_array_mode(mode: bool) -> None: - """ - Set the mode of returning shape as ivy.Array to the given mode instance. - - Parameter - --------- - mode - boolean whether to return shape as ivy.Array - - Examples - -------- - >>> ivy.set_shape_array_mode(False) - >>> ivy.shape_array_mode - False + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([-1, 0, 1], dtype=int32) + } - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True + >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), + ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([[-1., 0., 1.], + [-1., 0., 1.], + [1., 0., -1.]], dtype=float32), + b: array([[-1, 0, 0], + [1, 0, 1], + [1, 1, 1]], dtype=int32) + } """ - global shape_array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - shape_array_mode_stack.append(mode) - ivy.__setattr__("shape_array_mode", mode, True) + return current_backend(x).to_numpy(x, copy=copy) @handle_exceptions -def unset_shape_array_mode() -> None: - """ - Reset the mode of returning shape as ivy.Array to the previous state. - - Examples - -------- - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True - - >>> ivy.unset_shape_array_mode() - >>> ivy.shape_array_mode - False - """ - global shape_array_mode_stack - if shape_array_mode_stack: - shape_array_mode_stack.pop(-1) - mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False - ivy.__setattr__("shape_array_mode", mode, True) - - @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function @handle_device_shifting -def get_num_dims( - x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False -) -> int: +def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: """ - Return the number of dimensions of the array x. + Convert an array with a single element into a scalar. Parameters ---------- x - Input array to infer the number of dimensions for. - as_array - Whether to return the shape as a array, default False. + Input array with a single element. Returns ------- ret - Shape of the array + a scalar copying the element of the array ``x``. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - Examples - -------- + Functional Examples + ------------------- + With :class:`ivy.Array` input: - >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) - >>> b = ivy.get_num_dims(a, as_array=False) - >>> print(b) + >>> x = ivy.array([3]) + >>> y = ivy.to_scalar(x) + >>> print(y) 3 - With :class:`ivy.Container` input: + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) - >>> print(ivy.get_num_dims(a)) + >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) + >>> y = ivy.to_scalar(x) + >>> print(y) { - b: 2 + a: -1, + b: 3 } - >>> b = ivy.get_num_dims(a, as_array=True) - >>> print(b) + >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), + ... c=ivy.array([-1])) + >>> y = ivy.to_scalar(x) + >>> print(y) { - b: ivy.array(2) + a: 1, + b: 0, + c: -1 } """ - return current_backend(x).get_num_dims(x, as_array=as_array) + return current_backend(x).to_scalar(x) @handle_exceptions -def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): +@handle_nestable +def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: """ - Return the index and `inspect.Parameter` representation of the specified argument. - In the form of a dict with keys "idx" and "param". + Try and return the function, otherwise return None if an exception was raised during + function execution. Parameters ---------- fn - The function to retrieve the argument information for - name - The name of the argument - idx - the index of the argument in the inputs + Function to try and call and return. + args + list of arguments. + kwargs + dictionay of keyword arguments Returns ------- - ret - a `dict` containing the idx, and the `inspect.Parameter` for the argument, - which itself contains the parameter name, type, and other helpful information. - """ - ivy.utils.assertions.check_all_or_any_fn( - name, - idx, - fn=ivy.exists, - type="any", - limit=[1], - message="exactly one of the keyword arguments name or idx must be provided", - as_array=False, - ) - params = inspect.signature(fn).parameters - if ivy.exists(name): - return {"idx": list(params).index(name), "param": params[name]} - return {"idx": idx, "param": list(params.values())[idx]} + Either the function itself or None if an exception was raised + during function execution. + Examples + -------- + with a function that is executed without any exception: -def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): - attr_list = () - if hasattr(fn, other_attr_name): - attr_list = getattr(fn, other_attr_name) - if isinstance(attr_list, dict): - attr_list = attr_list.get(backend, ()) - ivy.utils.assertions.check_false( - dnd_dict and attr_list, - f"Cannot specify both {first_attr_name} and {other_attr_name} " - "cannot both be defined for the same function", - ) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> z = ivy.try_else_none(ivy.add, x, y) + >>> print(z.__name__) + add + with a function that is executed with an exception: -def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: - fn_unsupported_dnd = {} - fn_supported_dnd = {} - backend = ivy.current_backend_str() - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): - fn_supported_dnd = fn_supported_dnd.get(backend, {}) + >>> x = ivy.array([1, 2, 3]) + >>> y = 'hemant' + >>> z = ivy.try_else_none(ivy.add,x, y) + >>> print(z) + None + """ + try: + _ = fn(*args, **kwargs) + return fn + except Exception: + return None - ivy.utils.assertions.check_false( - fn_unsupported_dnd and fn_supported_dnd, - "unsupported_device_and_dtype and supported_device_and_dtype cannot" - " both be defined for the same function", - ) - us = "unsupported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") +@handle_exceptions +def unset_array_mode() -> None: + """ + Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back + to ivy.Array to the previous state. - ss = "supported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - return True + >>> ivy.unset_shape_array_mode() + >>> ivy.array_mode + True + """ + global array_mode_stack + if array_mode_stack: + array_mode_stack.pop(-1) + mode = array_mode_stack[-1] if array_mode_stack else True + ivy.__setattr__("array_mode", mode, True) -def _all_dnd_combinations(): - all_comb = {} - for device in ivy.all_devices: - all_comb[device] = ivy.all_dtypes - return all_comb +@handle_exceptions +def unset_exception_trace_mode() -> None: + """ + Reset the trace mode to the previously set mode. + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' -def _dnd_dict_intersection(a, b): - res = {} - for device in a: - if device in b: - intersection = set.intersection(set(a[device]), set(b[device])) - if intersection: - res[device] = tuple(intersection) - return res + >>> ivy.unset_exception_trace_mode() + >>> ivy.exception_trace_mode + 'full' + """ + global exception_trace_mode_stack + if exception_trace_mode_stack: + exception_trace_mode_stack.pop(-1) + mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" + ivy.__setattr__("exception_trace_mode", mode, True) -def _dnd_dict_difference(a, b): - res = a - for device in list(a): - if device in b: - difference = set.difference(set(a[device]), set(b[device])) - if difference: - res[device] = tuple(difference) - else: - del res[device] - return res +@handle_exceptions +def unset_inplace_mode() -> None: + """ + Reset the memory management behavior for in-place updates in Ivy to the previous + state. + Examples + -------- + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' -def _dnd_dict_union(a, b): - res = {} - for device in set(list(a) + list(b)): - u1 = set(a.get(device, ())) - u2 = set(b.get(device, ())) - res[device] = tuple(set.union(u1, u2)) + >>> unset_inplace_mode() + >>> ivy.inplace_mode + 'lenient' + """ + global inplace_mode_stack + if inplace_mode_stack: + inplace_mode_stack.pop(-1) + mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" + ivy.__setattr__("inplace_mode", mode, True) - return res +@handle_exceptions +def unset_min_base() -> None: + """ + Reset the global minimum base used by ivy for numerically stable power raising to + the previous value. -def _get_devices_and_dtypes(fn, recurse=False, complement=True): - supported_devices = ivy.function_supported_devices(fn, recurse=recurse) - supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) + Examples + -------- + >>> ivy.set_min_base(1e-07) + >>> y = ivy.min_base + >>> print(y) + 1e-07 - if hasattr(fn, "partial_mixed_handler"): - supported_devices = supported_devices["primary"] - supported_dtypes = supported_dtypes["primary"] + >>> ivy.unset_min_base() + >>> ivy.min_base + 1e-05 + """ + global min_base_stack + if min_base_stack: + min_base_stack.pop(-1) + val = min_base_stack[-1] if min_base_stack else 1e-05 + ivy.__setattr__("min_base", val, True) - supported = {} - # Generate a base supported set from other attributes - for device in supported_devices: - supported[device] = supported_dtypes - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported +@handle_exceptions +def unset_min_denominator() -> None: + """ + Reset the global minimum denominator used by ivy for numerically stable division to + the previous value. - backend = ivy.current_backend_str() + Examples + -------- + >>> ivy.set_min_denominator(1e-10) + >>> y = ivy.min_denominator + >>> print(y) + 1e-10 - # Their values are formatted like either - # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype.__get__() + >>> ivy.unset_min_denominator() + >>> ivy.min_denominator + 1e-12 + """ + global min_denominator_stack + if min_denominator_stack: + min_denominator_stack.pop(-1) + val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 + ivy.__setattr__("min_denominator", val, True) - if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): - fn_supported_dnd = fn_supported_dnd.get(backend, supported) - ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) - # dict intersection - supported = _dnd_dict_intersection(supported, fn_supported_dnd) +@handle_exceptions +def unset_nestable_mode() -> None: + """ + Reset the mode of whether to check if function inputs are ivy.Container to the + previous state. - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() + Examples + -------- + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False - if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) + >>> ivy.unset_nestable_mode() + >>> ivy.nestable_mode + True + """ + global nestable_mode_stack + if nestable_mode_stack: + nestable_mode_stack.pop(-1) + mode = nestable_mode_stack[-1] if nestable_mode_stack else True + ivy.__setattr__("nestable_mode", mode, True) - ivy.utils.assertions.check_isinstance( - list(fn_unsupported_dnd.values())[0], tuple - ) - # dict difference - supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - if complement: - # dict difference - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported +@handle_exceptions +def unset_precise_mode() -> None: + """ + Reset the mode of whether to use a promotion table that avoids any precision loss or + a compute effecient table that avoids most wider-than-necessary promotions. + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.unset_precise_mode() + >>> ivy.precise_mode + True + """ + global precise_mode_stack + if precise_mode_stack: + precise_mode_stack.pop(-1) + mode = precise_mode_stack[-1] if precise_mode_stack else True + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) @handle_exceptions -@handle_nestable -def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: +def unset_queue_timeout() -> None: """ - Return the supported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the supported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + Reset the global queue timeout value (in seconds) to the previous state. - Parameters - ---------- - fn - The function to check for the supported device and dtype attribute - recurse - Whether to recurse into used ivy functions. - Default is ``True``. + Examples + -------- + >>> ivy.set_queue_timeout(10.0) + >>> y = ivy.queue_timeout + >>> print(y) + 10.0 - Returns - ------- - ret - Tuple or dict containing the supported devices and dtypes of the function + >>> ivy.unset_queue_timeout() + >>> ivy.queue_timeout + 15.0 """ - ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", - ) + global queue_timeout_stack + if queue_timeout_stack: + queue_timeout_stack.pop(-1) + timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 + ivy.__setattr__("queue_timeout", timeout, True) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=False), - } - else: - supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) - if recurse: - supported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - supported_devices_dtypes, - _dnd_dict_intersection, - function_supported_devices_and_dtypes, - wrapper=lambda x: x, - ) - return supported_devices_dtypes +@handle_exceptions +def unset_shape_array_mode() -> None: + """ + Reset the mode of returning shape as ivy.Array to the previous state. + + Examples + -------- + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True + + >>> ivy.unset_shape_array_mode() + >>> ivy.shape_array_mode + False + """ + global shape_array_mode_stack + if shape_array_mode_stack: + shape_array_mode_stack.pop(-1) + mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False + ivy.__setattr__("shape_array_mode", mode, True) + + +@handle_exceptions +def unset_show_func_wrapper_trace_mode() -> None: + """ + Reset the mode of whether to show the full stack trace with function wrapping + traces. + + Examples + -------- + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False + + >>> ivy.unset_show_func_wrapper_trace_mode() + >>> ivy.show_func_wrapper_trace_mode + True + """ + global show_func_wrapper_trace_mode_stack + if show_func_wrapper_trace_mode_stack: + show_func_wrapper_trace_mode_stack.pop(-1) + mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True + ) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + + +@handle_exceptions +def unset_tmp_dir() -> None: + """ + Reset the directory for saving temporary files to the previous value. + + Examples + -------- + >>> ivy.set_tmp_dir("/my_dir") + >>> y = ivy.tmp_dir + >>> print(y) + /my_dir + + >>> ivy.unset_tmp_dir() + >>> ivy.tmp_dir + /tmp + """ + global tmp_dir_stack + if tmp_dir_stack: + tmp_dir_stack.pop(-1) + tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions @handle_nestable -def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def value_is_nan( + x: Union[ivy.Array, ivy.NativeArray, Number], + /, + *, + include_infs: bool = True, +) -> bool: """ - Return the unsupported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the unsupported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + Determine whether the single valued array or scalar is of nan type. Parameters ---------- - fn - The function to check for the unsupported device and dtype attribute - recurse - Whether to recurse into used ivy functions. + x + The input to check Input array. + include_infs + Whether to include infs and -infs in the check. Default is ``True``. Returns ------- ret - Tuple or dict containing the unsupported devices and dtypes of the function + Boolean as to whether the input value is a nan or not. + + Examples + -------- + >>> x = ivy.array([451]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False + + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + False + + >>> x = ivy.array([float('nan')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + True + + >>> x = ivy.array([0]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False """ - ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=True), - } - else: - unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) - if recurse: - unsupported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - unsupported_devices_dtypes, - _dnd_dict_union, - function_unsupported_devices_and_dtypes, - wrapper=lambda x: x, - ) - return unsupported_devices_dtypes + x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x + if not x_scalar == x: + return True + if include_infs and (x_scalar == INF or x_scalar == -INF): + return True + return False @handle_exceptions @@ -4115,142 +4187,48 @@ def vmap( return current_backend().vmap(func, in_axes, out_axes) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@to_native_arrays_and_back -@handle_device_shifting -def isin( - elements: Union[ivy.Array, ivy.NativeArray], - test_elements: Union[ivy.Array, ivy.NativeArray], - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> ivy.Array: - """ - Test if each element of elements is in test_elements. - - Parameters - ---------- - elements - input array - test_elements - values against which to test for each input element - assume_unique - If True, assumes both elements and test_elements contain unique elements, - which can speed up the calculation. Default value is False. - invert - If True, inverts the boolean return array, resulting in True values for - elements not in test_elements. Default value is False. - - Returns - ------- - ret - output a boolean array of the same shape as elements that is True for elements - in test_elements and False otherwise. - - Examples - -------- - >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y) - ivy.array([[False, False, False], [ True, True, True]]) - - >>> x = ivy.array([3, 2, 1, 0]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y, invert=True) - ivy.array([False, False, False, True]) - """ - return ivy.current_backend(elements, test_elements).isin( - elements, test_elements, assume_unique=assume_unique, invert=invert - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def itemsize( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> int: - """ - Return the size of the input array's elements. - - Parameters - ---------- - x - The input array. - - Returns - ------- - ret - An integer specifying the element size in bytes. - - Examples - -------- - >>> x = ivy.array([1,2,3], dtype=ivy.float64) - >>> ivy.itemsize(x) - 8 - - >>> x = ivy.array([1,2,3], dtype=ivy.complex128) - >>> ivy.itemsize(x) - 16 - """ - return ivy.current_backend(x).itemsize(x) - - -@handle_exceptions -@handle_nestable -@handle_device_shifting -def strides( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> Tuple[int]: - """ - Return the input array's strides across each dimension. - - Parameters - ---------- - x - The input array. - - Returns - ------- - ret - A tuple containing the strides. - - Examples - -------- - >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) - >>> ivy.strides(x) - (4, 8) - """ - if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): - return ivy.to_numpy(x).strides - # if x is an ivy array with a base, - # convert it to a numpy array with the same base: - ret = ivy.to_numpy(x.base) - ivy_numpy = ivy.with_backend("numpy") - for fn, args, kwargs, index in x._manipulation_stack: - ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) - ret = ret[index] if ivy.exists(index) else ret - return ret.to_native().strides - - -def is_ivy_nested_array(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Nested Array. - - Parameters - ---------- - x - The input to check - Returns - ------- - ret - Boolean, whether or not x is an ivy nested array. - """ - return isinstance(x, ivy.NestedArray) +trace_mode_dict["frontend"] = "ivy/functional/frontends" +trace_mode_dict["ivy"] = "ivy/" +trace_mode_dict["full"] = "" +trace_mode_dict["none"] = "" +ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True +ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True +ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True +ivy.exception_trace_mode = ( + exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" +) +ivy.show_func_wrapper_trace_mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True +) +# IMPORTANT: assign attribute directly to function instead of wrapper here +einops_reduce.unsupported_dtypes = { + "torch": ("float16",), + "tensorflow": ("complex",), + "paddle": ("complex", "uint8", "int8", "int16", "float16"), +} +ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 +ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 +stable_pow.unsupported_dtypes = ("bfloat16",) +ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 +ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" +get_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +set_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} +ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" +ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False diff --git a/ivy/functional/ivy/gradients.py b/ivy/functional/ivy/gradients.py index fc0c9b2a1f411..72e3961ffc7e9 100644 --- a/ivy/functional/ivy/gradients.py +++ b/ivy/functional/ivy/gradients.py @@ -22,18 +22,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # - - -def _get_duplicate_index_chains(xs): - """Generate a list of duplicate index chains for a given nested structure.""" - duplicate_index_chains = () - if isinstance(xs, ivy.Container): - duplicate_index_chains = xs.cont_duplicate_array_keychains() - elif isinstance(xs, (list, tuple, dict)): - duplicate_index_chains = ivy.duplicate_array_index_chains(xs) - return duplicate_index_chains +# --- Helpers --- # +# --------------- # def _arrays_to_float_variables(xs, xs_grad_idxs=None): @@ -60,6 +50,89 @@ def inner_fn(x): return ivy.nested_map(xs, map_fn, include_derived=True, shallow=False) +def _get_duplicate_index_chains(xs): + """Generate a list of duplicate index chains for a given nested structure.""" + duplicate_index_chains = () + if isinstance(xs, ivy.Container): + duplicate_index_chains = xs.cont_duplicate_array_keychains() + elif isinstance(xs, (list, tuple, dict)): + duplicate_index_chains = ivy.duplicate_array_index_chains(xs) + return duplicate_index_chains + + +def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): + """Extract all relevant results from the output nested structure of a function.""" + + def map_fn(x_): + if ivy.is_array(x_): + x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ + if create_var: + x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ + if len(x_.shape) == 0: + return ivy.to_native(x_) + if reshape: + if x_.size == 1: + if reshape: + return ivy.to_native(ivy.reshape(x_, [])) + return ivy.to_native(x_) + else: + return ivy.to_ivy(x_) + else: + return ivy.to_native(x_) + return x_ + + if ivy.is_array(x): + return [], map_fn(x) + + x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) + arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) + if _check_if_empty(arr_idxs): + return arr_idxs, [] + else: + if idxs is not None: + arr_idxs = [ + arr_idx + for arr_idx in arr_idxs + if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) + ] + arr_values = ivy.multi_index_nest(x, arr_idxs) + arr_idxs = _idxs_to_str(arr_idxs) + return arr_idxs, arr_values + + +def _get_native_y(y): + """Convert all outputs to native arrays.""" + array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) + y_final = [] + if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: + y_final = ivy.multi_index_nest(y, array_idxs) + return y_final + + +def _get_required_float_variables(xs, xs_grad_idxs): + """ + Convert all required arrays to float variables for gradient calculation. + + Also, returns a list of duplicate index chains for the nested + structure. + """ + if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: + xs_grad_idxs = None + duplicate_index_chains = _get_duplicate_index_chains(xs) + xs = _to_ivy(xs) + xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) + xs = _set_duplicates(xs, duplicate_index_chains) + xs_required = _get_required_native_variables(xs, xs_grad_idxs) + required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) + return ( + xs, + xs_grad_idxs, + xs_required, + required_duplicate_index_chains, + duplicate_index_chains, + ) + + def _get_required_native_variables(xs, xs_grad_idxs): """Extract all required native variables from a nested structure.""" # To make sure that only the required arrays are converted to native arrays @@ -102,68 +175,46 @@ def map_fn(x): return xs -def _get_required_float_variables(xs, xs_grad_idxs): - """ - Convert all required arrays to float variables for gradient calculation. - - Also, returns a list of duplicate index chains for the nested - structure. - """ - if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: - xs_grad_idxs = None - duplicate_index_chains = _get_duplicate_index_chains(xs) - xs = _to_ivy(xs) - xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) - xs = _set_duplicates(xs, duplicate_index_chains) - xs_required = _get_required_native_variables(xs, xs_grad_idxs) - required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) - return ( - xs, - xs_grad_idxs, - xs_required, - required_duplicate_index_chains, - duplicate_index_chains, +def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): + """Get the relevant outputs from the function return value.""" + if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ + [0] + ]: + ret_grad_idxs = None + ret_idxs, ret_values = _get_native_variables_and_indices( + func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape ) + if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): + return func_ret, {} + if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: + y = ret_values[0] + else: + y = ret_values + return ret_grad_idxs, y, ret_idxs -def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): - """Extract all relevant results from the output nested structure of a function.""" +def _is_variable(x, exclusive=False, to_ignore=None) -> bool: + x = ivy.to_native(x, nested=True, to_ignore=to_ignore) + return ivy.nested_map( + x, + lambda x: current_backend(x).is_variable(x, exclusive=exclusive), + include_derived=True, + shallow=False, + to_ignore=to_ignore, + ) - def map_fn(x_): - if ivy.is_array(x_): - x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ - if create_var: - x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ - if len(x_.shape) == 0: - return ivy.to_native(x_) - if reshape: - if x_.size == 1: - if reshape: - return ivy.to_native(ivy.reshape(x_, [])) - return ivy.to_native(x_) - else: - return ivy.to_ivy(x_) - else: - return ivy.to_native(x_) - return x_ - if ivy.is_array(x): - return [], map_fn(x) +def _process_func_ret_and_grads(func_ret, grads, retain_grads): + """ + Stop gradients propagation. - x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) - arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) - if _check_if_empty(arr_idxs): - return arr_idxs, [] - else: - if idxs is not None: - arr_idxs = [ - arr_idx - for arr_idx in arr_idxs - if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) - ] - arr_values = ivy.multi_index_nest(x, arr_idxs) - arr_idxs = _idxs_to_str(arr_idxs) - return arr_idxs, arr_values + Set the gradients of non-finite values to zero, and stopping + gradient propagation of the function results. + """ + grads = _non_finite_to_zero(grads) + func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) + grads = _to_ivy(grads) + return func_ret, grads def _set_duplicates(xs, duplicate_index_chains): @@ -192,33 +243,6 @@ def _set_duplicates(xs, duplicate_index_chains): return xs -def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): - """Get the relevant outputs from the function return value.""" - if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ - [0] - ]: - ret_grad_idxs = None - ret_idxs, ret_values = _get_native_variables_and_indices( - func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape - ) - if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): - return func_ret, {} - if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: - y = ret_values[0] - else: - y = ret_values - return ret_grad_idxs, y, ret_idxs - - -def _get_native_y(y): - """Convert all outputs to native arrays.""" - array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) - y_final = [] - if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: - y_final = ivy.multi_index_nest(y, array_idxs) - return y_final - - def _stop_grad_and_index(func_ret, retain_grads, grads): """Stop gradient propagation of the function results.""" if not retain_grads: @@ -232,46 +256,6 @@ def _stop_grad_and_index(func_ret, retain_grads, grads): return func_ret, grads -def _process_func_ret_and_grads(func_ret, grads, retain_grads): - """ - Stop gradients propagation. - - Set the gradients of non-finite values to zero, and stopping - gradient propagation of the function results. - """ - grads = _non_finite_to_zero(grads) - func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) - grads = _to_ivy(grads) - return func_ret, grads - - -_check_if_empty = ( - lambda idxs: not isinstance(idxs, list) - or np.asarray(idxs, dtype="object").size == 0 -) - - -_idxs_to_str = lambda idxs: [ - "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) -] - - -_to_ivy = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) - - -_non_finite_to_zero = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) - - # Private Variable Helpers # # -------------------------# @@ -284,17 +268,6 @@ def _variable(x): return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) -def _is_variable(x, exclusive=False, to_ignore=None) -> bool: - x = ivy.to_native(x, nested=True, to_ignore=to_ignore) - return ivy.nested_map( - x, - lambda x: current_backend(x).is_variable(x, exclusive=exclusive), - include_derived=True, - shallow=False, - to_ignore=to_ignore, - ) - - def _variable_data( x: Union[ivy.Array, ivy.NativeArray] ) -> Union[ivy.Array, ivy.NativeArray]: @@ -318,128 +291,367 @@ def _variable_data( return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) +# --- Main --- # +# ------------ # + + +# Optimizer Steps # + + @handle_exceptions -@handle_backend_invalid -@handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def stop_gradient( - x: Union[ivy.Array, ivy.NativeArray], +def adam_step( + dcdw: Union[ivy.Array, ivy.NativeArray], + mw: Union[ivy.Array, ivy.NativeArray], + vw: Union[ivy.Array, ivy.NativeArray], + step: Union[int, float], /, *, - preserve_type: bool = True, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Stop gradient computation. + Compute adam step delta, given the derivatives of some cost c with respect to + weights ws, using ADAM update. `[reference] + + `_ Parameters ---------- - x - Array for which to stop the gradient. - preserve_type - Whether to preserve gradient computation on ivy.Array instances. Default is - True. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + mw + running average of the gradients + vw + running average of second moments of the gradients + step + training step + beta1 + gradient forgetting factor (Default value = 0.9) + beta2 + second moment of gradient forgetting factor (Default value = 0.999) + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7) out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the effective grad of adam_step to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The same array x, but with no gradient information. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The adam step delta. Examples -------- With :class:`ivy.Array` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.stop_gradient(x, preserve_type=True) - >>> print(y) - ivy.array([1., 2., 3.]) + >>> dcdw = ivy.array([1, 2, 3]) + >>> mw = ivy.ones(3) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) + >>> print(adam_step_delta) + (ivy.array([0.2020105 , 0.22187898, 0.24144873]), + ivy.array([0.99999998, 1.09999998, 1.19999998]), + ivy.array([1.00000001, 1.00300001, 1.00800001])) - >>> x = ivy.zeros((2, 3)) - >>> ivy.stop_gradient(x, preserve_type=False, out=x) - >>> print(x) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) + >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) + >>> mw = ivy.zeros((2,3)) + >>> vw = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.86 + >>> beta2 = 0.95 + >>> epsilon = 1e-6 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + (ivy.array([[ 1., 1., -1.], + [ 1., 1., 1.]]), + ivy.array([[ 0.14, 0.56, -0.42], + [ 0.28, 0.42, 0.07]]), + ivy.array([[0.05 , 0.8 , 0.45 ], + [0.2 , 0.45 , 0.0125]])) - With one :class:`ivy.Container` inputs: + >>> dcdw = ivy.array([0.1, -0.7, 2]) + >>> mw = ivy.ones(1) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3.6) + >>> out = ivy.zeros_like(dcdw) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) + >>> print(out) + ivy.array([0.17294501, 0.15770318, 0.20863818]) - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.stop_gradient(x, preserve_type=False) - >>> print(y) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + With one :class:`ivy.Container` input: + + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.array([1., 4., 9.]) + >>> vw = ivy.array([0.,]) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), + b: ivy.array([2.02, 4.82, 8.17]) + }, { + a: ivy.array([0.87, 3.61, 8.09]), + b: ivy.array([1.26, 4., 8.48]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> ivy.stop_gradient(x, preserve_type=True, out=x) - >>> print(x) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> vw = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([0., 0.626, 0.626]), + b: ivy.array([0.626, 0.626, 0.626]) + }, { + a: ivy.array([0., 0.13, 0.26]), + b: ivy.array([0.39, 0.52, 0.65]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) """ - return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) - - -# AutoGrad # + step = float(step) + mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) + dcdw_sqrd = dcdw**2 + vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) + vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 + beta1_pow = beta1**step + beta2_pow = beta2**step + alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) + return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw @handle_exceptions -@handle_device_shifting -def execute_with_gradients( - func, - xs: Union[ivy.Array, ivy.NativeArray], +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def adam_update( + w: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], + lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, - retain_grads: bool = False, - xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], - ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], -) -> Tuple[ivy.Array, ivy.Array]: + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, + stop_gradients: bool = True, + out: Optional[ivy.Array] = None, +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Call function func with input of xs variables, and return the function result - func_ret and the gradients of each output variable w.r.t each input variable, + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, using ADAM update. `[reference] + + `_ Parameters ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - xs - Variables for which to compute the function gradients with respective to. This - can be a single array or an arbitrary nest of arrays. - retain_grads - Whether to retain the gradients of the returned values. (Default value = False) - xs_grad_idxs - Indices of the input arrays to compute gradients with respect to. If None, - gradients are returned with respect to all input arrays. If ``xs`` is an - ``ivy.Array`` or ``ivy.Container``, the default value is ``None``, otherwise the - default value is ``[[0]]``. - ret_grad_idxs - Indices of the returned arrays for which to return computed gradients. If None, - gradients are returned for all returned arrays. If the returned object from the - ``func`` is an ``ivy.Array`` or ``ivy.Container``, the default value is ``None`` - otherwise the default value is ``[[0]]``. + w + Weights of the function to be updated. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. + out + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - the function result func_ret and a dictionary of gradients of each output - variable w.r.t each input variable. + The new function weights ws_new, and also new mw and vw, following the adam + updates. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = 1 + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(updated_weights) + (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), + ivy.array([0.05, 0.02, 0.01]), + ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) + + >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) + >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = 2 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> out = ivy.zeros_like(w) + >>> stop_gradients = True + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, beta2=beta2, + ... epsilon=epsilon, out=out, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ( + ivy.array([[0.92558873, 1.92558754, 2.92558718], + [3.92558694, 1.92558682, 3.92558861], + [5.92558861, 3.92558694, 1.92558718]]), + ivy.array([[0.01, 0.02, 0.03], + [0.04, 0.05, 0.01], + [0.01, 0.05, 0.03]]), + ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], + [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], + [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) + ) + + With one :class:`ivy.Container` input: + + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([0.5, 0.2, 0.4]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(0.01) + >>> step = 2 + >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(updated_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), + ... b=ivy.array([0.3,0.2,0.2])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = 3 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> stop_gradients = False + >>> lr = ivy.array(0.001) + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ({ + a: ivy.array([0.99936122, 1.99936116, 2.99936128]), + b: ivy.array([3.99936128, 4.99936104, 5.99936104]) + }, { + a: ivy.array([0.01, 0.03, 0.03]), + b: ivy.array([0.03, 0.02, 0.02]) + }, { + a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), + b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) + }) + """ + effective_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + return ( + ivy.optimizer_update( + w, effective_grads, lr, stop_gradients=stop_gradients, out=out + ), + mw, + vw, + ) + + +# AutoGrad # + + +@handle_exceptions +@handle_device_shifting +def execute_with_gradients( + func, + xs: Union[ivy.Array, ivy.NativeArray], + /, + *, + retain_grads: bool = False, + xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], + ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], +) -> Tuple[ivy.Array, ivy.Array]: + """ + Call function func with input of xs variables, and return the function result + func_ret and the gradients of each output variable w.r.t each input variable, + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + xs + Variables for which to compute the function gradients with respective to. This + can be a single array or an arbitrary nest of arrays. + retain_grads + Whether to retain the gradients of the returned values. (Default value = False) + xs_grad_idxs + Indices of the input arrays to compute gradients with respect to. If None, + gradients are returned with respect to all input arrays. If ``xs`` is an + ``ivy.Array`` or ``ivy.Container``, the default value is ``None``, otherwise the + default value is ``[[0]]``. + ret_grad_idxs + Indices of the returned arrays for which to return computed gradients. If None, + gradients are returned for all returned arrays. If the returned object from the + ``func`` is an ``ivy.Array`` or ``ivy.Container``, the default value is ``None`` + otherwise the default value is ``[[0]]``. + + Returns + ------- + ret + the function result func_ret and a dictionary of gradients of each output + variable w.r.t each input variable. Examples -------- @@ -483,77 +695,6 @@ def execute_with_gradients( ) -execute_with_gradients.computes_gradients = True - - -@handle_exceptions -def value_and_grad(func: Callable) -> Callable: - """ - Create a function that evaluates both func and the gradient of func. - - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - - Returns - ------- - ret - A function that returns both func and the gradient of func. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> grad_fn = ivy.value_and_grad(func) - >>> value_grad = grad_fn(x) - >>> print(value_grad) - (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], - [0.93333334, 0.43333334, 2.0666666 ]])) - """ - return current_backend(None).value_and_grad(func) - - -value_and_grad.computes_gradients = True - - -@handle_exceptions -def jac(func: Callable) -> Callable: - """ - Call function func, and return func's Jacobian partial derivatives. - - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - - Returns - ------- - ret - the Jacobian function - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> jac_fn = ivy.jac(func) - >>> jacobian = jac_fn(x) - >>> print(jacobian) - ivy.array([[1.53 , 0.7 , 1.67 ], - ... [0.933, 0.433, 2.07 ]]) - """ - return current_backend(None).jac(func) - - -jac.computes_gradients = True - - @handle_exceptions def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: """ @@ -585,380 +726,297 @@ def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: return current_backend(None).grad(func, argnums=argnums) -grad.computes_gradients = True - - -# Optimizer Steps # - - @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def adam_step( +def gradient_descent_update( + w: Union[ivy.Array, ivy.NativeArray], dcdw: Union[ivy.Array, ivy.NativeArray], - mw: Union[ivy.Array, ivy.NativeArray], - vw: Union[ivy.Array, ivy.NativeArray], - step: Union[int, float], + lr: Union[float, ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, + stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Compute adam step delta, given the derivatives of some cost c with respect to - weights ws, using ADAM update. `[reference] - - `_ + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws]. Parameters ---------- + w + Weights of the function to be updated. dcdw Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - mw - running average of the gradients - vw - running average of second moments of the gradients - step - training step - beta1 - gradient forgetting factor (Default value = 0.9) - beta2 - second moment of gradient forgetting factor (Default value = 0.999) - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7) + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. out - optional output array, for writing the effective grad of adam_step to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The adam step delta. + The new weights, following the gradient descent updates. Examples -------- With :class:`ivy.Array` inputs: - >>> dcdw = ivy.array([1, 2, 3]) - >>> mw = ivy.ones(3) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) - >>> print(adam_step_delta) - (ivy.array([0.2020105 , 0.22187898, 0.24144873]), - ivy.array([0.99999998, 1.09999998, 1.19999998]), - ivy.array([1.00000001, 1.00300001, 1.00800001])) - - >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) - >>> mw = ivy.zeros((2,3)) - >>> vw = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.86 - >>> beta2 = 0.95 - >>> epsilon = 1e-6 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - (ivy.array([[ 1., 1., -1.], - [ 1., 1., 1.]]), - ivy.array([[ 0.14, 0.56, -0.42], - [ 0.28, 0.42, 0.07]]), - ivy.array([[0.05 , 0.8 , 0.45 ], - [0.2 , 0.45 , 0.0125]])) + >>> w = ivy.array([[1., 2, 3], + ... [4, 6, 1], + ... [1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1], + ... [0.3, 0.6, 0.4], + ... [0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) + >>> print(new_weights) + ivy.array([[ 0.95, 1.98, 2.99], + ... [ 3.97, 5.94, 0.96], + ... [ 0.96, -0.07, 6.98]]) - >>> dcdw = ivy.array([0.1, -0.7, 2]) - >>> mw = ivy.ones(1) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3.6) - >>> out = ivy.zeros_like(dcdw) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) + >>> w = ivy.array([1., 2., 3.]) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> out = ivy.zeros_like(w) + >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) >>> print(out) - ivy.array([0.17294501, 0.15770318, 0.20863818]) + ivy.array([0.85, 1.94, 2.97]) - With one :class:`ivy.Container` input: + With one :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.array([1., 4., 9.]) - >>> vw = ivy.array([0.,]) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), - b: ivy.array([2.02, 4.82, 8.17]) - }, { - a: ivy.array([0.87, 3.61, 8.09]), - b: ivy.array([1.26, 4., 8.48]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([3.33, 5.66, 1.95]) + } With multiple :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> vw = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([0., 0.626, 0.626]), - b: ivy.array([0.626, 0.626, 0.626]) - }, { - a: ivy.array([0., 0.13, 0.26]), - b: ivy.array([0.39, 0.52, 0.65]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), + ... b=ivy.array([2., 3.42, 1.69])) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([2.88, 4.69, 1.47]) + } """ - step = float(step) - mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) - dcdw_sqrd = dcdw**2 - vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) - vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 - beta1_pow = beta1**step - beta2_pow = beta2**step - alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) - return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw + return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) -adam_step.out_index = 0 +@handle_exceptions +def jac(func: Callable) -> Callable: + """ + Call function func, and return func's Jacobian partial derivatives. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + Returns + ------- + ret + the Jacobian function -# Optimizer Updates # + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> jac_fn = ivy.jac(func) + >>> jacobian = jac_fn(x) + >>> print(jacobian) + ivy.array([[1.53 , 0.7 , 1.67 ], + ... [0.933, 0.433, 2.07 ]]) + """ + return current_backend(None).jac(func) @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def optimizer_update( +def lamb_update( w: Union[ivy.Array, ivy.NativeArray], - effective_grad: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, + max_trust_ratio: Union[int, float] = 10, + decay_lambda: float = 0, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Update weights ws of some function, given the true or effective derivatives of some - cost c with respect to ws, [dc/dw for w in ws]. + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws], by applying LAMB method. Parameters ---------- w Weights of the function to be updated. - effective_grad - Effective gradients of the cost c with respect to the weights ws, - [dc/dw for w in ws]. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). + max_trust_ratio + The maximum value for the trust ratio. (Default value = 10) + decay_lambda + The factor used for weight decay. (Default value = 0). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the optimizer updates. + The new function weights ws_new, following the LAMB updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - ivy.array([1., 2., 3.]) - - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... out=None, stop_gradients=True) - >>> print(ws_new) - ivy.array([1., 2., 3.]) + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(new_weights) + (ivy.array([0.784, 1.78 , 2.78 ]), + ... ivy.array([0.05, 0.02, 0.01]), + ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) - >>> w = ivy.array([[1., 2.], [4., 5.]]) + >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) - >>> lr = ivy.array([3e-4, 1e-2]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) - >>> print(out) - ivy.array([[0.999, 1.95], - [4., 4.92]]) - - >>> w = ivy.array([1., 2., 3.]) - >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([4., 5., 6.]) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False, out=out) + >>> stop_gradients = True + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, out=out, + ... stop_gradients=stop_gradients) >>> print(out) - ivy.array([0.999, 2. , 3. ]) + ivy.array([[ 0.639, 1.64 , 2.64 ], + ... [ 3.64 , 5.64 , 0.639], + ... [ 0.639, -0.361, 6.64 ]]) - With one :class:`ivy.Container` input: + With one :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.array([0., 0., 0.]) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([3., 4., 5.]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(1.) + >>> step = ivy.array([2]) + >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(new_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) - >>> print(w) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - """ - deltas = effective_grad * lr - w = ivy.subtract(w, deltas, out=out) - if stop_gradients: - return ivy.stop_gradient(w, preserve_type=True, out=out) - return w - - -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def gradient_descent_update( - w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - stop_gradients: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws]. - - Parameters - ---------- - w - Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The new weights, following the gradient descent updates. - - Examples - -------- - With :class:`ivy.Array` inputs: + >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), + ... b=ivy.array([3.,4.,2.])) + >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), + ... b=ivy.array([0.6,0.4,0.7])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) - >>> w = ivy.array([[1., 2, 3], - ... [4, 6, 1], - ... [1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1], - ... [0.3, 0.6, 0.4], - ... [0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 + >>> stop_gradients = True + >>> lr = ivy.array(0.5) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, + ... stop_gradients=stop_gradients) >>> print(new_weights) - ivy.array([[ 0.95, 1.98, 2.99], - ... [ 3.97, 5.94, 0.96], - ... [ 0.96, -0.07, 6.98]]) - - >>> w = ivy.array([1., 2., 3.]) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> out = ivy.zeros_like(w) - >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) - >>> print(out) - ivy.array([0.85, 1.94, 2.97]) - - With one :class:`ivy.Container` inputs: - - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([3.33, 5.66, 1.95]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), - ... b=ivy.array([2., 3.42, 1.69])) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([2.88, 4.69, 1.47]) - } + ({ + a: ivy.array([-0.708, 1.29, 3.29]), + b: ivy.array([1.45, 2.45, 0.445]) + }, { + a: ivy.array([0.02, 0.03, 0.06]), + b: ivy.array([0.06, 0.04, 0.07]) + }, { + a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), + b: ivy.array([0.00036, 0.00016, 0.00049]) + }) """ - return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) + r1 = ivy.vector_norm(w) + eff_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + if decay_lambda > 0: + r2 = ivy.vector_norm(eff_grads + decay_lambda * w) + else: + r2 = ivy.vector_norm(eff_grads) + r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) + lr = r * lr + return ( + ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), + mw, + vw, + ) @handle_exceptions @@ -1012,338 +1070,264 @@ def lars_update( ) +# Optimizer Updates # + + @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def adam_update( +def optimizer_update( w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], + effective_grad: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, using ADAM update. `[reference] - - `_ + Update weights ws of some function, given the true or effective derivatives of some + cost c with respect to ws, [dc/dw for w in ws]. Parameters ---------- w Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + effective_grad + Effective gradients of the cost c with respect to the weights ws, + [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new function weights ws_new, and also new mw and vw, following the adam - updates. + The new function weights ws_new, following the optimizer updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = 1 - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(updated_weights) - (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), - ivy.array([0.05, 0.02, 0.01]), - ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) - >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = 2 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... out=None, stop_gradients=True) + >>> print(ws_new) + ivy.array([1., 2., 3.]) + + >>> w = ivy.array([[1., 2.], [4., 5.]]) >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, beta2=beta2, - ... epsilon=epsilon, out=out, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ( - ivy.array([[0.92558873, 1.92558754, 2.92558718], - [3.92558694, 1.92558682, 3.92558861], - [5.92558861, 3.92558694, 1.92558718]]), - ivy.array([[0.01, 0.02, 0.03], - [0.04, 0.05, 0.01], - [0.01, 0.05, 0.03]]), - ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], - [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], - [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) - ) + >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) + >>> lr = ivy.array([3e-4, 1e-2]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) + >>> print(out) + ivy.array([[0.999, 1.95], + [4., 4.92]]) + + >>> w = ivy.array([1., 2., 3.]) + >>> out = ivy.zeros_like(w) + >>> effective_grad = ivy.array([4., 5., 6.]) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False, out=out) + >>> print(out) + ivy.array([0.999, 2. , 3. ]) With one :class:`ivy.Container` input: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([0.5, 0.2, 0.4]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(0.01) - >>> step = 2 - >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(updated_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.array([0., 0., 0.]) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), ... b=ivy.array([3., 4., 5.])) - >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), - ... b=ivy.array([0.3,0.2,0.2])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = 3 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> stop_gradients = False - >>> lr = ivy.array(0.001) - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ({ - a: ivy.array([0.99936122, 1.99936116, 2.99936128]), - b: ivy.array([3.99936128, 4.99936104, 5.99936104]) - }, { - a: ivy.array([0.01, 0.03, 0.03]), - b: ivy.array([0.03, 0.02, 0.02]) - }, { - a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), - b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) - }) - """ - effective_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - return ( - ivy.optimizer_update( - w, effective_grads, lr, stop_gradients=stop_gradients, out=out - ), - mw, - vw, - ) - + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) + >>> print(w) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } -adam_update.out_index = 0 + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + """ + deltas = effective_grad * lr + w = ivy.subtract(w, deltas, out=out) + if stop_gradients: + return ivy.stop_gradient(w, preserve_type=True, out=out) + return w @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def lamb_update( - w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, +@handle_device_shifting +def stop_gradient( + x: Union[ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, - max_trust_ratio: Union[int, float] = 10, - decay_lambda: float = 0, - stop_gradients: bool = True, + preserve_type: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws], by applying LAMB method. + Stop gradient computation. Parameters ---------- - w - Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). - max_trust_ratio - The maximum value for the trust ratio. (Default value = 10) - decay_lambda - The factor used for weight decay. (Default value = 0). - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. + x + Array for which to stop the gradient. + preserve_type + Whether to preserve gradient computation on ivy.Array instances. Default is + True. out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the LAMB updates. + The same array x, but with no gradient information. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(new_weights) - (ivy.array([0.784, 1.78 , 2.78 ]), - ... ivy.array([0.05, 0.02, 0.01]), - ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.stop_gradient(x, preserve_type=True) + >>> print(y) + ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 - >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, out=out, - ... stop_gradients=stop_gradients) - >>> print(out) - ivy.array([[ 0.639, 1.64 , 2.64 ], - ... [ 3.64 , 5.64 , 0.639], - ... [ 0.639, -0.361, 6.64 ]]) + >>> x = ivy.zeros((2, 3)) + >>> ivy.stop_gradient(x, preserve_type=False, out=x) + >>> print(x) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) With one :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([3., 4., 5.]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(1.) - >>> step = ivy.array([2]) - >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(new_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.stop_gradient(x, preserve_type=False) + >>> print(y) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), - ... b=ivy.array([3.,4.,2.])) - >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), - ... b=ivy.array([0.6,0.4,0.7])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> ivy.stop_gradient(x, preserve_type=True, out=x) + >>> print(x) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + """ + return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 - >>> stop_gradients = True - >>> lr = ivy.array(0.5) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, - ... stop_gradients=stop_gradients) - >>> print(new_weights) - ({ - a: ivy.array([-0.708, 1.29, 3.29]), - b: ivy.array([1.45, 2.45, 0.445]) - }, { - a: ivy.array([0.02, 0.03, 0.06]), - b: ivy.array([0.06, 0.04, 0.07]) - }, { - a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), - b: ivy.array([0.00036, 0.00016, 0.00049]) - }) + +@handle_exceptions +def value_and_grad(func: Callable) -> Callable: """ - r1 = ivy.vector_norm(w) - eff_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - if decay_lambda > 0: - r2 = ivy.vector_norm(eff_grads + decay_lambda * w) - else: - r2 = ivy.vector_norm(eff_grads) - r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) - lr = r * lr - return ( - ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), - mw, - vw, - ) + Create a function that evaluates both func and the gradient of func. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + + Returns + ------- + ret + A function that returns both func and the gradient of func. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> grad_fn = ivy.value_and_grad(func) + >>> value_grad = grad_fn(x) + >>> print(value_grad) + (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], + [0.93333334, 0.43333334, 2.0666666 ]])) + """ + return current_backend(None).value_and_grad(func) +execute_with_gradients.computes_gradients = True +value_and_grad.computes_gradients = True +jac.computes_gradients = True +grad.computes_gradients = True +adam_step.out_index = 0 +adam_update.out_index = 0 lamb_update.out_index = 0 +_check_if_empty = ( + lambda idxs: not isinstance(idxs, list) + or np.asarray(idxs, dtype="object").size == 0 +) +_idxs_to_str = lambda idxs: [ + "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) +] +_to_ivy = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) +_non_finite_to_zero = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) diff --git a/ivy/functional/ivy/layers.py b/ivy/functional/ivy/layers.py index c8096d4d1fe26..1a239557e8b53 100644 --- a/ivy/functional/ivy/layers.py +++ b/ivy/functional/ivy/layers.py @@ -20,6 +20,111 @@ ) from ivy.utils.exceptions import handle_exceptions + +# --- Helpers --- # +# --------------- # + + +def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): + kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) + if padding == "SAME": + dim_size = dim_size * stride_size + else: + dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) + return dim_size + + +def _depth_max_pooling_helper( + x_shape, kernel, strides, dims, data_format="channel_last" +): + # Determine depth pooling. + # We assume that the kernel and the data have the same data_format. + depth_pooling = False + CHANNEL_LAST = "channel_last" + channel_idx = -1 if data_format == CHANNEL_LAST else 1 + if len(kernel) == dims + 2: + spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] + if kernel[channel_idx] != 1: + depth_pooling = True + if any(i != 1 for i in spatial_kernel): + raise NotImplementedError( + "MaxPooling supports exactly one of pooling across" + " depth or pooling across width/height." + ) + if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to equal the depth" + " stride" + ) + if x_shape[channel_idx] % kernel[channel_idx] != 0: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to evenly divide" + " the input depth" + ) + kernel = [kernel[channel_idx], *[1] * (dims - 1)] + strides = [strides[channel_idx], *[1] * (dims - 1)] + else: + kernel = spatial_kernel + if len(strides) == dims + 2: + strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] + return kernel, strides, depth_pooling + + +def _get_num_padded_values(i, p, n, k, s): + """ + Get number of padded values in a specific window. + + Parameters + ---------- + i window index + p total amount of padding + n input size + k kernel size + s stride + + Returns + ------- + number of padded values in a particular window represented by i + """ + current_index = s * i + left_padding = p // 2 + return max(0, left_padding - current_index) + max( + 0, current_index + k - n - left_padding + ) + + +def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): + if dims == 1: + if data_format == "channel_first": + return "NCW" + else: + return "NWC" + if dims == 2: + if data_format == "channel_first": + return "NCHW" + else: + return "NHWC" + elif dims == 3: + if data_format == "channel_first": + return "NCDHW" + else: + return "NDHWC" + + +# Helpers # + + +def _handle_padding(x, strides, filters, padding): + if isinstance(padding, str) and padding.upper() == "SAME": + if x % strides == 0: + pad = max(filters - strides, 0) + else: + pad = max(filters - (x % strides), 0) + else: + pad = 0 + return pad + + # Extra # # ------# @@ -72,34 +177,134 @@ def _in_projection( ) -# Linear # +def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): + if isinstance(kernel, int): + kernel = (kernel,) * dims + elif len(kernel) == 1: + kernel = (kernel[0],) * dims + elif (len(kernel) != dims) and (len(kernel) != dims + 2): + raise ValueError( + "The kernel should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" + ) + + if isinstance(strides, int): + strides = (strides,) * dims + elif len(strides) == 1: + strides = (strides[0],) * dims + elif (len(strides) != dims) and (len(strides) != dims + 2): + raise ValueError( + "The stride should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" + ) + + if isinstance(padding, int): + padding = [(padding,) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == 1: + padding = [(padding[0],) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == dims: + padding = [(padding[i],) * 2 for i in range(dims)] + elif isinstance(padding, list) and len(padding) == dims: + if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): + raise ValueError("Explicit padding must be a list of tuple of two integers") + if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: + raise ValueError( + f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" + ) + + if isinstance(dilation, int): + dilation = (dilation,) * dims + elif len(dilation) == 1: + dilation = (dilation[0],) * dims + elif len(dilation) != dims: + raise ValueError( + f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" + ) + if min(dilation) < 1: + raise ValueError("All values of `dilation` must be positive") + + # Other errors + if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: + raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") + assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" + + # Account for dilation when padding > kernel/2. Not the case in torch by default. + new_kernel = tuple( + [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] + ) + if isinstance(padding, list) and len(padding) == len(new_kernel): + ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) + + return kernel, strides, padding, dilation + + +# --- Main --- # +# ------------ # + + @handle_exceptions -@handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@inputs_to_native_shapes @handle_array_function -def linear( +def conv( x: Union[ivy.Array, ivy.NativeArray], - weight: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], /, *, + transpose: bool = False, + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. - The operation also supports batching of the weight matrices. This is useful if a - batch of different network parameters are to be represented. + """ + Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D + input x respectively and filters arrays. Parameters ---------- x - The input x to compute linear transformation on. - *[outer_batch_shape,inner_batch_shape,in_features]* - weight - The weight matrix. *[outer_batch_shape,out_features,in_features]* + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + transpose + True for computing transpose convolution, and False for dilated convolution. + When True, `x_dilations` must be 1 (the default). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output (Default value = None) + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) bias - The bias vector, default is ``None``. *[outer_batch_shape,out_features]* + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -107,179 +312,92 @@ def linear( Returns ------- ret - Result array of the linear transformation. - *[outer_batch_shape,inner_batch_shape,out_features]* - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 0., 0.]]) - >>> y = ivy.linear(x, w) - >>> print(y) - ivy.array([1.]) - - >>> x = ivy.array([[0.666, -0.4269, 1.911]]) - >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) - >>> y = ivy.zeros((1, 2)) - >>> ivy.linear(x, w, out=y) - >>> print(y) - ivy.array([[0.666, 1.91 ]]) - - >>> x = ivy.array([[1.546, 5.234, 6.487], - ... [0.157, 5.753, 4.52], - ... [5.165, 3.159, 7.101]]) - >>> w = ivy.array([[1.545, 2.547, 3.124], - ... [5.852, 8.753, 6.963]]) - >>> b = ivy.array([-1., 1.]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.linear(x, w, bias=b, out=y) - >>> print(y) - ivy.array([[ 34.98495483, 101.0293808 ], - [ 28.0159359 , 83.74752808], - [ 37.20942307, 108.3205719 ]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [4., 5., 6.]]), - ... b=ivy.array([1.1, 2.2, 3.3])) - >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [-1., 1., 2.]]), - ... b=ivy.array([[0., -1., 1.], - ... [0., 1., 1.]])) - >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) - >>> y = ivy.linear(x, w, bias=b) - >>> print(y) - { - a: ivy.array([[15., 6.], - [33., 12.]]), - b: ivy.array([2.1, 6.5]) - } - - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], - ... [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], - ... [7., 13., 17.]])) - >>> w = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.]]) - >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), - ... b=ivy.array([1., 1., 0.])) - >>> ivy.linear(x, w, bias=b, out=x) - >>> print(x) - { - a: ivy.array([[16.4, 35.2, 54.], - [155., 352., 549.]]), - b: ivy.array([[15.1, 32., 47.9], - [85., 196., 306.]]) - } - + The result of the transpose or dilated convolution operation. """ - outer_batch_shape = list(weight.shape[:-2]) - num_outer_batch_dims = len(outer_batch_shape) - inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) - num_inner_batch_dims = len(inner_batch_shape) - num_out_feats, num_in_feats = list(weight.shape[-2:]) - - # OBS x IBS x OF - y = ivy.matmul( - x, - ivy.swapaxes( - ivy.reshape( - weight, - outer_batch_shape - + [1] * max(num_inner_batch_dims - 1, 0) - + [num_out_feats, num_in_feats], - ), - -1, - -2, - ), - ) - - if ivy.exists(bias): - # OBS x [1]*len(IBS) x OF - bias_broadcast = ivy.reshape( - bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] + if transpose: + return conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, + ) + else: + return conv_general_dilated( + x, + filters, + strides, + padding, + dims=dims, + data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - - # OBS x IBS x OF - y = y + bias_broadcast - - if ivy.exists(out): - return ivy.inplace_update(out, y) - return y - - -linear.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -# Dropout # +# Convolutions # @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def dropout( +@handle_device_shifting +def conv1d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - scale: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, + data_format: str = "NWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int]] = 1, + dilations: Union[int, Tuple[int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly setting a fraction of input tensor to zeroes with probability. - - `prob` at each update during training time to prevent possible overfitting. - The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that - overall sum is unchanged at training time and inference time. + Compute a 1-D convolution given 3-D input x and filters arrays. Parameters ---------- x - The input array x to perform dropout on. - prob - The probability of zeroing out each array element, float between 0 and 1. - scale - Whether to scale the output by `1/(1-prob)`. Default is ``True``. - dtype - output array data type. If dtype is None, the output array data type - must be inferred from x. Default is ``None``. - training - Turn on dropout if training, turn off otherwise. Default is ``True``. - seed - Set a default seed for random number generating (for reproducibility). Default - is ``None``. - noise_shape - a sequence representing the shape of the binary dropout mask that will be - multiplied with the input. A shape dimension set to None means that a different - mask value will be applied to each element of the input across that dimension. A - dimension set to 1 means the same mask value will be applied to all elements of - the input across that dimension. + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + filters + Convolution filters *[fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + data_format + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -287,7 +405,7 @@ def dropout( Returns ------- ret - Result array after dropout is performed. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -297,175 +415,95 @@ def dropout( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - ivy.array([[ 1.42857146, 2.85714293, 4.28571415], - [ 0. , 7.14285755, 8.5714283 ], - [10. , 11.4285717 , 0. ], - [14.2857151 , 0. , 17.1428566 ]]) - - - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - ivy.array([[ 0. , 5.19999981], - [ 0. , 0. ], - [ 0. , 17.39999962]]) + >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC + >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO + >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) + >>> print(result) + ivy.array([[[0.], [3.], [0.]]]) - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3,scale=False) - >>> print(y) - ivy.array([[ 1., 2., 3.], - [ 4., 5., 0.], - [ 7., 0., 9.], - [10., 11., 0.]]) + With :class:`ivy.NativeArray` input: - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5,scale=False) - >>> print(y) - ivy.array([[0., 2.6], - [0., 0. ], - [0., 8.7]]) + >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) + >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) + >>> result = ivy.conv1d(x, filters, (2,),'VALID') + >>> print(result) + ivy.array([[[3., 1.], + ... [7., 5.]]]) - With :class:`ivy.Container` input: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], + ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), + ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) + >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) + >>> result = ivy.conv1d(x, filters, 3, 'VALID') + >>> print(result) { - a: ivy.array([[0., 0., 4.28571415], - [5.71428585, 7.14285755, 0.]]), - b: ivy.array([0., 11.4285717, 12.8571434]) - } - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - { - a: ivy.array([[0., 4.4000001, 6.5999999], - [22., 44., 0.]]), - b: ivy.array([[2.49000001, 0.55599999, 8.21000004], - [14., 0., 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - { - a: ivy.array([[0., 0., 3.], - [4., 5., 0.]]), - b: ivy.array([0., 8., 9.]) - } - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - { - a: ivy.array([[0., 2.2, 3.3], - [11., 22., 0.]]), - b: ivy.array([[1.245, 0.278, 4.105], - [7., 0., 0.]]) + a: ivy.array([[[6., 7.9, 1.2], + ... [15.6, 11.7, 6.1]]]), + ... b: ivy.array([[[15.4, 14.3, 8.8]]]) } """ - if prob == 0 or not training: - if dtype is not None: - x = ivy.astype(x, dtype) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) - if noise_shape is None: - noise_shape = x.shape - else: - noise_shape = list(noise_shape) - for i, v in enumerate(noise_shape): - if v is None: - noise_shape[i] = x.shape[i] - mask = ivy.where( - ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) - < prob, - 0.0, - 1.0, + return current_backend(x).conv1d( + x, + filters, + strides, + padding, + data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - x = x * mask - if scale: - x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) - - -dropout.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -# Attention # @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back @handle_array_function -def scaled_dot_product_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Union[ivy.Array, ivy.NativeArray], - value: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def conv1d_transpose( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int]], + padding: str, /, *, - scale: Optional[float] = None, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - dropout_p: Optional[float] = 0.0, - is_causal: Optional[bool] = False, - training: Optional[bool] = False, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NWC", + dilations: Union[int, Tuple[int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply scaled dot product attention to inputs x using optional mask. + Compute a 1-D transpose convolution given 3-D input x and filters arrays. Parameters ---------- - query - The queries input array. The shape of queries input array should be in - *[batch_shape,num_queries,feat_dim]*. The queries input array should have the - same size as keys and values. - key - The keys input array. The shape of keys input array should be in - *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same - size as queries and values. - value - The values input array. The shape of values input should be in - *[batch_shape,num_keys,feat_dim]*. The values input array should have the same - size as queries and keys. - scale - The scale float value. - The scale float value is used to scale the query-key pairs before softmax. - mask - The mask input array. The mask to apply to the query-key values. Default is - None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. - dropout_p - Specifies the dropout probablity, if greater than 0.0, dropout is applied - is_causal - If true, assumes causal attention masking - and errors if both `mask` and `is_causal` are set. - training - If True, dropout is used, otherwise dropout is not activated. + x + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + filters + Convolution filters *[fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) + data_format + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -473,9 +511,7 @@ def scaled_dot_product_attention( Returns ------- ret - The output following application of scaled dot-product attention. - The output array is the weighted sum produced by the attention score and value. - The shape of output array is *[batch_shape,num_queries,feat_dim]* . + The result of the transpose convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -485,403 +521,227 @@ def scaled_dot_product_attention( -------- With :class:`ivy.Array` input: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) - - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 6) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 64) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) + >>> y = ivy.zeros((1, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 32) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + With :class:`ivy.NativeArray` input: - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) + >>> filters = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 512, 32) - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) + With one :class:`ivy.Container` input: - ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) + >>> x = ivy.full((1, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) + >>> print(y.shape) + { + a: ivy.Shape(1, 10, 1), + b: ivy.Shape(1, 10, 1) + } - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + With multiple :class:`ivy.Container` inputs: - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) - - With :class:`ivy.Container` input: - - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a: ivy.array([[[5.19999981, 1.], - ... [2.59249449, 2.68226194], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[0.2, 1.], - ... [2.19603825, 2.9960382], - ... [4.4000001, 5.5999999]]]) + a: { + c: ivy.Shape(1, 28, 3), + d: ivy.Shape(1, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 3), + d: ivy.Shape(1, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + } } + """ + return current_backend(x).conv1d_transpose( + x, + filters, + strides, + padding, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + bias=bias, + out=out, + ) - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> mask = ivy.Container( - ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), - ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) - ... ) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) - { - a: ivy.array([[[4.26894283, 5.40236187], - ... [4.39999437, 5.59999037], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[4.35046196, 5.54282808], - ... [4.39989519, 5.5998764], - ... [4.4000001, 5.5999999]]]) - } - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def conv2d( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int]] = 1, + dilations: Union[int, Tuple[int, int]] = 1, + bias: Optional[ivy.Array] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D convolution given 4-D input x and filters arrays. - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + filters + Convolution filters *[fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + data_format + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", + input data formats, while "channel_last" corresponds to "HWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the convolution operation. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]]) + >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], + ... [[[0.]],[[1.]], [[0.]]], + ... [[[0.]],[[1.]], [[0.]]]]) + >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) >>> print(result) + ivy.array([[ + [[2.],[4.],[6.]], + [[3.],[6.],[9.]], + [[2.],[4.],[6.]] + ]]) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + With one :class:`ivy.Container` input: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) - >>> print(out) - ivy.array([[[4.03946018, 5.0280633 ], - ... [4.29981947, 5.29981089], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]])) + >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) + >>> print(result) + { + a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) + } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With multiple :class:`ivy.Container` inputs: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), + ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[1, 1, 1], + ... [0, 1, 1], + ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') >>> print(result) { - a: ivy.array([[[0.40000001, 1.29999995], - ... [2.06345534, 2.9634552], - ... [4.30000019, 5.30000019]]]), - b: ivy.array([[[0.40000001, 1.29999995], - ... [2.19336844, 3.09336829], - ... [4.30000019, 5.30000019]]]) + a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), + b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), + c:ivy.array([[[[2.], [0.], [0.]], + [[1.], [3.], [0.]], + [[0.], [1.], [2.]] + ]]) } + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... mask=mask, - ... dropout_p=0.1, - ... training=True) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[2, 0, 1], + ... [1, 3, 1], + ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') >>> print(result) { - a: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]), - b: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) + a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), + b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) } """ - ivy.assertions.check_all( - (not is_causal) or (is_causal and mask is None), - "is_causal and attn_mask cannot be set at the same time", - ) - embed_dim = query.shape[-1] - scale = 1 / (embed_dim**0.5) if not scale else scale - sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale - sim = ivy.dropout(sim, dropout_p, training=training) - if ivy.exists(mask): - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, - ) - elif is_causal: - L = query.shape[-2] # Source sequence length - S = key.shape[-2] # Target sequence length - mask = ivy.tril(ivy.ones((L, S)), k=0) - mask = ivy.astype(mask, ivy.bool) - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, - ) - attn = ivy.softmax(sim, axis=-1) - result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) - return result if not ivy.exists(out) else ivy.inplace_update(out, result) - - -@handle_exceptions -@handle_nestable -@handle_out_argument -# @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def multi_head_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - /, - *, - num_heads: Optional[int] = 8, - scale: Optional[float] = None, - attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - is_causal: Optional[bool] = False, - return_attention_weights: Optional[bool] = False, - average_attention_weights: Optional[bool] = True, - dropout: Optional[float] = 0.0, - training: Optional[bool] = False, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Apply multi-head attention to inputs x. This is an implementation of multi-headed - attention as described in the paper "Attention is all you Need" (Vaswani et al., - 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each - timestep in `query` attends to the corresponding sequence in `key`, and returns a - fixed-width vector. This layer first projects `query`, `key` and `value`. These are - (effectively) a list of tensors of length `num_attention_heads`, where the - corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, - , key_dim)`, `(batch_size, , - value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are - softmaxed to obtain attention probabilities. The value tensors are then interpolated - by these probabilities, then concatenated back to a single tensor. Finally, the - result tensor with the last dimension as value_dim can take an linear projection and - return. - - Parameters - ---------- - query - query embeddings *[batch_shape,num_queries,query_dim]*. - key - key embeddings *[batch_shape,num_queries,key_dim]*. - value - value embeddings *[batch_shape,num_queries,value_dim]*. - num_heads - The number of attention heads to use. - scale - The value by which to scale the query-key similarity measure before softmax. - attention_mask - The mask to apply to the query-key values. Default is ``None``. - *[batch_shape,num_queries,num_keys]*. - in_proj_weights - The weights used to project query, key and value *[3*E, E]. - q_proj_weights - The weights used to project query if in_proj_weights is None *[new_E, E]. - k_proj_weights - The weights used to project key if in_proj_weights is None *[new_E, E]. - v_proj_weights - The weights used to project value if in_proj_weights is None *[new_E, E]. - out_proj_weights - The weights used to project the output. - in_proj_bias - The bias used when projecting with query, key and value. - out_proj_bias - The bias used when projecting the output. - is_causal - If True, Uses a causal attention mask and ignores provided attention_mask. - return_attention_weights - If True, returns attention_weights alongside the output - as a tuple (output, attenion_weights). Defaults to `False`. - average_attention_weights - If true, indicates that the returned ``attention_weights`` should be averaged - across heads. Otherwise, ``attention_weights`` are provided separately per head. - Note that this flag only has an effect when ``return_attention_weights=True``. - Default: ``True`` (i.e. average weights across heads) - dropout - Specifies the dropout probablity, dropout is applied to attention_weights. - training - If True, dropout is used, otherwise dropout is not activated. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The output following application of multi-head attention. - *[batch_shape,num_queries,out_feat_dim]* if input is batched - otherwise *[num_queries, out_feat_dim] - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - num_dims = query.ndim - ivy.assertions.check_all( - num_dims > 1 and num_dims < 4, - "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" - f" input), got {num_dims}", - ) - if key is None and value is None: - key = value = query - if num_dims == 2: - query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] - if ivy.exists(in_proj_weights): - q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) - elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): - if ivy.exists(in_proj_bias): - b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) - else: - b_q = b_k = b_v = None - q, k, v = ( - ivy.linear(query, q_proj_weights, bias=b_q), - ivy.linear(key, k_proj_weights, bias=b_k), - ivy.linear(value, v_proj_weights, bias=b_v), - ) - else: - q, k, v = query, key, value - batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] - k_seq_length = k.shape[1] - ivy.assertions.check_true( - emb_dim % num_heads == 0, "features must be divisible by number of heads" - ) - dims_per_head = emb_dim // num_heads - # isolate heads - q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 3, 1) - ) - v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - # perform bmm - attn_scores = ivy.matmul(q, k) - # scale - scale = 1 / (dims_per_head**0.5) if not scale else scale - attn_scores *= scale - # apply attention mask - if ivy.exists(attention_mask) or is_causal: - if is_causal: - # create causal mask - attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) - attention_mask = attention_mask.astype("bool") - attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) - # perform softmax - attn_weights = ivy.softmax(attn_scores, axis=-1) - # perform dropout - attn_weights = ivy.dropout(attn_weights, dropout, training=training) - # bmm with values - attention_out = ivy.matmul(attn_weights, v) - attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( - (batch_size, q_seq_length, -1) + return current_backend(x).conv2d( + x, + filters, + strides, + padding, + data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - # proj out if out_proj_weight exists - if ivy.exists(out_proj_weights): - attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) - # if input was unbatched, unbatchify the output - if num_dims == 2: - attention_out = attention_out.squeeze(axis=0) - if return_attention_weights: - if average_attention_weights: - attn_weights = attn_weights.mean(axis=1) - if num_dims == 2: - attn_weights = attn_weights.squeeze(axis=0) - return attention_out, attn_weights - else: - return attention_out - - -# Convolutions # @handle_exceptions @@ -889,46 +749,47 @@ def multi_head_attention( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv1d( +def conv2d_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int]], + padding: str, /, *, - data_format: str = "NWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int]] = 1, - dilations: Union[int, Tuple[int]] = 1, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D convolution given 3-D input x and filters arrays. + Compute a 2-D transpose convolution given 4-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. filters - Convolution filters *[fw,d_in,d_out]*. + Convolution filters *[fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". - x_dilations + Either "channel_first" or "channel_last". "channel_first" corresponds to + "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . + x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations The dilation factor for each dimension of input. (Default value = 1) @@ -941,7 +802,7 @@ def conv1d( Returns ------- ret - The result of the convolution operation. + The result of the transpose convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -950,44 +811,72 @@ def conv1d( Examples -------- With :class:`ivy.Array` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 56, 6) - >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC - >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO - >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) - >>> print(result) - ivy.array([[[0.], [3.], [0.]]]) - - With :class:`ivy.NativeArray` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) + >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 128, 64) - >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) - >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) - >>> result = ivy.conv1d(x, filters, (2,),'VALID') - >>> print(result) - ivy.array([[[3., 1.], - ... [7., 5.]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) + >>> y = ivy.zeros((1, 258, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) + >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 258, 32) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With one :class:`ivy.Container` inputs: + >>> x = ivy.full((1, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) + >>> print(y.shape) + { + a: ivy.Shape(1, 10, 10, 1), + b: ivy.Shape(1, 10, 10, 1) + } - >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], - ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), - ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) - >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) - >>> result = ivy.conv1d(x, filters, 3, 'VALID') - >>> print(result) + With multiple :class:`ivy.Container` inputs: + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) { - a: ivy.array([[[6., 7.9, 1.2], - ... [15.6, 11.7, 6.1]]]), - ... b: ivy.array([[[15.4, 14.3, 8.8]]]) + a: { + c: ivy.Shape(1, 28, 28, 3), + d: ivy.Shape(1, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 56, 3), + d: ivy.Shape(1, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + } } """ - return current_backend(x).conv1d( + return current_backend(x).conv2d_transpose( x, filters, strides, padding, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -999,47 +888,52 @@ def conv1d( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv1d_transpose( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: str, +def conv3d( + x: Union[ivy.Array, ivy.NativeArray, ivy.Container], + filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], + strides: Union[int, Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NWC", - dilations: Union[int, Tuple[int]] = 1, + data_format: str = "NDHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int, int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D transpose convolution given 3-D input x and filters arrays. + Compute a 3-D convolution given 5-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds + to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) dilations The dilation factor for each dimension of input. (Default value = 1) bias - Bias array of shape *[d_out]*. + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1047,7 +941,7 @@ def conv1d_transpose( Returns ------- ret - The result of the transpose convolution operation. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1057,139 +951,106 @@ def conv1d_transpose( -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 6) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 64) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) - >>> y = ivy.zeros((1, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 32) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) - >>> filters = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 512, 32) + >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) + >>> filters = ivy.array([[[0.,1.,0.], + ... [0.,1.,0.], + ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) + >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) + >>> print(result) + ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) With one :class:`ivy.Container` input: - >>> x = ivy.full((1, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) - >>> print(y.shape) + >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) + >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 2, 'SAME') + >>> print(result) { - a: ivy.Shape(1, 10, 1), - b: ivy.Shape(1, 10, 1) + a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) } - With multiple :class:`ivy.Container` inputs: + With multiple :class:`ivy.Container` input: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) + >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 3, 5, 5, 1]), + ... b = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 5, 32 ,32, 1]), + ... c = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 32, 32, 32, 1])) + >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 1, 'SAME') + >>> print(result.cont_shapes) { - a: { - c: ivy.Shape(1, 28, 3), - d: ivy.Shape(1, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 3), - d: ivy.Shape(1, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - } + a: ivy.Shape(1, 3, 5, 5, 3), + b: ivy.Shape(1, 5, 32, 32, 3), + c: ivy.Shape(1, 32, 32, 32, 3) } """ - return current_backend(x).conv1d_transpose( + return current_backend(x).conv3d( x, filters, strides, padding, - output_shape=output_shape, data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, ) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv2d( +def conv3d_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int, int]], + padding: str, /, *, - data_format: str = "NHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int]] = 1, - dilations: Union[int, Tuple[int, int]] = 1, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NDHWC", + dilations: Union[int, Tuple[int, int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D convolution given 4-D input x and filters arrays. + Compute a 3-D transpose convolution given 5-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", - input data formats, while "channel_last" corresponds to "HWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). dilations The dilation factor for each dimension of input. (Default value = 1) bias - Bias array of shape *[d_out]*. + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1197,83 +1058,85 @@ def conv2d( Returns ------- ret - The result of the convolution operation. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The result of the transpose convolution operation. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]]) - >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], - ... [[[0.]],[[1.]], [[0.]]], - ... [[[0.]],[[1.]], [[0.]]]]) - >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) - >>> print(result) - ivy.array([[ - [[2.],[4.],[6.]], - [[3.],[6.],[9.]], - [[2.],[4.],[6.]] - ]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 6, 56, 56, 6) - With one :class:`ivy.Container` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) + >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') + >>> print(y.shape) + ivy.Shape(1, 9, 258, 258, 32) - >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]])) - >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) - >>> print(result) + With :class:`ivy.Container` inputs: + + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) + a: { + c: ivy.Shape(1, 6, 28, 28, 3), + d: ivy.Shape(1, 6, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 6, 56, 56, 3), + d: ivy.Shape(1, 6, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + } } - With multiple :class:`ivy.Container` inputs: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), - ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[1, 1, 1], - ... [0, 1, 1], - ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) { - a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), - b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), - c:ivy.array([[[[2.], [0.], [0.]], - [[1.], [3.], [0.]], - [[0.], [1.], [2.]] - ]]) + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[2, 0, 1], - ... [1, 3, 1], - ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) + >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) { - a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), - b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) } """ - return current_backend(x).conv2d( + return current_backend(x).conv3d_transpose( x, filters, strides, padding, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1285,50 +1148,58 @@ def conv2d( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv2d_transpose( +def conv_general_dilated( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: str, + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - bias: Optional[ivy.Array] = None, + dims: int = 2, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D transpose convolution given 4-D input x and filters arrays. + Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively + and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to - "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations - The dilation factor for each dimension of input. (Default value = 1) + The dilation factor for each dimension of filter. (Default value = 1) bias Bias array of shape *[d_out]*. out @@ -1339,80 +1210,17 @@ def conv2d_transpose( ------- ret The result of the transpose convolution operation. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 56, 6) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) - >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 128, 64) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) - >>> y = ivy.zeros((1, 258, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) - >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 258, 32) - - With one :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) - >>> print(y.shape) - { - a: ivy.Shape(1, 10, 10, 1), - b: ivy.Shape(1, 10, 10, 1) - } - - With multiple :class:`ivy.Container` inputs: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) - { - a: { - c: ivy.Shape(1, 28, 28, 3), - d: ivy.Shape(1, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 56, 3), - d: ivy.Shape(1, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - } - } """ - return current_backend(x).conv2d_transpose( + return current_backend(x).conv_general_dilated( x, filters, strides, padding, - output_shape=output_shape, + dims=dims, data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1424,29 +1232,106 @@ def conv2d_transpose( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def depthwise_conv2d( +def conv_general_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: str, /, *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + feature_group_count: int = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D depthwise convolution given 4-D input ``x`` and filters arrays. + Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x + respectively and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in]*. (d_in must be the same as d from x) + Convolution filters *[fd,fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output. + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + dilations + The dilation factor for each dimension of input. (Default value = 1) + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + bias + Bias array of shape *[d_out]*. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the transpose convolution operation. + """ + return current_backend(x).conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, + ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def depthwise_conv2d( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D depthwise convolution given 4-D input ``x`` and filters arrays. + + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + filters + Convolution filters *[fh,fw,d_in]*. (d_in must be the same as d from x) strides The stride of the sliding window for each dimension of input. padding @@ -1558,57 +1443,56 @@ def depthwise_conv2d( ) +# Dropout # + + @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv3d( - x: Union[ivy.Array, ivy.NativeArray, ivy.Container], - filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], - strides: Union[int, Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], +def dropout( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, /, *, - data_format: str = "NDHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int, int, int]] = 1, - bias: Optional[ivy.Array] = None, + scale: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 3-D convolution given 5-D input x and filters arrays. + Randomly setting a fraction of input tensor to zeroes with probability. + + `prob` at each update during training time to prevent possible overfitting. + The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that + overall sum is unchanged at training time and inference time. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds - to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]* + The input array x to perform dropout on. + prob + The probability of zeroing out each array element, float between 0 and 1. + scale + Whether to scale the output by `1/(1-prob)`. Default is ``True``. + dtype + output array data type. If dtype is None, the output array data type + must be inferred from x. Default is ``None``. + training + Turn on dropout if training, turn off otherwise. Default is ``True``. + seed + Set a default seed for random number generating (for reproducibility). Default + is ``None``. + noise_shape + a sequence representing the shape of the binary dropout mask that will be + multiplied with the input. A shape dimension set to None means that a different + mask value will be applied to each element of the input across that dimension. A + dimension set to 1 means the same mask value will be applied to all elements of + the input across that dimension. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1616,7 +1500,7 @@ def conv3d( Returns ------- ret - The result of the convolution operation. + Result array after dropout is performed. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1626,106 +1510,142 @@ def conv3d( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) - >>> filters = ivy.array([[[0.,1.,0.], - ... [0.,1.,0.], - ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) - >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) - >>> print(result) - ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3) + >>> print(y) + ivy.array([[ 1.42857146, 2.85714293, 4.28571415], + [ 0. , 7.14285755, 8.5714283 ], + [10. , 11.4285717 , 0. ], + [14.2857151 , 0. , 17.1428566 ]]) - With one :class:`ivy.Container` input: - >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) - >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + ivy.array([[ 0. , 5.19999981], + [ 0. , 0. ], + [ 0. , 17.39999962]]) + + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3,scale=False) + >>> print(y) + ivy.array([[ 1., 2., 3.], + [ 4., 5., 0.], + [ 7., 0., 9.], + [10., 11., 0.]]) + + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5,scale=False) + >>> print(y) + ivy.array([[0., 2.6], + [0., 0. ], + [0., 8.7]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) { - a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) + a: ivy.array([[0., 0., 4.28571415], + [5.71428585, 7.14285755, 0.]]), + b: ivy.array([0., 11.4285717, 12.8571434]) } - With multiple :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 4.4000001, 6.5999999], + [22., 44., 0.]]), + b: ivy.array([[2.49000001, 0.55599999, 8.21000004], + [14., 0., 0.]]) + } - >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 3, 5, 5, 1]), - ... b = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 5, 32 ,32, 1]), - ... c = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 32, 32, 32, 1])) - >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 1, 'SAME') - >>> print(result.cont_shapes) + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) { - a: ivy.Shape(1, 3, 5, 5, 3), - b: ivy.Shape(1, 5, 32, 32, 3), - c: ivy.Shape(1, 32, 32, 32, 3) + a: ivy.array([[0., 0., 3.], + [4., 5., 0.]]), + b: ivy.array([0., 8., 9.]) } - """ - return current_backend(x).conv3d( - x, - filters, - strides, - padding, - data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, - ) + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 2.2, 3.3], + [11., 22., 0.]]), + b: ivy.array([[1.245, 0.278, 4.105], + [7., 0., 0.]]) + } + """ + if prob == 0 or not training: + if dtype is not None: + x = ivy.astype(x, dtype) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + if noise_shape is None: + noise_shape = x.shape + else: + noise_shape = list(noise_shape) + for i, v in enumerate(noise_shape): + if v is None: + noise_shape[i] = x.shape[i] + mask = ivy.where( + ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) + < prob, + 0.0, + 1.0, + ) + x = x * mask + if scale: + x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + +# Linear # @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv3d_transpose( +def linear( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int, int]], - padding: str, + weight: Union[ivy.Array, ivy.NativeArray], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NDHWC", - dilations: Union[int, Tuple[int, int, int]] = 1, - bias: Optional[ivy.Array] = None, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute a 3-D transpose convolution given 5-D input x and filters arrays. + """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. + The operation also supports batching of the weight matrices. This is useful if a + batch of different network parameters are to be represented. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) - data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - dilations - The dilation factor for each dimension of input. (Default value = 1) + The input x to compute linear transformation on. + *[outer_batch_shape,inner_batch_shape,in_features]* + weight + The weight matrix. *[outer_batch_shape,out_features,in_features]* bias - Bias array of shape *[d_out]* + The bias vector, default is ``None``. *[outer_batch_shape,out_features]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1733,228 +1653,303 @@ def conv3d_transpose( Returns ------- ret - The result of the transpose convolution operation. + Result array of the linear transformation. + *[outer_batch_shape,inner_batch_shape,out_features]* + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 6, 56, 56, 6) + >>> x = ivy.array([1., 2., 3.]) + >>> w = ivy.array([[1., 0., 0.]]) + >>> y = ivy.linear(x, w) + >>> print(y) + ivy.array([1.]) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) - >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') - >>> print(y.shape) - ivy.Shape(1, 9, 258, 258, 32) + >>> x = ivy.array([[0.666, -0.4269, 1.911]]) + >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) + >>> y = ivy.zeros((1, 2)) + >>> ivy.linear(x, w, out=y) + >>> print(y) + ivy.array([[0.666, 1.91 ]]) - With :class:`ivy.Container` inputs: + >>> x = ivy.array([[1.546, 5.234, 6.487], + ... [0.157, 5.753, 4.52], + ... [5.165, 3.159, 7.101]]) + >>> w = ivy.array([[1.545, 2.547, 3.124], + ... [5.852, 8.753, 6.963]]) + >>> b = ivy.array([-1., 1.]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.linear(x, w, bias=b, out=y) + >>> print(y) + ivy.array([[ 34.98495483, 101.0293808 ], + [ 28.0159359 , 83.74752808], + [ 37.20942307, 108.3205719 ]]) - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [4., 5., 6.]]), + ... b=ivy.array([1.1, 2.2, 3.3])) + >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [-1., 1., 2.]]), + ... b=ivy.array([[0., -1., 1.], + ... [0., 1., 1.]])) + >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) + >>> y = ivy.linear(x, w, bias=b) + >>> print(y) { - a: { - c: ivy.Shape(1, 6, 28, 28, 3), - d: ivy.Shape(1, 6, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 6, 56, 56, 3), - d: ivy.Shape(1, 6, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - } + a: ivy.array([[15., 6.], + [33., 12.]]), + b: ivy.array([2.1, 6.5]) } With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], + ... [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], + ... [7., 13., 17.]])) + >>> w = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.]]) + >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), + ... b=ivy.array([1., 1., 0.])) + >>> ivy.linear(x, w, bias=b, out=x) + >>> print(x) { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) + a: ivy.array([[16.4, 35.2, 54.], + [155., 352., 549.]]), + b: ivy.array([[15.1, 32., 47.9], + [85., 196., 306.]]) } - - >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) - >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) - { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) - } """ - return current_backend(x).conv3d_transpose( + outer_batch_shape = list(weight.shape[:-2]) + num_outer_batch_dims = len(outer_batch_shape) + inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) + num_inner_batch_dims = len(inner_batch_shape) + num_out_feats, num_in_feats = list(weight.shape[-2:]) + + # OBS x IBS x OF + y = ivy.matmul( x, - filters, - strides, - padding, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - bias=bias, - out=out, + ivy.swapaxes( + ivy.reshape( + weight, + outer_batch_shape + + [1] * max(num_inner_batch_dims - 1, 0) + + [num_out_feats, num_in_feats], + ), + -1, + -2, + ), ) + if ivy.exists(bias): + # OBS x [1]*len(IBS) x OF + bias_broadcast = ivy.reshape( + bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] + ) + + # OBS x IBS x OF + y = y + bias_broadcast + + if ivy.exists(out): + return ivy.inplace_update(out, y) + return y + + +# LSTM # + @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv_general_dilated( +def lstm_update( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + init_h: Union[ivy.Array, ivy.NativeArray], + init_c: Union[ivy.Array, ivy.NativeArray], + kernel: Union[ivy.Array, ivy.NativeArray], + recurrent_kernel: Union[ivy.Array, ivy.NativeArray], /, *, - dims: int = 2, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively - and filters arrays. + Perform long-short term memory update by unrolling time dimension of input array. Parameters ---------- x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of filter. (Default value = 1) + input tensor of LSTM layer *[batch_shape, t, in]*. + init_h + initial state tensor for the cell output *[batch_shape, out]*. + init_c + initial state tensor for the cell hidden state *[batch_shape, out]*. + kernel + weights for cell kernel *[in, 4 x out]*. + recurrent_kernel + weights for cell recurrent kernel *[out, 4 x out]*. bias - Bias array of shape *[d_out]*. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + bias for cell kernel *[4 x out]*. (Default value = None) + recurrent_bias + bias for cell recurrent kernel *[4 x out]*. (Default value = None) Returns ------- ret - The result of the transpose convolution operation. + hidden state for all timesteps *[batch_shape,t,out]* and cell state for last + timestep *[batch_shape,out]* """ - return current_backend(x).conv_general_dilated( - x, - filters, - strides, - padding, - dims=dims, - data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, + # get shapes + x_shape = list(x.shape) + batch_shape = x_shape[:-2] + timesteps = x_shape[-2] + input_channels = x_shape[-1] + x_flat = ivy.reshape(x, (-1, input_channels)) + + # input kernel + Wi = kernel + Wi_x = ivy.reshape( + ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), + batch_shape + [timesteps, -1], ) + Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) + + # recurrent kernel + Wh = recurrent_kernel + + # lstm states + ht = init_h + ct = init_c + + # lstm outputs + hts_list = list() + + # unrolled time dimension with lstm steps + for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( + ivy.unstack(Wii_x, axis=-2), + ivy.unstack(Wif_x, axis=-2), + ivy.unstack(Wig_x, axis=-2), + ivy.unstack(Wio_x, axis=-2), + ): + htm1 = ht + ctm1 = ct + + Wh_htm1 = ivy.matmul(htm1, Wh) + ( + recurrent_bias if recurrent_bias is not None else 0 + ) + Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( + Wh_htm1, num_or_size_splits=4, axis=-1 + ) + + it = ivy.sigmoid(Wii_xt + Whi_htm1) + ft = ivy.sigmoid(Wif_xt + Whf_htm1) + gt = ivy.tanh(Wig_xt + Whg_htm1) + ot = ivy.sigmoid(Wio_xt + Who_htm1) + ct = ft * ctm1 + it * gt + ht = ot * ivy.tanh(ct) + + hts_list.append(ivy.expand_dims(ht, axis=-2)) + + return ivy.concat(hts_list, axis=-2), ct @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back +# @handle_array_like_without_promotion +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv_general_transpose( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: str, +def multi_head_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, /, *, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - feature_group_count: int = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + num_heads: Optional[int] = 8, + scale: Optional[float] = None, + attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + is_causal: Optional[bool] = False, + return_attention_weights: Optional[bool] = False, + average_attention_weights: Optional[bool] = True, + dropout: Optional[float] = 0.0, + training: Optional[bool] = False, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x - respectively and filters arrays. + Apply multi-head attention to inputs x. This is an implementation of multi-headed + attention as described in the paper "Attention is all you Need" (Vaswani et al., + 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each + timestep in `query` attends to the corresponding sequence in `key`, and returns a + fixed-width vector. This layer first projects `query`, `key` and `value`. These are + (effectively) a list of tensors of length `num_attention_heads`, where the + corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, + , key_dim)`, `(batch_size, , + value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are + softmaxed to obtain attention probabilities. The value tensors are then interpolated + by these probabilities, then concatenated back to a single tensor. Finally, the + result tensor with the last dimension as value_dim can take an linear projection and + return. Parameters ---------- - x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output. - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - dilations - The dilation factor for each dimension of input. (Default value = 1) - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - bias - Bias array of shape *[d_out]*. + query + query embeddings *[batch_shape,num_queries,query_dim]*. + key + key embeddings *[batch_shape,num_queries,key_dim]*. + value + value embeddings *[batch_shape,num_queries,value_dim]*. + num_heads + The number of attention heads to use. + scale + The value by which to scale the query-key similarity measure before softmax. + attention_mask + The mask to apply to the query-key values. Default is ``None``. + *[batch_shape,num_queries,num_keys]*. + in_proj_weights + The weights used to project query, key and value *[3*E, E]. + q_proj_weights + The weights used to project query if in_proj_weights is None *[new_E, E]. + k_proj_weights + The weights used to project key if in_proj_weights is None *[new_E, E]. + v_proj_weights + The weights used to project value if in_proj_weights is None *[new_E, E]. + out_proj_weights + The weights used to project the output. + in_proj_bias + The bias used when projecting with query, key and value. + out_proj_bias + The bias used when projecting the output. + is_causal + If True, Uses a causal attention mask and ignores provided attention_mask. + return_attention_weights + If True, returns attention_weights alongside the output + as a tuple (output, attenion_weights). Defaults to `False`. + average_attention_weights + If true, indicates that the returned ``attention_weights`` should be averaged + across heads. Otherwise, ``attention_weights`` are provided separately per head. + Note that this flag only has an effect when ``return_attention_weights=True``. + Default: ``True`` (i.e. average weights across heads) + dropout + Specifies the dropout probablity, dropout is applied to attention_weights. + training + If True, dropout is used, otherwise dropout is not activated. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1962,86 +1957,140 @@ def conv_general_transpose( Returns ------- ret - The result of the transpose convolution operation. - """ - return current_backend(x).conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - + The output following application of multi-head attention. + *[batch_shape,num_queries,out_feat_dim]* if input is batched + otherwise *[num_queries, out_feat_dim] + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + num_dims = query.ndim + ivy.assertions.check_all( + num_dims > 1 and num_dims < 4, + "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" + f" input), got {num_dims}", + ) + if key is None and value is None: + key = value = query + if num_dims == 2: + query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] + if ivy.exists(in_proj_weights): + q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) + elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): + if ivy.exists(in_proj_bias): + b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) + else: + b_q = b_k = b_v = None + q, k, v = ( + ivy.linear(query, q_proj_weights, bias=b_q), + ivy.linear(key, k_proj_weights, bias=b_k), + ivy.linear(value, v_proj_weights, bias=b_v), + ) + else: + q, k, v = query, key, value + batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] + k_seq_length = k.shape[1] + ivy.assertions.check_true( + emb_dim % num_heads == 0, "features must be divisible by number of heads" + ) + dims_per_head = emb_dim // num_heads + # isolate heads + q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 3, 1) + ) + v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + # perform bmm + attn_scores = ivy.matmul(q, k) + # scale + scale = 1 / (dims_per_head**0.5) if not scale else scale + attn_scores *= scale + # apply attention mask + if ivy.exists(attention_mask) or is_causal: + if is_causal: + # create causal mask + attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) + attention_mask = attention_mask.astype("bool") + attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) + # perform softmax + attn_weights = ivy.softmax(attn_scores, axis=-1) + # perform dropout + attn_weights = ivy.dropout(attn_weights, dropout, training=training) + # bmm with values + attention_out = ivy.matmul(attn_weights, v) + attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( + (batch_size, q_seq_length, -1) + ) + # proj out if out_proj_weight exists + if ivy.exists(out_proj_weights): + attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) + # if input was unbatched, unbatchify the output + if num_dims == 2: + attention_out = attention_out.squeeze(axis=0) + if return_attention_weights: + if average_attention_weights: + attn_weights = attn_weights.mean(axis=1) + if num_dims == 2: + attn_weights = attn_weights.squeeze(axis=0) + return attention_out, attn_weights + else: + return attention_out + + +# Attention # + @handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes @handle_array_function -def conv( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], +def scaled_dot_product_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Union[ivy.Array, ivy.NativeArray], + value: Union[ivy.Array, ivy.NativeArray], /, *, - transpose: bool = False, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + scale: Optional[float] = None, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + dropout_p: Optional[float] = 0.0, + is_causal: Optional[bool] = False, + training: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D - input x respectively and filters arrays. + Apply scaled dot product attention to inputs x using optional mask. Parameters ---------- - x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - transpose - True for computing transpose convolution, and False for dilated convolution. - When True, `x_dilations` must be 1 (the default). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output (Default value = None) - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]*. + query + The queries input array. The shape of queries input array should be in + *[batch_shape,num_queries,feat_dim]*. The queries input array should have the + same size as keys and values. + key + The keys input array. The shape of keys input array should be in + *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same + size as queries and values. + value + The values input array. The shape of values input should be in + *[batch_shape,num_keys,feat_dim]*. The values input array should have the same + size as queries and keys. + scale + The scale float value. + The scale float value is used to scale the query-key pairs before softmax. + mask + The mask input array. The mask to apply to the query-key values. Default is + None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. + dropout_p + Specifies the dropout probablity, if greater than 0.0, dropout is applied + is_causal + If true, assumes causal attention masking + and errors if both `mask` and `is_causal` are set. + training + If True, dropout is used, otherwise dropout is not activated. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -2049,294 +2098,252 @@ def conv( Returns ------- ret - The result of the transpose or dilated convolution operation. - """ - if transpose: - return conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - else: - return conv_general_dilated( - x, - filters, - strides, - padding, - dims=dims, - data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, - ) + The output following application of scaled dot-product attention. + The output array is the weighted sum produced by the attention score and value. + The shape of output array is *[batch_shape,num_queries,feat_dim]* . + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. -# LSTM # + Examples + -------- + With :class:`ivy.Array` input: + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def lstm_update( - x: Union[ivy.Array, ivy.NativeArray], - init_h: Union[ivy.Array, ivy.NativeArray], - init_c: Union[ivy.Array, ivy.NativeArray], - kernel: Union[ivy.Array, ivy.NativeArray], - recurrent_kernel: Union[ivy.Array, ivy.NativeArray], - /, - *, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[ivy.Array, ivy.Array]: - """ - Perform long-short term memory update by unrolling time dimension of input array. + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - Parameters - ---------- - x - input tensor of LSTM layer *[batch_shape, t, in]*. - init_h - initial state tensor for the cell output *[batch_shape, out]*. - init_c - initial state tensor for the cell hidden state *[batch_shape, out]*. - kernel - weights for cell kernel *[in, 4 x out]*. - recurrent_kernel - weights for cell recurrent kernel *[out, 4 x out]*. - bias - bias for cell kernel *[4 x out]*. (Default value = None) - recurrent_bias - bias for cell recurrent kernel *[4 x out]*. (Default value = None) - - Returns - ------- - ret - hidden state for all timesteps *[batch_shape,t,out]* and cell state for last - timestep *[batch_shape,out]* - """ - # get shapes - x_shape = list(x.shape) - batch_shape = x_shape[:-2] - timesteps = x_shape[-2] - input_channels = x_shape[-1] - x_flat = ivy.reshape(x, (-1, input_channels)) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - # input kernel - Wi = kernel - Wi_x = ivy.reshape( - ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), - batch_shape + [timesteps, -1], - ) - Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - # recurrent kernel - Wh = recurrent_kernel + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) - # lstm states - ht = init_h - ct = init_c + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - # lstm outputs - hts_list = list() + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - # unrolled time dimension with lstm steps - for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( - ivy.unstack(Wii_x, axis=-2), - ivy.unstack(Wif_x, axis=-2), - ivy.unstack(Wig_x, axis=-2), - ivy.unstack(Wio_x, axis=-2), - ): - htm1 = ht - ctm1 = ct + ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) - Wh_htm1 = ivy.matmul(htm1, Wh) + ( - recurrent_bias if recurrent_bias is not None else 0 - ) - Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( - Wh_htm1, num_or_size_splits=4, axis=-1 - ) + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) - it = ivy.sigmoid(Wii_xt + Whi_htm1) - ft = ivy.sigmoid(Wif_xt + Whf_htm1) - gt = ivy.tanh(Wig_xt + Whg_htm1) - ot = ivy.sigmoid(Wio_xt + Who_htm1) - ct = ft * ctm1 + it * gt - ht = ot * ivy.tanh(ct) + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - hts_list.append(ivy.expand_dims(ht, axis=-2)) + With :class:`ivy.Container` input: - return ivy.concat(hts_list, axis=-2), ct + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) + { + a: ivy.array([[[5.19999981, 1.], + ... [2.59249449, 2.68226194], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[0.2, 1.], + ... [2.19603825, 2.9960382], + ... [4.4000001, 5.5999999]]]) + } + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> mask = ivy.Container( + ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), + ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) + ... ) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) + { + a: ivy.array([[[4.26894283, 5.40236187], + ... [4.39999437, 5.59999037], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[4.35046196, 5.54282808], + ... [4.39989519, 5.5998764], + ... [4.4000001, 5.5999999]]]) + } -# Helpers # + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) -def _handle_padding(x, strides, filters, padding): - if isinstance(padding, str) and padding.upper() == "SAME": - if x % strides == 0: - pad = max(filters - strides, 0) - else: - pad = max(filters - (x % strides), 0) - else: - pad = 0 - return pad + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) + >>> print(out) + ivy.array([[[4.03946018, 5.0280633 ], + ... [4.29981947, 5.29981089], + ... [4.30000019, 5.30000019]]]) -def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): - if isinstance(kernel, int): - kernel = (kernel,) * dims - elif len(kernel) == 1: - kernel = (kernel[0],) * dims - elif (len(kernel) != dims) and (len(kernel) != dims + 2): - raise ValueError( - "The kernel should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - if isinstance(strides, int): - strides = (strides,) * dims - elif len(strides) == 1: - strides = (strides[0],) * dims - elif (len(strides) != dims) and (len(strides) != dims + 2): - raise ValueError( - "The stride should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) + >>> print(result) + { + a: ivy.array([[[0.40000001, 1.29999995], + ... [2.06345534, 2.9634552], + ... [4.30000019, 5.30000019]]]), + b: ivy.array([[[0.40000001, 1.29999995], + ... [2.19336844, 3.09336829], + ... [4.30000019, 5.30000019]]]) + } + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - if isinstance(padding, int): - padding = [(padding,) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == 1: - padding = [(padding[0],) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == dims: - padding = [(padding[i],) * 2 for i in range(dims)] - elif isinstance(padding, list) and len(padding) == dims: - if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): - raise ValueError("Explicit padding must be a list of tuple of two integers") - if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: - raise ValueError( - f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... mask=mask, + ... dropout_p=0.1, + ... training=True) + >>> print(result) + { + a: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]), + b: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) + } + """ + ivy.assertions.check_all( + (not is_causal) or (is_causal and mask is None), + "is_causal and attn_mask cannot be set at the same time", + ) + embed_dim = query.shape[-1] + scale = 1 / (embed_dim**0.5) if not scale else scale + sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale + sim = ivy.dropout(sim, dropout_p, training=training) + if ivy.exists(mask): + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, ) - - if isinstance(dilation, int): - dilation = (dilation,) * dims - elif len(dilation) == 1: - dilation = (dilation[0],) * dims - elif len(dilation) != dims: - raise ValueError( - f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" + elif is_causal: + L = query.shape[-2] # Source sequence length + S = key.shape[-2] # Target sequence length + mask = ivy.tril(ivy.ones((L, S)), k=0) + mask = ivy.astype(mask, ivy.bool) + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, ) - if min(dilation) < 1: - raise ValueError("All values of `dilation` must be positive") - - # Other errors - if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: - raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") - assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" - - # Account for dilation when padding > kernel/2. Not the case in torch by default. - new_kernel = tuple( - [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] - ) - if isinstance(padding, list) and len(padding) == len(new_kernel): - ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) - - return kernel, strides, padding, dilation - - -def _depth_max_pooling_helper( - x_shape, kernel, strides, dims, data_format="channel_last" -): - # Determine depth pooling. - # We assume that the kernel and the data have the same data_format. - depth_pooling = False - CHANNEL_LAST = "channel_last" - channel_idx = -1 if data_format == CHANNEL_LAST else 1 - if len(kernel) == dims + 2: - spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] - if kernel[channel_idx] != 1: - depth_pooling = True - if any(i != 1 for i in spatial_kernel): - raise NotImplementedError( - "MaxPooling supports exactly one of pooling across" - " depth or pooling across width/height." - ) - if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to equal the depth" - " stride" - ) - if x_shape[channel_idx] % kernel[channel_idx] != 0: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to evenly divide" - " the input depth" - ) - kernel = [kernel[channel_idx], *[1] * (dims - 1)] - strides = [strides[channel_idx], *[1] * (dims - 1)] - else: - kernel = spatial_kernel - if len(strides) == dims + 2: - strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] - return kernel, strides, depth_pooling - - -def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): - kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) - if padding == "SAME": - dim_size = dim_size * stride_size - else: - dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) - return dim_size - - -def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): - if dims == 1: - if data_format == "channel_first": - return "NCW" - else: - return "NWC" - if dims == 2: - if data_format == "channel_first": - return "NCHW" - else: - return "NHWC" - elif dims == 3: - if data_format == "channel_first": - return "NCDHW" - else: - return "NDHWC" - - -def _get_num_padded_values(i, p, n, k, s): - """ - Get number of padded values in a specific window. + attn = ivy.softmax(sim, axis=-1) + result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) + return result if not ivy.exists(out) else ivy.inplace_update(out, result) - Parameters - ---------- - i window index - p total amount of padding - n input size - k kernel size - s stride - Returns - ------- - number of padded values in a particular window represented by i - """ - current_index = s * i - left_padding = p // 2 - return max(0, left_padding - current_index) + max( - 0, current_index + k - n - left_padding - ) +linear.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +dropout.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} diff --git a/ivy/functional/ivy/linear_algebra.py b/ivy/functional/ivy/linear_algebra.py index 612315f0b987a..8190c0d4b57d5 100644 --- a/ivy/functional/ivy/linear_algebra.py +++ b/ivy/functional/ivy/linear_algebra.py @@ -345,6 +345,95 @@ def det( return current_backend(x).det(x, out=out) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def diag( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the specified diagonals of the input array, or an array with the input + array's elements as diagonals. + + Parameters + ---------- + x + An array with rank >= 1. + k + An integer that controls which diagonal to consider. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + If x is a 1-D array, the function returns a 2-D square array with the elements + of input as diagonals. + If x is a 2-D array, the function returns a 1-D array with the diagonal elements + of x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x) + ivy.array([0, 4, 8]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=1) + ivy.array([1, 5]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=-1) + ivy.array([3, 7]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(ivy.diag(x)) + ivy.array([[0, 0, 0], + [0, 4, 0], + [0, 0, 8]]) + """ + return current_backend(x).diag(x, k=k, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -929,6 +1018,41 @@ def inv( return current_backend(x).inv(x, adjoint=adjoint, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lu_factor( + A: Union[ivy.Array, ivy.NativeArray], + /, + *, + pivot: bool = True, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Parameters + ---------- + A + tensor of shape (*, m, n) where * is zero or more batch dimensions. + + pivot + Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. pivot = False not supported on CPU. Default: True. + + out + tuple of two tensors to write the output to. Ignored if None. Default: None. + + Returns + ------- + ret + A named tuple (LU, pivots). + """ + return current_backend(A).lu_factor(A, pivot=pivot, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2315,6 +2439,51 @@ def tensordot( return current_backend(x1, x2).tensordot(x1, x2, axes=axes, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def tensorsolve( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = 2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + ndim1 = ivy.get_num_dims(x1) + ndim2 = ivy.get_num_dims(x2) + + if axes is not None: + allaxes = list(range(0, ndim1)) + for k in axes: + allaxes.remove(k) + allaxes.insert(ndim1, k) + + x1 = ivy.matrix_transpose(x1, allaxes) + + old_shape = x1.shape[-(ndim1 - ndim2) :] + + prod = 1 + for k in old_shape: + prod *= k + + if ivy.shape(ivy.flatten(x1))[0] != prod**2: + raise ivy.utils.exceptions.IvyException( + "Input arrays must satisfy the requirement " + "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" + ) + + x1 = ivy.reshape(x1, (prod, prod)) + x2 = ivy.flatten(x2) + res = ivy.solve(x1, x2) + res = ivy.reshape(res, old_shape) + return res + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2443,6 +2612,80 @@ def trace( return current_backend(x).trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def vander( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a Vandermonde matrix. The columns of the output matrix are elementwise + powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, + the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a + geometric progression in each row is named for Alexandre-Theophile Vandermonde. + + Parameters + ---------- + x + 1-D input array. + N + Number of columns in the output. If N is not specified, + a square array is returned (N = len(x)) + increasing + Order of the powers of the columns. If True, the powers increase + from left to right, if False (the default) they are reversed. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Vandermonde matrix. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x) + ivy.array( + [[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3) + ivy.array( + [[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3, increasing=True) + ivy.array( + [[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]] + ) + """ + return current_backend(x).vander(x, N=N, increasing=increasing, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2657,169 +2900,6 @@ def vector_norm( ) -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def diag( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the specified diagonals of the input array, or an array with the input - array's elements as diagonals. - - Parameters - ---------- - x - An array with rank >= 1. - k - An integer that controls which diagonal to consider. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - If x is a 1-D array, the function returns a 2-D square array with the elements - of input as diagonals. - If x is a 2-D array, the function returns a 1-D array with the diagonal elements - of x. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x) - ivy.array([0, 4, 8]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=1) - ivy.array([1, 5]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=-1) - ivy.array([3, 7]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(ivy.diag(x)) - ivy.array([[0, 0, 0], - [0, 4, 0], - [0, 0, 8]]) - """ - return current_backend(x).diag(x, k=k, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def vander( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Generate a Vandermonde matrix. The columns of the output matrix are elementwise - powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, - the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a - geometric progression in each row is named for Alexandre-Theophile Vandermonde. - - Parameters - ---------- - x - 1-D input array. - N - Number of columns in the output. If N is not specified, - a square array is returned (N = len(x)) - increasing - Order of the powers of the columns. If True, the powers increase - from left to right, if False (the default) they are reversed. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Vandermonde matrix. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x) - ivy.array( - [[ 1, 1, 1, 1], - [ 8, 4, 2, 1], - [ 27, 9, 3, 1], - [125, 25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3) - ivy.array( - [[ 1, 1, 1], - [ 4, 2, 1], - [ 9, 3, 1], - [25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3, increasing=True) - ivy.array( - [[ 1, 1, 1], - [ 1, 2, 4], - [ 1, 3, 9], - [ 1, 5, 25]] - ) - """ - return current_backend(x).vander(x, N=N, increasing=increasing, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2854,84 +2934,3 @@ def vector_to_skew_symmetric_matrix( instances in place of any of the arguments. """ return current_backend(vector).vector_to_skew_symmetric_matrix(vector, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lu_factor( - A: Union[ivy.Array, ivy.NativeArray], - /, - *, - pivot: bool = True, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Parameters - ---------- - A - tensor of shape (*, m, n) where * is zero or more batch dimensions. - - pivot - Whether to compute the LU decomposition with partial pivoting, or the regular LU - decomposition. pivot = False not supported on CPU. Default: True. - - out - tuple of two tensors to write the output to. Ignored if None. Default: None. - - Returns - ------- - ret - A named tuple (LU, pivots). - """ - return current_backend(A).lu_factor(A, pivot=pivot, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def tensorsolve( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - ndim1 = ivy.get_num_dims(x1) - ndim2 = ivy.get_num_dims(x2) - - if axes is not None: - allaxes = list(range(0, ndim1)) - for k in axes: - allaxes.remove(k) - allaxes.insert(ndim1, k) - - x1 = ivy.matrix_transpose(x1, allaxes) - - old_shape = x1.shape[-(ndim1 - ndim2) :] - - prod = 1 - for k in old_shape: - prod *= k - - if ivy.shape(ivy.flatten(x1))[0] != prod**2: - raise ivy.utils.exceptions.IvyException( - "Input arrays must satisfy the requirement " - "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" - ) - - x1 = ivy.reshape(x1, (prod, prod)) - x2 = ivy.flatten(x2) - res = ivy.solve(x1, x2) - res = ivy.reshape(res, old_shape) - return res - # return current_backend(x1, x2).tensorsolve(x1, x2, axes=axes, out=out) diff --git a/ivy/functional/ivy/losses.py b/ivy/functional/ivy/losses.py index d866cd8bc2487..d803eb6fca839 100644 --- a/ivy/functional/ivy/losses.py +++ b/ivy/functional/ivy/losses.py @@ -12,8 +12,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # +# --- Helpers --- # +# --------------- # def _reduce_loss(red, loss, axis, out): @@ -25,64 +25,8 @@ def _reduce_loss(red, loss, axis, out): return ivy.negative(loss, out=out) -# Extra # -# ------# - - -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def cross_entropy( - true: Union[ivy.Array, ivy.NativeArray], - pred: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - epsilon: float = 1e-7, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute cross-entropy between predicted and true discrete distributions. - - Parameters - ---------- - true - input array containing true labels. - pred - input array containing the predicted labels. - axis - the axis along which to compute the cross-entropy. If axis is ``-1``, - the cross-entropy will be computed along the last dimension. Default: ``-1``. - epsilon - a float in [0.0, 1.0] specifying the amount of smoothing when calculating - the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - The cross-entropy loss between the given distributions - - Examples - -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.cross_entropy(x, y)) - ivy.array(1.3862944) - - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.cross_entropy(x, z)) - ivy.array(0.35667497) - """ - ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) - pred = ivy.clip(pred, epsilon, 1 - epsilon) - log_pred = ivy.log(pred) - return _reduce_loss(reduction, log_pred * true, axis, out) +# --- Main --- # +# ------------ # @handle_exceptions @@ -268,6 +212,66 @@ def binary_cross_entropy( return _reduce_loss(reduction, loss, axis, out) +# Extra # +# ------# + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def cross_entropy( + true: Union[ivy.Array, ivy.NativeArray], + pred: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + epsilon: float = 1e-7, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute cross-entropy between predicted and true discrete distributions. + + Parameters + ---------- + true + input array containing true labels. + pred + input array containing the predicted labels. + axis + the axis along which to compute the cross-entropy. If axis is ``-1``, + the cross-entropy will be computed along the last dimension. Default: ``-1``. + epsilon + a float in [0.0, 1.0] specifying the amount of smoothing when calculating + the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + The cross-entropy loss between the given distributions + + Examples + -------- + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.cross_entropy(x, y)) + ivy.array(1.3862944) + + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.cross_entropy(x, z)) + ivy.array(0.35667497) + """ + ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) + pred = ivy.clip(pred, epsilon, 1 - epsilon) + log_pred = ivy.log(pred) + return _reduce_loss(reduction, log_pred * true, axis, out) + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/manipulation.py b/ivy/functional/ivy/manipulation.py index 71d5e65b61c89..af74e3523eb5f 100644 --- a/ivy/functional/ivy/manipulation.py +++ b/ivy/functional/ivy/manipulation.py @@ -20,6 +20,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _calculate_out_shape(axis, array_shape): if type(axis) not in (tuple, list): axis = (axis,) @@ -33,6 +37,141 @@ def _calculate_out_shape(axis, array_shape): return out_shape +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def clip( + x: Union[ivy.Array, ivy.NativeArray], + x_min: Union[Number, ivy.Array, ivy.NativeArray], + x_max: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Clips (limits) the values in an array. + + Given an interval, values outside the interval are clipped to the interval edges + (element-wise). For example, if an interval of [0, 1] is specified, values smaller + than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller + or equal to maximum value to return correct results. + + Parameters + ---------- + x + Input array containing elements to clip. + x_min + Minimum value. + x_max + Maximum value. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + An array with the elements of x, but where values < x_min are replaced with + x_min, and those > x_max with x_max. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.zeros_like(x) + >>> ivy.clip(x, 2., 7., out=y) + >>> print(y) + ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) + >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + { + a: ivy.array([1., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> x_min = ivy.Container(a=0, b=-3) + >>> x_max = ivy.Container(a=1, b=-1) + >>> y = ivy.clip(x, x_min,x_max) + >>> print(y) + { + a: ivy.array([0., 1., 1.]), + b: ivy.array([-1., -1., -1.]) + } + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 0., 1]) + >>> x_max = ivy.array([5., 4., 3.]) + >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> z = ivy.clip(y, x_min, x_max) + >>> print(z) + { + a: ivy.array([3., 1., 2.]), + b: ivy.array([3., 4., 3.]) + } + """ + return current_backend(x).clip(x, x_min, x_max, out=out) + + # Array API Standard # # -------------------# @@ -97,6 +236,103 @@ def concat( return current_backend(xs[0]).concat(xs, axis=axis, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def constant_pad( + x: Union[ivy.Array, ivy.NativeArray], + /, + pad_width: Iterable[Tuple[int]], + *, + value: Number = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Pad an array with a constant value. + + Parameters + ---------- + x + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of + axes of x. + value + The constant value to pad the array with. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Padded array of rank equal to x with shape increased according to pad_width. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5]) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) + >>> print(y) + ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[2.], [3.]]) + >>> y = ivy.zeros((4, 3)) + >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) + >>> print(y) + ivy.array([[5., 5., 5.], + [5., 2., 5.], + [5., 3., 5.], + [5., 5., 5.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), + ... b = ivy.array([3., 4., 5.])) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) + >>> print(y) + { + a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), + b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) + } + """ + return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -435,40 +671,113 @@ def permute_dims( @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def reshape( +def repeat( x: Union[ivy.Array, ivy.NativeArray], /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + repeats: Union[int, Iterable[int]], *, - copy: Optional[bool] = None, - order: str = "C", - allowzero: bool = True, + axis: int = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Give a new shape to an array without changing its data. + Repeat values along a given dimension. Parameters ---------- x - Input array to be reshaped. - shape - a new shape compatible with the original shape. One shape dimension - can be -1. In this case, the value is inferred from the length of the array and - remaining dimensions. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - order - Read the elements of x using this index order, and place the elements into - the reshaped array using this index order. + Input array. + repeats + The number of repetitions for each element. repeats is broadcast to fit the + shape of the given axis. + axis + The axis along which to repeat values. By default, use the flattened input + array, and return a flat output array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The repeated output array. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.repeat(x, 2) + >>> print(y) + ivy.array([3, 3, 4, 4, 5, 5]) + + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> y = ivy.repeat(x, [1, 2], axis=0) + >>> print(y) + ivy.array([[1, 2, 3], + [4, 5, 6], + [4, 5, 6]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([0., 1., 2.])) + >>> y = ivy.repeat(x, 2, axis=0) + >>> print(y) + { + a: ivy.array([0., 0., 1., 1., 2., 2.]), + b: ivy.array([0., 0., 1., 1., 2., 2.]) + } + """ + return current_backend(x).repeat(x, repeats, axis=axis, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def reshape( + x: Union[ivy.Array, ivy.NativeArray], + /, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + *, + copy: Optional[bool] = None, + order: str = "C", + allowzero: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Give a new shape to an array without changing its data. + + Parameters + ---------- + x + Input array to be reshaped. + shape + a new shape compatible with the original shape. One shape dimension + can be -1. In this case, the value is inferred from the length of the array and + remaining dimensions. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + order + Read the elements of x using this index order, and place the elements into + the reshaped array using this index order. ``C`` means to read / write the elements using C-like index order, with the last axis index changing fastest, back to the first axis index changing slowest. @@ -671,6 +980,93 @@ def roll( return current_backend(x).roll(x, shift, axis=axis, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def split( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], ivy.Array, ivy.NativeArray] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> List[ivy.Array]: + """ + Split an array into multiple sub-arrays. + + Parameters + ---------- + x + array to be divided into sub-arrays. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + num_or_size_splits + Number of equal arrays to divide the array into along the given axis if an + integer. The size of each split element if a sequence of integers or 1-D array. + Default is to divide into as many 1-dimensional arrays as the axis dimension. + axis + The axis along which to split, default is ``0``. + with_remainder + If the tensor does not split evenly, then store the last remainder entry. + Default is ``False``. + + Returns + ------- + ret + A list of sub-arrays. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.split(x) + >>> print(y) + [ivy.array([1]),ivy.array([2]),ivy.array([3])] + + >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) + >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) + >>> print(y) + [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] + + >>> x = ivy.array([4, 6, 5, 3]) + >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) + >>> print(y) + ivy.array([[4], [6, 5, 3]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([10, 45, 2])) + >>> y = ivy.split(x) + >>> print(y) + { + a:(list[3],shape=[1]) + } + """ + return current_backend(x).split( + x, + copy=copy, + num_or_size_splits=num_or_size_splits, + axis=axis, + with_remainder=with_remainder, + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -851,394 +1247,6 @@ def stack( return res -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def clip( - x: Union[ivy.Array, ivy.NativeArray], - x_min: Union[Number, ivy.Array, ivy.NativeArray], - x_max: Union[Number, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Clips (limits) the values in an array. - - Given an interval, values outside the interval are clipped to the interval edges - (element-wise). For example, if an interval of [0, 1] is specified, values smaller - than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller - or equal to maximum value to return correct results. - - Parameters - ---------- - x - Input array containing elements to clip. - x_min - Minimum value. - x_max - Maximum value. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - An array with the elements of x, but where values < x_min are replaced with - x_min, and those > x_max with x_max. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.zeros_like(x) - >>> ivy.clip(x, 2., 7., out=y) - >>> print(y) - ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) - >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - { - a: ivy.array([1., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> x_min = ivy.Container(a=0, b=-3) - >>> x_max = ivy.Container(a=1, b=-1) - >>> y = ivy.clip(x, x_min,x_max) - >>> print(y) - { - a: ivy.array([0., 1., 1.]), - b: ivy.array([-1., -1., -1.]) - } - - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 0., 1]) - >>> x_max = ivy.array([5., 4., 3.]) - >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> z = ivy.clip(y, x_min, x_max) - >>> print(z) - { - a: ivy.array([3., 1., 2.]), - b: ivy.array([3., 4., 3.]) - } - """ - return current_backend(x).clip(x, x_min, x_max, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def constant_pad( - x: Union[ivy.Array, ivy.NativeArray], - /, - pad_width: Iterable[Tuple[int]], - *, - value: Number = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Pad an array with a constant value. - - Parameters - ---------- - x - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of - axes of x. - value - The constant value to pad the array with. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Padded array of rank equal to x with shape increased according to pad_width. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3, 4, 5]) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) - >>> print(y) - ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[2.], [3.]]) - >>> y = ivy.zeros((4, 3)) - >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) - >>> print(y) - ivy.array([[5., 5., 5.], - [5., 2., 5.], - [5., 3., 5.], - [5., 5., 5.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), - ... b = ivy.array([3., 4., 5.])) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) - >>> print(y) - { - a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), - b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) - } - """ - return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def repeat( - x: Union[ivy.Array, ivy.NativeArray], - /, - repeats: Union[int, Iterable[int]], - *, - axis: int = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Repeat values along a given dimension. - - Parameters - ---------- - x - Input array. - repeats - The number of repetitions for each element. repeats is broadcast to fit the - shape of the given axis. - axis - The axis along which to repeat values. By default, use the flattened input - array, and return a flat output array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The repeated output array. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.repeat(x, 2) - >>> print(y) - ivy.array([3, 3, 4, 4, 5, 5]) - - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> y = ivy.repeat(x, [1, 2], axis=0) - >>> print(y) - ivy.array([[1, 2, 3], - [4, 5, 6], - [4, 5, 6]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([0., 1., 2.])) - >>> y = ivy.repeat(x, 2, axis=0) - >>> print(y) - { - a: ivy.array([0., 0., 1., 1., 2., 2.]), - b: ivy.array([0., 0., 1., 1., 2., 2.]) - } - """ - return current_backend(x).repeat(x, repeats, axis=axis, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def split( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], ivy.Array, ivy.NativeArray] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> List[ivy.Array]: - """ - Split an array into multiple sub-arrays. - - Parameters - ---------- - x - array to be divided into sub-arrays. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - num_or_size_splits - Number of equal arrays to divide the array into along the given axis if an - integer. The size of each split element if a sequence of integers or 1-D array. - Default is to divide into as many 1-dimensional arrays as the axis dimension. - axis - The axis along which to split, default is ``0``. - with_remainder - If the tensor does not split evenly, then store the last remainder entry. - Default is ``False``. - - Returns - ------- - ret - A list of sub-arrays. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.split(x) - >>> print(y) - [ivy.array([1]),ivy.array([2]),ivy.array([3])] - - >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) - >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) - >>> print(y) - [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] - - >>> x = ivy.array([4, 6, 5, 3]) - >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) - >>> print(y) - ivy.array([[4], [6, 5, 3]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([10, 45, 2])) - >>> y = ivy.split(x) - >>> print(y) - { - a:(list[3],shape=[1]) - } - """ - return current_backend(x).split( - x, - copy=copy, - num_or_size_splits=num_or_size_splits, - axis=axis, - with_remainder=with_remainder, - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable diff --git a/ivy/functional/ivy/meta.py b/ivy/functional/ivy/meta.py index ab885bb3b207c..66d5452e48832 100644 --- a/ivy/functional/ivy/meta.py +++ b/ivy/functional/ivy/meta.py @@ -7,6 +7,11 @@ # local from typing import Optional, Union, Callable, Tuple, Any + +# --- Helpers --- # +# --------------- # + + # Extra # # ------# @@ -194,6 +199,70 @@ def _train_task( return final_cost, variables, all_grads +def _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + batched, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, +): + if batched: + return _train_tasks_batched( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) + return _train_tasks_with_for_loop( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) + + def _train_tasks_batched( batch, inner_batch_fn, @@ -336,68 +405,8 @@ def _train_tasks_with_for_loop( return total_cost / num_tasks -def _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - batched, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, -): - if batched: - return _train_tasks_batched( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) - return _train_tasks_with_for_loop( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) +# --- Main --- # +# ------------ # # Public # @@ -526,7 +535,136 @@ def fomaml_step( return cost, grads -fomaml_step.computes_gradients = True +# Second Order + + +@handle_exceptions +@handle_array_function +def maml_step( + batch: ivy.Container, + inner_cost_fn: Callable, + outer_cost_fn: Callable, + variables: ivy.Container, + inner_grad_steps: int, + inner_learning_rate: float, + /, + *, + inner_optimization_step: Callable = gradient_descent_update, + inner_batch_fn: Optional[Callable] = None, + outer_batch_fn: Optional[Callable] = None, + average_across_steps: bool = False, + batched: bool = True, + inner_v: Optional[ivy.Container] = None, + keep_inner_v: bool = True, + outer_v: Optional[ivy.Container] = None, + keep_outer_v: bool = True, + return_inner_v: Union[str, bool] = False, + num_tasks: Optional[int] = None, + stop_gradients: bool = True, +) -> Tuple[ivy.Array, ivy.Container, Any]: + """ + Perform step of vanilla second order MAML. + + Parameters + ---------- + batch + The input batch + inner_cost_fn + callable for the inner loop cost function, receiving sub-batch, inner vars and + outer vars + outer_cost_fn + callable for the outer loop cost function, receving task-specific sub-batch, + inner vars and outer vars. If None, the cost from the inner loop will also be + optimized in the outer loop. + variables + Variables to be optimized during the meta step + inner_grad_steps + Number of gradient steps to perform during the inner loop. + inner_learning_rate + The learning rate of the inner loop. + inner_optimization_step + The function used for the inner loop optimization. + Default is ivy.gradient_descent_update. + inner_batch_fn + Function to apply to the task sub-batch, before passing to the inner_cost_fn. + Default is ``None``. + outer_batch_fn + Function to apply to the task sub-batch, before passing to the outer_cost_fn. + Default is ``None``. + average_across_steps + Whether to average the inner loop steps for the outer loop update. + Default is ``False``. + batched + Whether to batch along the time dimension, and run the meta steps in batch. + Default is ``True``. + inner_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_inner_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + outer_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_outer_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + return_inner_v + Either 'first', 'all', or False. 'first' means the variables for the first task + inner loop will also be returned. variables for all tasks will be returned with + 'all'. Default is ``False``. + num_tasks + Number of unique tasks to inner-loop optimize for the meta step. Determined from + batch by default. + stop_gradients + Whether to stop the gradients of the cost. Default is ``True``. + + Returns + ------- + ret + The cost and the gradients with respect to the outer loop variables. + """ + if num_tasks is None: + num_tasks = batch.cont_shape[0] + unique_outer = outer_v is not None + func_ret, grads = ivy.execute_with_gradients( + lambda v: _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables.cont_set_at_key_chains(v) if unique_outer else v, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + 2, + average_across_steps, + batched, + inner_v, + keep_inner_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + False, + ), + ( + variables.cont_at_key_chains(outer_v, ignore_none=True) + if keep_outer_v + else variables.cont_prune_key_chains(outer_v, ignore_none=True) + ), + ) + if isinstance(func_ret, tuple): + grads = grads["0"] if "0" in grads else grads + cost = func_ret[0] + rest = func_ret[1] + else: + cost = func_ret + rest = () + if stop_gradients: + cost = ivy.stop_gradient(cost, preserve_type=False) + return cost, grads.sum(axis=0), rest @handle_exceptions @@ -665,139 +803,6 @@ def reptile_step( return cost, grads +fomaml_step.computes_gradients = True reptile_step.computes_gradients = True - - -# Second Order - - -@handle_exceptions -@handle_array_function -def maml_step( - batch: ivy.Container, - inner_cost_fn: Callable, - outer_cost_fn: Callable, - variables: ivy.Container, - inner_grad_steps: int, - inner_learning_rate: float, - /, - *, - inner_optimization_step: Callable = gradient_descent_update, - inner_batch_fn: Optional[Callable] = None, - outer_batch_fn: Optional[Callable] = None, - average_across_steps: bool = False, - batched: bool = True, - inner_v: Optional[ivy.Container] = None, - keep_inner_v: bool = True, - outer_v: Optional[ivy.Container] = None, - keep_outer_v: bool = True, - return_inner_v: Union[str, bool] = False, - num_tasks: Optional[int] = None, - stop_gradients: bool = True, -) -> Tuple[ivy.Array, ivy.Container, Any]: - """ - Perform step of vanilla second order MAML. - - Parameters - ---------- - batch - The input batch - inner_cost_fn - callable for the inner loop cost function, receiving sub-batch, inner vars and - outer vars - outer_cost_fn - callable for the outer loop cost function, receving task-specific sub-batch, - inner vars and outer vars. If None, the cost from the inner loop will also be - optimized in the outer loop. - variables - Variables to be optimized during the meta step - inner_grad_steps - Number of gradient steps to perform during the inner loop. - inner_learning_rate - The learning rate of the inner loop. - inner_optimization_step - The function used for the inner loop optimization. - Default is ivy.gradient_descent_update. - inner_batch_fn - Function to apply to the task sub-batch, before passing to the inner_cost_fn. - Default is ``None``. - outer_batch_fn - Function to apply to the task sub-batch, before passing to the outer_cost_fn. - Default is ``None``. - average_across_steps - Whether to average the inner loop steps for the outer loop update. - Default is ``False``. - batched - Whether to batch along the time dimension, and run the meta steps in batch. - Default is ``True``. - inner_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_inner_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - outer_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_outer_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - return_inner_v - Either 'first', 'all', or False. 'first' means the variables for the first task - inner loop will also be returned. variables for all tasks will be returned with - 'all'. Default is ``False``. - num_tasks - Number of unique tasks to inner-loop optimize for the meta step. Determined from - batch by default. - stop_gradients - Whether to stop the gradients of the cost. Default is ``True``. - - Returns - ------- - ret - The cost and the gradients with respect to the outer loop variables. - """ - if num_tasks is None: - num_tasks = batch.cont_shape[0] - unique_outer = outer_v is not None - func_ret, grads = ivy.execute_with_gradients( - lambda v: _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables.cont_set_at_key_chains(v) if unique_outer else v, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - 2, - average_across_steps, - batched, - inner_v, - keep_inner_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - False, - ), - ( - variables.cont_at_key_chains(outer_v, ignore_none=True) - if keep_outer_v - else variables.cont_prune_key_chains(outer_v, ignore_none=True) - ), - ) - if isinstance(func_ret, tuple): - grads = grads["0"] if "0" in grads else grads - cost = func_ret[0] - rest = func_ret[1] - else: - cost = func_ret - rest = () - if stop_gradients: - cost = ivy.stop_gradient(cost, preserve_type=False) - return cost, grads.sum(axis=0), rest - - maml_step.computes_gradients = True diff --git a/ivy/functional/ivy/nest.py b/ivy/functional/ivy/nest.py index 05e47e2358582..c7a052fd7cdb6 100644 --- a/ivy/functional/ivy/nest.py +++ b/ivy/functional/ivy/nest.py @@ -11,6 +11,246 @@ from ivy.utils.exceptions import handle_exceptions +@handle_exceptions +def all_nested_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, + /, + include_nests: bool = False, + _index: Optional[Union[int, Sequence[int]]] = None, + _base: bool = True, + extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return indices of all the elements in nest. + + Parameters + ---------- + nest + The nest to check the leaves of. + include_nests + Whether to also include indices of the nests themselves, not only + leaves. Default is ``False``. + _index + The indices detected so far. None at the beginning. Used internally, + do not set manually. + _base + Whether the current function call is the first function call in the + recursive stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + out + Optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + A set of indices of all elements in nest + + Both the description and the type hints above assumes an array input + for simplicity, but this function is nestable, and therefore also + accepts :class:ivy.Container instances in place of the arguments. + + Examples + -------- + With :class:`Dict` input: + + >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} + >>> y = ivy.all_nested_indices(x) + >>> print(y) + [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] + + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4.]) + >>> y = ivy.all_nested_indices(x, False, out=x) + >>> print(y) + [[]] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.all_nested_indices(x, True) + >>> print(y) + [['a'], ['b']] + """ + _index = list() if _index is None else _index + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + ind = ivy.argwhere(ivy.ones_like(nest)) + indices = list() + for i in range(len(ind)): + indices.append(_index + ind.to_list()[i]) + return indices + _indices = [ + all_nested_indices( + item, include_nests, _index + [i], False, extra_nest_types + ) + for i, item in enumerate(nest) + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + elif isinstance(nest, dict): + _indices = [ + all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) + for k, v in nest.items() + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + else: + return [_index] + return [index for index in _indices if index] + + +@handle_exceptions +def copy_nest( + nest: Union[ivy.Array, ivy.NativeArray, Iterable], + /, + include_derived: bool = False, + to_mutable: bool = False, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[ivy.Array, ivy.NativeArray, Iterable]: + """ + Copy a nest deeply, but without copying leaves of the nest, only the nest lists, + tuples and dicts are copied. + + Parameters + ---------- + nest + The nest to copy. + include_derived + Whether to also recursive for classes derived from tuple, list and dict. + Default is ``False``. + to_mutable + Whether to convert the nest to a mutable form, changing all tuples to lists. + Default is ``False``. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + The copied nest. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + ivy.array([[1., 2., 3.], + [7., 8., 9.]]) + + With :code:`Iterable` input: + + >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + >>> copied_nest = ivy.copy_nest(nest, include_derived = True) + >>> print(copied_nest) + [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + + >>> nest = ([23, 25, 1337], [63, 98, 6]) + >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) + >>> print(copied_nest) + [[23, 25, 1337], [63, 98, 6]] + + >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + class_instance = type(nest) + check_fn = ( + (lambda x_, t: isinstance(nest, t)) + if include_derived + else (lambda x_, t: type(nest) is t) + ) + if check_fn(nest, tuple): + ret_list = [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + if to_mutable: + return ret_list + if hasattr(nest, "_fields"): + return class_instance(**dict(zip(nest._fields, ret_list))) + return class_instance(tuple(ret_list)) + elif check_fn(nest, list) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + return copy.deepcopy(nest) + return class_instance( + [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + ) + elif check_fn(nest, dict): + class_instance = type(nest) + dict_ = { + k: copy_nest( + v, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for k, v in nest.items() + } + if isinstance(nest, OrderedDict): + return class_instance(**dict_) + return class_instance(dict_) + return nest + + +@handle_exceptions +def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): + """ + Group all unique index chains in a nest. This function is useful for finding all + unique index chains in a nest, and then duplicating the values at those index chains + for functional frameworks. + + Parameters + ---------- + nest + nest to get duplicate index chains for. + + Returns + ------- + list of index chains to duplicate. + """ + all_index_chains = ivy.nested_argwhere(nest, lambda _: True) + duplicates = [] + duplicate_index_chains = {} + for index_chain in all_index_chains: + val = ivy.index_nest(nest, index_chain) + if ivy.is_array(val): + for i in range(len(duplicates)): + if val is duplicates[i]: + duplicate_index_chains[i].append(index_chain) + break + else: + duplicates.append(val) + duplicate_index_chains[len(duplicates) - 1] = [index_chain] + return list(duplicate_index_chains.values()) + + # Extra # # ------# @@ -92,140 +332,146 @@ def index_nest( @handle_exceptions -def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: +def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: + if len(index) == 1: + idx = index[0] + if isinstance(nest, list): + nest.insert(idx, value) + else: + nest[index[0]] = value + else: + insert_into_nest_at_index(nest[index[0]], index[1:], value) + + +@handle_exceptions +def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: """ - Prune a nested object at a specified index. + Insert a value into the nested item at specified indices with specified values. Parameters ---------- nest - The nested object to prune. - index - A tuple of indices for the index at which to prune. + The nested object to insert into. + indices + A tuple of tuples of indices for the indices at which to insert + values. + values + The new values for inserting. """ - if len(index) == 1: - del nest[index[0]] - else: - prune_nest_at_index(nest[index[0]], index[1:]) + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + [ + insert_into_nest_at_index(nest, index, value) + for index, value in zip(indices, values) + ] + + +# noinspection PyShadowingBuiltins @handle_exceptions -def set_nest_at_index( - nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], - index: Sequence[Union[str, int]], - value: Any, - /, - shallow: bool = True, - _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: +def map( + fn: Callable, + constant: Optional[Dict[str, Any]] = None, + unique: Optional[Dict[str, Iterable[Any]]] = None, + mean: bool = False, +) -> List: """ - Set the value of a nested item at a specified index. + Apply a function on each item of an iterable x. Parameters ---------- - nest - The nested object to update. - index - A tuple of indices for the index at which to update. - value - The new value for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - _result - Placeholder for the result of the update. do not set this paramter. + fn + The function to map onto x. + constant + keyword arguments which remain constant between each function call. + Default is ``None``. + unique + keyword arguments which are unique for each function call. Default is ``None``. + mean + Whether to compute the mean across the return values, and return this mean. + Default is ``False``. Returns ------- ret - nest with changed value at the given index. + x following the application of fn to each of its iterated items. Examples -------- - With :class:`ivy.Array` inputs: + With :code:`int` inputs: - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = (1, 1) - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([[1., 2.], [3., 5.]]) + >>> def special_square(x : float) -> float : return np.square(x) + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x' : [1,2,3]}, + ... mean = False) + >>> print(results) + [1, 4, 9] - >>> x = ivy.array([1., 2., 3., 4.]) - >>> y = [1] - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([1., 5., 3., 4.]) + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x':[0,1,2]}, + ... mean = True) + >>> print(results) + 1.6666666666666667 - With :code:`Dict` input: + >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = False) + >>> print(results) + [array([1,1]), + array([1,2]), + array([1,3])] - >>> x = {1 : [1, [2, 3]], 2: (4, 5)} - >>> y = (1, 1) - >>> z = 2 - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - {1: [1, 2], 2: (4, 5)} + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = True) + >>> print(results) + [1. 2.] - With :code:`List` inputs: + With float inputs: - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> y = (2, 1, 0) - >>> z = 'H' - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - [['a','b','c'],['d','e','f'],['g',['H','i']]] + >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':10., 'b':1.}, + ... unique = {'x':[0.,1.,2.]}, + ... mean = False) + >>> print(results) + [1.0, 11.0, 21.0] - With :class:`ivy.Container` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) - >>> y = ('b',) - >>> z = ivy.array([3., 4.]) - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - { - a: ivy.array([1., 2.]), - b: ivy.array([3., 4.]) - } + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = False) + >>> print(results) + [ivy.array([0., 10., 100.]), + ivy.array([1., 10., 101.])] + + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = True) + >>> print(results) + ivy.array([ 0.5, 10. , 100. ]) """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if _result is None: - if shallow: - _result = nest_type(nest) - else: - _result = copy_nest(nest, include_derived=True) - _result = list(_result) if is_tuple else _result - if len(index) == 1: - if shallow: - try: - nest[index[0]] = value - except TypeError: - pass - _result[index[0]] = value - else: - _result[index[0]] = set_nest_at_index( - nest[index[0]], index[1:], value, shallow, _result[index[0]] + c = ivy.default(constant, {}) + u = ivy.default(unique, {}) + rets = [ + r + for r in _map( + lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() ) - try: - _result = nest_type(_result) - except TypeError: - _result = nest_type(*_result) - return _result - + ] + if mean: + rets = sum(rets) / len(rets) -@handle_exceptions -def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: - if len(index) == 1: - idx = index[0] - if isinstance(nest, list): - nest.insert(idx, value) - else: - nest[index[0]] = value - else: - insert_into_nest_at_index(nest[index[0]], index[1:], value) + return rets @handle_exceptions @@ -337,162 +583,71 @@ def map_nest_at_index( @handle_exceptions -def multi_index_nest( - nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], - indices: Iterable[Iterable[int]], +def map_nest_at_indices( + nest: Iterable, + indices: Tuple, + fn: Callable, /, -) -> Iterable[Any]: + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: """ - Repeatedly index a nested object, using a tuple of tuples of indices or keys in the - case of dicts. + Map a function to the values of a nested item at the specified indices. Parameters ---------- nest - The nested object to slice. + The nested object to update. indices - A tuple of tuples of indices to apply. + A tuple of tuples of indices for the indices at which to update. + fn + The function to perform on the nest at the given index. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. Returns ------- ret - The result elements through indexing the nested object. + nest with applicable of fn on given indices. Examples -------- - With :code:`Tuple` inputs: - - >>> x = (1, 2) - >>> y = [[0]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [1] - - With :class:`ivy.Array` inputs: + With :code:`List` inputs: - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> y = [[0],[1]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [ivy.array([1., 2.], ivy.array([3., 4.])] + >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] + >>> indices = [[0, 4], [1, 5]] + >>> function = lambda x : x + 'b' + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] - With :class:`ivy.Container` input: + With :code:`Tuple` inputs: - >>> x = ivy.Container(a=ivy.array([1,2]), - ... b=[30,40]) - >>> y = ('a', ('b', 0)) - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [ivy.array([1, 2]), 30] + >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) + >>> indices = ((0, 2),(1, 0),(1, 2)) + >>> function = abs + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + ([-9, 8, 27], [9, -4, 5, 7]) With :code:`Dict` input: - >>> x = {'a': 0, 'b': [1, [2, 3]], 'c': (4, 5)} - >>> y = (('b', 1), 'a') - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [[2, 3], 0] - - With :code:`List` inputs: - - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> y = [[2, 1, 0], [0, 1]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - ['h', 'b'] - """ - return [index_nest(nest, index) for index in indices] - - -@handle_exceptions -def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: - """ - Prune a nested object at specified indices. - - Parameters - ---------- - nest - The nested object to prune. - indices - A tuple of tuples of indices for the indices at which to prune. - """ - # Delete first deeper elements and elements with larger index - indices_sorted = sorted( - indices, - key=str, - reverse=True, - ) - [prune_nest_at_index(nest, index) for index in indices_sorted] - - -@handle_exceptions -def set_nest_at_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], - indices: Union[List[int], Tuple[int], Iterable[int]], - values: Union[List[int], Tuple[int], Iterable[int]], - /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: - """ - Set the value of a nested item at specified indices with specified values. - - Parameters - ---------- - nest - The nested object to update. - indices - A tuple of tuples of indices for the indices at which to update. - values - The new values for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - - Returns - ------- - ret - nest with updated values at the given indices. - - Examples - -------- - With :code:`List` inputs: - - >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] - >>> indices = [[0, 4], [1, 3]] - >>> values = [111, 'x'] - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] - - With :code:`Tuple` inputs: - - >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] - >>> indices = ((0, 1),(1, 2)) - >>> values = ('ivy', 'x') - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) - - With :code:`Dict` input: - - >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} - >>> indices = (('a', 1), ('b', 2), ('c', 0)) - >>> values = (11., 22., 33.) - >>> ivy.set_nest_at_indices(nest, indices, values) + >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} + >>> indices = (('a', 2), ('b', 0), ('c', 1)) + >>> function = lambda x : x + 1 + >>> ivy.map_nest_at_indices(nest, indices, function) >>> print(nest) - {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} + {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} With :class:`ivy.Array` inputs: - >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) - >>> indices = ((0, 1),(1, 2)) - >>> values = (11., 22.) - >>> ivy.set_nest_at_indices(nest, indices, values) + >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) + >>> indices = ((0, 1),(1, 1),(1, 2)) + >>> function = lambda x : x ** 2 + >>> ivy.map_nest_at_indices(nest, indices, function) >>> print(nest) - ivy.array([[1., 11., 3.], [4., 5., 22.]]) + ivy.array([[ -9., 64., -17.], + [ 11., 9., 25.]]) """ is_tuple = isinstance(nest, tuple) nest_type = type(nest) if is_tuple else lambda x: x @@ -501,10 +656,8 @@ def set_nest_at_indices( else: result = copy_nest(nest, include_derived=True) result = list(result) if is_tuple else result - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - for index, value in zip(indices, values): - result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) + for i, index in enumerate(indices): + result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) try: result = nest_type(result) except TypeError: @@ -513,125 +666,87 @@ def set_nest_at_indices( @handle_exceptions -def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: - """ - Insert a value into the nested item at specified indices with specified values. - - Parameters - ---------- - nest - The nested object to insert into. - indices - A tuple of tuples of indices for the indices at which to insert - values. - values - The new values for inserting. - """ - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - [ - insert_into_nest_at_index(nest, index, value) - for index, value in zip(indices, values) - ] - - -@handle_exceptions -def map_nest_at_indices( - nest: Iterable, - indices: Tuple, - fn: Callable, +def multi_index_nest( + nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], + indices: Iterable[Iterable[int]], /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: +) -> Iterable[Any]: """ - Map a function to the values of a nested item at the specified indices. + Repeatedly index a nested object, using a tuple of tuples of indices or keys in the + case of dicts. Parameters ---------- nest - The nested object to update. + The nested object to slice. indices - A tuple of tuples of indices for the indices at which to update. - fn - The function to perform on the nest at the given index. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. + A tuple of tuples of indices to apply. Returns ------- ret - nest with applicable of fn on given indices. + The result elements through indexing the nested object. Examples -------- - With :code:`List` inputs: + With :code:`Tuple` inputs: - >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] - >>> indices = [[0, 4], [1, 5]] - >>> function = lambda x : x + 'b' - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] + >>> x = (1, 2) + >>> y = [[0]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [1] - With :code:`Tuple` inputs: + With :class:`ivy.Array` inputs: - >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) - >>> indices = ((0, 2),(1, 0),(1, 2)) - >>> function = abs - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ([-9, 8, 27], [9, -4, 5, 7]) + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> y = [[0],[1]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [ivy.array([1., 2.], ivy.array([3., 4.])] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1,2]), + ... b=[30,40]) + >>> y = ('a', ('b', 0)) + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [ivy.array([1, 2]), 30] With :code:`Dict` input: - >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} - >>> indices = (('a', 2), ('b', 0), ('c', 1)) - >>> function = lambda x : x + 1 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} + >>> x = {'a': 0, 'b': [1, [2, 3]], 'c': (4, 5)} + >>> y = (('b', 1), 'a') + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [[2, 3], 0] - With :class:`ivy.Array` inputs: + With :code:`List` inputs: - >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) - >>> indices = ((0, 1),(1, 1),(1, 2)) - >>> function = lambda x : x ** 2 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ivy.array([[ -9., 64., -17.], - [ 11., 9., 25.]]) + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> y = [[2, 1, 0], [0, 1]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + ['h', 'b'] """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if shallow: - result = nest_type(nest) - else: - result = copy_nest(nest, include_derived=True) - result = list(result) if is_tuple else result - for i, index in enumerate(indices): - result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) - try: - result = nest_type(result) - except TypeError: - result = nest_type(*result) - return result + return [index_nest(nest, index) for index in indices] @handle_exceptions -def nested_argwhere( +def nested_any( nest: Iterable, fn: Callable, check_nests: bool = False, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - _index: Optional[List] = None, _base: bool = True, - stop_after_n_found: Optional[int] = None, extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[Iterable, bool]: +) -> bool: """ - Check the leaf nodes of nested x via function fn, and returns all nest indices where - the method evaluates as True. + Check the leaf nodes of nest x via function fn, and returns True if any evaluate to + True, else False. Parameters ---------- @@ -642,9 +757,66 @@ def nested_argwhere( check_nests Whether to also check the nests for the condition, not only nest leaves. Default is ``False``. - to_ignore - Types to ignore when deciding whether to go deeper into the nest or not - _index + _base + Whether the current function call is the first function call in the recursive + stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + A boolean, whether the function evaluates to true for any leaf node. + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + if ivy.any(fn(nest)): + return True + for i, item in enumerate(nest): + if nested_any(item, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif isinstance(nest, dict): + for k, v in nest.items(): + if nested_any(v, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif fn(nest): + return True + return False + + +@handle_exceptions +def nested_argwhere( + nest: Iterable, + fn: Callable, + check_nests: bool = False, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + _index: Optional[List] = None, + _base: bool = True, + stop_after_n_found: Optional[int] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[Iterable, bool]: + """ + Check the leaf nodes of nested x via function fn, and returns all nest indices where + the method evaluates as True. + + Parameters + ---------- + nest + The nest to check the leaves of. + fn + The conditon function, returning True or False. + check_nests + Whether to also check the nests for the condition, not only nest leaves. + Default is ``False``. + to_ignore + Types to ignore when deciding whether to go deeper into the nest or not + _index The indices detected so far. None at the beginning. Used internally, do not set manually. _base @@ -799,330 +971,125 @@ def nested_argwhere( @handle_exceptions -def all_nested_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, +def nested_map( + x: Union[ivy.Array, ivy.NativeArray, Iterable], /, - include_nests: bool = False, - _index: Optional[Union[int, Sequence[int]]] = None, - _base: bool = True, - extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + fn: Callable, + include_derived: Optional[Union[Dict[type, bool], bool]] = None, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + to_mutable: bool = False, + max_depth: Optional[int] = None, + _depth: int = 0, + _tuple_check_fn: Optional[Callable] = None, + _list_check_fn: Optional[Callable] = None, + _dict_check_fn: Optional[Callable] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: """ - Return indices of all the elements in nest. + Apply a function on x in a nested manner, whereby all dicts, lists and tuples are + traversed to their lowest leaves before applying the method and returning x. If x is + not nested, the method is applied to x directly. Parameters ---------- - nest - The nest to check the leaves of. - include_nests - Whether to also include indices of the nests themselves, not only - leaves. Default is ``False``. - _index - The indices detected so far. None at the beginning. Used internally, - do not set manually. - _base - Whether the current function call is the first function call in the - recursive stack. Used internally, do not set manually. + x + The item to apply the mapped function to. + fn + The function to map onto x. + include_derived + Whether to also recursive for classes derived from tuple, list and dict. + Default is ``False``. + to_ignore + Types to ignore when deciding whether to go deeper into the nest or not + to_mutable + Whether to convert the nest to a mutable form, changing all tuples to lists. + Default is ``False``. + max_depth + The maximum nested depth to reach. Default is 1. Increase this if the nest is + deeper. + _depth + Placeholder for tracking the recursive depth, do not set this parameter. + _tuple_check_fn + Placeholder for the tuple check function, do not set this parameter. + _list_check_fn + Placeholder for the list check function, do not set this parameter. + _dict_check_fn + Placeholder for the dict check function, do not set this parameter. extra_nest_types Types to recursively check when deciding whether to go deeper into the nest or not - out - Optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. Returns ------- ret - A set of indices of all elements in nest - - Both the description and the type hints above assumes an array input - for simplicity, but this function is nestable, and therefore also - accepts :class:ivy.Container instances in place of the arguments. + x following the applicable of fn to it's nested leaves, or x itself if x is not + nested. Examples -------- - With :class:`Dict` input: - - >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} - >>> y = ivy.all_nested_indices(x) - >>> print(y) - [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] + With :class:`Tuple` inputs: - With :class:`ivy.Array` input: + >>> x = ([[1., 2.], [3., 4.]]) + >>> function = lambda a : a * 2 + >>> ivy.nested_map(x, function) + [[2.0, 4.0], [6.0, 8.0]] + >>> print(x) + [[2.0, 4.0], [6.0, 8.0]] - >>> x = ivy.array([0., 1., 2., 3., 4.]) - >>> y = ivy.all_nested_indices(x, False, out=x) - >>> print(y) - [[]] + With :code:`Dict` input: - With :class:`ivy.Container` input: + >>> x = {1 : [1, [2, 3]], 2: (4, 5)} + >>> function = lambda a : a + 1 + >>> ivy.nested_map(x, function) + {1 : [2, [3, 4]], 2: (5, 6)} + >>> print(x) + {1 : [2, [3, 4]], 2: (5, 6)} - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.all_nested_indices(x, True) - >>> print(y) - [['a'], ['b']] - """ - _index = list() if _index is None else _index - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - ind = ivy.argwhere(ivy.ones_like(nest)) - indices = list() - for i in range(len(ind)): - indices.append(_index + ind.to_list()[i]) - return indices - _indices = [ - all_nested_indices( - item, include_nests, _index + [i], False, extra_nest_types - ) - for i, item in enumerate(nest) - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - elif isinstance(nest, dict): - _indices = [ - all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) - for k, v in nest.items() - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - else: - return [_index] - return [index for index in _indices if index] + With :code:`List` inputs: + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> function = lambda a: a + 'H' + >>> ivy.nested_map(x, function) + [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] + >>> print(x) + [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] -# noinspection PyShadowingBuiltins + With :class:`ivy.Container` input: + >>> x = ivy.Container( + ... a=ivy.array([[1, 2, 3], [9, 8, 7]]) , b=ivy.array([[4, 5, 6], [12, 13, 14]]) + ... ) + >>> function = lambda a : a + 1 + >>> ivy.nested_map(x, function) + { + a: ivy.array([[2, 3, 4], + [10, 9, 8]]), + b: ivy.array([[5, 6, 7], + [13, 14, 15]]) + } + >>> print(x) + { + a: ivy.array([[2, 3, 4], + [10, 9, 8]]), + b: ivy.array([[5, 6, 7], + [13, 14, 15]]) + } -@handle_exceptions -def map( - fn: Callable, - constant: Optional[Dict[str, Any]] = None, - unique: Optional[Dict[str, Iterable[Any]]] = None, - mean: bool = False, -) -> List: - """ - Apply a function on each item of an iterable x. + >>> nest = ([1, 2], [3, 4], [5, 6], {"a": 1, "b": 2, "c": 3}) + >>> function = lambda a : a * 2 + >>> ivy.nested_map(nest, function, to_ignore=list) + ([1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6], {'a': 2, 'b': 4, 'c': 6}) - Parameters - ---------- - fn - The function to map onto x. - constant - keyword arguments which remain constant between each function call. - Default is ``None``. - unique - keyword arguments which are unique for each function call. Default is ``None``. - mean - Whether to compute the mean across the return values, and return this mean. - Default is ``False``. - - Returns - ------- - ret - x following the application of fn to each of its iterated items. - - Examples - -------- - With :code:`int` inputs: - - >>> def special_square(x : float) -> float : return np.square(x) - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x' : [1,2,3]}, - ... mean = False) - >>> print(results) - [1, 4, 9] - - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x':[0,1,2]}, - ... mean = True) - >>> print(results) - 1.6666666666666667 - - >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = False) - >>> print(results) - [array([1,1]), - array([1,2]), - array([1,3])] - - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = True) - >>> print(results) - [1. 2.] - - With float inputs: - - >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':10., 'b':1.}, - ... unique = {'x':[0.,1.,2.]}, - ... mean = False) - >>> print(results) - [1.0, 11.0, 21.0] - - With :class:`ivy.Array` inputs: - - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = False) - >>> print(results) - [ivy.array([0., 10., 100.]), - ivy.array([1., 10., 101.])] - - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = True) - >>> print(results) - ivy.array([ 0.5, 10. , 100. ]) - """ - c = ivy.default(constant, {}) - u = ivy.default(unique, {}) - rets = [ - r - for r in _map( - lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() - ) - ] - if mean: - rets = sum(rets) / len(rets) - - return rets - - -@handle_exceptions -def nested_map( - x: Union[ivy.Array, ivy.NativeArray, Iterable], - /, - fn: Callable, - include_derived: Optional[Union[Dict[type, bool], bool]] = None, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - to_mutable: bool = False, - max_depth: Optional[int] = None, - _depth: int = 0, - _tuple_check_fn: Optional[Callable] = None, - _list_check_fn: Optional[Callable] = None, - _dict_check_fn: Optional[Callable] = None, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: - """ - Apply a function on x in a nested manner, whereby all dicts, lists and tuples are - traversed to their lowest leaves before applying the method and returning x. If x is - not nested, the method is applied to x directly. - - Parameters - ---------- - x - The item to apply the mapped function to. - fn - The function to map onto x. - include_derived - Whether to also recursive for classes derived from tuple, list and dict. - Default is ``False``. - to_ignore - Types to ignore when deciding whether to go deeper into the nest or not - to_mutable - Whether to convert the nest to a mutable form, changing all tuples to lists. - Default is ``False``. - max_depth - The maximum nested depth to reach. Default is 1. Increase this if the nest is - deeper. - _depth - Placeholder for tracking the recursive depth, do not set this parameter. - _tuple_check_fn - Placeholder for the tuple check function, do not set this parameter. - _list_check_fn - Placeholder for the list check function, do not set this parameter. - _dict_check_fn - Placeholder for the dict check function, do not set this parameter. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - - Returns - ------- - ret - x following the applicable of fn to it's nested leaves, or x itself if x is not - nested. - - Examples - -------- - With :class:`Tuple` inputs: - - >>> x = ([[1., 2.], [3., 4.]]) - >>> function = lambda a : a * 2 - >>> ivy.nested_map(x, function) - [[2.0, 4.0], [6.0, 8.0]] - >>> print(x) - [[2.0, 4.0], [6.0, 8.0]] - - With :code:`Dict` input: - - >>> x = {1 : [1, [2, 3]], 2: (4, 5)} - >>> function = lambda a : a + 1 - >>> ivy.nested_map(x, function) - {1 : [2, [3, 4]], 2: (5, 6)} - >>> print(x) - {1 : [2, [3, 4]], 2: (5, 6)} - - With :code:`List` inputs: - - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> function = lambda a: a + 'H' - >>> ivy.nested_map(x, function) - [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] - >>> print(x) - [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] - - With :class:`ivy.Container` input: - - >>> x = ivy.Container( - ... a=ivy.array([[1, 2, 3], [9, 8, 7]]) , b=ivy.array([[4, 5, 6], [12, 13, 14]]) - ... ) - >>> function = lambda a : a + 1 - >>> ivy.nested_map(x, function) - { - a: ivy.array([[2, 3, 4], - [10, 9, 8]]), - b: ivy.array([[5, 6, 7], - [13, 14, 15]]) - } - >>> print(x) - { - a: ivy.array([[2, 3, 4], - [10, 9, 8]]), - b: ivy.array([[5, 6, 7], - [13, 14, 15]]) - } - - >>> nest = ([1, 2], [3, 4], [5, 6], {"a": 1, "b": 2, "c": 3}) - >>> function = lambda a : a * 2 - >>> ivy.nested_map(nest, function, to_ignore=list) - ([1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6], {'a': 2, 'b': 4, 'c': 6}) - - >>> nest = [[1, 2], [3, [4, 5]], [[6], [7, 8, [9, 10]]]] - >>> function = lambda a : a * 2 - >>> ivy.nested_map(nest, function, max_depth = 3) - [[2, 4], [6, [8, 10]], [[12], [14, 16, [9, 10]]]] + >>> nest = [[1, 2], [3, [4, 5]], [[6], [7, 8, [9, 10]]]] + >>> function = lambda a : a * 2 + >>> ivy.nested_map(nest, function, max_depth = 3) + [[2, 4], [6, [8, 10]], [[12], [14, 16, [9, 10]]]] >>> nest = ([23, 25, 1337], [63, 98, 6]) >>> function = lambda a : a + 1 @@ -1261,211 +1228,46 @@ def nested_map( @handle_exceptions -def nested_any( - nest: Iterable, - fn: Callable, - check_nests: bool = False, - _base: bool = True, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> bool: +def nested_multi_map( + func: Callable, + nests: List[Iterable], + index_chains=None, + to_apply=True, + prune_unapplied=False, + index_chain="", + config=None, + to_ivy=True, +): """ - Check the leaf nodes of nest x via function fn, and returns True if any evaluate to - True, else False. + Apply function to all array values from a collection of identically structured ivy + arrays. Parameters ---------- + func + Function to apply to each nest entry. nest - The nest to check the leaves of. - fn - The conditon function, returning True or False. - check_nests - Whether to also check the nests for the condition, not only nest leaves. - Default is ``False``. - _base - Whether the current function call is the first function call in the recursive - stack. Used internally, do not set manually. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - + nests to map. + index_chains + The key-chains to apply or not apply the method to. Default is ``None``. + to_apply + If True, the method will be applied to index_chains, otherwise index_chains will + be skipped. Default is ``True``. + prune_unapplied + Whether to prune index_chains for which the function was not applied, + otherwise the leftmost nest value is used. Default is ``False``. + index_chain + Chain of keys for this dict entry (Default value = '') + config + The configuration for the nests. Default is the same as nest0. + to_ivy + convert the output to ivy_arrays. Default is ``True`` Returns ------- - ret - A boolean, whether the function evaluates to true for any leaf node. - """ - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - if ivy.any(fn(nest)): - return True - for i, item in enumerate(nest): - if nested_any(item, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif isinstance(nest, dict): - for k, v in nest.items(): - if nested_any(v, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif fn(nest): - return True - return False - - -@handle_exceptions -def copy_nest( - nest: Union[ivy.Array, ivy.NativeArray, Iterable], - /, - include_derived: bool = False, - to_mutable: bool = False, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[ivy.Array, ivy.NativeArray, Iterable]: - """ - Copy a nest deeply, but without copying leaves of the nest, only the nest lists, - tuples and dicts are copied. - - Parameters - ---------- - nest - The nest to copy. - include_derived - Whether to also recursive for classes derived from tuple, list and dict. - Default is ``False``. - to_mutable - Whether to convert the nest to a mutable form, changing all tuples to lists. - Default is ``False``. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - - Returns - ------- - ret - The copied nest. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - ivy.array([[1., 2., 3.], - [7., 8., 9.]]) - - With :code:`Iterable` input: - - >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - >>> copied_nest = ivy.copy_nest(nest, include_derived = True) - >>> print(copied_nest) - [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - - >>> nest = ([23, 25, 1337], [63, 98, 6]) - >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) - >>> print(copied_nest) - [[23, 25, 1337], [63, 98, 6]] - - >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} - """ - extra_nest_types = ivy.default(extra_nest_types, ()) - class_instance = type(nest) - check_fn = ( - (lambda x_, t: isinstance(nest, t)) - if include_derived - else (lambda x_, t: type(nest) is t) - ) - if check_fn(nest, tuple): - ret_list = [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - if to_mutable: - return ret_list - if hasattr(nest, "_fields"): - return class_instance(**dict(zip(nest._fields, ret_list))) - return class_instance(tuple(ret_list)) - elif check_fn(nest, list) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - return copy.deepcopy(nest) - return class_instance( - [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - ) - elif check_fn(nest, dict): - class_instance = type(nest) - dict_ = { - k: copy_nest( - v, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for k, v in nest.items() - } - if isinstance(nest, OrderedDict): - return class_instance(**dict_) - return class_instance(dict_) - return nest - - -@handle_exceptions -def nested_multi_map( - func: Callable, - nests: List[Iterable], - index_chains=None, - to_apply=True, - prune_unapplied=False, - index_chain="", - config=None, - to_ivy=True, -): - """ - Apply function to all array values from a collection of identically structured ivy - arrays. - - Parameters - ---------- - func - Function to apply to each nest entry. - nest - nests to map. - index_chains - The key-chains to apply or not apply the method to. Default is ``None``. - to_apply - If True, the method will be applied to index_chains, otherwise index_chains will - be skipped. Default is ``True``. - prune_unapplied - Whether to prune index_chains for which the function was not applied, - otherwise the leftmost nest value is used. Default is ``False``. - index_chain - Chain of keys for this dict entry (Default value = '') - config - The configuration for the nests. Default is the same as nest0. - to_ivy - convert the output to ivy_arrays. Default is ``True`` - Returns - ------- - nest containing the result of the function. The structure of the output is the - same as the input with the result of the function applied to each applicable - leaf and the value at that leaf in the first nest for a non-applicable leaf if - prune_unapplied is False else unapplied leaves are pruned. + nest containing the result of the function. The structure of the output is the + same as the input with the result of the function applied to each applicable + leaf and the value at that leaf in the first nest for a non-applicable leaf if + prune_unapplied is False else unapplied leaves are pruned. """ nest0 = None for nest in nests: @@ -1583,38 +1385,6 @@ def _found_in_index_chains(this_index_chain, index_chains): ) -@handle_exceptions -def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): - """ - Group all unique index chains in a nest. This function is useful for finding all - unique index chains in a nest, and then duplicating the values at those index chains - for functional frameworks. - - Parameters - ---------- - nest - nest to get duplicate index chains for. - - Returns - ------- - list of index chains to duplicate. - """ - all_index_chains = ivy.nested_argwhere(nest, lambda _: True) - duplicates = [] - duplicate_index_chains = {} - for index_chain in all_index_chains: - val = ivy.index_nest(nest, index_chain) - if ivy.is_array(val): - for i in range(len(duplicates)): - if val is duplicates[i]: - duplicate_index_chains[i].append(index_chain) - break - else: - duplicates.append(val) - duplicate_index_chains[len(duplicates) - 1] = [index_chain] - return list(duplicate_index_chains.values()) - - def prune_empty(nest): """ Prune empty nests from a nest. @@ -1650,3 +1420,233 @@ def prune_empty(nest): if not valid and not (ivy.is_array(nest) or isinstance(nest, (int, float, str))): return None return nest + + +@handle_exceptions +def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: + """ + Prune a nested object at a specified index. + + Parameters + ---------- + nest + The nested object to prune. + index + A tuple of indices for the index at which to prune. + """ + if len(index) == 1: + del nest[index[0]] + else: + prune_nest_at_index(nest[index[0]], index[1:]) + + +@handle_exceptions +def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: + """ + Prune a nested object at specified indices. + + Parameters + ---------- + nest + The nested object to prune. + indices + A tuple of tuples of indices for the indices at which to prune. + """ + # Delete first deeper elements and elements with larger index + indices_sorted = sorted( + indices, + key=str, + reverse=True, + ) + [prune_nest_at_index(nest, index) for index in indices_sorted] + + +@handle_exceptions +def set_nest_at_index( + nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], + index: Sequence[Union[str, int]], + value: Any, + /, + shallow: bool = True, + _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: + """ + Set the value of a nested item at a specified index. + + Parameters + ---------- + nest + The nested object to update. + index + A tuple of indices for the index at which to update. + value + The new value for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + _result + Placeholder for the result of the update. do not set this paramter. + + Returns + ------- + ret + nest with changed value at the given index. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = (1, 1) + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([[1., 2.], [3., 5.]]) + + >>> x = ivy.array([1., 2., 3., 4.]) + >>> y = [1] + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([1., 5., 3., 4.]) + + With :code:`Dict` input: + + >>> x = {1 : [1, [2, 3]], 2: (4, 5)} + >>> y = (1, 1) + >>> z = 2 + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + {1: [1, 2], 2: (4, 5)} + + With :code:`List` inputs: + + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> y = (2, 1, 0) + >>> z = 'H' + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + [['a','b','c'],['d','e','f'],['g',['H','i']]] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) + >>> y = ('b',) + >>> z = ivy.array([3., 4.]) + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + { + a: ivy.array([1., 2.]), + b: ivy.array([3., 4.]) + } + """ + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if _result is None: + if shallow: + _result = nest_type(nest) + else: + _result = copy_nest(nest, include_derived=True) + _result = list(_result) if is_tuple else _result + if len(index) == 1: + if shallow: + try: + nest[index[0]] = value + except TypeError: + pass + _result[index[0]] = value + else: + _result[index[0]] = set_nest_at_index( + nest[index[0]], index[1:], value, shallow, _result[index[0]] + ) + try: + _result = nest_type(_result) + except TypeError: + _result = nest_type(*_result) + return _result + + +@handle_exceptions +def set_nest_at_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], + indices: Union[List[int], Tuple[int], Iterable[int]], + values: Union[List[int], Tuple[int], Iterable[int]], + /, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: + """ + Set the value of a nested item at specified indices with specified values. + + Parameters + ---------- + nest + The nested object to update. + indices + A tuple of tuples of indices for the indices at which to update. + values + The new values for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + + Returns + ------- + ret + nest with updated values at the given indices. + + Examples + -------- + With :code:`List` inputs: + + >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] + >>> indices = [[0, 4], [1, 3]] + >>> values = [111, 'x'] + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] + + With :code:`Tuple` inputs: + + >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] + >>> indices = ((0, 1),(1, 2)) + >>> values = ('ivy', 'x') + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) + + With :code:`Dict` input: + + >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} + >>> indices = (('a', 1), ('b', 2), ('c', 0)) + >>> values = (11., 22., 33.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} + + With :class:`ivy.Array` inputs: + + >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) + >>> indices = ((0, 1),(1, 2)) + >>> values = (11., 22.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + ivy.array([[1., 11., 3.], [4., 5., 22.]]) + """ + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if shallow: + result = nest_type(nest) + else: + result = copy_nest(nest, include_derived=True) + result = list(result) if is_tuple else result + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + for index, value in zip(indices, values): + result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) + try: + result = nest_type(result) + except TypeError: + result = nest_type(*result) + return result diff --git a/ivy/functional/ivy/norms.py b/ivy/functional/ivy/norms.py index e0804ca4dc668..48d0cce6435e0 100644 --- a/ivy/functional/ivy/norms.py +++ b/ivy/functional/ivy/norms.py @@ -13,6 +13,18 @@ from ivy.utils.exceptions import handle_exceptions +layer_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + + # Extra # # ------# @@ -138,15 +150,3 @@ def layer_norm( return ivy.multiply(ivy.multiply(x, scale), new_std, out=out) return ivy.multiply(x, new_std, out=out) - - -layer_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} diff --git a/ivy/functional/ivy/random.py b/ivy/functional/ivy/random.py index 7ed96ae869d1c..70e8c21a5db45 100644 --- a/ivy/functional/ivy/random.py +++ b/ivy/functional/ivy/random.py @@ -20,8 +20,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # +# --- Helpers --- # +# --------------- # def _check_bounds_and_get_shape(low, high, shape): @@ -51,6 +51,17 @@ def _check_bounds_and_get_shape(low, high, shape): return ivy.Shape(()) +def _check_shapes_broadcastable(out, inp): + if out is not None: + ivy.utils.assertions.check_shapes_broadcastable(out, inp) + + +def _check_valid_scale(std): + ivy.utils.assertions.check_greater( + std, 0, allow_equal=True, message="std must be non-negative" + ) + + def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_all_or_any_fn( low, @@ -73,66 +84,51 @@ def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_less(low, high) -def _check_valid_scale(std): - ivy.utils.assertions.check_greater( - std, 0, allow_equal=True, message="std must be non-negative" - ) - - -def _check_shapes_broadcastable(out, inp): - if out is not None: - ivy.utils.assertions.check_shapes_broadcastable(out, inp) - - -# Extra # -# ------# +# --- Main --- # +# ------------ # @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function -@infer_dtype @handle_device_shifting @infer_device -def random_uniform( +def multinomial( + population_size: int, + num_samples: int, + /, *, - low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, - high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, - shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, + batch_size: int = 1, + probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + replace: bool = True, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a uniform distribution. Samples are uniformly distributed over - the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In - other words, any value within the given interval is equally likely to be drawn by - uniform. + Draws samples from a multinomial distribution. Specifically, returns a tensor where + each row contains num_samples indices sampled from the multinomial probability + distribution located in the corresponding row of tensor input. Parameters ---------- - low - Lower boundary of the output interval. All values generated will be greater than - or equal to ``low``. If array, must have same shape as ``high``. - high - Upper boundary of the output interval. All the values generated will be less - than ``high``. If array, must have same shape as ``low``. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. - Can only be specified when ``low`` and ``high`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. + population_size + The size of the population from which to draw samples. + num_samples + Number of independent samples to draw from the population. + batch_size + Number of tensors to generate. Default is 1. + probs + The unnormalized probabilities for all elements in population, + default is uniform *[batch_shape, population_size]* + replace + Whether to replace samples once they've been drawn. Default is ``True``. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` + (Default value = None) seed A python integer. Used to create a random seed distribution out @@ -142,67 +138,145 @@ def random_uniform( Returns ------- ret - Drawn samples from the parameterized uniform distribution. + Drawn samples indices from the multinomial distribution. - Functional Examples - ------------------- + Examples + -------- + >>> y = ivy.multinomial(10, 5) + >>> print(y) + ivy.array([[1, 8, 7, 8, 3]]) - >>> ivy.random_uniform() - ivy.array(0.26431865) + >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) + >>> print(y) + ivy.array([[3, 9, 7, 5, 1], + [1, 0, 8, 6, 7]]) - >>> ivy.random_uniform(shape=3) - ivy.array([0.475, 0.878, 0.861]) + >>> y = ivy.multinomial(10, 5, replace=False) + >>> print(y) + ivy.array([[2, 6, 4, 7, 0]]) - >>> ivy.random_uniform(shape=(2,3)) - ivy.array([[0.929 , 0.545 , 0.789 ], - [0.519 , 0.0435, 0.381 ]]) + With :class:`ivy.Array` input: - >>> ivy.random_uniform(low=3.0, high=6.0) - ivy.array(3.4608004) + >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) + >>> print(y) + ivy.array([5, 2, 7, 6, 9]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) - ivy.array([[1.85], - [1.81]]) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) + >>> print(y) + ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) - >>> z = ivy.zeros(()) - >>> ivy.random_uniform(low=1.0, high=2.0, out=z) - ivy.array(1.8458502) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), + ... replace=False) + >>> print(y) + ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') - ivy.array([[1.81, 1.8 ], - [1.32, 1.43]]) + With :class:`ivy.NativeArray` input: - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', - ... dtype='int32') - ivy.array([[1, 1], - [1, 1]]) + >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) + >>> print(y) + ivy.array([5, 7, 4, 2, 1]) - >>> z = ivy.zeros((1,2)) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', - ... dtype='float64', out=z) - ivy.array([[1.34, 1.02]]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) + >>> print(y) + ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) - >>> x = ivy.array([4.8, 5.6]) - >>> y = ivy.array([9.8, 7.4]) - >>> ivy.random_uniform(low=x, high=y) - ivy.array([0.475, 0.878]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), + ... replace=False) + >>> print(y) + ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) + """ + return ivy.current_backend().multinomial( + population_size, + num_samples, + batch_size=batch_size, + probs=probs, + replace=replace, + device=device, + seed=seed, + out=out, + ) - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) - ivy.array([6.67270088, 7.31128597]) - >>> ivy.random_uniform(low=x, high=y, device='cpu') - ivy.array([6.88, 6.75]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +@infer_device +def randint( + low: Union[int, ivy.NativeArray, ivy.Array], + high: Union[int, ivy.NativeArray, ivy.Array], + /, + *, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array filled with random integers generated uniformly between low + (inclusive) and high (exclusive). - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') - ivy.array([8.62, 6.47]) + Parameters + ---------- + low + Lowest integer that can be drawn from the distribution. + high + One above the highest integer that can be drawn from the distribution. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn + Can only be specified when ``mean`` and ``std`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. + device + device on which to create the array. 'cuda:0', + 'cuda:1', 'cpu' etc. (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default integer data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) - ivy.array([5. , 7.3]) + Returns + ------- + ret + Returns an array with the given shape filled with integers from + the uniform distribution in the “half-open” interval [low, high) + + Examples + -------- + >>> y = ivy.randint(0, 9, shape=(1,1)) + >>> print(y) + ivy.array([[5]]) + + >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) + >>> print(y) + ivy.array([[ 8, 16], + [12, 9]]) + + >>> x = ivy.array([1, 2, 3]) + >>> ivy.randint(0, 10, shape=(3,), out=x) + >>> print(x) + ivy.array([2, 6, 7]) + + >>> y = ivy.zeros((3, 3)) + >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) + >>> print(y) + ivy.array([[ 7, 7, 5], + [12, 8, 8], + [ 8, 11, 3]]) """ - return ivy.current_backend().random_uniform( - low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed + return ivy.current_backend().randint( + low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out ) @@ -320,47 +394,55 @@ def random_normal( ) +# Extra # +# ------# + + @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function +@infer_dtype @handle_device_shifting @infer_device -def multinomial( - population_size: int, - num_samples: int, - /, +def random_uniform( *, - batch_size: int = 1, - probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - replace: bool = True, + low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, + high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, + shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a multinomial distribution. Specifically, returns a tensor where - each row contains num_samples indices sampled from the multinomial probability - distribution located in the corresponding row of tensor input. + Draws samples from a uniform distribution. Samples are uniformly distributed over + the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In + other words, any value within the given interval is equally likely to be drawn by + uniform. Parameters ---------- - population_size - The size of the population from which to draw samples. - num_samples - Number of independent samples to draw from the population. - batch_size - Number of tensors to generate. Default is 1. - probs - The unnormalized probabilities for all elements in population, - default is uniform *[batch_shape, population_size]* - replace - Whether to replace samples once they've been drawn. Default is ``True``. + low + Lower boundary of the output interval. All values generated will be greater than + or equal to ``low``. If array, must have same shape as ``high``. + high + Upper boundary of the output interval. All the values generated will be less + than ``high``. If array, must have same shape as ``low``. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. + Can only be specified when ``low`` and ``high`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None) + (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out @@ -370,145 +452,67 @@ def multinomial( Returns ------- ret - Drawn samples indices from the multinomial distribution. - - Examples - -------- - >>> y = ivy.multinomial(10, 5) - >>> print(y) - ivy.array([[1, 8, 7, 8, 3]]) - - >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) - >>> print(y) - ivy.array([[3, 9, 7, 5, 1], - [1, 0, 8, 6, 7]]) - - >>> y = ivy.multinomial(10, 5, replace=False) - >>> print(y) - ivy.array([[2, 6, 4, 7, 0]]) - - With :class:`ivy.Array` input: + Drawn samples from the parameterized uniform distribution. - >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) - >>> print(y) - ivy.array([5, 2, 7, 6, 9]) + Functional Examples + ------------------- - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) - >>> print(y) - ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) + >>> ivy.random_uniform() + ivy.array(0.26431865) - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), - ... replace=False) - >>> print(y) - ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) + >>> ivy.random_uniform(shape=3) + ivy.array([0.475, 0.878, 0.861]) - With :class:`ivy.NativeArray` input: + >>> ivy.random_uniform(shape=(2,3)) + ivy.array([[0.929 , 0.545 , 0.789 ], + [0.519 , 0.0435, 0.381 ]]) - >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) - >>> print(y) - ivy.array([5, 7, 4, 2, 1]) + >>> ivy.random_uniform(low=3.0, high=6.0) + ivy.array(3.4608004) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) - >>> print(y) - ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) + ivy.array([[1.85], + [1.81]]) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), - ... replace=False) - >>> print(y) - ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) - """ - return ivy.current_backend().multinomial( - population_size, - num_samples, - batch_size=batch_size, - probs=probs, - replace=replace, - device=device, - seed=seed, - out=out, - ) + >>> z = ivy.zeros(()) + >>> ivy.random_uniform(low=1.0, high=2.0, out=z) + ivy.array(1.8458502) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') + ivy.array([[1.81, 1.8 ], + [1.32, 1.43]]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -@infer_device -def randint( - low: Union[int, ivy.NativeArray, ivy.Array], - high: Union[int, ivy.NativeArray, ivy.Array], - /, - *, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array filled with random integers generated uniformly between low - (inclusive) and high (exclusive). + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', + ... dtype='int32') + ivy.array([[1, 1], + [1, 1]]) - Parameters - ---------- - low - Lowest integer that can be drawn from the distribution. - high - One above the highest integer that can be drawn from the distribution. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn - Can only be specified when ``mean`` and ``std`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. - device - device on which to create the array. 'cuda:0', - 'cuda:1', 'cpu' etc. (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default integer data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. + >>> z = ivy.zeros((1,2)) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', + ... dtype='float64', out=z) + ivy.array([[1.34, 1.02]]) - Returns - ------- - ret - Returns an array with the given shape filled with integers from - the uniform distribution in the “half-open” interval [low, high) + >>> x = ivy.array([4.8, 5.6]) + >>> y = ivy.array([9.8, 7.4]) + >>> ivy.random_uniform(low=x, high=y) + ivy.array([0.475, 0.878]) - Examples - -------- - >>> y = ivy.randint(0, 9, shape=(1,1)) - >>> print(y) - ivy.array([[5]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) + ivy.array([6.67270088, 7.31128597]) - >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) - >>> print(y) - ivy.array([[ 8, 16], - [12, 9]]) + >>> ivy.random_uniform(low=x, high=y, device='cpu') + ivy.array([6.88, 6.75]) - >>> x = ivy.array([1, 2, 3]) - >>> ivy.randint(0, 10, shape=(3,), out=x) - >>> print(x) - ivy.array([2, 6, 7]) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') + ivy.array([8.62, 6.47]) - >>> y = ivy.zeros((3, 3)) - >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) - >>> print(y) - ivy.array([[ 7, 7, 5], - [12, 8, 8], - [ 8, 11, 3]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) + ivy.array([5. , 7.3]) """ - return ivy.current_backend().randint( - low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out + return ivy.current_backend().random_uniform( + low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed ) diff --git a/ivy/functional/ivy/searching.py b/ivy/functional/ivy/searching.py index 573c7fc1068ed..c5391a424a431 100644 --- a/ivy/functional/ivy/searching.py +++ b/ivy/functional/ivy/searching.py @@ -236,6 +236,81 @@ def argmin( ) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def argwhere( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the indices of all non-zero elements of the input array. + + Parameters + ---------- + x + input array, for which indices are desired. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Indices of non-zero elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> y = ivy.zeros((3, 2), dtype=ivy.int64) + >>> res = ivy.argwhere(x, out=y) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + With a :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0], [1]]), + b: ivy.array([[0], [1]]) + } + + >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0]]), + b: ivy.array([[0], [1]]) + } + """ + return current_backend(x).argwhere(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -454,78 +529,3 @@ def where( } """ return current_backend(x1).where(condition, x1, x2, out=out) - - -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def argwhere( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the indices of all non-zero elements of the input array. - - Parameters - ---------- - x - input array, for which indices are desired. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Indices of non-zero elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> y = ivy.zeros((3, 2), dtype=ivy.int64) - >>> res = ivy.argwhere(x, out=y) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - With a :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0], [1]]), - b: ivy.array([[0], [1]]) - } - - >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0]]), - b: ivy.array([[0], [1]]) - } - """ - return current_backend(x).argwhere(x, out=out) diff --git a/ivy/functional/ivy/set.py b/ivy/functional/ivy/set.py index 1a3e69831ede7..2c27798531508 100644 --- a/ivy/functional/ivy/set.py +++ b/ivy/functional/ivy/set.py @@ -148,6 +148,111 @@ def unique_all( return ivy.current_backend(x).unique_all(x, axis=axis, by_value=by_value) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def unique_counts( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Return the unique elements of an input array ``x`` and the corresponding counts for + each unique element in ``x``. + + .. admonition:: Data-dependent output shape + :class: important + + The shapes of two of the output arrays for this function depend on the data + values in the input array; hence, array libraries which build computation graphs + (e.g., JAX, Dask, etc.) may find this function difficult to implement without + knowing array values. Accordingly, such libraries may choose to omit this + function. See :ref:`data-dependent-output-shapes` section for more details. + + .. note:: + Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). + For input arrays having floating-point data types, value-based equality implies + the following behavior. + + - As ``nan`` values compare as ``False``, ``nan`` values should be considered + distinct. + - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be + considered distinct, and the corresponding unique element will be + implementation-dependent (e.g., an implementation could choose to return + ``-0`` if ``-0`` occurs before ``+0``). + + Parameters + ---------- + x + input array. If ``x`` has more than one dimension, the function must flatten + ``x`` and return the unique elements of the flattened array. + + Returns + ------- + ret + a namedtuple ``(values, counts)`` whose + + - first element must have the field name ``values`` and must be an + array containing the unique elements of ``x``. + The array must have the same data type as ``x``. + - second element must have the field name ``counts`` and must be an array + containing the number of times each unique element occurs in ``x``. + The returned array must have same shape as ``values`` and must + have the default array index data type. + + .. note:: + The order of unique elements is not specified and may vary between + implementations. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2,1,3,4,1,3]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) + + >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) + + >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, + 2.29999995]), + counts=ivy.array([3, 1, 1, 1, 1])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), + ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) + >>> y = ivy.unique_counts(x) + >>> print(y) + { + a: (list[2],shape=[4]), + b: (list[2],shape=[4]) + } + """ + return ivy.current_backend(x).unique_counts(x) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -348,108 +453,3 @@ def unique_values( array([0., 1., 2., 3., 4., 5., nan, -0.]) """ return ivy.current_backend(x).unique_values(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def unique_counts( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Return the unique elements of an input array ``x`` and the corresponding counts for - each unique element in ``x``. - - .. admonition:: Data-dependent output shape - :class: important - - The shapes of two of the output arrays for this function depend on the data - values in the input array; hence, array libraries which build computation graphs - (e.g., JAX, Dask, etc.) may find this function difficult to implement without - knowing array values. Accordingly, such libraries may choose to omit this - function. See :ref:`data-dependent-output-shapes` section for more details. - - .. note:: - Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). - For input arrays having floating-point data types, value-based equality implies - the following behavior. - - - As ``nan`` values compare as ``False``, ``nan`` values should be considered - distinct. - - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be - considered distinct, and the corresponding unique element will be - implementation-dependent (e.g., an implementation could choose to return - ``-0`` if ``-0`` occurs before ``+0``). - - Parameters - ---------- - x - input array. If ``x`` has more than one dimension, the function must flatten - ``x`` and return the unique elements of the flattened array. - - Returns - ------- - ret - a namedtuple ``(values, counts)`` whose - - - first element must have the field name ``values`` and must be an - array containing the unique elements of ``x``. - The array must have the same data type as ``x``. - - second element must have the field name ``counts`` and must be an array - containing the number of times each unique element occurs in ``x``. - The returned array must have same shape as ``values`` and must - have the default array index data type. - - .. note:: - The order of unique elements is not specified and may vary between - implementations. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2,1,3,4,1,3]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) - - >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) - - >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, - 2.29999995]), - counts=ivy.array([3, 1, 1, 1, 1])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), - ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) - >>> y = ivy.unique_counts(x) - >>> print(y) - { - a: (list[2],shape=[4]), - b: (list[2],shape=[4]) - } - """ - return ivy.current_backend(x).unique_counts(x) diff --git a/ivy/functional/ivy/sorting.py b/ivy/functional/ivy/sorting.py index a3423a308e41a..085db2f176276 100644 --- a/ivy/functional/ivy/sorting.py +++ b/ivy/functional/ivy/sorting.py @@ -140,113 +140,6 @@ def argsort( ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def sort( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a sorted copy of an array. - - Parameters - ---------- - x - input array - axis - axis along which to sort. If set to ``-1``, the function must sort along the - last axis. Default: ``-1``. - descending - direction The direction in which to sort the values - stable - sort stability. If ``True``, - the returned indices must maintain the relative order of ``x`` values which - compare as equal. If ``False``, the returned indices may or may not maintain the - relative order of ``x`` values which compare as equal (i.e., the relative order - of ``x`` values which compare as equal is implementation-dependent). - Default: ``True``. - out - optional output array, for writing the result to. It must have the same shape - as ``x``. - - Returns - ------- - ret - An array with the same dtype and shape as ``x``, with the elements sorted - along the given `axis`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([7, 8, 6]) - >>> y = ivy.sort(x) - >>> print(y) - ivy.array([6, 7, 8]) - - >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) - >>> y = ivy.sort(x, axis=1, descending=True, stable=False) - >>> print(y) - ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) - - >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) - >>> y = ivy.zeros(5) - >>> ivy.sort(x, descending=True, stable=False, out=y) - >>> print(y) - ivy.array([3.2, 2.5, 1.5, 0.7]) - - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.sort(x, out=x) - >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-6.6, -5.5, -4.4]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) - >>> y = ivy.sort(x, descending=True) - >>> print(y) - { - a: ivy.array([8, 6, 6]), - b: ivy.array([[9., 0.7], [0.4, 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) - >>> y = ivy.sort(x, descending=False, stable=False) - >>> print(y) - { - a: ivy.array([0.7, 1., 3.]), - b: ivy.array([[0.9, 4.], [0.2, 0.6]]) - } - """ - return ivy.current_backend(x).sort( - x, axis=axis, descending=descending, stable=stable, out=out - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -369,3 +262,110 @@ def searchsorted( out=out, ret_dtype=ret_dtype, ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def sort( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a sorted copy of an array. + + Parameters + ---------- + x + input array + axis + axis along which to sort. If set to ``-1``, the function must sort along the + last axis. Default: ``-1``. + descending + direction The direction in which to sort the values + stable + sort stability. If ``True``, + the returned indices must maintain the relative order of ``x`` values which + compare as equal. If ``False``, the returned indices may or may not maintain the + relative order of ``x`` values which compare as equal (i.e., the relative order + of ``x`` values which compare as equal is implementation-dependent). + Default: ``True``. + out + optional output array, for writing the result to. It must have the same shape + as ``x``. + + Returns + ------- + ret + An array with the same dtype and shape as ``x``, with the elements sorted + along the given `axis`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([7, 8, 6]) + >>> y = ivy.sort(x) + >>> print(y) + ivy.array([6, 7, 8]) + + >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) + >>> y = ivy.sort(x, axis=1, descending=True, stable=False) + >>> print(y) + ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) + + >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) + >>> y = ivy.zeros(5) + >>> ivy.sort(x, descending=True, stable=False, out=y) + >>> print(y) + ivy.array([3.2, 2.5, 1.5, 0.7]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.sort(x, out=x) + >>> print(x) + ivy.array([[ 1.1, 2.2, 3.3], + [-6.6, -5.5, -4.4]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) + >>> y = ivy.sort(x, descending=True) + >>> print(y) + { + a: ivy.array([8, 6, 6]), + b: ivy.array([[9., 0.7], [0.4, 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) + >>> y = ivy.sort(x, descending=False, stable=False) + >>> print(y) + { + a: ivy.array([0.7, 1., 3.]), + b: ivy.array([[0.9, 4.], [0.2, 0.6]]) + } + """ + return ivy.current_backend(x).sort( + x, axis=axis, descending=descending, stable=stable, out=out + ) diff --git a/ivy/functional/ivy/statistical.py b/ivy/functional/ivy/statistical.py index aaea746d500d0..81d7c3b8f17d9 100644 --- a/ivy/functional/ivy/statistical.py +++ b/ivy/functional/ivy/statistical.py @@ -16,8 +16,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# --------# +# --- Helpers --- # +# --------------- # def _get_promoted_type_of_operands(operands): @@ -31,10 +31,11 @@ def _get_promoted_type_of_operands(operands): return ivy.as_native_dtype(dtype) -# Array API Standard # -# -------------------# +# --- Main --- # +# ------------ # +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -42,105 +43,138 @@ def _get_promoted_type_of_operands(operands): @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def min( +def cumprod( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the minimum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the minimum value is zero, the - minimum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the maximum possible value - for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, - return ``+infinity``). - - **Special Cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` - (i.e., ``NaN`` values propagate). + Return the cumulative product of the elements along a given axis. Parameters ---------- x - Input array. Should have a real-valued data type. + Input array. axis - axis or axes along which minimum values must be computed. By default, the - minimum value must be computed over the entire array. If a tuple of integers, - minimum values must be computed over multiple axes. Default: ``None``. - - keepdims - optional boolean, if ``True``, the reduced axes (dimensions) must be included - in the result as singleton dimensions, and, accordingly, the result must be - compatible with the input array (see :ref:`broadcasting`). Otherwise, - if ``False``, the reduced axes (dimensions) must not be included in the result. - Default: ``False``. + int , axis along which the cumulative product is computed. By default 0. + exclusive + optional bool, Whether to perform the cumprod exclusively. Defaults is False. + reverse + Whether to perform the cumprod from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - if the minimum value was computed over the entire array, a zero-dimensional - array containing the minimum value; otherwise, a non-zero-dimensional array - containing the minimum values. The returned array must have the same data type - as ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Input array with cumulatively multiplied elements along axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.min(x) - >>> print(z) - ivy.array(1) - - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array([0, 0, 0]) - >>> y = ivy.min(x, out=z) - >>> print(z) - ivy.array(0) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x) + >>> print(y) + ivy.array([2, 6, 24]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x, axis=0, keepdims=True) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x, exclusive=True) >>> print(y) - ivy.array([[0, 1, 2]]) + ivy.array([1, 2, 6]) - >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x) + >>> x = ivy.array([[2, 3], + [5, 7], + [11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) >>> print(y) - ivy.array(0) + ivy.array([[ 1., 2.], + [ 1., 5.], + [ 1., 11.]]) + + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) + >>> print(x) + ivy.array([[1, 1], + [2, 3], + [10, 21]]) + + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> x.cumprod(axis=0, exclusive=True, out=y) + >>> print(x) + ivy.array([[1., 1.], + [2., 3.], + [10., 21.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) - >>> z = ivy.min(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x) + >>> print(y) { - a: ivy.array(1), - b: ivy.array(2) + a: ivy.array([2, 6, 24]), + b: ivy.array([3, 12, 60]) + } + + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x, exclusive=True) + >>> print(y) + { + a: ivy.array([1, 2, 6]), + b: ivy.array([1, 3, 12]) + } + + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) + >>> print(y) + { + a: ivy.array([[1, 2], + [1, 5], + [1, 11]]), + b: ivy.array([[1, 3], + [1, 4], + [1, 5]]) + } + + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> x.cumprod(axis=0, exclusive=True, out=x) + >>> print(x) + { + a: ivy.array([[1, 1], + [2, 3], + [10, 21]]), + b: ivy.array([[1, 1], + [3, 4], + [15, 42]]) } """ - return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).cumprod( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + ) + + +# Extra # +# ------# @handle_exceptions @@ -151,107 +185,141 @@ def min( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def max( +def cumsum( x: Union[ivy.Array, ivy.NativeArray], - /, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the maximum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the maximum value is zero, the - maximum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the minimum possible - value for the input array ``x`` data type (e.g., if ``x`` is a floating-point - array, return ``-infinity``). - - **Special Cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values - propagate). + Return the cumulative sum of the elements along a given axis. Parameters ---------- x - input array. Should have a numeric data type. + Input array. axis - axis or axes along which maximum values must be computed. By default, the - maximum value must be computed over the entire array. If a tuple of integers, - maximum values must be computed over multiple axes. Default: ``None``. - keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + Axis along which the cumulative sum is computed. Default is ``0``. + exclusive + Whether to perform cumsum exclusively. Default is ``False``. + reverse + Whether to perform the cumsum from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - if the maximum value was computed over the entire array, a zero-dimensional - array containing the maximum value; otherwise, a non-zero-dimensional array - containing the maximum values. The returned array must have the same data type - as ``x``. - - - This method conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Array which holds the result of applying cumsum at each + original array elements along the specified axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.max(x) - >>> print(z) - ivy.array(3) + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cumsum(x, exclusive= True, reverse=False) + >>> print(y) + ivy.array([0, 1, 6, 8]) - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array(0) - >>> y = ivy.max(x, out=z) - >>> print(z) - ivy.array(2) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) + >>> print(y) + ivy.array([[7, 7, 2], + [1, 3, 0]]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.max(x, axis=0, keepdims=True) + >>> x = ivy.array([[1, 5, 2], + ... [4, 3, 0]]) + >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) >>> print(y) - ivy.array([[4, 6, 10]]) + ivy.array([[4, 3, 0], + [0, 0, 0]]) + + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[11, 9, 5], + [14, 11, 5], + [14, 13, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.max(x) + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cumsum(x, axis= 0) >>> print(y) { - a: ivy.array(2.), - b: ivy.array(5.) + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) } - >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), - ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) - >>> z = ivy.max(x, axis=1) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cumsum(x,axis=1,reverse=True, out=y) + >>> print(y) { - a: ivy.array([3, 2]), - b: ivy.array([4, 2]) + a: ivy.array([[8, 7, 4]]), + b: ivy.array([[16, 13, 8], + [16, 11, 5]]), + c: ivy.array([[7, 5, 1], + [18, 15, 9], + [5, 5, 3]]) + } + + >>> x = ivy.Container(a=ivy.array([[0], + ... [5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cumsum(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [5]]), + b: ivy.array([[6, 8, 7], + [10, 10, 10]]), + c: ivy.array([[1, 2], + [4, 6], + [10, 10]]) } """ - return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) @handle_exceptions @@ -262,108 +330,126 @@ def max( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def mean( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, +def einsum( + equation: str, + *operands: Union[ivy.Array, ivy.NativeArray], out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the arithmetic mean of the input array ``x``. - - **Special Cases** - - Let ``N`` equal the number of elements over which to compute the arithmetic mean. - - If ``N`` is ``0``, the arithmetic mean is ``NaN``. - - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values - propagate). + Sum the product of the elements of the input operands along dimensions specified + using a notation based on the Einstein summation convention. Parameters ---------- - x - input array. Should have a floating-point data type. - axis - axis or axes along which arithmetic means must be computed. By default, the mean - must be computed over the entire array. If a Sequence of integers, arithmetic - means must be computed over multiple axes. Default: ``None``. - keepdims - bool, if ``True``, the reduced axes (dimensions) must be included in the result - as singleton dimensions, and, accordingly, the result must be compatible with - the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced - axes (dimensions) must not be included in the result. Default: ``False``. + equation + A str describing the contraction, in the same format as numpy.einsum. + operands + seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes + should be consistent with equation. out optional output array, for writing the result to. Returns ------- ret - array, if the arithmetic mean was computed over the entire array, a - zero-dimensional array containing the arithmetic mean; otherwise, a - non-zero-dimensional array containing the arithmetic means. The returned - array must have the same data type as ``x``. - .. note:: - While this specification recommends that this function only accept input - arrays having a floating-point data type, specification-compliant array - libraries may choose to accept input arrays having an integer data type. - While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have the - default floating-point data type. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ in the standard. + The array with sums computed. - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Functional Examples + ------------------- - Examples - -------- With :class:`ivy.Array` input: - >>> x = ivy.array([3., 4., 5.]) - >>> y = ivy.mean(x) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> y = ivy.einsum('ii', x) >>> print(y) - ivy.array(4.) + ivy.array(12) - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.array(0.) - >>> ivy.mean(x, out=y) - >>> print(y) - ivy.array(1.) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> z = ivy.einsum('ij -> j', x) + >>> print(z) + ivy.array([ 9, 12, 15]) - >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) - >>> y = ivy.array([0., 0.]) - >>> ivy.mean(x, axis=1, out=y) - >>> print(y) - ivy.array([-1.4, 1.4]) + >>> A = ivy.array([0, 1, 2]) + >>> B = ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]) + >>> C = ivy.einsum('i,ij->i', A, B) + >>> print(C) + ivy.array([ 0, 22, 76]) + >>> A = ivy.array([[1, 1, 1], + ... [2, 2, 2], + ... [5, 5, 5]]) + >>> B = ivy.array([[0, 1, 0], + ... [1, 1, 0], + ... [1, 1, 1]]) + >>> C = ivy.einsum('ij,jk->ik', A, B) + >>> print(C) + ivy.array([[ 2, 3, 1], + [ 4, 6, 2], + [10, 15, 5]]) - With :class:`ivy.Container` input: + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i->', A) + >>> print(C) + ivy.array(45) - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.mean(x) - >>> print(y) + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->i', A, B) + >>> print(C) + ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' + >>> print(C) + ivy.array(510) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,j->ij', A, B) + >>> print(C) + ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], + [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], + [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], + [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], + [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], + [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], + [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], + [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], + [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]), + ... b=ivy.array([[ 0, 1, 2], + ... [ 4, 5, 6], + ... [ 8, 9, 10]])) + >>> z = ivy.einsum('i,ij->i', x, y) + >>> print(z) { - a: ivy.array(0.), - b: ivy.array(0.90000004) + a: ivy.array([0, 22, 76]), + b: ivy.array([0, 15, 54]) } - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), - ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) - >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) - >>> ivy.mean(x, axis=0, out=y) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), + ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) + >>> y = ivy.einsum('ii', x) >>> print(y) { - a: ivy.array([1.5, 2.5, 3.5]), - b: ivy.array([4.5, 5.5, 6.5]) + a: ivy.array(2), + b: ivy.array(15) } """ - return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(operands[0]).einsum(equation, *operands, out=out) @handle_exceptions @@ -374,72 +460,61 @@ def mean( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def prod( +def max( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the product of input array x elements. + Calculate the maximum value of the input array ``x``. - **Special Cases** + .. note:: + When the number of elements over which to compute the maximum value is zero, the + maximum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the minimum possible + value for the input array ``x`` data type (e.g., if ``x`` is a floating-point + array, return ``-infinity``). - Let ``N`` equal the number of elements over which to compute the product. + **Special Cases** - - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). + For floating-point operands, - For both both real-valued and complex floating-point operands, special - cases must be handled as the operation is implemented by successive application - of :func:`ivy.multiply`: + - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- x input array. Should have a numeric data type. axis - axis or axes along which products must be computed. By default, the product must - be computed over the entire array. If a tuple of integers, products must be - computed over multiple axes. Default: ``None``. + axis or axes along which maximum values must be computed. By default, the + maximum value must be computed over the entire array. If a tuple of integers, + maximum values must be computed over multiple axes. Default: ``None``. keepdims - bool, if True, the reduced axes (dimensions) must be included in the result as + if ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the - input array (see Broadcasting). Otherwise, if False, the reduced axes + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. - dtype - data type of the returned array. If None, - if the default data type corresponding to the data type “kind” (integer or - floating-point) of x has a smaller range of values than the data type of x - (e.g., x has data type int64 and the default data type is int32, or x has data - type uint64 and the default data type is int64), the returned array must have - the same data type as x. if x has a floating-point data type, the returned array - must have the default floating-point data type. if x has a signed integer data - type (e.g., int16), the returned array must have the default integer data type. - if x has an unsigned integer data type (e.g., uint16), the returned array must - have an unsigned integer data type having the same number of bits as the default - integer data type (e.g., if the default integer data type is int32, the returned - array must have a uint32 data type). If the data type (either specified or - resolved) differs from the data type of x, the input array should be cast to the - specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - array, if the product was computed over the entire array, a zero-dimensional - array containing the product; otherwise, a non-zero-dimensional array containing - the products. The returned array must have a data type as described by the dtype - parameter above. + if the maximum value was computed over the entire array, a zero-dimensional + array containing the maximum value; otherwise, a non-zero-dimensional array + containing the maximum values. The returned array must have the same data type + as ``x``. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.max.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -451,55 +526,41 @@ def prod( With :class:`ivy.Array` input: >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.prod(x) + >>> z = ivy.max(x) >>> print(z) - ivy.array(6) + ivy.array(3) - >>> x = ivy.array([1, 0, 3]) - >>> z = ivy.prod(x) + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array(0) + >>> y = ivy.max(x, out=z) >>> print(z) - ivy.array(0) - - >>> x = ivy.array([[3., 4., 5.]]) - >>> y = ivy.prod(x, keepdims=True) - >>> print(y) - ivy.array([60.]) - - >>> x = ivy.array([2., 1.]) - >>> y = ivy.array(0.) - >>> ivy.prod(x, out=y) - >>> print(y) - ivy.array(2.) + ivy.array(2) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.prod(x, axis=1) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.max(x, axis=0, keepdims=True) >>> print(y) - ivy.array([2., 9.]) + ivy.array([[4, 6, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.prod(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.max(x) >>> print(y) { - a: ivy.array(-0.), - b: ivy.array(0.30800003) + a: ivy.array(2.), + b: ivy.array(5.) } - >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), - ... b=ivy.array([[ 4., 5.], [5., 6.]])) - >>> y = ivy.prod(x, axis=1, keepdims=True) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), + ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) + >>> z = ivy.max(x, axis=1) + >>> print(z) { - a: ivy.array([[2.], - [12.]]), - b: ivy.array([[20.], - [30.]]) + a: ivy.array([3, 2]), + b: ivy.array([4, 2]) } """ - return current_backend(x).prod( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out - ) + return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -510,138 +571,114 @@ def prod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def std( +def mean( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the standard deviation of the input array ``x``. + Calculate the arithmetic mean of the input array ``x``. **Special Cases** - Let ``N`` equal the number of elements over which to compute the standard deviation. - - - If ``N - correction`` is less than or equal to ``0``, - the standard deviation is ``NaN``. - - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` - (i.e., ``NaN`` values propagate). + Let ``N`` equal the number of elements over which to compute the arithmetic mean. + - If ``N`` is ``0``, the arithmetic mean is ``NaN``. + - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- x - input array. + input array. Should have a floating-point data type. axis - axis or axes along which standard deviations must be computed. By default, the - standard deviation must be computed over the entire array. If a tuple of - integers, standard deviations must be computed over multiple axes. - Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other - than ``0`` has the effect of adjusting the divisor during the calculation of the - standard deviation according to ``N-c`` where ``N`` corresponds to the total - number of elements over which the standard deviation is computed and ``c`` - corresponds to the provided degrees of freedom adjustment. When computing the - standard deviation of a population, setting this parameter to ``0`` is the - standard choice (i.e., the provided array contains data constituting an - entire population). When computing the corrected sample standard deviation, - setting this parameter to ``1`` is the standard choice (i.e., the provided array - contains data sampled from a larger population; this is commonly referred to as - Bessel's correction). - Default: ``0``. + axis or axes along which arithmetic means must be computed. By default, the mean + must be computed over the entire array. If a Sequence of integers, arithmetic + means must be computed over multiple axes. Default: ``None``. keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + bool, if ``True``, the reduced axes (dimensions) must be included in the result + as singleton dimensions, and, accordingly, the result must be compatible with + the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced + axes (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - if the standard deviation was computed over the entire array, a zero-dimensional - array containing the standard deviation; otherwise, a non-zero-dimensional array - containing the standard deviations. The returned array must have the same data - type as ``x``. - + array, if the arithmetic mean was computed over the entire array, a + zero-dimensional array containing the arithmetic mean; otherwise, a + non-zero-dimensional array containing the arithmetic means. The returned + array must have the same data type as ``x``. .. note:: While this specification recommends that this function only accept input - arrays having a real-valued floating-point data type, specification-compliant - array libraries may choose to accept input arrays having an integer data - type. While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have - the default real-valued floating-point data type. + arrays having a floating-point data type, specification-compliant array + libraries may choose to accept input arrays having an integer data type. + While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have the + default floating-point data type. + This function conforms to the `Array API Standard `_. This docstring is an extension of the - `docstring `_ - in the standard. + `docstring `_ in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.std(x) - >>> print(y) - ivy.array(0.81649661) - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = ivy.std(x, correction=1) - >>> print(z) - ivy.array(1.) + With :class:`ivy.Array` input: - >>> x = ivy.array([[0., 4.]]) - >>> y = ivy.std(x, keepdims=True) + >>> x = ivy.array([3., 4., 5.]) + >>> y = ivy.mean(x) >>> print(y) - ivy.array([[2.]]) + ivy.array(4.) - >>> x = ivy.array([2., 1.]) + >>> x = ivy.array([0., 1., 2.]) >>> y = ivy.array(0.) - >>> ivy.std(x, out=y) + >>> ivy.mean(x, out=y) >>> print(y) - ivy.array(0.5) + ivy.array(1.) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.std(x, axis=1) + >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) + >>> y = ivy.array([0., 0.]) + >>> ivy.mean(x, axis=1, out=y) >>> print(y) - ivy.array([0.5, 0. ]) + ivy.array([-1.4, 1.4]) + With :class:`ivy.Container` input: >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = x.std() + >>> y = ivy.mean(x) >>> print(y) { - a: ivy.array(0.81649661), - b: ivy.array(0.509902) + a: ivy.array(0.), + b: ivy.array(0.90000004) } - >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), - ... b=ivy.array([[ 4., 2.], [2., 1.]])) - >>> y = x.std(axis=1, keepdims=True) + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), + ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) + >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) + >>> ivy.mean(x, axis=0, out=y) >>> print(y) { - a: ivy.array([[1.], - [1.5]]), - b: ivy.array([[1.], - [0.5]]) + a: ivy.array([1.5, 2.5, 3.5]), + b: ivy.array([4.5, 5.5, 6.5]) } """ - return current_backend(x).std( - x, axis=axis, correction=correction, keepdims=keepdims, out=out - ) + return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) + + +# Array API Standard # +# -------------------# -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -649,81 +686,63 @@ def std( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sum( +def min( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: Optional[bool] = False, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the sum of the input array x. + Calculate the minimum value of the input array ``x``. + + .. note:: + When the number of elements over which to compute the minimum value is zero, the + minimum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the maximum possible value + for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, + return ``+infinity``). **Special Cases** - Let ``N`` equal the number of elements over which to compute the sum. - - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). - For floating-point operands, - - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). - For both real-valued and complex floating-point operands, special cases must - be handled as if the operation is implemented by successive application of - :func:`ivy.add`: + - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - Input array. Should have a numeric data type. + Input array. Should have a real-valued data type. axis - Axis or axes along which sums must be computed. By default, the sum must be - computed over the entire array. If a tuple of integers, sums must be computed - over multiple axes. Default: ``None``. - dtype - Data type of the returned array. If ``None``, - If the default data type corresponding to the data type "kind" (integer or - floating-point) of ``x`` has a smaller range of values than the data type of - ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is - ``int32``, or ``x`` has data type ``uint64`` and the default data type is - ``int64``), the returned array must have the same data type as ``x``. - If ``x`` has a floating-point data type, the returned array must have the - default floating-point data type. - If ``x`` has a signed integer data type (e.g., ``int16``), the returned - array must have the default integer data type. - If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned - array must have an unsigned integer data type having the same number of bits - as the default integer data type (e.g., if the default integer data type is - ``int32``, the returned array must have a ``uint32`` data type). - - If the data type (either specified or resolved) differs from the data type of - ``x``, the input array should be cast to the specified data type before - computing the sum. Default: ``None``. - - .. note:: - keyword argument is intended to help prevent data type overflows. + axis or axes along which minimum values must be computed. By default, the + minimum value must be computed over the entire array. If a tuple of integers, + minimum values must be computed over multiple axes. Default: ``None``. keepdims - If ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + optional boolean, if ``True``, the reduced axes (dimensions) must be included + in the result as singleton dimensions, and, accordingly, the result must be + compatible with the input array (see :ref:`broadcasting`). Otherwise, + if ``False``, the reduced axes (dimensions) must not be included in the result. + Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - If the sum was computed over the entire array, a zero-dimensional array - containing the sum; otherwise, an array containing the sums. The returned array - must have a data type as described by the ``dtype`` parameter above. + if the minimum value was computed over the entire array, a zero-dimensional + array containing the minimum value; otherwise, a non-zero-dimensional array + containing the minimum values. The returned array must have the same data type + as ``x``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.min.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -734,52 +753,38 @@ def sum( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.41, 0.89]) - >>> y = ivy.sum(x) - >>> print(y) - ivy.array(1.3) - - >>> x = ivy.array([0.5, 0.7, 2.4]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) - >>> print(y) - ivy.array(3.6) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.min(x) + >>> print(z) + ivy.array(1) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.sum(x, axis = 1, keepdims = False) - >>> print(y) - ivy.array([3, 20]) + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array([0, 0, 0]) + >>> y = ivy.min(x, out=z) + >>> print(z) + ivy.array(0) >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.array([0,0,0]) - >>> ivy.sum(x, axis = 0, keepdims = False, out = y) - >>> print(y) - ivy.array([4, 7, 12]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.sum(x) + >>> y = ivy.min(x, axis=0, keepdims=True) >>> print(y) - ivy.array(1.9) + ivy.array([[0, 1, 2]]) - >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) + >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.min(x) >>> print(y) - ivy.array(8.) + ivy.array(0) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.sum(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) + >>> z = ivy.min(x) + >>> print(z) { - a: ivy.array(3.), - b: ivy.array(12.) + a: ivy.array(1), + b: ivy.array(2) } """ - return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -790,65 +795,72 @@ def sum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def var( +def prod( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the variance of the input array x. + Calculate the product of input array x elements. **Special Cases** - Let N equal the number of elements over which to compute the variance. + Let ``N`` equal the number of elements over which to compute the product. - If N - correction is less than or equal to 0, the variance is NaN. + - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). - If x_i is NaN, the variance is NaN (i.e., NaN values propagate). + For both both real-valued and complex floating-point operands, special + cases must be handled as the operation is implemented by successive application + of :func:`ivy.multiply`: Parameters ---------- x - input array. Should have a floating-point data type. + input array. Should have a numeric data type. axis - axis or axes along which variances must be computed. By default, the variance - must be computed over the entire array. If a tuple of integers, variances must - be computed over multiple axes. Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other than 0 - has the effect of adjusting the divisor during the calculation of the variance - according to N-c where N corresponds to the total number of elements over which - the variance is computed and c corresponds to the provided degrees of freedom - adjustment. When computing the variance of a population, setting this parameter - to 0 is the standard choice (i.e., the provided array contains data constituting - an entire population). When computing the unbiased sample variance, setting this - parameter to 1 is the standard choice (i.e., the provided array contains data - sampled from a larger population; this is commonly referred to as Bessel's - correction). Default: ``0``. + axis or axes along which products must be computed. By default, the product must + be computed over the entire array. If a tuple of integers, products must be + computed over multiple axes. Default: ``None``. keepdims - if True, the reduced axes (dimensions) must be included in the result as + bool, if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see Broadcasting). Otherwise, if False, the reduced axes (dimensions) must not be included in the result. Default: ``False``. + dtype + data type of the returned array. If None, + if the default data type corresponding to the data type “kind” (integer or + floating-point) of x has a smaller range of values than the data type of x + (e.g., x has data type int64 and the default data type is int32, or x has data + type uint64 and the default data type is int64), the returned array must have + the same data type as x. if x has a floating-point data type, the returned array + must have the default floating-point data type. if x has a signed integer data + type (e.g., int16), the returned array must have the default integer data type. + if x has an unsigned integer data type (e.g., uint16), the returned array must + have an unsigned integer data type having the same number of bits as the default + integer data type (e.g., if the default integer data type is int32, the returned + array must have a uint32 data type). If the data type (either specified or + resolved) differs from the data type of x, the input array should be cast to the + specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - if the variance was computed over the entire array, a zero-dimensional array - containing the variance; otherwise, a non-zero-dimensional array containing the - variances. The returned array must have the same data type as x. + array, if the product was computed over the entire array, a zero-dimensional + array containing the product; otherwise, a non-zero-dimensional array containing + the products. The returned array must have a data type as described by the dtype + parameter above. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.prod.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -859,46 +871,58 @@ def var( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.var(x) - >>> print(y) - ivy.array(0.07472222) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.prod(x) + >>> print(z) + ivy.array(6) - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.array(0.0) - >>> ivy.var(x, out=y) + >>> x = ivy.array([1, 0, 3]) + >>> z = ivy.prod(x) + >>> print(z) + ivy.array(0) + + >>> x = ivy.array([[3., 4., 5.]]) + >>> y = ivy.prod(x, keepdims=True) >>> print(y) - ivy.array(0.07472222) + ivy.array([60.]) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> print(ivy.var(x, axis=1, keepdims=True)) - ivy.array([[0.00666667], - [0.11555555]]) + >>> x = ivy.array([2., 1.]) + >>> y = ivy.array(0.) + >>> ivy.prod(x, out=y) + >>> print(y) + ivy.array(2.) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> y = ivy.var(x, correction=1) + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.prod(x, axis=1) >>> print(y) - ivy.array(0.08966666) + ivy.array([2., 9.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), - ... b=ivy.array([0.7, 0.1, 0.9])) - >>> y = ivy.var(x) + + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = ivy.prod(x) >>> print(y) { - a: ivy.array(0.12666667), - b: ivy.array(0.11555555) + a: ivy.array(-0.), + b: ivy.array(0.30800003) + } + + >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), + ... b=ivy.array([[ 4., 5.], [5., 6.]])) + >>> y = ivy.prod(x, axis=1, keepdims=True) + >>> print(y) + { + a: ivy.array([[2.], + [12.]]), + b: ivy.array([[20.], + [30.]]) } """ - return current_backend(x).var( - x, axis=axis, correction=correction, keepdims=keepdims, out=out + return current_backend(x).prod( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out ) -# Extra # -# ------# - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -907,141 +931,135 @@ def var( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumsum( +def std( x: Union[ivy.Array, ivy.NativeArray], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, + /, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative sum of the elements along a given axis. + Calculate the standard deviation of the input array ``x``. + + **Special Cases** + + Let ``N`` equal the number of elements over which to compute the standard deviation. + + - If ``N - correction`` is less than or equal to ``0``, + the standard deviation is ``NaN``. + - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - Input array. + input array. axis - Axis along which the cumulative sum is computed. Default is ``0``. - exclusive - Whether to perform cumsum exclusively. Default is ``False``. - reverse - Whether to perform the cumsum from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. + axis or axes along which standard deviations must be computed. By default, the + standard deviation must be computed over the entire array. If a tuple of + integers, standard deviations must be computed over multiple axes. + Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other + than ``0`` has the effect of adjusting the divisor during the calculation of the + standard deviation according to ``N-c`` where ``N`` corresponds to the total + number of elements over which the standard deviation is computed and ``c`` + corresponds to the provided degrees of freedom adjustment. When computing the + standard deviation of a population, setting this parameter to ``0`` is the + standard choice (i.e., the provided array contains data constituting an + entire population). When computing the corrected sample standard deviation, + setting this parameter to ``1`` is the standard choice (i.e., the provided array + contains data sampled from a larger population; this is commonly referred to as + Bessel's correction). + Default: ``0``. + keepdims + if ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cumsum at each - original array elements along the specified axis. + if the standard deviation was computed over the entire array, a zero-dimensional + array containing the standard deviation; otherwise, a non-zero-dimensional array + containing the standard deviations. The returned array must have the same data + type as ``x``. + + .. note:: + While this specification recommends that this function only accept input + arrays having a real-valued floating-point data type, specification-compliant + array libraries may choose to accept input arrays having an integer data + type. While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have + the default real-valued floating-point data type. + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cumsum(x, exclusive= True, reverse=False) + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.std(x) >>> print(y) - ivy.array([0, 1, 6, 8]) + ivy.array(0.81649661) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) + >>> x = ivy.array([-1., 0., 1.]) + >>> z = ivy.std(x, correction=1) + >>> print(z) + ivy.array(1.) + + >>> x = ivy.array([[0., 4.]]) + >>> y = ivy.std(x, keepdims=True) >>> print(y) - ivy.array([[7, 7, 2], - [1, 3, 0]]) + ivy.array([[2.]]) - >>> x = ivy.array([[1, 5, 2], - ... [4, 3, 0]]) - >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) + >>> x = ivy.array([2., 1.]) + >>> y = ivy.array(0.) + >>> ivy.std(x, out=y) >>> print(y) - ivy.array([[4, 3, 0], - [0, 0, 0]]) + ivy.array(0.5) - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[11, 9, 5], - [14, 11, 5], - [14, 13, 10]]) + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.std(x, axis=1) + >>> print(y) + ivy.array([0.5, 0. ]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cumsum(x, axis= 0) + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = x.std() >>> print(y) { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) + a: ivy.array(0.81649661), + b: ivy.array(0.509902) } - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cumsum(x,axis=1,reverse=True, out=y) + >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), + ... b=ivy.array([[ 4., 2.], [2., 1.]])) + >>> y = x.std(axis=1, keepdims=True) >>> print(y) { - a: ivy.array([[8, 7, 4]]), - b: ivy.array([[16, 13, 8], - [16, 11, 5]]), - c: ivy.array([[7, 5, 1], - [18, 15, 9], - [5, 5, 3]]) - } - - >>> x = ivy.Container(a=ivy.array([[0], - ... [5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cumsum(x,axis=0,out=x) - >>> print(x) - { - a: ivy.array([[0], - [5]]), - b: ivy.array([[6, 8, 7], - [10, 10, 10]]), - c: ivy.array([[1, 2], - [4, 6], - [10, 10]]) + a: ivy.array([[1.], + [1.5]]), + b: ivy.array([[1.], + [0.5]]) } """ - return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) + return current_backend(x).std( + x, axis=axis, correction=correction, keepdims=keepdims, out=out + ) @handle_exceptions @@ -1052,134 +1070,137 @@ def cumsum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumprod( +def sum( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, + axis: Optional[Union[int, Sequence[int]]] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative product of the elements along a given axis. + Calculate the sum of the input array x. + + **Special Cases** + + Let ``N`` equal the number of elements over which to compute the sum. + - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). + + For floating-point operands, + - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). + + For both real-valued and complex floating-point operands, special cases must + be handled as if the operation is implemented by successive application of + :func:`ivy.add`: Parameters ---------- x - Input array. + Input array. Should have a numeric data type. axis - int , axis along which the cumulative product is computed. By default 0. - exclusive - optional bool, Whether to perform the cumprod exclusively. Defaults is False. - reverse - Whether to perform the cumprod from last to first element in the selected - axis. Default is ``False`` (from first to last element) + Axis or axes along which sums must be computed. By default, the sum must be + computed over the entire array. If a tuple of integers, sums must be computed + over multiple axes. Default: ``None``. + dtype + Data type of the returned array. If ``None``, + If the default data type corresponding to the data type "kind" (integer or + floating-point) of ``x`` has a smaller range of values than the data type of + ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is + ``int32``, or ``x`` has data type ``uint64`` and the default data type is + ``int64``), the returned array must have the same data type as ``x``. + If ``x`` has a floating-point data type, the returned array must have the + default floating-point data type. + If ``x`` has a signed integer data type (e.g., ``int16``), the returned + array must have the default integer data type. + If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned + array must have an unsigned integer data type having the same number of bits + as the default integer data type (e.g., if the default integer data type is + ``int32``, the returned array must have a ``uint32`` data type). + + If the data type (either specified or resolved) differs from the data type of + ``x``, the input array should be cast to the specified data type before + computing the sum. Default: ``None``. + + .. note:: + keyword argument is intended to help prevent data type overflows. + + keepdims + If ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Input array with cumulatively multiplied elements along axis. + If the sum was computed over the entire array, a zero-dimensional array + containing the sum; otherwise, an array containing the sums. The returned array + must have a data type as described by the ``dtype`` parameter above. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x) + >>> x = ivy.array([0.41, 0.89]) + >>> y = ivy.sum(x) >>> print(y) - ivy.array([2, 6, 24]) + ivy.array(1.3) - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x, exclusive=True) + >>> x = ivy.array([0.5, 0.7, 2.4]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) >>> print(y) - ivy.array([1, 2, 6]) + ivy.array(3.6) - >>> x = ivy.array([[2, 3], - [5, 7], - [11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.sum(x, axis = 1, keepdims = False) >>> print(y) - ivy.array([[ 1., 2.], - [ 1., 5.], - [ 1., 11.]]) - - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) - >>> print(x) - ivy.array([[1, 1], - [2, 3], - [10, 21]]) + ivy.array([3, 20]) - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> x.cumprod(axis=0, exclusive=True, out=y) - >>> print(x) - ivy.array([[1., 1.], - [2., 3.], - [10., 21.]]) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.array([0,0,0]) + >>> ivy.sum(x, axis = 0, keepdims = False, out = y) + >>> print(y) + ivy.array([4, 7, 12]) - With :class:`ivy.Container` input: + With :class:`ivy.NativeArray` input: - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x) + >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.sum(x) >>> print(y) - { - a: ivy.array([2, 6, 24]), - b: ivy.array([3, 12, 60]) - } + ivy.array(1.9) - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x, exclusive=True) + >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) >>> print(y) - { - a: ivy.array([1, 2, 6]), - b: ivy.array([1, 3, 12]) - } + ivy.array(8.) - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) - >>> print(y) - { - a: ivy.array([[1, 2], - [1, 5], - [1, 11]]), - b: ivy.array([[1, 3], - [1, 4], - [1, 5]]) - } + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> x.cumprod(axis=0, exclusive=True, out=x) - >>> print(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.sum(x) + >>> print(y) { - a: ivy.array([[1, 1], - [2, 3], - [10, 21]]), - b: ivy.array([[1, 1], - [3, 4], - [15, 42]]) + a: ivy.array(3.), + b: ivy.array(12.) } """ - return current_backend(x).cumprod( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out - ) + return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) @handle_exceptions @@ -1190,123 +1211,106 @@ def cumprod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def einsum( - equation: str, - *operands: Union[ivy.Array, ivy.NativeArray], +def var( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Sum the product of the elements of the input operands along dimensions specified - using a notation based on the Einstein summation convention. + Calculate the variance of the input array x. + + **Special Cases** + + Let N equal the number of elements over which to compute the variance. + + If N - correction is less than or equal to 0, the variance is NaN. + + If x_i is NaN, the variance is NaN (i.e., NaN values propagate). Parameters ---------- - equation - A str describing the contraction, in the same format as numpy.einsum. - operands - seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes - should be consistent with equation. + x + input array. Should have a floating-point data type. + axis + axis or axes along which variances must be computed. By default, the variance + must be computed over the entire array. If a tuple of integers, variances must + be computed over multiple axes. Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other than 0 + has the effect of adjusting the divisor during the calculation of the variance + according to N-c where N corresponds to the total number of elements over which + the variance is computed and c corresponds to the provided degrees of freedom + adjustment. When computing the variance of a population, setting this parameter + to 0 is the standard choice (i.e., the provided array contains data constituting + an entire population). When computing the unbiased sample variance, setting this + parameter to 1 is the standard choice (i.e., the provided array contains data + sampled from a larger population; this is commonly referred to as Bessel's + correction). Default: ``0``. + keepdims + if True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see Broadcasting). Otherwise, if False, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - The array with sums computed. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: + if the variance was computed over the entire array, a zero-dimensional array + containing the variance; otherwise, a non-zero-dimensional array containing the + variances. The returned array must have the same data type as x. - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> y = ivy.einsum('ii', x) - >>> print(y) - ivy.array(12) - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> z = ivy.einsum('ij -> j', x) - >>> print(z) - ivy.array([ 9, 12, 15]) - - >>> A = ivy.array([0, 1, 2]) - >>> B = ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]) - >>> C = ivy.einsum('i,ij->i', A, B) - >>> print(C) - ivy.array([ 0, 22, 76]) - - >>> A = ivy.array([[1, 1, 1], - ... [2, 2, 2], - ... [5, 5, 5]]) - >>> B = ivy.array([[0, 1, 0], - ... [1, 1, 0], - ... [1, 1, 1]]) - >>> C = ivy.einsum('ij,jk->ik', A, B) - >>> print(C) - ivy.array([[ 2, 3, 1], - [ 4, 6, 2], - [10, 15, 5]]) + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i->', A) - >>> print(C) - ivy.array(45) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->i', A, B) - >>> print(C) - ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + Examples + -------- + With :class:`ivy.Array` input: - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' - >>> print(C) - ivy.array(510) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.var(x) + >>> print(y) + ivy.array(0.07472222) - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,j->ij', A, B) - >>> print(C) - ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], - [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], - [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], - [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], - [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], - [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], - [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], - [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], - [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.array(0.0) + >>> ivy.var(x, out=y) + >>> print(y) + ivy.array(0.07472222) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> print(ivy.var(x, axis=1, keepdims=True)) + ivy.array([[0.00666667], + [0.11555555]]) - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]), - ... b=ivy.array([[ 0, 1, 2], - ... [ 4, 5, 6], - ... [ 8, 9, 10]])) - >>> z = ivy.einsum('i,ij->i', x, y) - >>> print(z) - { - a: ivy.array([0, 22, 76]), - b: ivy.array([0, 15, 54]) - } + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> y = ivy.var(x, correction=1) + >>> print(y) + ivy.array(0.08966666) With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), - ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) - >>> y = ivy.einsum('ii', x) + >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), + ... b=ivy.array([0.7, 0.1, 0.9])) + >>> y = ivy.var(x) >>> print(y) { - a: ivy.array(2), - b: ivy.array(15) + a: ivy.array(0.12666667), + b: ivy.array(0.11555555) } """ - return current_backend(operands[0]).einsum(equation, *operands, out=out) + return current_backend(x).var( + x, axis=axis, correction=correction, keepdims=keepdims, out=out + ) diff --git a/ivy/functional/ivy/utility.py b/ivy/functional/ivy/utility.py index 09bd432b60291..cefd9d446c711 100644 --- a/ivy/functional/ivy/utility.py +++ b/ivy/functional/ivy/utility.py @@ -237,6 +237,19 @@ def any( return ivy.current_backend(x).any(x, axis=axis, keepdims=keepdims, out=out) +@staticmethod +def load(filepath, format=None, type="module"): + if type == "module": + return ivy.Module.load(filepath) + elif type == "container": + if format is not None: + return ivy.Container.cont_load(filepath, format=format) + else: + return ivy.Container.cont_load(filepath) + else: + raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.") + + # Extra # # ----- # @@ -251,16 +264,3 @@ def save(item, filepath, format=None): item.save(filepath) else: raise ivy.utils.exceptions.IvyException("Unsupported item type for saving.") - - -@staticmethod -def load(filepath, format=None, type="module"): - if type == "module": - return ivy.Module.load(filepath) - elif type == "container": - if format is not None: - return ivy.Container.cont_load(filepath, format=format) - else: - return ivy.Container.cont_load(filepath) - else: - raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.") From e795c0044500f36e6ba2cacae8aaa14237e41425 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Thu, 31 Aug 2023 17:05:49 +0530 Subject: [PATCH 2/3] Revert "refactor(All APIs): Added new pre-commit hook for ordering functions" This reverts commit 6515e43972e40d996653e086ff1c637973966898. --- .pre-commit-config.yaml | 2 +- ivy/functional/backends/jax/activations.py | 32 +- ivy/functional/backends/jax/creation.py | 138 +- ivy/functional/backends/jax/data_type.py | 168 +- ivy/functional/backends/jax/device.py | 146 +- ivy/functional/backends/jax/elementwise.py | 256 +- .../backends/jax/experimental/activations.py | 40 +- .../backends/jax/experimental/creation.py | 77 +- .../backends/jax/experimental/elementwise.py | 284 +- .../backends/jax/experimental/layers.py | 772 ++- .../jax/experimental/linear_algebra.py | 82 +- .../backends/jax/experimental/manipulation.py | 354 +- .../backends/jax/experimental/random.py | 65 +- .../backends/jax/experimental/statistical.py | 416 +- ivy/functional/backends/jax/general.py | 252 +- ivy/functional/backends/jax/gradients.py | 68 +- ivy/functional/backends/jax/layers.py | 156 +- ivy/functional/backends/jax/linear_algebra.py | 92 +- ivy/functional/backends/jax/manipulation.py | 200 +- ivy/functional/backends/jax/random.py | 98 +- ivy/functional/backends/jax/searching.py | 26 +- ivy/functional/backends/jax/sorting.py | 30 +- ivy/functional/backends/jax/statistical.py | 197 +- ivy/functional/backends/mxnet/activations.py | 16 +- ivy/functional/backends/mxnet/creation.py | 80 +- ivy/functional/backends/mxnet/data_type.py | 127 +- ivy/functional/backends/mxnet/device.py | 64 +- ivy/functional/backends/mxnet/elementwise.py | 240 +- .../mxnet/experimental/activations.py | 12 +- .../backends/mxnet/experimental/creation.py | 40 +- .../mxnet/experimental/elementwise.py | 172 +- .../backends/mxnet/experimental/layers.py | 152 +- .../mxnet/experimental/linear_algebra.py | 100 +- .../mxnet/experimental/manipulation.py | 154 +- .../backends/mxnet/experimental/norms.py | 20 +- .../backends/mxnet/experimental/random.py | 38 +- .../mxnet/experimental/statistical.py | 86 +- ivy/functional/backends/mxnet/general.py | 40 +- ivy/functional/backends/mxnet/gradients.py | 30 +- ivy/functional/backends/mxnet/layers.py | 28 +- .../backends/mxnet/linear_algebra.py | 42 +- ivy/functional/backends/mxnet/manipulation.py | 88 +- ivy/functional/backends/mxnet/random.py | 48 +- ivy/functional/backends/mxnet/searching.py | 18 +- ivy/functional/backends/mxnet/sorting.py | 24 +- ivy/functional/backends/mxnet/statistical.py | 72 +- ivy/functional/backends/numpy/activations.py | 116 +- ivy/functional/backends/numpy/creation.py | 140 +- ivy/functional/backends/numpy/data_type.py | 166 +- ivy/functional/backends/numpy/device.py | 100 +- ivy/functional/backends/numpy/elementwise.py | 560 +- .../numpy/experimental/activations.py | 64 +- .../backends/numpy/experimental/creation.py | 125 +- .../numpy/experimental/elementwise.py | 574 +- .../backends/numpy/experimental/layers.py | 1462 +++-- .../numpy/experimental/linear_algebra.py | 118 +- .../numpy/experimental/manipulation.py | 434 +- .../backends/numpy/experimental/random.py | 60 +- .../backends/numpy/experimental/searching.py | 6 +- .../backends/numpy/experimental/sorting.py | 6 +- .../numpy/experimental/statistical.py | 618 +- ivy/functional/backends/numpy/general.py | 188 +- ivy/functional/backends/numpy/gradients.py | 58 +- ivy/functional/backends/numpy/helpers.py | 50 +- ivy/functional/backends/numpy/layers.py | 120 +- .../backends/numpy/linear_algebra.py | 98 +- ivy/functional/backends/numpy/manipulation.py | 158 +- ivy/functional/backends/numpy/random.py | 73 +- ivy/functional/backends/numpy/searching.py | 26 +- ivy/functional/backends/numpy/sorting.py | 38 +- ivy/functional/backends/numpy/statistical.py | 216 +- ivy/functional/backends/numpy/utility.py | 10 +- ivy/functional/backends/paddle/activations.py | 128 +- ivy/functional/backends/paddle/creation.py | 453 +- ivy/functional/backends/paddle/data_type.py | 139 +- ivy/functional/backends/paddle/device.py | 88 +- ivy/functional/backends/paddle/elementwise.py | 1846 +++--- .../paddle/experimental/activations.py | 72 +- .../backends/paddle/experimental/creation.py | 143 +- .../paddle/experimental/elementwise.py | 546 +- .../backends/paddle/experimental/layers.py | 499 +- .../paddle/experimental/linear_algebra.py | 74 +- .../backends/paddle/experimental/losses.py | 24 +- .../paddle/experimental/manipulation.py | 615 +- .../backends/paddle/experimental/norms.py | 56 +- .../backends/paddle/experimental/random.py | 127 +- .../paddle/experimental/statistical.py | 714 +- ivy/functional/backends/paddle/general.py | 202 +- ivy/functional/backends/paddle/gradients.py | 222 +- ivy/functional/backends/paddle/layers.py | 72 +- .../backends/paddle/linear_algebra.py | 142 +- .../backends/paddle/manipulation.py | 336 +- ivy/functional/backends/paddle/random.py | 111 +- ivy/functional/backends/paddle/searching.py | 50 +- ivy/functional/backends/paddle/sorting.py | 48 +- ivy/functional/backends/paddle/statistical.py | 359 +- .../backends/tensorflow/activations.py | 44 +- .../backends/tensorflow/control_flow_ops.py | 77 +- .../backends/tensorflow/creation.py | 160 +- .../backends/tensorflow/data_type.py | 176 +- ivy/functional/backends/tensorflow/device.py | 126 +- .../backends/tensorflow/elementwise.py | 482 +- .../tensorflow/experimental/activations.py | 44 +- .../tensorflow/experimental/creation.py | 120 +- .../tensorflow/experimental/elementwise.py | 362 +- .../tensorflow/experimental/layers.py | 1492 +++-- .../tensorflow/experimental/linear_algebra.py | 216 +- .../tensorflow/experimental/manipulation.py | 360 +- .../backends/tensorflow/experimental/norms.py | 48 +- .../tensorflow/experimental/random.py | 90 +- .../tensorflow/experimental/statistical.py | 660 +- ivy/functional/backends/tensorflow/general.py | 270 +- .../backends/tensorflow/gradients.py | 180 +- ivy/functional/backends/tensorflow/layers.py | 92 +- .../backends/tensorflow/linear_algebra.py | 60 +- .../backends/tensorflow/manipulation.py | 226 +- ivy/functional/backends/tensorflow/random.py | 82 +- .../backends/tensorflow/searching.py | 46 +- ivy/functional/backends/tensorflow/sorting.py | 46 +- .../backends/tensorflow/statistical.py | 159 +- ivy/functional/backends/torch/activations.py | 114 +- ivy/functional/backends/torch/creation.py | 258 +- ivy/functional/backends/torch/data_type.py | 135 +- ivy/functional/backends/torch/device.py | 107 +- ivy/functional/backends/torch/elementwise.py | 1128 ++-- .../torch/experimental/activations.py | 44 +- .../backends/torch/experimental/creation.py | 149 +- .../torch/experimental/elementwise.py | 340 +- .../backends/torch/experimental/layers.py | 1062 ++- .../torch/experimental/linear_algebra.py | 112 +- .../backends/torch/experimental/losses.py | 35 +- .../torch/experimental/manipulation.py | 396 +- .../backends/torch/experimental/norms.py | 116 +- .../backends/torch/experimental/random.py | 104 +- .../backends/torch/experimental/searching.py | 6 +- .../backends/torch/experimental/sorting.py | 6 +- .../torch/experimental/statistical.py | 756 +-- ivy/functional/backends/torch/general.py | 334 +- ivy/functional/backends/torch/gradients.py | 146 +- ivy/functional/backends/torch/layers.py | 124 +- .../backends/torch/linear_algebra.py | 198 +- ivy/functional/backends/torch/manipulation.py | 264 +- ivy/functional/backends/torch/random.py | 97 +- ivy/functional/backends/torch/searching.py | 26 +- ivy/functional/backends/torch/sorting.py | 50 +- ivy/functional/backends/torch/statistical.py | 270 +- ivy/functional/backends/torch/utility.py | 10 +- .../frontends/torch/comparison_ops.py | 2 +- ivy/functional/ivy/activations.py | 368 +- ivy/functional/ivy/constants.py | 81 +- ivy/functional/ivy/control_flow_ops.py | 202 +- ivy/functional/ivy/creation.py | 2453 ++++--- ivy/functional/ivy/data_type.py | 2070 +++--- ivy/functional/ivy/device.py | 1395 ++-- ivy/functional/ivy/elementwise.py | 3898 +++++------ .../ivy/experimental/activations.py | 340 +- ivy/functional/ivy/experimental/creation.py | 994 ++- .../ivy/experimental/elementwise.py | 1238 ++-- ivy/functional/ivy/experimental/general.py | 8 - ivy/functional/ivy/experimental/layers.py | 3953 ++++++------ .../ivy/experimental/linear_algebra.py | 1773 +++-- ivy/functional/ivy/experimental/losses.py | 178 +- .../ivy/experimental/manipulation.py | 3534 +++++----- ivy/functional/ivy/experimental/norms.py | 332 +- ivy/functional/ivy/experimental/random.py | 204 +- .../ivy/experimental/sparse_array.py | 1396 ++-- .../ivy/experimental/statistical.py | 1238 ++-- ivy/functional/ivy/general.py | 5724 +++++++++-------- ivy/functional/ivy/gradients.py | 1658 ++--- ivy/functional/ivy/layers.py | 3219 +++++---- ivy/functional/ivy/linear_algebra.py | 487 +- ivy/functional/ivy/losses.py | 124 +- ivy/functional/ivy/manipulation.py | 798 ++- ivy/functional/ivy/meta.py | 399 +- ivy/functional/ivy/nest.py | 1732 ++--- ivy/functional/ivy/norms.py | 24 +- ivy/functional/ivy/random.py | 466 +- ivy/functional/ivy/searching.py | 150 +- ivy/functional/ivy/set.py | 210 +- ivy/functional/ivy/sorting.py | 214 +- ivy/functional/ivy/statistical.py | 1558 +++-- ivy/functional/ivy/utility.py | 26 +- 182 files changed, 36430 insertions(+), 36313 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 34c2ff59c9c14..f4c4c298876fb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,6 +26,6 @@ repos: # Exclude everything in frontends except __init__.py, and func_wrapper.py exclude: 'ivy/functional/(frontends|backends)/(?!.*/func_wrapper\.py$).*(?!__init__\.py$)' - repo: https://github.com/unifyai/lint-hook - rev: 5abf78187bb7a5f839174ce4febd4e63f5d758f6 + rev: 27646397c5390f644a645f439535b1061b9c0105 hooks: - id: ivy-lint diff --git a/ivy/functional/backends/jax/activations.py b/ivy/functional/backends/jax/activations.py index 7732873e41b6f..6d980aaf6022e 100644 --- a/ivy/functional/backends/jax/activations.py +++ b/ivy/functional/backends/jax/activations.py @@ -23,10 +23,6 @@ def gelu( return jax.nn.gelu(x, approximate) -def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.hard_swish(x) - - def leaky_relu( x: JaxArray, /, @@ -38,18 +34,6 @@ def leaky_relu( return jnp.asarray(jnp.where(x > 0, x, jnp.multiply(x, alpha)), x.dtype) -def log_softmax( - x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None -): - if axis is None: - axis = -1 - return jax.nn.log_softmax(x, axis) - - -def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): - return x * jnp.tanh(jax.nn.softplus(x)) - - def relu( x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None ) -> JaxArray: @@ -94,3 +78,19 @@ def softplus( if threshold is not None: return jnp.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) + + +def log_softmax( + x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None +): + if axis is None: + axis = -1 + return jax.nn.log_softmax(x, axis) + + +def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): + return x * jnp.tanh(jax.nn.softplus(x)) + + +def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.hard_swish(x) diff --git a/ivy/functional/backends/jax/creation.py b/ivy/functional/backends/jax/creation.py index 45b8ad89b441f..4fbc39f493756 100644 --- a/ivy/functional/backends/jax/creation.py +++ b/ivy/functional/backends/jax/creation.py @@ -79,19 +79,6 @@ def asarray( return jnp.asarray(obj, dtype=dtype) -def copy_array( - x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - x = ( - jax.core.ShapedArray(x.shape, x.dtype) - if isinstance(x, jax.core.ShapedArray) - else jnp.array(x) - ) - if to_ivy_array: - return ivy.to_ivy(x) - return x - - def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -140,15 +127,6 @@ def from_dlpack(x, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jax.dlpack.from_dlpack(capsule) -def frombuffer( - buffer: bytes, - dtype: Optional[jnp.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> JaxArray: - return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) - - def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -256,42 +234,6 @@ def meshgrid( return jnp.meshgrid(*arrays, sparse=sparse, indexing=indexing) -def one_hot( - indices: JaxArray, - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - device: jaxlib.xla_extension.Device, - out: Optional[JaxArray] = None, -) -> JaxArray: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = jnp.float32 - else: - if not on_none: - dtype = jnp.array(on_value).dtype - elif not off_none: - dtype = jnp.array(off_value).dtype - - res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] - res = res.reshape(list(indices.shape) + [depth]) - - if not on_none and not off_none: - res = jnp.where(res == 1, on_value, off_value) - - if axis is not None: - res = jnp.moveaxis(res, -1, axis) - - return res - - def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -321,17 +263,6 @@ def triu(x: JaxArray, /, *, k: int = 0, out: Optional[JaxArray] = None) -> JaxAr return jnp.triu(x, k) -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: jaxlib.xla_extension.Device, -) -> Tuple[JaxArray]: - return jnp.triu_indices(n=n_rows, k=k, m=n_cols) - - def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -358,3 +289,72 @@ def zeros_like( array = asarray + + +def copy_array( + x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + x = ( + jax.core.ShapedArray(x.shape, x.dtype) + if isinstance(x, jax.core.ShapedArray) + else jnp.array(x) + ) + if to_ivy_array: + return ivy.to_ivy(x) + return x + + +def one_hot( + indices: JaxArray, + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[jnp.dtype] = None, + device: jaxlib.xla_extension.Device, + out: Optional[JaxArray] = None, +) -> JaxArray: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = jnp.float32 + else: + if not on_none: + dtype = jnp.array(on_value).dtype + elif not off_none: + dtype = jnp.array(off_value).dtype + + res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] + res = res.reshape(list(indices.shape) + [depth]) + + if not on_none and not off_none: + res = jnp.where(res == 1, on_value, off_value) + + if axis is not None: + res = jnp.moveaxis(res, -1, axis) + + return res + + +def frombuffer( + buffer: bytes, + dtype: Optional[jnp.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> JaxArray: + return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + + +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: jaxlib.xla_extension.Device, +) -> Tuple[JaxArray]: + return jnp.triu_indices(n=n_rows, k=k, m=n_cols) diff --git a/ivy/functional/backends/jax/data_type.py b/ivy/functional/backends/jax/data_type.py index 9a53044c244f3..e25e54e745bbe 100644 --- a/ivy/functional/backends/jax/data_type.py +++ b/ivy/functional/backends/jax/data_type.py @@ -9,26 +9,6 @@ from ivy.functional.backends.jax import JaxArray from ivy.functional.ivy.data_type import _handle_nestable_dtype_info -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} ivy_dtype_dict = { jnp.dtype("int8"): "int8", jnp.dtype("int16"): "int16", @@ -61,6 +41,7 @@ jnp.complex128: "complex128", jnp.bool_: "bool", } + native_dtype_dict = { "int8": jnp.dtype("int8"), "int16": jnp.dtype("int16"), @@ -79,6 +60,27 @@ "bool": jnp.dtype("bool"), } +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} + class Finfo: def __init__(self, jnp_finfo: jnp.finfo): @@ -108,6 +110,69 @@ def smallest_normal(self): return float(self._jnp_finfo.tiny) +# Array API Standard # +# -------------------# + + +def astype( + x: JaxArray, + dtype: jnp.dtype, + /, + *, + copy: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + ivy.utils.assertions._check_jax_x64_flag(dtype) + if x.dtype == dtype: + return jnp.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: + try: + return jnp.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: JaxArray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return jnp.broadcast_to(x.reshape(-1), shape) + return jnp.broadcast_to(x, shape) + + +@_handle_nestable_dtype_info +def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return Finfo(jnp.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return jnp.iinfo(ivy.as_native_dtype(type)) + + +def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return jnp.result_type(arrays_and_dtypes) + + result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = jnp.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) + + # Extra # # ------# @@ -180,45 +245,6 @@ def as_native_dtype( ) -# Array API Standard # -# -------------------# - - -def astype( - x: JaxArray, - dtype: jnp.dtype, - /, - *, - copy: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - ivy.utils.assertions._check_jax_x64_flag(dtype) - if x.dtype == dtype: - return jnp.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: - try: - return jnp.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: JaxArray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return jnp.broadcast_to(x.reshape(-1), shape) - return jnp.broadcast_to(x, shape) - - def dtype(x: Union[JaxArray, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -238,20 +264,6 @@ def dtype_bits(dtype_in: Union[jnp.dtype, str, np.dtype], /) -> int: ) -@_handle_nestable_dtype_info -def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return Finfo(jnp.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return jnp.iinfo(ivy.as_native_dtype(type)) - - def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -259,13 +271,3 @@ def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: return True else: return False - - -def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return jnp.result_type(arrays_and_dtypes) - - result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = jnp.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) diff --git a/ivy/functional/backends/jax/device.py b/ivy/functional/backends/jax/device.py index 9f41e96c4be8c..5158c9a0809e5 100644 --- a/ivy/functional/backends/jax/device.py +++ b/ivy/functional/backends/jax/device.py @@ -15,35 +15,8 @@ ) -# noinspection PyMethodMayBeStatic -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._save_dir = os.path.join(self._save_dir, "profile") - - def start(self): - jax.profiler.start_trace(self._save_dir) - - def stop(self): - jax.profiler.stop_trace() - - def __enter__(self): - self.start() - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - -# --- Helpers --- # -# --------------- # - - -def _dev_is_available(base_dev): - try: - jax.devices(base_dev) - return True - except RuntimeError: - return False +# Helpers # +# --------# def _to_array(x): @@ -56,6 +29,43 @@ def _to_array(x): return x +# API # +# ----# + + +def dev( + x: JaxArray, + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, jaxlib.xla_extension.Device]: + if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): + return "" + try: + dv = _to_array(x).device_buffer.device + dv = dv() + except Exception: + dv = jax.devices()[0] + if as_native: + return dv + return as_ivy_dev(dv) + + +def to_device( + x: JaxArray, + device: jaxlib.xla_extension.Device, + /, + *, + stream: Optional[int] = None, + out: Optional[JaxArray] = None, +): + if device is not None: + cur_dev = as_native_dev(dev(x)) + if cur_dev != device: + x = jax.device_put(x, as_native_dev(device)) + return x + + # this is a non-wrapped function used to place JAX arrays on respective devices, # since if we use to_device, it will return ivy.array which is not desirable def _to_device(x, device=None): @@ -66,10 +76,6 @@ def _to_device(x, device=None): return x -# --- Main --- # -# ------------ # - - def as_ivy_dev(device, /): if isinstance(device, str): return ivy.Device(device) @@ -93,44 +99,30 @@ def as_native_dev(device, /): return jax.devices(device)[idx] -def clear_cached_mem_on_dev(device: str, /): - return None +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( + *args, device_shifting_dev=device_shifting_dev, **kwargs + ) + with jax.default_device(device_shifting_dev): + return fn(*args, **kwargs) -# API # -# ----# +def clear_cached_mem_on_dev(device: str, /): + return None -def dev( - x: JaxArray, - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, jaxlib.xla_extension.Device]: - if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): - return "" +def _dev_is_available(base_dev): try: - dv = _to_array(x).device_buffer.device - dv = dv() - except Exception: - dv = jax.devices()[0] - if as_native: - return dv - return as_ivy_dev(dv) + jax.devices(base_dev) + return True + except RuntimeError: + return False def gpu_is_available() -> bool: return _dev_is_available("gpu") -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( - *args, device_shifting_dev=device_shifting_dev, **kwargs - ) - with jax.default_device(device_shifting_dev): - return fn(*args, **kwargs) - - def num_gpus() -> int: try: return len(jax.devices("gpu")) @@ -138,20 +130,24 @@ def num_gpus() -> int: return 0 -def to_device( - x: JaxArray, - device: jaxlib.xla_extension.Device, - /, - *, - stream: Optional[int] = None, - out: Optional[JaxArray] = None, -): - if device is not None: - cur_dev = as_native_dev(dev(x)) - if cur_dev != device: - x = jax.device_put(x, as_native_dev(device)) - return x - - def tpu_is_available() -> bool: return _dev_is_available("tpu") + + +# noinspection PyMethodMayBeStatic +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._save_dir = os.path.join(self._save_dir, "profile") + + def start(self): + jax.profiler.start_trace(self._save_dir) + + def stop(self): + jax.profiler.stop_trace() + + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/ivy/functional/backends/jax/elementwise.py b/ivy/functional/backends/jax/elementwise.py index 77f304a382779..6de81f65a16eb 100644 --- a/ivy/functional/backends/jax/elementwise.py +++ b/ivy/functional/backends/jax/elementwise.py @@ -16,18 +16,6 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - -def _abs_variant_sign(x): - return jnp.where(x != 0, x / jnp.abs(x), 0) - - -# --- Main --- # -# ------------ # - - def abs( x: Union[float, JaxArray], /, @@ -63,16 +51,6 @@ def add( return jnp.add(x1, x2) -def angle( - z: JaxArray, - /, - *, - deg: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.angle(z, deg=deg) - - def asin(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.arcsin(x) @@ -178,10 +156,6 @@ def cosh(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.cosh(x) -def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.deg2rad(x) - - def divide( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -209,28 +183,10 @@ def equal( return jnp.equal(x1, x2) -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.scipy.special.erf(x) - - def exp(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.exp(x) -def exp2( - x: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.power(2, x) - - def expm1(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.expm1(x) @@ -265,29 +221,6 @@ def fmin( return jnp.fmin(x1, x2) -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def fmod( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmod(x1, x2) - - -def gcd( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.gcd(x1, x2) - - def greater( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -310,15 +243,6 @@ def greater_equal( return jnp.greater_equal(x1, x2) -def imag( - val: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.imag(val) - - def isfinite(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isfinite(x) @@ -344,10 +268,6 @@ def isnan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isnan(x) -def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.isreal(x) - - def lcm(x1: JaxArray, x2: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: x1, x2 = promote_types_of_inputs(x1, x2) return jnp.lcm(x1, x2) @@ -433,34 +353,6 @@ def logical_xor( return jnp.logical_xor(x1, x2) -def maximum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 >= x2, x1, x2) - return jnp.maximum(x1, x2) - - -def minimum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 <= x2, x1, x2) - return jnp.minimum(x1, x2) - - def multiply( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -519,20 +411,6 @@ def pow( return jnp.power(x1, x2) -def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.rad2deg(x) - - -def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.real(x) - - -def reciprocal( - x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.reciprocal(x) - - @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def remainder( x1: Union[float, JaxArray], @@ -564,6 +442,10 @@ def round( return ret +def _abs_variant_sign(x): + return jnp.where(x != 0, x / jnp.abs(x), 0) + + def sign( x: JaxArray, /, *, np_variant: Optional[bool] = True, out: Optional[JaxArray] = None ) -> JaxArray: @@ -604,16 +486,6 @@ def subtract( return jnp.subtract(x1, x2) -def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.tan(x) - - -def tanh( - x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.tanh(x) - - def trapz( y: JaxArray, /, @@ -626,9 +498,129 @@ def trapz( return jnp.trapz(y, x=x, dx=dx, axis=axis) +def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.tan(x) + + +def tanh( + x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.tanh(x) + + @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def trunc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: if "int" in str(x.dtype): return x else: return jnp.trunc(x) + + +def exp2( + x: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.power(2, x) + + +def imag( + val: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.imag(val) + + +def angle( + z: JaxArray, + /, + *, + deg: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.angle(z, deg=deg) + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.scipy.special.erf(x) + + +def maximum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 >= x2, x1, x2) + return jnp.maximum(x1, x2) + + +def minimum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 <= x2, x1, x2) + return jnp.minimum(x1, x2) + + +def reciprocal( + x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.reciprocal(x) + + +def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.deg2rad(x) + + +def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.rad2deg(x) + + +def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.isreal(x) + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def fmod( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmod(x1, x2) + + +def gcd( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.gcd(x1, x2) + + +def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.real(x) diff --git a/ivy/functional/backends/jax/experimental/activations.py b/ivy/functional/backends/jax/experimental/activations.py index 8d270c6c8458a..788126f4b44cf 100644 --- a/ivy/functional/backends/jax/experimental/activations.py +++ b/ivy/functional/backends/jax/experimental/activations.py @@ -8,15 +8,6 @@ import ivy -def elu( - x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None -) -> JaxArray: - ret = jax.nn.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret - - def logit( x: JaxArray, /, @@ -31,10 +22,6 @@ def logit( return jnp.log(x / (1 - x)) -def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.log_sigmoid(input) - - def relu6(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: relu6_func = jax.nn.relu6 @@ -51,6 +38,20 @@ def custom_grad_func(x_and_grad, one): return new_func(x).astype(x.dtype) +def thresholded_relu( + x: JaxArray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.where(x > threshold, x, 0).astype(x.dtype) + + +def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.log_sigmoid(input) + + def selu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: ret = jax.nn.selu(x).astype(x.dtype) if ivy.exists(out): @@ -65,11 +66,10 @@ def silu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return ret -def thresholded_relu( - x: JaxArray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[JaxArray] = None, +def elu( + x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None ) -> JaxArray: - return jnp.where(x > threshold, x, 0).astype(x.dtype) + ret = jax.nn.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret diff --git a/ivy/functional/backends/jax/experimental/creation.py b/ivy/functional/backends/jax/experimental/creation.py index 343a8dfaf9b49..2bbc0b1b895e6 100644 --- a/ivy/functional/backends/jax/experimental/creation.py +++ b/ivy/functional/backends/jax/experimental/creation.py @@ -9,23 +9,27 @@ from ivy.functional.backends.jax import JaxArray import ivy +# Array API Standard # +# ------------------ # -def blackman_window( - size: int, - /, + +def vorbis_window( + window_length: JaxArray, *, - periodic: bool = True, - dtype: Optional[jnp.dtype] = None, + dtype: jnp.dtype = jnp.float32, out: Optional[JaxArray] = None, ) -> JaxArray: - if size < 2: - return jnp.ones([size], dtype=dtype) - if periodic: - count = jnp.arange(size) / size - else: - count = jnp.linspace(start=0, stop=size, num=size) - return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( - 0.08 * jnp.cos(2 * jnp.pi * 2 * count) + return jnp.array( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, ) @@ -73,14 +77,6 @@ def tril_indices( return jnp.tril_indices(n=n_rows, k=k, m=n_cols) -def trilu( - x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - if upper: - return jnp.triu(x, k) - return jnp.tril(x, k) - - def unsorted_segment_min( data: JaxArray, segment_ids: JaxArray, @@ -108,25 +104,28 @@ def unsorted_segment_sum( return jax.ops.segment_sum(data, segment_ids, num_segments) -# Array API Standard # -# ------------------ # - - -def vorbis_window( - window_length: JaxArray, +def blackman_window( + size: int, + /, *, - dtype: jnp.dtype = jnp.float32, + periodic: bool = True, + dtype: Optional[jnp.dtype] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.array( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, + if size < 2: + return jnp.ones([size], dtype=dtype) + if periodic: + count = jnp.arange(size) / size + else: + count = jnp.linspace(start=0, stop=size, num=size) + return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( + 0.08 * jnp.cos(2 * jnp.pi * 2 * count) ) + + +def trilu( + x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + if upper: + return jnp.triu(x, k) + return jnp.tril(x, k) diff --git a/ivy/functional/backends/jax/experimental/elementwise.py b/ivy/functional/backends/jax/experimental/elementwise.py index c05fecdf8c937..f6517ca26c95d 100644 --- a/ivy/functional/backends/jax/experimental/elementwise.py +++ b/ivy/functional/backends/jax/experimental/elementwise.py @@ -19,65 +19,41 @@ jax_ArrayLike = Union[JaxArray, Number] -# --- Helpers --- # -# --------------- # - - -# def gradient( -# x: JaxArray, -# /, -# *, -# spacing: Optional[Union[int, list, tuple]] = 1, -# axis: Optional[Union[int, list, tuple]] = None, -# edge_order: Optional[int] = 1, -# ) -> Union[JaxArray, List[JaxArray]]: -# if type(spacing) == int: -# return jnp.gradient(x, spacing, axis=axis) -# return jnp.gradient(x, *spacing, axis=axis) - - -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim - - -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.sinc(x) -# --- Main --- # -# ------------ # +@with_supported_dtypes( + {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version +) +def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jlax.lgamma(x) -def allclose( +def fmax( x1: JaxArray, x2: JaxArray, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[JaxArray] = None, -) -> bool: - return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmax(x1, x2) -def conj( - x: JaxArray, +def float_power( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.conj(x) + x1, x2 = promote_types_of_inputs(x1, x2) + if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): + out_dtype = jnp.complex128 + else: + out_dtype = jnp.float64 + return jnp.float_power(x1, x2).astype(out_dtype) def copysign( @@ -110,6 +86,65 @@ def count_nonzero( return jnp.array(jnp.count_nonzero(a, axis=axis, keepdims=keepdims), dtype=dtype) +def nansum( + x: JaxArray, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[jnp.dtype] = None, + keepdims: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + if isinstance(axis, list): + axis = tuple(axis) + return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +def isclose( + a: JaxArray, + b: JaxArray, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + +def signbit( + x: Union[JaxArray, float, int, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.signbit(x) + + +def hypot( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.hypot(x1, x2) + + +def allclose( + x1: JaxArray, + x2: JaxArray, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[JaxArray] = None, +) -> bool: + return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) + + def diff( x: JaxArray, /, @@ -128,15 +163,6 @@ def diff( return jnp.diff(x, n=n, axis=axis, prepend=prepend, append=append) -def digamma( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return js.special.digamma(x) - - def fix( x: JaxArray, /, @@ -146,36 +172,65 @@ def fix( return jnp.fix(x, out=out) -def float_power( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], +def nextafter( + x1: JaxArray, + x2: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): - out_dtype = jnp.complex128 - else: - out_dtype = jnp.float64 - return jnp.float_power(x1, x2).astype(out_dtype) + return jnp.nextafter(x1, x2) -def fmax( - x1: JaxArray, - x2: JaxArray, +def zeta( + x: JaxArray, + q: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmax(x1, x2) + temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) + inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) + temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) + ret = js.special.zeta(x, q) + ret = ret.at[nan_indices].set(jnp.nan) + ret = ret.at[inf_indices].set(jnp.inf) + return ret -def frexp( - x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None -) -> Tuple[JaxArray, JaxArray]: - return jnp.frexp(x) +# def gradient( +# x: JaxArray, +# /, +# *, +# spacing: Optional[Union[int, list, tuple]] = 1, +# axis: Optional[Union[int, list, tuple]] = None, +# edge_order: Optional[int] = 1, +# ) -> Union[JaxArray, List[JaxArray]]: +# if type(spacing) == int: +# return jnp.gradient(x, spacing, axis=axis) +# return jnp.gradient(x, *spacing, axis=axis) + + +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim + + +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis def gradient( @@ -367,27 +422,18 @@ def gradient( return outvals -def hypot( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.hypot(x1, x2) +def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + x, y = promote_types_of_inputs(x, y) + return js.special.xlogy(x, y) -def isclose( - a: JaxArray, - b: JaxArray, +def conj( + x: JaxArray, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + return jnp.conj(x) def ldexp( @@ -396,11 +442,10 @@ def ldexp( return jnp.ldexp(x1, x2) -@with_supported_dtypes( - {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version -) -def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jlax.lgamma(x) +def frexp( + x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None +) -> Tuple[JaxArray, JaxArray]: + return jnp.frexp(x) def modf( @@ -412,63 +457,10 @@ def modf( return jnp.modf(x) -def nansum( - x: JaxArray, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[jnp.dtype] = None, - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - if isinstance(axis, list): - axis = tuple(axis) - return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -def nextafter( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.nextafter(x1, x2) - - -def signbit( - x: Union[JaxArray, float, int, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.signbit(x) - - -def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.sinc(x) - - -def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - x, y = promote_types_of_inputs(x, y) - return js.special.xlogy(x, y) - - -def zeta( +def digamma( x: JaxArray, - q: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) - inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) - temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) - ret = js.special.zeta(x, q) - ret = ret.at[nan_indices].set(jnp.nan) - ret = ret.at[inf_indices].set(jnp.inf) - return ret + return js.special.digamma(x) diff --git a/ivy/functional/backends/jax/experimental/layers.py b/ivy/functional/backends/jax/experimental/layers.py index 6ea8306098c5d..e65256730b950 100644 --- a/ivy/functional/backends/jax/experimental/layers.py +++ b/ivy/functional/backends/jax/experimental/layers.py @@ -23,10 +23,6 @@ from ivy.functional.backends.jax.experimental.manipulation import _to_nested_tuple -# --- Helpers --- # -# --------------- # - - def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # determine depth pooling _, _, depth_pooling = _depth_max_pooling_helper( @@ -52,8 +48,239 @@ def _pad_str_to_list(inputs, dims, padding, strides, new_window_shape): return pad_list -# --- Main --- # -# ------------ # +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + # This function assumes that param validation is already done + window_shape = tuple(window_shape) + strides = (1,) + strides + (1,) if len(strides) == dim else strides + dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape + if isinstance(dilation, int): + dilation = (1,) + (dilation,) * dim + (1,) + else: + dilation = (1,) + tuple(dilation) + (1,) + + is_single_input = False + if inputs.ndim == len(dims) - 1: + # add singleton batch dimension because lax.reduce_window always + # needs a batch dimension. + inputs = inputs[None] + is_single_input = True + + assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" + + # shape of window after dilation + new_window_shape = tuple( + [ + window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) + for i in range(1, len(dims) - 1) + ] + ) + inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( + inputs, window_shape, strides, dim, data_format="channel_last" + ) + if not depth_pooling: + # manually creating padding list + if isinstance(padding, str): + pad_list = _pad_str_to_list( + inputs, dims, padding, strides, new_window_shape + ) + else: + if isinstance(padding, int): + padding = [(padding,) * 2] * dim + pad_list = [(0, 0)] + list(padding) + [(0, 0)] + + if ceil_mode: + c = [] + for i in range(len(dims) - 2): + pad_list[i + 1], ceil = _padding_ceil_mode( + inputs.shape[i + 1], + new_window_shape[i], + pad_list[i + 1], + strides[i + 1], + True, + ) + c.append(ceil) + + if count_include_pad: + # manually pad inputs with 0 if ceil_mode is True + # because they're not counted in average calculation + if ceil_mode: + ceil = [(0, c[i]) for i in range(len(dims) - 2)] + for i in range(len(dims) - 2): + pad_list[i + 1] = ( + pad_list[i + 1][0], + pad_list[i + 1][1] - ceil[i][1], + ) + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + inputs = jnp.pad( + inputs, + [(0, 0)] + ceil + [(0, 0)], + mode="constant", + constant_values=0.0, + ) + else: + # manually pad inputs with 1s + # because they are counted in average calculation + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + pad_list = [(0, 0)] * len(pad_list) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + pad_list = [(0, 0)] * (dim + 2) + + if not ivy.is_array(inputs): + # if dtype is not set here, jax casts it to float64 + inputs = jnp.array(inputs, dtype=jnp.float32) + if not ivy.is_array(init): + init = jnp.array(init, dtype=jnp.float32) + promoted_type = jnp.promote_types(inputs.dtype, init.dtype) + inputs = inputs.astype(promoted_type) + init = init.astype(promoted_type) + y = jlax.reduce_window( + inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation + ) + if is_single_input: + y = jnp.squeeze(y, axis=0) + return y + + +def max_pool1d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCW": + x = jnp.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCW": + res = jnp.transpose(res, (0, 2, 1)) + return res + + +def max_pool2d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCHW": + x = jnp.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCHW": + return jnp.transpose(res, (0, 3, 1, 2)) + + return res + + +def max_pool3d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NCDHW": + x = jnp.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCDHW": + res = jnp.transpose(res, (0, 4, 1, 2, 3)) + + return res def avg_pool1d( @@ -285,27 +512,74 @@ def dct( return dct_out -def dropout1d( +def idct( x: JaxArray, - prob: float, /, *, - training: bool = True, - data_format: str = "NWC", + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - if training: - x_shape = x.shape - is_batched = len(x_shape) == 3 - if data_format == "NCW": - perm = (0, 2, 1) if is_batched else (1, 0) - x = jnp.transpose(x, perm) - x_shape = x.shape - _, rng_input = jax.random.split(RNG.key) - mask = jax.random.bernoulli(rng_input, 1 - prob, x_shape) - res = jnp.where(mask, x / (1 - prob), 0) - if data_format == "NCW": - res = jnp.transpose(res, perm) + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def fft( + x: JaxArray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return jnp.fft.fft(x, n, dim, norm) + + +def dropout1d( + x: JaxArray, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[JaxArray] = None, +) -> JaxArray: + if training: + x_shape = x.shape + is_batched = len(x_shape) == 3 + if data_format == "NCW": + perm = (0, 2, 1) if is_batched else (1, 0) + x = jnp.transpose(x, perm) + x_shape = x.shape + _, rng_input = jax.random.split(RNG.key) + mask = jax.random.bernoulli(rng_input, 1 - prob, x_shape) + res = jnp.where(mask, x / (1 - prob), 0) + if data_format == "NCW": + res = jnp.transpose(res, perm) else: res = x return res @@ -365,225 +639,6 @@ def dropout3d( return res -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version -) -def embedding( - weights: JaxArray, - indices: JaxArray, - /, - *, - max_norm: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - - embeddings = jnp.take(weights, indices, axis=0) - if max_norm is not None: - norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = jnp.where( - norms > max_norm, embeddings * max_norm / norms, embeddings - ) - embeddings = jnp.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings - ) - return embeddings - - -def fft( - x: JaxArray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return jnp.fft.fft(x, n, dim, norm) - - -def fft2( - x: JaxArray, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_elem_in_list( - norm, - ["backward", "ortho", "forward"], - message=f"Unrecognized normalization mode {norm}", - ) - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" - ) - return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) - - -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - # This function assumes that param validation is already done - window_shape = tuple(window_shape) - strides = (1,) + strides + (1,) if len(strides) == dim else strides - dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape - if isinstance(dilation, int): - dilation = (1,) + (dilation,) * dim + (1,) - else: - dilation = (1,) + tuple(dilation) + (1,) - - is_single_input = False - if inputs.ndim == len(dims) - 1: - # add singleton batch dimension because lax.reduce_window always - # needs a batch dimension. - inputs = inputs[None] - is_single_input = True - - assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" - - # shape of window after dilation - new_window_shape = tuple( - [ - window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) - for i in range(1, len(dims) - 1) - ] - ) - inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( - inputs, window_shape, strides, dim, data_format="channel_last" - ) - if not depth_pooling: - # manually creating padding list - if isinstance(padding, str): - pad_list = _pad_str_to_list( - inputs, dims, padding, strides, new_window_shape - ) - else: - if isinstance(padding, int): - padding = [(padding,) * 2] * dim - pad_list = [(0, 0)] + list(padding) + [(0, 0)] - - if ceil_mode: - c = [] - for i in range(len(dims) - 2): - pad_list[i + 1], ceil = _padding_ceil_mode( - inputs.shape[i + 1], - new_window_shape[i], - pad_list[i + 1], - strides[i + 1], - True, - ) - c.append(ceil) - - if count_include_pad: - # manually pad inputs with 0 if ceil_mode is True - # because they're not counted in average calculation - if ceil_mode: - ceil = [(0, c[i]) for i in range(len(dims) - 2)] - for i in range(len(dims) - 2): - pad_list[i + 1] = ( - pad_list[i + 1][0], - pad_list[i + 1][1] - ceil[i][1], - ) - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - inputs = jnp.pad( - inputs, - [(0, 0)] + ceil + [(0, 0)], - mode="constant", - constant_values=0.0, - ) - else: - # manually pad inputs with 1s - # because they are counted in average calculation - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - pad_list = [(0, 0)] * len(pad_list) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - pad_list = [(0, 0)] * (dim + 2) - - if not ivy.is_array(inputs): - # if dtype is not set here, jax casts it to float64 - inputs = jnp.array(inputs, dtype=jnp.float32) - if not ivy.is_array(init): - init = jnp.array(init, dtype=jnp.float32) - promoted_type = jnp.promote_types(inputs.dtype, init.dtype) - inputs = inputs.astype(promoted_type) - init = init.astype(promoted_type) - y = jlax.reduce_window( - inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation - ) - if is_single_input: - y = jnp.squeeze(y, axis=0) - return y - - -def idct( - x: JaxArray, - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - def ifft( x: JaxArray, dim: int, @@ -616,17 +671,6 @@ def ifft( return jnp.fft.ifft(x, n, dim, norm) -def ifftn( - x: JaxArray, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.fft.ifftn(x, s, axes, norm) - - def interpolate( x: JaxArray, size: Union[Sequence[int], int], @@ -669,129 +713,19 @@ def interpolate( ) -def max_pool1d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCW": - x = jnp.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCW": - res = jnp.transpose(res, (0, 2, 1)) - return res - - -def max_pool2d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCHW": - x = jnp.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCHW": - return jnp.transpose(res, (0, 3, 1, 2)) - - return res - - -def max_pool3d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - if data_format == "NCDHW": - x = jnp.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCDHW": - res = jnp.transpose(res, (0, 4, 1, 2, 3)) - - return res +interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (align_corners is None or not align_corners) + and mode + not in [ + "area", + "nearest", + "nd", + "tf_area", + "mitchellcubic", + "gaussian", + "bicubic", + ] +) def reduce_window( @@ -829,6 +763,79 @@ def reduce_window( ) +def fft2( + x: JaxArray, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_elem_in_list( + norm, + ["backward", "ortho", "forward"], + message=f"Unrecognized normalization mode {norm}", + ) + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " + ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) + + +def ifftn( + x: JaxArray, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.fft.ifftn(x, s, axes, norm) + + +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version +) +def embedding( + weights: JaxArray, + indices: JaxArray, + /, + *, + max_norm: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + + embeddings = jnp.take(weights, indices, axis=0) + if max_norm is not None: + norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = jnp.where( + norms > max_norm, embeddings * max_norm / norms, embeddings + ) + embeddings = jnp.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings + ) + return embeddings + + @with_unsupported_dtypes({"0.4.14 and below": ("float16", "complex")}, backend_version) def rfftn( x: JaxArray, @@ -860,18 +867,3 @@ def rfftn( if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") return jnp.fft.rfftn(x, s, axes, norm).astype(jnp.complex128) - - -interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (align_corners is None or not align_corners) - and mode - not in [ - "area", - "nearest", - "nd", - "tf_area", - "mitchellcubic", - "gaussian", - "bicubic", - ] -) diff --git a/ivy/functional/backends/jax/experimental/linear_algebra.py b/ivy/functional/backends/jax/experimental/linear_algebra.py index 214010733ee24..897254757e7ba 100644 --- a/ivy/functional/backends/jax/experimental/linear_algebra.py +++ b/ivy/functional/backends/jax/experimental/linear_algebra.py @@ -10,31 +10,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -dot.support_native_out = True - - -def adjoint( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return jnp.conjugate(jnp.transpose(x, axes=axes)) - - -def cond( - x: JaxArray, - /, - *, - p: Optional[Union[int, str, None]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.cond(x, p=p) - - def diagflat( x: JaxArray, /, @@ -111,14 +86,23 @@ def diagflat( return ret -def dot( +def kron( a: JaxArray, b: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dot(a, b) + return jnp.kron(a, b) + + +def matrix_exp( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jla.expm(x) def eig( @@ -136,39 +120,55 @@ def eigvals(x: JaxArray, /) -> JaxArray: return jnp.linalg.eigvals(x) -def kron( - a: JaxArray, - b: JaxArray, +def adjoint( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.kron(a, b) + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return jnp.conjugate(jnp.transpose(x, axes=axes)) -def lu_factor( - x: JaxArray, +def multi_dot( + x: Sequence[JaxArray], /, *, - pivot: Optional[bool] = True, out: Optional[JaxArray] = None, -) -> Tuple[JaxArray]: - raise IvyNotImplementedException() +) -> JaxArray: + return jnp.linalg.multi_dot(x) -def matrix_exp( +def cond( x: JaxArray, /, *, + p: Optional[Union[int, str, None]] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jla.expm(x) + return jnp.linalg.cond(x, p=p) -def multi_dot( - x: Sequence[JaxArray], +def lu_factor( + x: JaxArray, + /, + *, + pivot: Optional[bool] = True, + out: Optional[JaxArray] = None, +) -> Tuple[JaxArray]: + raise IvyNotImplementedException() + + +def dot( + a: JaxArray, + b: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.linalg.multi_dot(x) + return jnp.dot(a, b) + + +dot.support_native_out = True diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index e38a8343ba23d..7258a7ae7ebe7 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -20,180 +20,102 @@ from ivy.functional.backends.jax import JaxArray -# --- Helpers --- # -# --------------- # - - -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x - - -def _to_nested_tuple(nested_list): - ret = () - if hasattr(nested_list, "__iter__"): - for inner_list in nested_list: - if hasattr(inner_list, "__iter__"): - ret += (tuple(inner_list),) - else: - ret += (inner_list,) - return ret - if ret == (): - return nested_list - - -# --- Main --- # -# ------------ # - - -def atleast_1d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_1d(*arys) - - -def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: - return jnp.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_3d(*arys) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return jnp.broadcast_shapes(*shapes) - - -def concat_from_sequence( - input_sequence: Union[Tuple[JaxArray], List[JaxArray]], +def moveaxis( + a: JaxArray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], /, *, - new_axis: int = 0, - axis: int = 0, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = jnp.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = jnp.stack(input_sequence, axis=axis) - return ret - - -def dsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + return jnp.moveaxis(a, source, destination) -def dstack( - arrays: Sequence[JaxArray], +def heaviside( + x1: JaxArray, + x2: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dstack(arrays) + return jnp.heaviside(x1, x2) -def expand( - x: JaxArray, - shape: Union[List[int], List[Tuple]], +def flipud( + m: JaxArray, /, *, copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - shape = list(shape) - if len(shape) > len(x.shape): - x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape[i] - return jnp.broadcast_to(x, tuple(shape)) + return jnp.flipud(m) -def fill_diagonal( - a: JaxArray, - v: Union[int, float], +def vstack( + arrays: Sequence[JaxArray], /, *, - wrap: bool = False, + out: Optional[JaxArray] = None, ) -> JaxArray: - shape = jnp.array(a.shape) - end = None - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (jnp.cumprod(shape[:-1])).sum() - a = jnp.reshape(a, (-1,)) - a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) - a = jnp.reshape(a, shape) - return a + return jnp.vstack(arrays) -def fliplr( - m: JaxArray, +def hstack( + arrays: Sequence[JaxArray], /, *, - copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.fliplr(m) + return jnp.hstack(arrays) -def flipud( +def rot90( m: JaxArray, /, *, copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.flipud(m) + if isinstance(axes, list): + axes = tuple(axes) + return jnp.rot90(m, k, axes) -def heaviside( - x1: JaxArray, - x2: JaxArray, +def top_k( + x: JaxArray, + k: int, /, *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.heaviside(x1, x2) + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[JaxArray, JaxArray]] = None, +) -> Tuple[JaxArray, JaxArray]: + k = min(k, x.shape[axis]) + if not largest: + indices = jnp.argsort(x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + else: + indices = jnp.argsort(-x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + if not sorted: + indices = jnp.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) + val = jnp.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) -def hsplit( - ary: JaxArray, - indices_or_sections: Union[int, Tuple[int, ...]], +def fliplr( + m: JaxArray, /, *, copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def hstack( - arrays: Sequence[JaxArray], - /, - *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.hstack(arrays) + return jnp.fliplr(m) def i0( @@ -205,16 +127,21 @@ def i0( return jnp.i0(x) -def moveaxis( - a: JaxArray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.moveaxis(a, source, destination) +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x + + +def _to_nested_tuple(nested_list): + ret = () + if hasattr(nested_list, "__iter__"): + for inner_list in nested_list: + if hasattr(inner_list, "__iter__"): + ret += (tuple(inner_list),) + else: + ret += (inner_list,) + return ret + if ret == (): + return nested_list def pad( @@ -301,18 +228,57 @@ def pad( return ret -def rot90( - m: JaxArray, +def vsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), +) -> List[JaxArray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def dsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], + /, + *, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def atleast_1d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_1d(*arys) + + +def dstack( + arrays: Sequence[JaxArray], + /, + *, out: Optional[JaxArray] = None, ) -> JaxArray: - if isinstance(axes, list): - axes = tuple(axes) - return jnp.rot90(m, k, axes) + return jnp.dstack(arrays) + + +def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: + return jnp.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_3d(*arys) def take_along_axis( @@ -332,28 +298,56 @@ def take_along_axis( return jnp.take_along_axis(arr, indices, axis, mode=mode) -def top_k( +def hsplit( + ary: JaxArray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return jnp.broadcast_shapes(*shapes) + + +def expand( x: JaxArray, - k: int, + shape: Union[List[int], List[Tuple]], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[JaxArray, JaxArray]] = None, -) -> Tuple[JaxArray, JaxArray]: - k = min(k, x.shape[axis]) - if not largest: - indices = jnp.argsort(x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - else: - indices = jnp.argsort(-x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - if not sorted: - indices = jnp.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) - val = jnp.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + copy: Optional[bool] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + shape = list(shape) + if len(shape) > len(x.shape): + x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape[i] + return jnp.broadcast_to(x, tuple(shape)) + + +def concat_from_sequence( + input_sequence: Union[Tuple[JaxArray], List[JaxArray]], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = jnp.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = jnp.stack(input_sequence, axis=axis) + return ret def unique_consecutive( @@ -399,24 +393,22 @@ def unique_consecutive( ) -def vsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def vstack( - arrays: Sequence[JaxArray], +def fill_diagonal( + a: JaxArray, + v: Union[int, float], /, *, - out: Optional[JaxArray] = None, + wrap: bool = False, ) -> JaxArray: - return jnp.vstack(arrays) + shape = jnp.array(a.shape) + end = None + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (jnp.cumprod(shape[:-1])).sum() + a = jnp.reshape(a, (-1,)) + a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) + a = jnp.reshape(a, shape) + return a diff --git a/ivy/functional/backends/jax/experimental/random.py b/ivy/functional/backends/jax/experimental/random.py index 0b75f787a241c..cfec5e0f496ce 100644 --- a/ivy/functional/backends/jax/experimental/random.py +++ b/ivy/functional/backends/jax/experimental/random.py @@ -15,27 +15,26 @@ from ivy.func_wrapper import with_unsupported_dtypes from .. import backend_version +# Extra # +# ----- # -def bernoulli( - probs: Union[float, JaxArray], + +# dirichlet +def dirichlet( + alpha: Union[JaxArray, float, Sequence[float]], + /, *, - logits: Optional[Union[float, JaxArray]] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: Optional[jaxlib.xla_extension.Device] = None, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, dtype: Optional[jnp.dtype] = None, seed: Optional[int] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - if seed: + if seed is not None: rng_input = jax.random.PRNGKey(seed) else: RNG_, rng_input = jax.random.split(_getRNG()) _setRNG(RNG_) - if logits is not None: - probs = jax.nn.softmax(logits, axis=-1) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return jax.random.bernoulli(rng_input, probs, shape=shape) + return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) def beta( @@ -57,28 +56,6 @@ def beta( return jax.random.beta(rng_input, a, b, shape, dtype) -# Extra # -# ----- # - - -# dirichlet -def dirichlet( - alpha: Union[JaxArray, float, Sequence[float]], - /, - *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: Optional[jnp.dtype] = None, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if seed is not None: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) - - @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, JaxArray], @@ -128,3 +105,25 @@ def poisson( else: ret = jax.random.poisson(rng_input, lam, shape=list_shape).astype(dtype) return ret + + +def bernoulli( + probs: Union[float, JaxArray], + *, + logits: Optional[Union[float, JaxArray]] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: Optional[jaxlib.xla_extension.Device] = None, + dtype: Optional[jnp.dtype] = None, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + if logits is not None: + probs = jax.nn.softmax(logits, axis=-1) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return jax.random.bernoulli(rng_input, probs, shape=shape) diff --git a/ivy/functional/backends/jax/experimental/statistical.py b/ivy/functional/backends/jax/experimental/statistical.py index b391273e1b628..64f1bfc031973 100644 --- a/ivy/functional/backends/jax/experimental/statistical.py +++ b/ivy/functional/backends/jax/experimental/statistical.py @@ -9,198 +9,6 @@ from ..statistical import _infer_dtype -# --- Helpers --- # -# --------------- # - - -def __find_cummax_indices( - x: JaxArray, - axis: int = 0, -) -> JaxArray: - n, indice, indices = 0, [], [] - - if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - else: - z_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - else: - indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) - else: - n, indices = 0, [] - for idx, y in enumerate(x): - if idx == 0 or x[n] <= y: - n = idx - indices.append(n) - - return jnp.asarray(indices, dtype="int64") - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - -# --- Main --- # -# ------------ # - - -def bincount( - x: JaxArray, - /, - *, - weights: Optional[JaxArray] = None, - minlength: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - if weights is not None: - ret = jnp.bincount(x, weights=weights, minlength=minlength) - ret = ret.astype(weights.dtype) - else: - ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) - return ret - - -def corrcoef( - x: JaxArray, - /, - *, - y: Optional[JaxArray] = None, - rowvar: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.corrcoef(x, y=y, rowvar=rowvar) - - -def cov( - x1: JaxArray, - x2: JaxArray = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[JaxArray] = None, - aweights: Optional[JaxArray] = None, - dtype: Optional[jnp.dtype] = None, -) -> JaxArray: - if not dtype: - x1 = jnp.asarray(x1, dtype=jnp.float64) - - if jnp.ndim(x1) > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if jnp.ndim(x2) > 2: - raise ValueError("x2 has more than 2 dimensions") - - if fweights is not None: - fweights = jnp.asarray(fweights, dtype=jnp.int64) - - return jnp.cov( - m=x1, - y=x2, - rowvar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - ) - - -def cummax( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> Tuple[JaxArray, JaxArray]: - if x.dtype in (jnp.bool_, jnp.float16): - x = x.astype(jnp.float64) - elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): - x = x.astype(jnp.int64) - elif x.dtype in (jnp.complex128, jnp.complex64): - x = jnp.real(x).astype(jnp.float64) - - if exclusive or (reverse and exclusive): - if exclusive and reverse: - indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) - x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - x, indices = jnp.concatenate( - (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 - ), jnp.concatenate( - (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 - ) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - indices = __find_cummax_indices(x, axis=axis) - res = jlax.cummax(x, axis=axis) - return res, indices - - if reverse: - y = jnp.flip(x, axis=axis) - indices = __find_cummax_indices(y, axis=axis) - indices = jnp.flip(indices, axis=axis) - else: - indices = __find_cummax_indices(x, axis=axis) - return jlax.cummax(x, axis, reverse=reverse), indices - - -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cummin( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if axis < 0: - axis = axis + len(x.shape) - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - return jlax.cummin(x, axis, reverse=reverse).astype(dtype) - - @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16",)}, backend_version, @@ -312,16 +120,6 @@ def histogram( return ret -def igamma( - a: JaxArray, - /, - *, - x: JaxArray, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jlax.igamma(a=a, x=x) - - @with_unsupported_dtypes( {"0.4.14 and below": ("complex64", "complex128")}, backend_version ) @@ -364,6 +162,34 @@ def nanmean( return jnp.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) +def quantile( + a: JaxArray, + q: Union[float, JaxArray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + axis = tuple(axis) if isinstance(axis, list) else axis + interpolation = "nearest" if interpolation == "nearest_jax" else interpolation + return jnp.quantile( + a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ) + + +def corrcoef( + x: JaxArray, + /, + *, + y: Optional[JaxArray] = None, + rowvar: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.corrcoef(x, y=y, rowvar=rowvar) + + def nanmedian( input: JaxArray, /, @@ -393,18 +219,184 @@ def nanmedian( ) -def quantile( - a: JaxArray, - q: Union[float, JaxArray], +def bincount( + x: JaxArray, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, + weights: Optional[JaxArray] = None, + minlength: int = 0, out: Optional[JaxArray] = None, ) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - interpolation = "nearest" if interpolation == "nearest_jax" else interpolation - return jnp.quantile( - a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + if weights is not None: + ret = jnp.bincount(x, weights=weights, minlength=minlength) + ret = ret.astype(weights.dtype) + else: + ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) + return ret + + +def cov( + x1: JaxArray, + x2: JaxArray = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[JaxArray] = None, + aweights: Optional[JaxArray] = None, + dtype: Optional[jnp.dtype] = None, +) -> JaxArray: + if not dtype: + x1 = jnp.asarray(x1, dtype=jnp.float64) + + if jnp.ndim(x1) > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if jnp.ndim(x2) > 2: + raise ValueError("x2 has more than 2 dimensions") + + if fweights is not None: + fweights = jnp.asarray(fweights, dtype=jnp.int64) + + return jnp.cov( + m=x1, + y=x2, + rowvar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, ) + + +def cummax( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> Tuple[JaxArray, JaxArray]: + if x.dtype in (jnp.bool_, jnp.float16): + x = x.astype(jnp.float64) + elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): + x = x.astype(jnp.int64) + elif x.dtype in (jnp.complex128, jnp.complex64): + x = jnp.real(x).astype(jnp.float64) + + if exclusive or (reverse and exclusive): + if exclusive and reverse: + indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) + x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + x, indices = jnp.concatenate( + (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 + ), jnp.concatenate( + (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 + ) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + indices = __find_cummax_indices(x, axis=axis) + res = jlax.cummax(x, axis=axis) + return res, indices + + if reverse: + y = jnp.flip(x, axis=axis) + indices = __find_cummax_indices(y, axis=axis) + indices = jnp.flip(indices, axis=axis) + else: + indices = __find_cummax_indices(x, axis=axis) + return jlax.cummax(x, axis, reverse=reverse), indices + + +def __find_cummax_indices( + x: JaxArray, + axis: int = 0, +) -> JaxArray: + n, indice, indices = 0, [], [] + + if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + else: + z_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + else: + indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) + else: + n, indices = 0, [] + for idx, y in enumerate(x): + if idx == 0 or x[n] <= y: + n = idx + indices.append(n) + + return jnp.asarray(indices, dtype="int64") + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cummin( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if axis < 0: + axis = axis + len(x.shape) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + return jlax.cummin(x, axis, reverse=reverse).astype(dtype) + + +def igamma( + a: JaxArray, + /, + *, + x: JaxArray, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jlax.igamma(a=a, x=x) diff --git a/ivy/functional/backends/jax/general.py b/ivy/functional/backends/jax/general.py index 1315cef38980a..9221319cfeaed 100644 --- a/ivy/functional/backends/jax/general.py +++ b/ivy/functional/backends/jax/general.py @@ -22,8 +22,33 @@ from . import backend_version -# --- Helpers --- # -# --------------- # +def container_types(): + flat_mapping_spec = importlib.util.find_spec( + "FlatMapping", "haiku._src.data_structures" + ) + if not flat_mapping_spec: + from haiku._src.data_structures import FlatMapping + else: + FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) + return [FlatMapping] + + +def current_backend_str() -> str: + return "jax" + + +def is_native_array(x, /, *, exclusive=False): + if exclusive: + return isinstance(x, NativeArray) + return isinstance( + x, + ( + NativeArray, + jax.interpreters.ad.JVPTracer, + jax.core.ShapedArray, + jax.interpreters.partial_eval.DynamicJaxprTracer, + ), + ) def _mask_to_index(query, x): @@ -36,35 +61,63 @@ def _mask_to_index(query, x): return jnp.where(query), expected_shape -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view +def get_item( + x: JaxArray, + /, + query: Union[JaxArray, Tuple], + *, + copy: bool = None, +) -> JaxArray: + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return jnp.array([], dtype=x.dtype) + else: + return jnp.expand_dims(x, 0) + query, _ = _mask_to_index(query, x) + elif isinstance(query, list): + query = (query,) + return x.__getitem__(query) -# --- Main --- # -# ------------ # +def set_item( + x: JaxArray, + query: Union[JaxArray, Tuple], + val: JaxArray, + /, + *, + copy: Optional[bool] = False, +) -> JaxArray: + if ivy.is_array(query) and ivy.is_bool_dtype(query): + query, expected_shape = _mask_to_index(query, x) + val = _broadcast_to(val, expected_shape)._data + ret = x.at[query].set(val) + if copy: + return ret + return ivy.inplace_update(x, _to_device(ret)) def array_equal(x0: JaxArray, x1: JaxArray, /) -> bool: return bool(jnp.array_equal(x0, x1)) -def container_types(): - flat_mapping_spec = importlib.util.find_spec( - "FlatMapping", "haiku._src.data_structures" - ) - if not flat_mapping_spec: - from haiku._src.data_structures import FlatMapping +@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) +def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return np.array(_to_array(x)) else: - FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) - return [FlatMapping] + return np.asarray(_to_array(x)) -def current_backend_str() -> str: - return "jax" +def to_scalar(x: JaxArray, /) -> Number: + if isinstance(x, Number): + return x + else: + return _to_array(x).item() + + +def to_list(x: JaxArray, /) -> list: + return _to_array(x).tolist() def gather( @@ -99,36 +152,6 @@ def gather( return result -def gather_nd( - params: JaxArray, - indices: JaxArray, - /, - *, - batch_dims: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = jnp.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result - - def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -160,23 +183,34 @@ def gather_nd_helper(params, indices): return ret -def get_item( - x: JaxArray, +def gather_nd( + params: JaxArray, + indices: JaxArray, /, - query: Union[JaxArray, Tuple], *, - copy: bool = None, + batch_dims: int = 0, + out: Optional[JaxArray] = None, ) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return jnp.array([], dtype=x.dtype) + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] else: - return jnp.expand_dims(x, 0) - query, _ = _mask_to_index(query, x) - elif isinstance(query, list): - query = (query,) - return x.__getitem__(query) + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = jnp.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result def get_num_dims(x: JaxArray, /, *, as_array: bool = False) -> Union[JaxArray, int]: @@ -254,38 +288,16 @@ def inplace_update( return val -def inplace_variables_supported(): - return False - - -def is_native_array(x, /, *, exclusive=False): - if exclusive: - return isinstance(x, NativeArray) - return isinstance( - x, - ( - NativeArray, - jax.interpreters.ad.JVPTracer, - jax.core.ShapedArray, - jax.interpreters.partial_eval.DynamicJaxprTracer, - ), - ) - - -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def isin( - elements: JaxArray, - test_elements: JaxArray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> JaxArray: - return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view -def itemsize(x: JaxArray) -> int: - return x.itemsize +def inplace_variables_supported(): + return False def multiprocessing(context: Optional[str] = None): @@ -330,6 +342,9 @@ def scatter_flat( return target +scatter_flat.support_native_out = True + + def scatter_nd( indices: JaxArray, updates: JaxArray, @@ -377,21 +392,7 @@ def scatter_nd( return target -def set_item( - x: JaxArray, - query: Union[JaxArray, Tuple], - val: JaxArray, - /, - *, - copy: Optional[bool] = False, -) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - query, expected_shape = _mask_to_index(query, x) - val = _broadcast_to(val, expected_shape)._data - ret = x.at[query].set(val) - if copy: - return ret - return ivy.inplace_update(x, _to_device(ret)) +scatter_nd.support_native_out = True def shape( @@ -406,25 +407,6 @@ def shape( return ivy.Shape(x.shape) -def to_list(x: JaxArray, /) -> list: - return _to_array(x).tolist() - - -@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) -def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return np.array(_to_array(x)) - else: - return np.asarray(_to_array(x)) - - -def to_scalar(x: JaxArray, /) -> Number: - if isinstance(x, Number): - return x - else: - return _to_array(x).item() - - def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -436,5 +418,17 @@ def vmap( ) -scatter_flat.support_native_out = True -scatter_nd.support_native_out = True +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def isin( + elements: JaxArray, + test_elements: JaxArray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> JaxArray: + return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) + + +def itemsize(x: JaxArray) -> int: + return x.itemsize diff --git a/ivy/functional/backends/jax/gradients.py b/ivy/functional/backends/jax/gradients.py index dcd85bc2b395e..a5e6d789a070b 100644 --- a/ivy/functional/backends/jax/gradients.py +++ b/ivy/functional/backends/jax/gradients.py @@ -18,8 +18,19 @@ ) -# --- Helpers --- # -# --------------- # +# ToDo: modify these functions to track whether variable() has been called +def variable(x, /): + return x + + +def is_variable(x, /, *, exclusive=False): + if exclusive: + return False + return isinstance(x, NativeArray) + + +def variable_data(x: JaxArray, /) -> JaxArray: + return x def _forward_fn( @@ -57,10 +68,6 @@ def _forward_fn( return ret_values -# --- Main --- # -# ------------ # - - def execute_with_gradients( func, xs: JaxArray, @@ -120,18 +127,21 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): - grad_fn = lambda x_in: ivy.to_native(func(x_in)) - callback_fn = lambda x_in: ivy.to_ivy( - jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) - ) +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + value, grad = jax.value_and_grad(grad_fn)(xs) + return ivy.to_ivy(value), ivy.to_ivy(grad) + return callback_fn -def is_variable(x, /, *, exclusive=False): - if exclusive: - return False - return isinstance(x, NativeArray) +def stop_gradient( + x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + return jlax.stop_gradient(x) def jac(func: Callable): @@ -148,27 +158,9 @@ def jac(func: Callable): return callback_fn -def stop_gradient( - x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - return jlax.stop_gradient(x) - - -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - value, grad = jax.value_and_grad(grad_fn)(xs) - return ivy.to_ivy(value), ivy.to_ivy(grad) - +def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): + grad_fn = lambda x_in: ivy.to_native(func(x_in)) + callback_fn = lambda x_in: ivy.to_ivy( + jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) + ) return callback_fn - - -# ToDo: modify these functions to track whether variable() has been called -def variable(x, /): - return x - - -def variable_data(x: JaxArray, /) -> JaxArray: - return x diff --git a/ivy/functional/backends/jax/layers.py b/ivy/functional/backends/jax/layers.py index c27e92a2680eb..ff4c264439da7 100644 --- a/ivy/functional/backends/jax/layers.py +++ b/ivy/functional/backends/jax/layers.py @@ -16,18 +16,41 @@ ) -# --- Helpers --- # -# --------------- # +def _transpose_padding_helper(k, s, padding, dilation, diff=0): + k = (k - 1) * dilation + 1 + if padding == "SAME": + pad_len = k + s - 2 + pad_len -= diff + if s > k - 1: + pad_a = k - 1 + else: + pad_a = int(jnp.ceil(pad_len / 2)) + else: + pad_len = k + s - 2 + max(k - s, 0) + pad_a = k - 1 + pad_b = pad_len - pad_a + return pad_a, pad_b -def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): - first = True if filter_format == "channel_first" else False - if dims == 1: - return "WIO" if not first else "OIW" - if dims == 2: - return "HWIO" if not first else "OIHW" - elif dims == 3: - return "DHWIO" if not first else "OIDHW" +def _get_tranpose_padding( + x_shape, filter_shape, strides, padding, dims, dilations, output_shape +): + new_shape = [ + _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) + for i in range(dims) + ] + if output_shape is None: + output_shape = [x_shape[0], *new_shape, filter_shape[-1]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] + shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] + pad_list = [ + _transpose_padding_helper( + filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] + ) + for i in range(dims) + ] + return pad_list def _get_new_padding_before_conv( @@ -71,47 +94,6 @@ def _get_new_padding_before_conv( return padding -def _get_tranpose_padding( - x_shape, filter_shape, strides, padding, dims, dilations, output_shape -): - new_shape = [ - _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) - for i in range(dims) - ] - if output_shape is None: - output_shape = [x_shape[0], *new_shape, filter_shape[-1]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] - shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] - pad_list = [ - _transpose_padding_helper( - filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] - ) - for i in range(dims) - ] - return pad_list - - -def _transpose_padding_helper(k, s, padding, dilation, diff=0): - k = (k - 1) * dilation + 1 - if padding == "SAME": - pad_len = k + s - 2 - pad_len -= diff - if s > k - 1: - pad_a = k - 1 - else: - pad_a = int(jnp.ceil(pad_len / 2)) - else: - pad_len = k + s - 2 + max(k - s, 0) - pad_a = k - 1 - pad_b = pad_len - pad_a - return pad_a, pad_b - - -# --- Main --- # -# ------------ # - - def conv1d( x: JaxArray, filters: JaxArray, @@ -249,6 +231,37 @@ def conv2d_transpose( return res +def depthwise_conv2d( + x: JaxArray, + filters: JaxArray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[JaxArray] = None, +) -> JaxArray: + strides = [strides] * 2 if isinstance(strides, int) else strides + strides = [strides[1], strides[2]] if len(strides) == 4 else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters + cn = filters.shape[-1] + filters = jnp.expand_dims(filters, -2) + return jlax.conv_general_dilated( + x, + filters, + strides, + padding, + None, + dilations, + (data_format, "HWIO", data_format), + feature_group_count=cn, + ) + + def conv3d( x: JaxArray, filters: JaxArray, @@ -317,6 +330,16 @@ def conv3d_transpose( return res +def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): + first = True if filter_format == "channel_first" else False + if dims == 1: + return "WIO" if not first else "OIW" + if dims == 2: + return "HWIO" if not first else "OIHW" + elif dims == 3: + return "DHWIO" if not first else "OIDHW" + + def conv_general_dilated( x: JaxArray, filters: JaxArray, @@ -432,34 +455,3 @@ def conv_general_transpose( if data_format == "channel_first": return jnp.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res - - -def depthwise_conv2d( - x: JaxArray, - filters: JaxArray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[JaxArray] = None, -) -> JaxArray: - strides = [strides] * 2 if isinstance(strides, int) else strides - strides = [strides[1], strides[2]] if len(strides) == 4 else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters - cn = filters.shape[-1] - filters = jnp.expand_dims(filters, -2) - return jlax.conv_general_dilated( - x, - filters, - strides, - padding, - None, - dilations, - (data_format, "HWIO", data_format), - feature_group_count=cn, - ) diff --git a/ivy/functional/backends/jax/linear_algebra.py b/ivy/functional/backends/jax/linear_algebra.py index 64c80f2339090..383a9d032abc6 100644 --- a/ivy/functional/backends/jax/linear_algebra.py +++ b/ivy/functional/backends/jax/linear_algebra.py @@ -58,19 +58,13 @@ def det(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.linalg.det(x) -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def diag( - x: JaxArray, - /, - *, - k: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.diag(x, k=k) +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] + ) + eigenvalues, eigenvectors = jnp.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) @@ -98,13 +92,15 @@ def diagonal( return ret -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] - ) - eigenvalues, eigenvectors = jnp.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) +def tensorsolve( + x1: JaxArray, + x2: JaxArray, + /, + *, + axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.tensorsolve(x1, x2, axes) @with_unsupported_dtypes( @@ -392,17 +388,6 @@ def tensordot( return jnp.tensordot(x1, x2, axes) -def tensorsolve( - x1: JaxArray, - x2: JaxArray, - /, - *, - axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.tensorsolve(x1, x2, axes) - - @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -419,21 +404,6 @@ def trace( return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, - backend_version, -) -def vander( - x: JaxArray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def vecdot( x1: JaxArray, x2: JaxArray, /, *, axis: int = -1, out: Optional[JaxArray] = None @@ -468,6 +438,36 @@ def vector_norm( return jnp.sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def diag( + x: JaxArray, + /, + *, + k: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.diag(x, k=k) + + +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, + backend_version, +) +def vander( + x: JaxArray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes( { "0.4.14 and below": ( diff --git a/ivy/functional/backends/jax/manipulation.py b/ivy/functional/backends/jax/manipulation.py index 4ebc42a93b9c7..61e638b9509cb 100644 --- a/ivy/functional/backends/jax/manipulation.py +++ b/ivy/functional/backends/jax/manipulation.py @@ -12,61 +12,10 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x -# --- Main --- # -# ------------ # - - -def clip( - x: JaxArray, - x_min: Union[Number, JaxArray], - x_max: Union[Number, JaxArray], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - if ( - hasattr(x_min, "dtype") - and hasattr(x_max, "dtype") - and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) - ): - if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( - jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float32) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - elif ( - jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) - ) and ( - jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float64) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - else: - promoted_type = jnp.promote_types(x.dtype, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x.astype(promoted_type) - # jnp.clip isn't used because of inconsistent gradients - x = jnp.where(x > x_max, x_max, x) - return jnp.where(x < x_min, x_min, x) - - # Array API Standard # # -------------------# @@ -93,18 +42,6 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) -@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) -def constant_pad( - x: JaxArray, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) - - def expand_dims( x: JaxArray, /, @@ -142,17 +79,6 @@ def permute_dims( return jnp.transpose(x, axes) -def repeat( - x: JaxArray, - /, - repeats: Union[int, Iterable[int]], - *, - axis: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.repeat(x, repeats, axis) - - def reshape( x: JaxArray, /, @@ -188,6 +114,38 @@ def roll( return jnp.roll(x, shift, axis) +def squeeze( + x: JaxArray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + else: + ret = jnp.squeeze(x, axis=axis) + return ret + + +def stack( + arrays: Union[Tuple[JaxArray], List[JaxArray]], + /, + *, + axis: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + try: + return jnp.stack(arrays, axis=axis) + except ValueError as error: + raise ivy.utils.exceptions.IvyIndexError(error) + + # Extra # # ------# @@ -226,54 +184,76 @@ def split( return jnp.split(x, num_or_size_splits, axis) -def squeeze( +def repeat( x: JaxArray, /, + repeats: Union[int, Iterable[int]], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, + axis: Optional[int] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - else: - ret = jnp.squeeze(x, axis=axis) - return ret + return jnp.repeat(x, repeats, axis) -def stack( - arrays: Union[Tuple[JaxArray], List[JaxArray]], - /, - *, - axis: int = 0, - out: Optional[JaxArray] = None, +def tile( + x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None ) -> JaxArray: - try: - return jnp.stack(arrays, axis=axis) - except ValueError as error: - raise ivy.utils.exceptions.IvyIndexError(error) + return jnp.tile(x, repeats) -def swapaxes( +def clip( x: JaxArray, - axis0: int, - axis1: int, + x_min: Union[Number, JaxArray], + x_max: Union[Number, JaxArray], /, *, - copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.swapaxes(x, axis0, axis1) + if ( + hasattr(x_min, "dtype") + and hasattr(x_max, "dtype") + and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) + ): + if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( + jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float32) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + elif ( + jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) + ) and ( + jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float64) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + else: + promoted_type = jnp.promote_types(x.dtype, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x.astype(promoted_type) + # jnp.clip isn't used because of inconsistent gradients + x = jnp.where(x > x_max, x_max, x) + return jnp.where(x < x_min, x_min, x) -def tile( - x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None +@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) +def constant_pad( + x: JaxArray, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.tile(x, repeats) + return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) def unstack( @@ -298,3 +278,15 @@ def zero_pad( x: JaxArray, /, pad_width: List[List[int]], *, out: Optional[JaxArray] = None ): return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=0) + + +def swapaxes( + x: JaxArray, + axis0: int, + axis1: int, + /, + *, + copy: Optional[bool] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.swapaxes(x, axis0, axis1) diff --git a/ivy/functional/backends/jax/random.py b/ivy/functional/backends/jax/random.py index 38b599f8ef405..53aedb4e811e4 100644 --- a/ivy/functional/backends/jax/random.py +++ b/ivy/functional/backends/jax/random.py @@ -26,8 +26,12 @@ def __init__(self): self.key = jax.random.PRNGKey(0) -# --- Helpers --- # -# --------------- # +RNG = RNGWrapper() + + +def _setRNG(key): + global RNG + RNG.key = key def _getRNG(): @@ -35,13 +39,47 @@ def _getRNG(): return RNG.key -def _setRNG(key): - global RNG - RNG.key = key +def random_uniform( + *, + low: Union[float, JaxArray] = 0.0, + high: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + shape = _check_bounds_and_get_shape(low, high, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.uniform( + rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 + ).astype(dtype) -# --- Main --- # -# ------------ # +def random_normal( + *, + mean: Union[float, JaxArray] = 0.0, + std: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) @@ -114,49 +152,6 @@ def randint( return jax.random.randint(rng_input, shape, low, high, dtype) -def random_normal( - *, - mean: Union[float, JaxArray] = 0.0, - std: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean - - -def random_uniform( - *, - low: Union[float, JaxArray] = 0.0, - high: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - shape = _check_bounds_and_get_shape(low, high, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.uniform( - rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 - ).astype(dtype) - - def seed(*, seed_value: int = 0) -> None: _setRNG(jax.random.PRNGKey(seed_value)) return @@ -181,6 +176,3 @@ def shuffle( # jax.random.shuffle is deprecated; identical behaviour reproduced with # jax.random.permutation return jax.random.permutation(key=rng_input, x=x, axis=axis, independent=True) - - -RNG = RNGWrapper() diff --git a/ivy/functional/backends/jax/searching.py b/ivy/functional/backends/jax/searching.py index e5f6a320ef428..794436773c74a 100644 --- a/ivy/functional/backends/jax/searching.py +++ b/ivy/functional/backends/jax/searching.py @@ -64,19 +64,6 @@ def argmin( return ret -# Extra # -# ----- # - - -def argwhere( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.argwhere(x) - - def nonzero( x: JaxArray, /, @@ -103,3 +90,16 @@ def where( ) -> JaxArray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(jnp.where(condition, x1, x2), x1.dtype, copy=False) + + +# Extra # +# ----- # + + +def argwhere( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.argwhere(x) diff --git a/ivy/functional/backends/jax/sorting.py b/ivy/functional/backends/jax/sorting.py index 5ccd39bef98f6..5111eca178530 100644 --- a/ivy/functional/backends/jax/sorting.py +++ b/ivy/functional/backends/jax/sorting.py @@ -26,15 +26,20 @@ def argsort( ) -# msort -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def msort( - a: Union[JaxArray, list, tuple], +def sort( + x: JaxArray, /, *, + axis: int = -1, + descending: bool = False, + stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.sort(a, axis=0, kind="mergesort") + kind = "stable" if stable else "quicksort" + ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) + if descending: + ret = jnp.asarray(jnp.flip(ret, axis=axis)) + return ret def searchsorted( @@ -74,17 +79,12 @@ def searchsorted( return ret.astype(ret_dtype) -def sort( - x: JaxArray, +# msort +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def msort( + a: Union[JaxArray, list, tuple], /, *, - axis: int = -1, - descending: bool = False, - stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - kind = "stable" if stable else "quicksort" - ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) - if descending: - ret = jnp.asarray(jnp.flip(ret, axis=axis)) - return ret + return jnp.sort(a, axis=0, kind="mergesort") diff --git a/ivy/functional/backends/jax/statistical.py b/ivy/functional/backends/jax/statistical.py index b6af324bafde4..b03be2a367141 100644 --- a/ivy/functional/backends/jax/statistical.py +++ b/ivy/functional/backends/jax/statistical.py @@ -9,102 +9,20 @@ from ivy.functional.backends.jax import JaxArray from . import backend_version - -# --- Helpers --- # -# --------------- # - - -def _infer_dtype(dtype: jnp.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype - - -# --- Main --- # -# ------------ # - - -# Extra # -# ------# +# Array API Standard # +# -------------------# -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cumprod( +def min( x: JaxArray, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return jnp.cumprod(x, axis, dtype=dtype) - elif exclusive and reverse: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - return jnp.flip(x, axis=(axis,)) - - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumprod(x, -1, dtype=dtype) - return jnp.swapaxes(x, axis, -1) - else: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - return jnp.flip(x, axis=axis) - - -def cumsum( - x: JaxArray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[jnp.dtype] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - res = jnp.flip(x, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumsum(x, -1, dtype=dtype) - res = jnp.swapaxes(x, axis, -1) - elif reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - res = jnp.flip(x, axis=axis) - return res - return jnp.cumsum(x, axis, dtype=dtype) - - -def einsum( - equation: str, *operands: JaxArray, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.einsum(equation, *operands) + axis = tuple(axis) if isinstance(axis, list) else axis + return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) def max( @@ -131,20 +49,11 @@ def mean( return jnp.mean(x, axis=axis, keepdims=keepdims) -# Array API Standard # -# -------------------# - - -def min( - x: JaxArray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) +def _infer_dtype(dtype: jnp.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype def prod( @@ -224,3 +133,85 @@ def var( x.dtype, copy=False, ) + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cumprod( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return jnp.cumprod(x, axis, dtype=dtype) + elif exclusive and reverse: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + return jnp.flip(x, axis=(axis,)) + + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumprod(x, -1, dtype=dtype) + return jnp.swapaxes(x, axis, -1) + else: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + return jnp.flip(x, axis=axis) + + +def cumsum( + x: JaxArray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + res = jnp.flip(x, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumsum(x, -1, dtype=dtype) + res = jnp.swapaxes(x, axis, -1) + elif reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + res = jnp.flip(x, axis=axis) + return res + return jnp.cumsum(x, axis, dtype=dtype) + + +def einsum( + equation: str, *operands: JaxArray, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.einsum(equation, *operands) diff --git a/ivy/functional/backends/mxnet/activations.py b/ivy/functional/backends/mxnet/activations.py index a6f3fc0e9b713..8efd486936334 100644 --- a/ivy/functional/backends/mxnet/activations.py +++ b/ivy/functional/backends/mxnet/activations.py @@ -32,14 +32,6 @@ def leaky_relu( return mx.nd.LeakyReLU(x, slope=alpha) -def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): - raise IvyNotImplementedException() - - -def mish(x: None, /, *, out: Optional[None] = None) -> None: - raise IvyNotImplementedException() - - def relu(x: None, /, *, complex_mode="jax", out: Optional[None] = None) -> None: return mx.nd.relu(x) @@ -79,3 +71,11 @@ def softplus( if threshold is not None: return mx.nd.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) + + +def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): + raise IvyNotImplementedException() + + +def mish(x: None, /, *, out: Optional[None] = None) -> None: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/creation.py b/ivy/functional/backends/mxnet/creation.py index 417bea7a45a17..632b1daca8f9e 100644 --- a/ivy/functional/backends/mxnet/creation.py +++ b/ivy/functional/backends/mxnet/creation.py @@ -60,15 +60,7 @@ def asarray( return ret -def copy_array( - x: Union[(None, mx.ndarray.NDArray)], - *, - to_ivy_array: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() +array = asarray def empty( @@ -115,15 +107,6 @@ def from_dlpack( raise IvyNotImplementedException() -def frombuffer( - buffer: bytes, - dtype: Optional[None] = float, - count: Optional[int] = (-1), - offset: Optional[int] = 0, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def full( shape: Union[(ivy.NativeShape, Sequence[int])], fill_value: Union[(int, float, bool)], @@ -171,21 +154,6 @@ def meshgrid( raise IvyNotImplementedException() -def one_hot( - indices: Union[(None, mx.ndarray.NDArray)], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[None] = None, - device: str, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def ones( shape: Optional[ivy.NativeShape] = None, *, @@ -227,12 +195,6 @@ def triu( raise IvyNotImplementedException() -def triu_indices( - n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str -) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - def zeros( *size: Union[(int, Sequence[int])], shape: Optional[ivy.NativeShape] = None, @@ -258,4 +220,42 @@ def zeros_like( return ivy.to_device(ret, device) -array = asarray +def copy_array( + x: Union[(None, mx.ndarray.NDArray)], + *, + to_ivy_array: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() + + +def one_hot( + indices: Union[(None, mx.ndarray.NDArray)], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[None] = None, + device: str, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def frombuffer( + buffer: bytes, + dtype: Optional[None] = float, + count: Optional[int] = (-1), + offset: Optional[int] = 0, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def triu_indices( + n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str +) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/data_type.py b/ivy/functional/backends/mxnet/data_type.py index 6d15845c7c19f..1ec0e1174ceec 100644 --- a/ivy/functional/backends/mxnet/data_type.py +++ b/ivy/functional/backends/mxnet/data_type.py @@ -5,18 +5,6 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from ivy.utils.exceptions import IvyNotImplementedException -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "u1": "uint8", -} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int32"): "int32", @@ -46,6 +34,19 @@ "bool": np.bool_, } +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "u1": "uint8", +} + class Finfo: def __init__(self, mx_finfo: mx.np.finfo): @@ -83,6 +84,57 @@ def __repr__(self): return repr(self._mx_finfo) +def astype( + x: Union[(None, mx.ndarray.NDArray)], + dtype: Union[(None, str)], + /, + *, + copy: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return mx.nd.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays( + *arrays: Union[(None, mx.ndarray.NDArray)] +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def broadcast_to( + x: Union[(None, mx.ndarray.NDArray)], + /, + shape: Union[(ivy.NativeShape, Sequence[int])], + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +@_handle_nestable_dtype_info +def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: + if isinstance(type, mx.ndarray.NDArray): + type = type.dtype + return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: + # using np.iinfo as mx use np dtypes and mxnet iinfo not provided + if isinstance(type, mx.ndarray.NDArray): + type = type.asnumpy().dtype + return np.iinfo(ivy.as_native_dtype(type)) + + +def result_type( + *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] +) -> ivy.Dtype: + raise IvyNotImplementedException() + + def as_ivy_dtype( dtype_in: Union[(str, int, float, complex, bool, np.dtype)], / ) -> ivy.Dtype: @@ -137,36 +189,6 @@ def as_native_dtype(dtype_in: Union[(None, str, bool, int, float, np.dtype)]) -> ) -def astype( - x: Union[(None, mx.ndarray.NDArray)], - dtype: Union[(None, str)], - /, - *, - copy: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return mx.nd.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays( - *arrays: Union[(None, mx.ndarray.NDArray)] -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_to( - x: Union[(None, mx.ndarray.NDArray)], - /, - shape: Union[(ivy.NativeShape, Sequence[int])], - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def dtype( x: Union[(None, mx.ndarray.NDArray, np.ndarray)], *, as_native: bool = False ) -> ivy.Dtype: @@ -179,26 +201,5 @@ def dtype_bits(dtype_in: Union[(None, str, np.dtype)], /) -> int: raise IvyNotImplementedException() -@_handle_nestable_dtype_info -def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: - if isinstance(type, mx.ndarray.NDArray): - type = type.dtype - return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: - # using np.iinfo as mx use np dtypes and mxnet iinfo not provided - if isinstance(type, mx.ndarray.NDArray): - type = type.asnumpy().dtype - return np.iinfo(ivy.as_native_dtype(type)) - - def is_native_dtype(dtype_in: Union[(None, str)], /) -> bool: raise IvyNotImplementedException() - - -def result_type( - *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] -) -> ivy.Dtype: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/device.py b/ivy/functional/backends/mxnet/device.py index 2a47dc7848916..309621811a71e 100644 --- a/ivy/functional/backends/mxnet/device.py +++ b/ivy/functional/backends/mxnet/device.py @@ -11,21 +11,23 @@ from ivy.utils.exceptions import IvyNotImplementedException -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - raise IvyNotImplementedException() - - def start(self): - raise IvyNotImplementedException() - - def stop(self): - raise IvyNotImplementedException() +def dev( + x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False +) -> Union[(ivy.Device, str)]: + if as_native: + return x.context + return as_ivy_dev(x.context) - def __enter__(self): - raise IvyNotImplementedException() - def __exit__(self, exc_type, exc_val, exc_tb): - raise IvyNotImplementedException() +def to_device( + x: Union[(None, mx.ndarray.NDArray)], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return x.as_in_context(as_native_dev(device)) def as_ivy_dev(device): @@ -60,12 +62,8 @@ def clear_cached_mem_on_dev(device: str, /): raise IvyNotImplementedException() -def dev( - x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False -) -> Union[(ivy.Device, str)]: - if as_native: - return x.context - return as_ivy_dev(x.context) +def num_gpus() -> int: + return mx.context.num_gpus() def gpu_is_available() -> bool: @@ -74,20 +72,22 @@ def gpu_is_available() -> bool: return False -def num_gpus() -> int: - return mx.context.num_gpus() +def tpu_is_available() -> bool: + return False -def to_device( - x: Union[(None, mx.ndarray.NDArray)], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return x.as_in_context(as_native_dev(device)) +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + raise IvyNotImplementedException() + def start(self): + raise IvyNotImplementedException() -def tpu_is_available() -> bool: - return False + def stop(self): + raise IvyNotImplementedException() + + def __enter__(self): + raise IvyNotImplementedException() + + def __exit__(self, exc_type, exc_val, exc_tb): + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/elementwise.py b/ivy/functional/backends/mxnet/elementwise.py index f3a98ade3f1cf..42a8a2788b605 100644 --- a/ivy/functional/backends/mxnet/elementwise.py +++ b/ivy/functional/backends/mxnet/elementwise.py @@ -46,16 +46,6 @@ def add( return mx.nd.add(mx.nd.multiply(x1, alpha), x2) -def angle( - input: Union[(None, mx.ndarray.NDArray)], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def asin( x: Union[(None, mx.ndarray.NDArray)], /, @@ -188,15 +178,6 @@ def cosh( return mx.nd.cosh(x) -def deg2rad( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def divide( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -222,15 +203,6 @@ def equal( raise IvyNotImplementedException() -def erf( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def exp( x: Union[(None, mx.ndarray.NDArray)], /, @@ -240,15 +212,6 @@ def exp( return mx.nd.exp(x) -def exp2( - x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def expm1( x: Union[(None, mx.ndarray.NDArray)], /, @@ -287,16 +250,6 @@ def fmin( raise IvyNotImplementedException() -def fmod( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def greater( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -317,15 +270,6 @@ def greater_equal( return mx.nd.greater_equal(x1, x2) -def imag( - val: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def isfinite( x: Union[(None, mx.ndarray.NDArray)], /, @@ -355,15 +299,6 @@ def isnan( raise IvyNotImplementedException() -def isreal( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def lcm( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -489,28 +424,6 @@ def logical_xor( return mx.nd.logical_xor(x1, x2) -def maximum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def minimum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def multiply( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -559,33 +472,6 @@ def pow( return mx.nd.power(x1, x2) -def rad2deg( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def real( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def reciprocal( - x: Union[(float, None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.reciprocal(x) - - def remainder( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -687,6 +573,18 @@ def subtract( return mx.nd.subtract(x1, x2) +def trapz( + y: Union[(None, mx.ndarray.NDArray)], + /, + *, + x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + dx: float = 1.0, + axis: int = (-1), + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def tan( x: Union[(None, mx.ndarray.NDArray)], /, @@ -706,19 +604,121 @@ def tanh( return mx.nd.tanh(x) -def trapz( - y: Union[(None, mx.ndarray.NDArray)], +def trunc( + x: Union[(None, mx.ndarray.NDArray)], /, *, - x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - dx: float = 1.0, - axis: int = (-1), out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def trunc( +def imag( + val: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def angle( + input: Union[(None, mx.ndarray.NDArray)], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def exp2( + x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def erf( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def maximum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def minimum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def reciprocal( + x: Union[(float, None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.reciprocal(x) + + +def deg2rad( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def rad2deg( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def isreal( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def fmod( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def real( x: Union[(None, mx.ndarray.NDArray)], /, *, diff --git a/ivy/functional/backends/mxnet/experimental/activations.py b/ivy/functional/backends/mxnet/experimental/activations.py index 74eac3b3ada9c..2ab8f7443e0af 100644 --- a/ivy/functional/backends/mxnet/experimental/activations.py +++ b/ivy/functional/backends/mxnet/experimental/activations.py @@ -14,7 +14,9 @@ def logit( raise IvyNotImplementedException() -def logsigmoid(input: None) -> None: +def thresholded_relu( + x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None +) -> None: raise IvyNotImplementedException() @@ -22,15 +24,13 @@ def relu6(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def selu(x: None, /, *, out: Optional[None] = None) -> None: +def logsigmoid(input: None) -> None: raise IvyNotImplementedException() -def silu(x: None, /, *, out: Optional[None] = None) -> None: +def selu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def thresholded_relu( - x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None -) -> None: +def silu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/creation.py b/ivy/functional/backends/mxnet/experimental/creation.py index 3942f38dae26c..633a0125504d0 100644 --- a/ivy/functional/backends/mxnet/experimental/creation.py +++ b/ivy/functional/backends/mxnet/experimental/creation.py @@ -5,44 +5,42 @@ from ivy.utils.exceptions import IvyNotImplementedException -def blackman_window( - size: int, - /, - *, +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def hann_window( - size: int, - /, - *, +def kaiser_bessel_derived_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def kaiser_bessel_derived_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def vorbis_window( + window_length: Union[(None, mx.ndarray.NDArray)], *, - dtype: Optional[None] = None, + dtype: None = np.float32, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: @@ -55,10 +53,12 @@ def tril_indices( raise IvyNotImplementedException() -def vorbis_window( - window_length: Union[(None, mx.ndarray.NDArray)], +def blackman_window( + size: int, + /, *, - dtype: None = np.float32, + periodic: bool = True, + dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/elementwise.py b/ivy/functional/backends/mxnet/experimental/elementwise.py index c6b93dab5e8a7..ce69fe9e91ed3 100644 --- a/ivy/functional/backends/mxnet/experimental/elementwise.py +++ b/ivy/functional/backends/mxnet/experimental/elementwise.py @@ -7,21 +7,41 @@ from .. import backend_version -def allclose( +@with_supported_dtypes( + {"1.9.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.np.log(mx.npx.gamma(x)) + + +def sinc( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def fmax( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> bool: +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def conj( - x: Union[(None, mx.ndarray.NDArray)], +def float_power( + x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -51,32 +71,20 @@ def count_nonzero( raise IvyNotImplementedException() -def diff( - x: Union[(None, mx.ndarray.NDArray, list, tuple)], - /, - *, - n: int = 1, - axis: int = (-1), - prepend: Optional[ - Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] - ] = None, - append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def fix( +def nansum( x: Union[(None, mx.ndarray.NDArray)], /, *, + axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, + dtype: Optional[None] = None, + keepdims: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def float_power( - x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def gcd( + x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, @@ -85,30 +93,21 @@ def float_power( raise IvyNotImplementedException() -def fmax( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def isclose( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def frexp( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[ - Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] - ] = None, -) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: - raise IvyNotImplementedException() - - -def gcd( - x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], - x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def signbit( + x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -116,17 +115,6 @@ def gcd( raise IvyNotImplementedException() -def gradient( - x: None, - /, - *, - spacing: Union[(int, list, tuple)] = 1, - axis: Optional[Union[(int, list, tuple)]] = None, - edge_order: int = 1, -) -> Union[(None, List[None])]: - raise IvyNotImplementedException() - - def hypot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -137,22 +125,21 @@ def hypot( raise IvyNotImplementedException() -def isclose( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def allclose( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> bool: raise IvyNotImplementedException() -def ldexp( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray, int)], +def fix( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -160,52 +147,55 @@ def ldexp( raise IvyNotImplementedException() -@with_supported_dtypes( - {"1.9.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( - x: Union[(None, mx.ndarray.NDArray)], +def nextafter( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.np.log(mx.npx.gamma(x)) + raise IvyNotImplementedException() -def modf( - x: Union[(None, mx.ndarray.NDArray)], +def diff( + x: Union[(None, mx.ndarray.NDArray, list, tuple)], /, *, + n: int = 1, + axis: int = (-1), + prepend: Optional[ + Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] + ] = None, + append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def nansum( +def zeta( x: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, mx.ndarray.NDArray)], /, *, - axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, - dtype: Optional[None] = None, - keepdims: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def nextafter( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def gradient( + x: None, /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: + spacing: Union[(int, list, tuple)] = 1, + axis: Optional[Union[(int, list, tuple)]] = None, + edge_order: int = 1, +) -> Union[(None, List[None])]: raise IvyNotImplementedException() -def signbit( - x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], +def xlogy( + x: Union[(None, mx.ndarray.NDArray)], + y: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -213,7 +203,7 @@ def signbit( raise IvyNotImplementedException() -def sinc( +def conj( x: Union[(None, mx.ndarray.NDArray)], /, *, @@ -222,9 +212,9 @@ def sinc( raise IvyNotImplementedException() -def xlogy( - x: Union[(None, mx.ndarray.NDArray)], - y: Union[(None, mx.ndarray.NDArray)], +def ldexp( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray, int)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -232,9 +222,19 @@ def xlogy( raise IvyNotImplementedException() -def zeta( +def frexp( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[ + Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] + ] = None, +) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: + raise IvyNotImplementedException() + + +def modf( x: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, diff --git a/ivy/functional/backends/mxnet/experimental/layers.py b/ivy/functional/backends/mxnet/experimental/layers.py index 4d9b6a781898e..84fab7a4ba311 100644 --- a/ivy/functional/backends/mxnet/experimental/layers.py +++ b/ivy/functional/backends/mxnet/experimental/layers.py @@ -7,6 +7,70 @@ from ivy.utils.exceptions import IvyNotImplementedException +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + raise IvyNotImplementedException() + + +def max_pool1d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: Union[str, int, Tuple[int]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool2d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: Union[str, int, Tuple[int], Tuple[int, int]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool3d( + x: mx.nd.NDArray, + kernel: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + strides: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + padding: Union[str, int, Tuple[int], Tuple[int, int, int]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + def avg_pool1d( x: mx.nd.NDArray, kernel: Union[int, Tuple[int]], @@ -67,69 +131,54 @@ def dct( raise IvyNotImplementedException() -def dropout1d( +def fft( x: mx.nd.NDArray, - prob: float, + dim: int, /, *, - training: bool = True, - data_format: str = "NWC", + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, out: Optional[mx.nd.NDArray] = None, ) -> mx.nd.NDArray: raise IvyNotImplementedException() -def dropout2d( +def dropout1d( x: mx.nd.NDArray, prob: float, /, *, training: bool = True, - data_format: str = "NHWC", + data_format: str = "NWC", out: Optional[mx.nd.NDArray] = None, ) -> mx.nd.NDArray: raise IvyNotImplementedException() -def dropout3d( +def dropout2d( x: mx.nd.NDArray, prob: float, /, *, training: bool = True, - data_format: str = "NDHWC", + data_format: str = "NHWC", out: Optional[mx.nd.NDArray] = None, ) -> mx.nd.NDArray: raise IvyNotImplementedException() -def fft( +def dropout3d( x: mx.nd.NDArray, - dim: int, + prob: float, /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + training: bool = True, + data_format: str = "NDHWC", out: Optional[mx.nd.NDArray] = None, ) -> mx.nd.NDArray: raise IvyNotImplementedException() -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - raise IvyNotImplementedException() - - def ifft( x: mx.nd.NDArray, dim: int, @@ -182,52 +231,3 @@ def interpolate( out: Optional[mx.nd.NDArray] = None, ): raise IvyNotImplementedException() - - -def max_pool1d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: Union[str, int, Tuple[int]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool2d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: Union[str, int, Tuple[int], Tuple[int, int]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool3d( - x: mx.nd.NDArray, - kernel: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - strides: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - padding: Union[str, int, Tuple[int], Tuple[int, int, int]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/linear_algebra.py b/ivy/functional/backends/mxnet/experimental/linear_algebra.py index a5da1d277cf8d..8212bc3c54f76 100644 --- a/ivy/functional/backends/mxnet/experimental/linear_algebra.py +++ b/ivy/functional/backends/mxnet/experimental/linear_algebra.py @@ -4,25 +4,24 @@ from ivy.utils.exceptions import IvyNotImplementedException -dot.support_native_out = True - - -def adjoint( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def cond( - x: Union[(None, mx.ndarray.NDArray)], +def eigh_tridiagonal( + alpha: Union[(None, mx.ndarray.NDArray)], + beta: Union[(None, mx.ndarray.NDArray)], /, *, - p: Optional[Union[(None, int, str)]] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] + ] = None, + tol: Optional[float] = None, +) -> Union[ + ( + None, + mx.ndarray.NDArray, + Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], + ) +]: raise IvyNotImplementedException() @@ -40,43 +39,31 @@ def diagflat( raise IvyNotImplementedException() -def dot( - a: mx.ndarray.NDArray, - b: mx.ndarray.NDArray, +def kron( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, - out: Optional[mx.ndarray.NDArray] = None, -) -> mx.ndarray.NDArray: - return mx.symbol.dot(a, b, out=out) + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() -def eig( +def matrix_exp( x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Tuple[None]: +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def eigh_tridiagonal( - alpha: Union[(None, mx.ndarray.NDArray)], - beta: Union[(None, mx.ndarray.NDArray)], +def eig( + x: Union[(None, mx.ndarray.NDArray)], /, *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] - ] = None, - tol: Optional[float] = None, -) -> Union[ - ( - None, - mx.ndarray.NDArray, - Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], - ) -]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Tuple[None]: raise IvyNotImplementedException() @@ -86,9 +73,8 @@ def eigvals( raise IvyNotImplementedException() -def kron( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def adjoint( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -96,19 +82,33 @@ def kron( raise IvyNotImplementedException() -def matrix_exp( - x: Union[(None, mx.ndarray.NDArray)], +def multi_dot( + x: Sequence[Union[(None, mx.ndarray.NDArray)]], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> None: raise IvyNotImplementedException() -def multi_dot( - x: Sequence[Union[(None, mx.ndarray.NDArray)]], +def cond( + x: Union[(None, mx.ndarray.NDArray)], /, *, + p: Optional[Union[(None, int, str)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def dot( + a: mx.ndarray.NDArray, + b: mx.ndarray.NDArray, + /, + *, + out: Optional[mx.ndarray.NDArray] = None, +) -> mx.ndarray.NDArray: + return mx.symbol.dot(a, b, out=out) + + +dot.support_native_out = True diff --git a/ivy/functional/backends/mxnet/experimental/manipulation.py b/ivy/functional/backends/mxnet/experimental/manipulation.py index aa8ac9077a4a0..66b46e813ae47 100644 --- a/ivy/functional/backends/mxnet/experimental/manipulation.py +++ b/ivy/functional/backends/mxnet/experimental/manipulation.py @@ -5,50 +5,39 @@ from ivy.utils.exceptions import IvyNotImplementedException -def atleast_1d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def atleast_2d( - *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def atleast_3d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: +def moveaxis( + a: Union[(None, mx.ndarray.NDArray)], + source: Union[(int, Sequence[int])], + destination: Union[(int, Sequence[int])], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def concat_from_sequence( - input_sequence: Union[(Tuple[None], List[None])], +def heaviside( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, - new_axis: int = 0, - axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def flipud( + m: Union[(None, mx.ndarray.NDArray)], /, *, copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dstack( +def vstack( arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, @@ -57,28 +46,41 @@ def dstack( raise IvyNotImplementedException() -def expand( - x: Union[(None, mx.ndarray.NDArray)], - shape: Union[(List[int], List[Tuple])], +def hstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def fliplr( +def rot90( m: Union[(None, mx.ndarray.NDArray)], /, *, copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + k: int = 1, + axes: Tuple[(int, int)] = (0, 1), + out: Union[(None, mx.ndarray.NDArray)] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def flipud( +def top_k( + x: None, + k: int, + /, + *, + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[(None, None)]] = None, +) -> Tuple[(None, None)]: + raise IvyNotImplementedException() + + +def fliplr( m: Union[(None, mx.ndarray.NDArray)], /, *, @@ -88,9 +90,8 @@ def flipud( raise IvyNotImplementedException() -def heaviside( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def i0( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -98,7 +99,7 @@ def heaviside( raise IvyNotImplementedException() -def hsplit( +def vsplit( ary: Union[(None, mx.ndarray.NDArray)], indices_or_sections: Union[(int, Tuple[(int, ...)])], /, @@ -108,45 +109,40 @@ def hsplit( raise IvyNotImplementedException() -def hstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def dsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: + copy: Optional[bool] = None, +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def i0( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +def atleast_1d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def moveaxis( - a: Union[(None, mx.ndarray.NDArray)], - source: Union[(int, Sequence[int])], - destination: Union[(int, Sequence[int])], +def dstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def rot90( - m: Union[(None, mx.ndarray.NDArray)], - /, - *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[(int, int)] = (0, 1), - out: Union[(None, mx.ndarray.NDArray)] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +def atleast_2d( + *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def atleast_3d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() @@ -162,33 +158,37 @@ def take_along_axis( raise IvyNotImplementedException() -def top_k( - x: None, - k: int, +def hsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[(None, None)]] = None, -) -> Tuple[(None, None)]: + copy: Optional[bool] = None, +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def vsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: + raise IvyNotImplementedException() + + +def expand( + x: Union[(None, mx.ndarray.NDArray)], + shape: Union[(List[int], List[Tuple])], /, *, copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def vstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def concat_from_sequence( + input_sequence: Union[(Tuple[None], List[None])], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/norms.py b/ivy/functional/backends/mxnet/experimental/norms.py index ee8cb935f70f0..19fea95289a8d 100644 --- a/ivy/functional/backends/mxnet/experimental/norms.py +++ b/ivy/functional/backends/mxnet/experimental/norms.py @@ -4,6 +4,16 @@ from ivy.utils.exceptions import IvyNotImplementedException +def l2_normalize( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[int] = None, + out: Optional[None] = None, +) -> None: + raise IvyNotImplementedException() + + def batch_norm( x: Union[(None, mx.ndarray.NDArray)], mean: Union[(None, mx.ndarray.NDArray)], @@ -48,16 +58,6 @@ def instance_norm( raise IvyNotImplementedException() -def l2_normalize( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: Optional[int] = None, - out: Optional[None] = None, -) -> None: - raise IvyNotImplementedException() - - def lp_normalize( x: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/experimental/random.py b/ivy/functional/backends/mxnet/experimental/random.py index 36c313891b167..8c3c4dfc0554d 100644 --- a/ivy/functional/backends/mxnet/experimental/random.py +++ b/ivy/functional/backends/mxnet/experimental/random.py @@ -5,15 +5,14 @@ from ivy.utils.exceptions import IvyNotImplementedException -def bernoulli( - probs: Union[(float, None, mx.ndarray.NDArray)], +def dirichlet( + alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], + /, *, - logits: Union[(float, None, mx.ndarray.NDArray)] = None, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - device: str, - dtype: None, - seed: Optional[int] = None, + size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + seed: Optional[int] = None, + dtype: Optional[None] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -32,18 +31,6 @@ def beta( raise IvyNotImplementedException() -def dirichlet( - alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], - /, - *, - size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - seed: Optional[int] = None, - dtype: Optional[None] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def gamma( alpha: Union[(float, None, mx.ndarray.NDArray)], beta: Union[(float, None, mx.ndarray.NDArray)], @@ -68,3 +55,16 @@ def poisson( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def bernoulli( + probs: Union[(float, None, mx.ndarray.NDArray)], + *, + logits: Union[(float, None, mx.ndarray.NDArray)] = None, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + device: str, + dtype: None, + seed: Optional[int] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/statistical.py b/ivy/functional/backends/mxnet/experimental/statistical.py index 985ab019ce0e1..252e27dd08e8c 100644 --- a/ivy/functional/backends/mxnet/experimental/statistical.py +++ b/ivy/functional/backends/mxnet/experimental/statistical.py @@ -4,43 +4,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def bincount( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - minlength: int = 0, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def corrcoef( - x: None, - /, - *, - y: None, - rowvar: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: - raise IvyNotImplementedException() - - -def cov( - x1: None, - x2: None = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[None] = None, - aweights: Optional[None] = None, - dtype: Optional[type] = None, -) -> None: - raise IvyNotImplementedException() - - def histogram( a: None, /, @@ -81,6 +44,30 @@ def nanmean( raise IvyNotImplementedException() +def quantile( + a: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, float)], + /, + *, + axis: Optional[Union[(int, Sequence[int])]] = None, + interpolation: str = "linear", + keepdims: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def corrcoef( + x: None, + /, + *, + y: None, + rowvar: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> None: + raise IvyNotImplementedException() + + def nanmedian( input: Union[(None, mx.ndarray.NDArray)], /, @@ -92,14 +79,27 @@ def nanmedian( raise IvyNotImplementedException() -def quantile( - a: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, float)], +def bincount( + x: Union[(None, mx.ndarray.NDArray)], /, *, - axis: Optional[Union[(int, Sequence[int])]] = None, - interpolation: str = "linear", - keepdims: bool = False, + weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + minlength: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def cov( + x1: None, + x2: None = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[None] = None, + aweights: Optional[None] = None, + dtype: Optional[type] = None, +) -> None: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/general.py b/ivy/functional/backends/mxnet/general.py index 34eebe8331057..c469b517c47c1 100644 --- a/ivy/functional/backends/mxnet/general.py +++ b/ivy/functional/backends/mxnet/general.py @@ -5,26 +5,10 @@ from ivy.utils.exceptions import IvyNotImplementedException -def container_types(): - return [] - - def current_backend_str() -> str: return "mxnet" -def gather( - x: mx.ndarray.NDArray, - indices: mx.ndarray.NDArray, - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[mx.ndarray.NDArray] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def is_native_array( x: Union[(None, mx.ndarray.NDArray)], /, @@ -37,10 +21,6 @@ def is_native_array( return isinstance(x, mx.ndarray.NDArray) or isinstance(x, np.ndarray) -def itemsize(x: mx.ndarray.NDArray, /) -> int: - return x.asnumpy().itemsize - - def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: if copy: if x.shape == (): @@ -49,3 +29,23 @@ def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: return x.copy().asnumpy() else: return x.asnumpy() + + +def itemsize(x: mx.ndarray.NDArray, /) -> int: + return x.asnumpy().itemsize + + +def container_types(): + return [] + + +def gather( + x: mx.ndarray.NDArray, + indices: mx.ndarray.NDArray, + /, + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[mx.ndarray.NDArray] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/gradients.py b/ivy/functional/backends/mxnet/gradients.py index 28189a4412a33..dd3e9041601be 100644 --- a/ivy/functional/backends/mxnet/gradients.py +++ b/ivy/functional/backends/mxnet/gradients.py @@ -8,6 +8,18 @@ from ivy.utils.exceptions import IvyNotImplementedException +def variable(x, /): + return x + + +def is_variable(x, /, *, exclusive=False): + return isinstance(x, mx.ndarray.NDArray) + + +def variable_data(x, /): + raise IvyNotImplementedException() + + def execute_with_gradients( func, xs, @@ -20,29 +32,17 @@ def execute_with_gradients( raise IvyNotImplementedException() -def grad(func, argnums=0): +def value_and_grad(func): raise IvyNotImplementedException() -def is_variable(x, /, *, exclusive=False): - return isinstance(x, mx.ndarray.NDArray) - - def jac(func): raise IvyNotImplementedException() -def stop_gradient(x, /, *, preserve_type=True, out=None): - raise IvyNotImplementedException() - - -def value_and_grad(func): +def grad(func, argnums=0): raise IvyNotImplementedException() -def variable(x, /): - return x - - -def variable_data(x, /): +def stop_gradient(x, /, *, preserve_type=True, out=None): raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/layers.py b/ivy/functional/backends/mxnet/layers.py index 3937442087024..1ae0560d0c356 100644 --- a/ivy/functional/backends/mxnet/layers.py +++ b/ivy/functional/backends/mxnet/layers.py @@ -66,6 +66,20 @@ def conv2d_transpose( raise IvyNotImplementedException() +def depthwise_conv2d( + x: Union[(None, mx.ndarray.NDArray)], + filters: Union[(None, mx.ndarray.NDArray)], + strides: Union[(int, Tuple[(int, int)])], + padding: Union[(str, int, Sequence[Tuple[(int, int)]])], + /, + *, + data_format: str = "NHWC", + dilations: Union[(int, Tuple[(int, int)])] = 1, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def conv3d( x: Union[(None, mx.ndarray.NDArray)], filters: Union[(None, mx.ndarray.NDArray)], @@ -131,17 +145,3 @@ def conv_general_transpose( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def depthwise_conv2d( - x: Union[(None, mx.ndarray.NDArray)], - filters: Union[(None, mx.ndarray.NDArray)], - strides: Union[(int, Tuple[(int, int)])], - padding: Union[(str, int, Sequence[Tuple[(int, int)]])], - /, - *, - data_format: str = "NHWC", - dilations: Union[(int, Tuple[(int, int)])] = 1, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/linear_algebra.py b/ivy/functional/backends/mxnet/linear_algebra.py index 52682499037c2..5077fd3a24bc2 100644 --- a/ivy/functional/backends/mxnet/linear_algebra.py +++ b/ivy/functional/backends/mxnet/linear_algebra.py @@ -38,16 +38,6 @@ def det( return mx.nd.linalg.det(x) -def diag( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - k: int = 0, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def diagonal( x: Union[(None, mx.ndarray.NDArray)], /, @@ -259,17 +249,6 @@ def trace( raise IvyNotImplementedException() -def vander( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def vecdot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -294,6 +273,27 @@ def vector_norm( raise IvyNotImplementedException() +def diag( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + k: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def vander( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def vector_to_skew_symmetric_matrix( vector: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/manipulation.py b/ivy/functional/backends/mxnet/manipulation.py index 49dd4aa56f818..feb262db2a152 100644 --- a/ivy/functional/backends/mxnet/manipulation.py +++ b/ivy/functional/backends/mxnet/manipulation.py @@ -6,17 +6,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def clip( - x: Union[(None, mx.ndarray.NDArray)], - x_min: Union[(Number, None, mx.ndarray.NDArray)], - x_max: Union[(Number, None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.clip(x, x_min, x_max) - - def concat( xs: Union[(Tuple[(None, ...)], List[None])], /, @@ -27,12 +16,6 @@ def concat( raise IvyNotImplementedException() -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): - raise IvyNotImplementedException() - - def expand_dims( x: Union[(None, mx.ndarray.NDArray)], /, @@ -66,17 +49,6 @@ def permute_dims( raise IvyNotImplementedException() -def repeat( - x: Union[(None, mx.ndarray.NDArray)], - /, - repeats: Union[(int, List[int])], - *, - axis: int = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def reshape( x: Union[(None, mx.ndarray.NDArray)], /, @@ -101,6 +73,27 @@ def roll( raise IvyNotImplementedException() +def squeeze( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.squeeze(x, axis=axis) + + +def stack( + arrays: Union[(Tuple[None], List[None])], + /, + *, + axis: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def split( x: Union[(None, mx.ndarray.NDArray)], /, @@ -113,24 +106,36 @@ def split( raise IvyNotImplementedException() -def squeeze( +def repeat( x: Union[(None, mx.ndarray.NDArray)], /, + repeats: Union[(int, List[int])], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, + axis: int = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.squeeze(x, axis=axis) + raise IvyNotImplementedException() -def stack( - arrays: Union[(Tuple[None], List[None])], +def tile( + x: Union[(None, mx.ndarray.NDArray)], /, + repeats: Sequence[int], *, - axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.tile(x, repeats) + + +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): + raise IvyNotImplementedException() + + +def zero_pad( + x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): raise IvyNotImplementedException() @@ -146,14 +151,15 @@ def swapaxes( raise IvyNotImplementedException() -def tile( +def clip( x: Union[(None, mx.ndarray.NDArray)], + x_min: Union[(Number, None, mx.ndarray.NDArray)], + x_max: Union[(Number, None, mx.ndarray.NDArray)], /, - repeats: Sequence[int], *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.tile(x, repeats) + return mx.nd.clip(x, x_min, x_max) def unstack( @@ -165,9 +171,3 @@ def unstack( keepdims: bool = False, ) -> List[None]: raise IvyNotImplementedException() - - -def zero_pad( - x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/random.py b/ivy/functional/backends/mxnet/random.py index 97a4c3959de99..4f1e25f4763d5 100644 --- a/ivy/functional/backends/mxnet/random.py +++ b/ivy/functional/backends/mxnet/random.py @@ -11,14 +11,12 @@ from ivy.utils.exceptions import IvyNotImplementedException -def multinomial( - population_size: int, - num_samples: int, - /, +def random_uniform( *, - batch_size: int = 1, - probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - replace: bool = True, + low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, + shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, + dtype: None, device: str, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -26,40 +24,42 @@ def multinomial( raise IvyNotImplementedException() -def randint( - low: Union[(float, None, mx.ndarray.NDArray)], - high: Union[(float, None, mx.ndarray.NDArray)], - /, +def random_normal( *, + mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - device: str, - dtype: Optional[Union[(None, ivy.Dtype)]] = None, + dtype: None, seed: Optional[int] = None, + device: str, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def random_normal( +def multinomial( + population_size: int, + num_samples: int, + /, *, - mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - dtype: None, - seed: Optional[int] = None, + batch_size: int = 1, + probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + replace: bool = True, device: str, + seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def random_uniform( +def randint( + low: Union[(float, None, mx.ndarray.NDArray)], + high: Union[(float, None, mx.ndarray.NDArray)], + /, *, - low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, - shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, - dtype: None, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, device: str, + dtype: Optional[Union[(None, ivy.Dtype)]] = None, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: diff --git a/ivy/functional/backends/mxnet/searching.py b/ivy/functional/backends/mxnet/searching.py index 8e758dae93f55..ed20dee3fccc9 100644 --- a/ivy/functional/backends/mxnet/searching.py +++ b/ivy/functional/backends/mxnet/searching.py @@ -33,15 +33,6 @@ def argmin( raise IvyNotImplementedException() -def argwhere( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def nonzero( x: Union[(None, mx.ndarray.NDArray)], /, @@ -62,3 +53,12 @@ def where( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def argwhere( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/sorting.py b/ivy/functional/backends/mxnet/sorting.py index 85954c0a9e91b..356c4b9414403 100644 --- a/ivy/functional/backends/mxnet/sorting.py +++ b/ivy/functional/backends/mxnet/sorting.py @@ -17,35 +17,35 @@ def argsort( raise IvyNotImplementedException() -def msort( - a: Union[(None, mx.ndarray.NDArray, list, tuple)], +def sort( + x: Union[(None, mx.ndarray.NDArray)], /, *, + axis: int = (-1), + descending: bool = False, + stable: bool = True, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def searchsorted( - x: Union[(None, mx.ndarray.NDArray)], - v: Union[(None, mx.ndarray.NDArray)], +def msort( + a: Union[(None, mx.ndarray.NDArray, list, tuple)], /, *, - side: Literal[("left", "right")] = "left", - sorter: Optional[Union[(ivy.Array, ivy.NativeArray, List[int])]] = None, - ret_dtype: None = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def sort( +def searchsorted( x: Union[(None, mx.ndarray.NDArray)], + v: Union[(None, mx.ndarray.NDArray)], /, *, - axis: int = (-1), - descending: bool = False, - stable: bool = True, + side: Literal[("left", "right")] = "left", + sorter: Optional[Union[(ivy.Array, ivy.NativeArray, List[int])]] = None, + ret_dtype: None = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/statistical.py b/ivy/functional/backends/mxnet/statistical.py index c688052b9fc2e..67a24d9922d61 100644 --- a/ivy/functional/backends/mxnet/statistical.py +++ b/ivy/functional/backends/mxnet/statistical.py @@ -7,40 +7,7 @@ import ivy -def cumprod( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def cumsum( - x: Union[(None, mx.ndarray.NDArray)], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def einsum( - equation: str, - *operands: Union[(None, mx.ndarray.NDArray)], - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def max( +def min( x: Union[(None, mx.ndarray.NDArray)], /, *, @@ -51,7 +18,7 @@ def max( raise IvyNotImplementedException() -def mean( +def max( x: Union[(None, mx.ndarray.NDArray)], /, *, @@ -62,7 +29,7 @@ def mean( raise IvyNotImplementedException() -def min( +def mean( x: Union[(None, mx.ndarray.NDArray)], /, *, @@ -143,3 +110,36 @@ def var( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def cumprod( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def cumsum( + x: Union[(None, mx.ndarray.NDArray)], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def einsum( + equation: str, + *operands: Union[(None, mx.ndarray.NDArray)], + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/numpy/activations.py b/ivy/functional/backends/numpy/activations.py index f7ef8d66048b9..2e230498b9574 100644 --- a/ivy/functional/backends/numpy/activations.py +++ b/ivy/functional/backends/numpy/activations.py @@ -9,34 +9,14 @@ from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array -relu.support_native_out = True -softmax.support_native_out = True -softplus.support_native_out = True -log_softmax.support_native_out = True -mish.support_native_out = True -hardswish.support_native_out = True - - @_scalar_output_to_0d_array -def gelu( - x: np.ndarray, - /, - *, - approximate: bool = False, - complex_mode="jax", - out: Optional[np.ndarray] = None, +def relu( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None ) -> np.ndarray: - if approximate: - ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) - else: - ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) - return ivy.astype(ret, x.dtype, copy=False) + return np.maximum(x, 0, out=out, dtype=x.dtype) -@_scalar_output_to_0d_array -def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) - return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) +relu.support_native_out = True def leaky_relu( @@ -51,36 +31,19 @@ def leaky_relu( @_scalar_output_to_0d_array -def log_softmax( - x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None -) -> np.ndarray: - if axis is None: - axis = -1 - x_max = np.max(x, axis=axis, keepdims=True) - if x_max.ndim > 0: - x_max[~np.isfinite(x_max)] = 0 - elif not np.isfinite(x_max): - x_max = 0 - exp_tmp = np.exp(x - x_max) - - with np.errstate(divide="ignore"): - s = np.sum(exp_tmp, axis=axis, keepdims=True) - ret = np.log(s) - - ret = x - x_max - ret - return ret - - -@_scalar_output_to_0d_array -def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return x * np.tanh(np.log1p(np.exp(x))) - - -@_scalar_output_to_0d_array -def relu( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +def gelu( + x: np.ndarray, + /, + *, + approximate: bool = False, + complex_mode="jax", + out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.maximum(x, 0, out=out, dtype=x.dtype) + if approximate: + ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) + else: + ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) + return ivy.astype(ret, x.dtype, copy=False) def sigmoid(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -97,6 +60,9 @@ def softmax( return np.divide(exp_x, np.sum(exp_x, axis=axis, keepdims=True), out=out) +softmax.support_native_out = True + + @_scalar_output_to_0d_array def softplus( x: np.ndarray, @@ -126,3 +92,47 @@ def softplus( if threshold is not None: return np.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) + + +softplus.support_native_out = True + + +@_scalar_output_to_0d_array +def log_softmax( + x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None +) -> np.ndarray: + if axis is None: + axis = -1 + x_max = np.max(x, axis=axis, keepdims=True) + if x_max.ndim > 0: + x_max[~np.isfinite(x_max)] = 0 + elif not np.isfinite(x_max): + x_max = 0 + exp_tmp = np.exp(x - x_max) + + with np.errstate(divide="ignore"): + s = np.sum(exp_tmp, axis=axis, keepdims=True) + ret = np.log(s) + + ret = x - x_max - ret + return ret + + +log_softmax.support_native_out = True + + +@_scalar_output_to_0d_array +def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return x * np.tanh(np.log1p(np.exp(x))) + + +mish.support_native_out = True + + +@_scalar_output_to_0d_array +def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) + return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) + + +hardswish.support_native_out = True diff --git a/ivy/functional/backends/numpy/creation.py b/ivy/functional/backends/numpy/creation.py index 0eb26a54646d9..7600f3cf9dba0 100644 --- a/ivy/functional/backends/numpy/creation.py +++ b/ivy/functional/backends/numpy/creation.py @@ -64,17 +64,6 @@ def asarray( return np.copy(ret) if copy else ret -def copy_array( - x: np.ndarray, - *, - to_ivy_array: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() - - def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -118,17 +107,6 @@ def from_dlpack(x, /, *, out: Optional[np.ndarray] = None): return np.from_dlpack(x) -def frombuffer( - buffer: bytes, - dtype: Optional[np.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> np.ndarray: - if isinstance(dtype, list): - dtype = np.dtype(dtype[0]) - return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) - - def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -192,6 +170,68 @@ def meshgrid( return np.meshgrid(*arrays, sparse=sparse, indexing=indexing) +def ones( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return _to_device(np.ones(shape, dtype), device=device) + + +def ones_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.ones_like(x, dtype=dtype), device=device) + + +def tril( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tril(x, k) + + +def triu( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.triu(x, k) + + +def zeros( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return _to_device(np.zeros(shape, dtype), device=device) + + +def zeros_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.zeros_like(x, dtype=dtype), device=device) + + +# Extra # +# ------# + + +array = asarray + + +def copy_array( + x: np.ndarray, + *, + to_ivy_array: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() + + def one_hot( indices: np.ndarray, depth: int, @@ -228,32 +268,15 @@ def one_hot( return res -def ones( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.ones(shape, dtype), device=device) - - -def ones_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.ones_like(x, dtype=dtype), device=device) - - -def tril( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tril(x, k) - - -def triu( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +def frombuffer( + buffer: bytes, + dtype: Optional[np.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, ) -> np.ndarray: - return np.triu(x, k) + if isinstance(dtype, list): + dtype = np.dtype(dtype[0]) + return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) def triu_indices( @@ -267,26 +290,3 @@ def triu_indices( return tuple( _to_device(np.asarray(np.triu_indices(n=n_rows, k=k, m=n_cols)), device=device) ) - - -def zeros( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.zeros(shape, dtype), device=device) - - -def zeros_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.zeros_like(x, dtype=dtype), device=device) - - -# Extra # -# ------# - - -array = asarray diff --git a/ivy/functional/backends/numpy/data_type.py b/ivy/functional/backends/numpy/data_type.py index 172989ff82968..317cc71c8608d 100644 --- a/ivy/functional/backends/numpy/data_type.py +++ b/ivy/functional/backends/numpy/data_type.py @@ -9,26 +9,6 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from . import backend_version -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int16"): "int16", @@ -60,6 +40,7 @@ np.complex128: "complex128", np.bool_: "bool", } + native_dtype_dict = { "int8": np.dtype("int8"), "int16": np.dtype("int16"), @@ -77,6 +58,27 @@ "bool": np.dtype("bool"), } +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} + class Finfo: def __init__(self, np_finfo: np.finfo): @@ -106,6 +108,68 @@ def smallest_normal(self): return float(self._np_finfo.tiny) +# Array API Standard # +# -------------------# + + +def astype( + x: np.ndarray, + dtype: np.dtype, + /, + *, + copy: bool = True, + out: Optional[ivy.Array] = None, +) -> np.ndarray: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return np.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: + try: + return np.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def broadcast_to( + x: np.ndarray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return np.broadcast_to(x.reshape([-1]), shape) + return np.broadcast_to(x, shape) + + +@_handle_nestable_dtype_info +def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype + return Finfo(np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype + return np.iinfo(ivy.as_native_dtype(type)) + + +def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return np.result_type(arrays_and_dtypes) + result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = np.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) + + # Extra # # ------# @@ -174,45 +238,6 @@ def as_native_dtype(dtype_in: Union[np.dtype, str, bool, int, float], /) -> np.d ) -# Array API Standard # -# -------------------# - - -def astype( - x: np.ndarray, - dtype: np.dtype, - /, - *, - copy: bool = True, - out: Optional[ivy.Array] = None, -) -> np.ndarray: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return np.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: - try: - return np.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def broadcast_to( - x: np.ndarray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return np.broadcast_to(x.reshape([-1]), shape) - return np.broadcast_to(x, shape) - - def dtype(x: np.ndarray, *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.to_native(x).dtype @@ -232,20 +257,6 @@ def dtype_bits(dtype_in: Union[np.dtype, str], /) -> int: ) -@_handle_nestable_dtype_info -def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype - return Finfo(np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype - return np.iinfo(ivy.as_native_dtype(type)) - - def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -253,12 +264,3 @@ def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: return True else: return False - - -def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return np.result_type(arrays_and_dtypes) - result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = np.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) diff --git a/ivy/functional/backends/numpy/device.py b/ivy/functional/backends/numpy/device.py index 4576b42831bc5..b10c8119f8cc1 100644 --- a/ivy/functional/backends/numpy/device.py +++ b/ivy/functional/backends/numpy/device.py @@ -11,30 +11,34 @@ from ivy.functional.ivy.device import Profiler as BaseProfiler -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper numpy profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None +def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: + if as_native: + return "cpu" + return as_ivy_dev("cpu") - def start(self): - self._start_time = time.perf_counter() - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) +def as_ivy_dev(device: str, /): + return ivy.Device("cpu") - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def as_native_dev(device: str, /): + return "cpu" -# --- Helpers --- # -# --------------- # +def clear_cached_mem_on_dev(device: str, /): + return None + + +def tpu_is_available() -> bool: + return False + + +def num_gpus() -> int: + return 0 + + +def gpu_is_available() -> bool: + return False # private version of to_device to be used in backend implementations @@ -56,40 +60,6 @@ def _to_device(x: np.ndarray, device=None) -> np.ndarray: return x -# --- Main --- # -# ------------ # - - -def as_ivy_dev(device: str, /): - return ivy.Device("cpu") - - -def as_native_dev(device: str, /): - return "cpu" - - -def clear_cached_mem_on_dev(device: str, /): - return None - - -def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: - if as_native: - return "cpu" - return as_ivy_dev("cpu") - - -def gpu_is_available() -> bool: - return False - - -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return fn(*args, **kwargs) - - -def num_gpus() -> int: - return 0 - - def to_device( x: np.ndarray, device: str, @@ -115,5 +85,27 @@ def to_device( return x -def tpu_is_available() -> bool: - return False +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return fn(*args, **kwargs) + + +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper numpy profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + + def start(self): + self._start_time = time.perf_counter() + + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/ivy/functional/backends/numpy/elementwise.py b/ivy/functional/backends/numpy/elementwise.py index 5d643eb77e7a3..6c4459bbaaf4f 100644 --- a/ivy/functional/backends/numpy/elementwise.py +++ b/ivy/functional/backends/numpy/elementwise.py @@ -10,18 +10,6 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - -def _abs_variant_sign(x): - return np.divide(x, np.abs(x), where=x != 0) - - -# --- Main --- # -# ------------ # - - @_scalar_output_to_0d_array def abs( x: Union[float, np.ndarray], @@ -32,16 +20,25 @@ def abs( return np.absolute(x, out=out) +abs.support_native_out = True + + @_scalar_output_to_0d_array def acos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccos(x, out=out) +acos.support_native_out = True + + @_scalar_output_to_0d_array def acosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccosh(x, out=out) +acosh.support_native_out = True + + @_scalar_output_to_0d_array def add( x1: Union[float, np.ndarray], @@ -58,14 +55,7 @@ def add( return np.add(x1, x2, out=out) -def angle( - z: np.ndarray, - /, - *, - deg: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.angle(z, deg=deg) +add.support_native_out = True @_scalar_output_to_0d_array @@ -73,16 +63,25 @@ def asin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsin(x, out=out) +asin.support_native_out = True + + @_scalar_output_to_0d_array def asinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsinh(x, out=out) +asinh.support_native_out = True + + @_scalar_output_to_0d_array def atan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctan(x, out=out) +atan.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def atan2( @@ -92,11 +91,17 @@ def atan2( return np.arctan2(x1, x2, out=out) +atan2.support_native_out = True + + @_scalar_output_to_0d_array def atanh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctanh(x, out=out) +atanh.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_and( @@ -110,6 +115,9 @@ def bitwise_and( return np.bitwise_and(x1, x2, out=out) +bitwise_and.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_invert( @@ -118,6 +126,9 @@ def bitwise_invert( return np.invert(x, out=out) +bitwise_invert.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_left_shift( @@ -131,6 +142,9 @@ def bitwise_left_shift( return np.left_shift(x1, x2, out=out) +bitwise_left_shift.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_or( @@ -144,6 +158,9 @@ def bitwise_or( return np.bitwise_or(x1, x2, out=out) +bitwise_or.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_right_shift( @@ -157,6 +174,9 @@ def bitwise_right_shift( return np.right_shift(x1, x2, out=out) +bitwise_right_shift.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_xor( @@ -170,6 +190,9 @@ def bitwise_xor( return np.bitwise_xor(x1, x2, out=out) +bitwise_xor.support_native_out = True + + @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) @_scalar_output_to_0d_array def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -182,21 +205,24 @@ def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret +ceil.support_native_out = True + + @_scalar_output_to_0d_array def cos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cos(x, out=out) +cos.support_native_out = True + + @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) @_scalar_output_to_0d_array def cosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cosh(x, out=out) -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.deg2rad(x, out=out) +cosh.support_native_out = True @_scalar_output_to_0d_array @@ -216,6 +242,9 @@ def divide( return ret +divide.support_native_out = True + + @_scalar_output_to_0d_array def equal( x1: Union[float, np.ndarray], @@ -228,31 +257,7 @@ def equal( return np.equal(x1, x2, out=out) -# Extra # -# ------# - - -@_scalar_output_to_0d_array -def erf(x, /, *, out: Optional[np.ndarray] = None): - a1 = 0.254829592 - a2 = -0.284496736 - a3 = 1.421413741 - a4 = -1.453152027 - a5 = 1.061405429 - p = 0.3275911 - - sign = np.sign(x) - x = np.abs(x) - - # A&S formula 7.1.26 - t = 1.0 / (1.0 + p * x) - y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) - ret = sign * y - if hasattr(x, "dtype"): - ret = np.asarray(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret +equal.support_native_out = True @_scalar_output_to_0d_array @@ -260,6 +265,9 @@ def exp(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.exp(x, out=out) +exp.support_native_out = True + + def exp2( x: Union[np.ndarray, float, list, tuple], /, @@ -269,11 +277,17 @@ def exp2( return np.exp2(x, out=out) +exp2.support_native_out = True + + @_scalar_output_to_0d_array def expm1(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.expm1(x, out=out) +expm1.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -286,6 +300,9 @@ def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret +floor.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor_divide( @@ -320,32 +337,7 @@ def fmin( ) -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def fmod( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmod( - x1, - x2, - out=None, - ) - - -def gcd( - x1: Union[np.ndarray, int, list, tuple], - x2: Union[np.ndarray, float, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.gcd(x1, x2, out=out) +fmin.support_native_out = True @_scalar_output_to_0d_array @@ -360,6 +352,9 @@ def greater( return np.greater(x1, x2, out=out) +greater.support_native_out = True + + @_scalar_output_to_0d_array def greater_equal( x1: Union[float, np.ndarray], @@ -372,13 +367,7 @@ def greater_equal( return np.greater_equal(x1, x2, out=out) -def imag( - val: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.imag(val) +greater_equal.support_native_out = True @_scalar_output_to_0d_array @@ -386,6 +375,9 @@ def isfinite(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarra return np.isfinite(x, out=out) +isfinite.support_native_out = True + + @_scalar_output_to_0d_array def isinf( x: np.ndarray, @@ -409,9 +401,7 @@ def isnan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.isnan(x, out=out) -@_scalar_output_to_0d_array -def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.isreal(x) +isnan.support_native_out = True @_scalar_output_to_0d_array @@ -430,6 +420,9 @@ def lcm( ) +lcm.support_native_out = True + + @_scalar_output_to_0d_array def less( x1: Union[float, np.ndarray], @@ -442,6 +435,9 @@ def less( return np.less(x1, x2, out=out) +less.support_native_out = True + + @_scalar_output_to_0d_array def less_equal( x1: Union[float, np.ndarray], @@ -454,26 +450,41 @@ def less_equal( return np.less_equal(x1, x2, out=out) +less_equal.support_native_out = True + + @_scalar_output_to_0d_array def log(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log(x, out=out) +log.support_native_out = True + + @_scalar_output_to_0d_array def log10(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log10(x, out=out) +log10.support_native_out = True + + @_scalar_output_to_0d_array def log1p(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log1p(x, out=out) +log1p.support_native_out = True + + @_scalar_output_to_0d_array def log2(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log2(x, out=out) +log2.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def logaddexp( @@ -483,6 +494,9 @@ def logaddexp( return np.logaddexp(x1, x2, out=out) +logaddexp.support_native_out = True + + def logaddexp2( x1: Union[np.ndarray, int, list, tuple], x2: Union[np.ndarray, int, list, tuple], @@ -497,6 +511,9 @@ def logaddexp2( return np.logaddexp2(x1, x2, out=out) +logaddexp2.support_native_out = True + + @_scalar_output_to_0d_array def logical_and( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -504,11 +521,17 @@ def logical_and( return np.logical_and(x1, x2, out=out) +logical_and.support_native_out = True + + @_scalar_output_to_0d_array def logical_not(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.logical_not(x, out=out) +logical_not.support_native_out = True + + @_scalar_output_to_0d_array def logical_or( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -516,6 +539,9 @@ def logical_or( return np.logical_or(x1, x2, out=out) +logical_or.support_native_out = True + + @_scalar_output_to_0d_array def logical_xor( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -523,40 +549,7 @@ def logical_xor( return np.logical_xor(x1, x2, out=out) -@_scalar_output_to_0d_array -def maximum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -): - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 >= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.maximum(x1, x2, out=out) - - -@_scalar_output_to_0d_array -def minimum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 <= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.minimum(x1, x2, out=out) +logical_xor.support_native_out = True @_scalar_output_to_0d_array @@ -571,17 +564,7 @@ def multiply( return np.multiply(x1, x2, out=out) -def nan_to_num( - x: np.ndarray, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) +multiply.support_native_out = True @_scalar_output_to_0d_array @@ -591,6 +574,9 @@ def negative( return np.negative(x, out=out) +negative.support_native_out = True + + @_scalar_output_to_0d_array def not_equal( x1: Union[float, np.ndarray], @@ -603,6 +589,9 @@ def not_equal( return np.not_equal(x1, x2, out=out) +not_equal.support_native_out = True + + @_scalar_output_to_0d_array def positive( x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None @@ -610,6 +599,9 @@ def positive( return np.positive(x, out=out) +positive.support_native_out = True + + @_scalar_output_to_0d_array def pow( x1: Union[float, np.ndarray], @@ -622,22 +614,7 @@ def pow( return np.power(x1, x2, out=out) -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.rad2deg(x, out=out) - - -def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.real(x) - - -@_scalar_output_to_0d_array -def reciprocal( - x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - numerator = np.ones_like(x) - return np.true_divide(numerator, x, out=out) +pow.support_native_out = True @_scalar_output_to_0d_array @@ -660,6 +637,9 @@ def remainder( return np.remainder(x1, x2, out=out) +remainder.support_native_out = True + + @_scalar_output_to_0d_array def round( x: np.ndarray, /, *, decimals: int = 0, out: Optional[np.ndarray] = None @@ -673,6 +653,13 @@ def round( return ret +round.support_native_out = True + + +def _abs_variant_sign(x): + return np.divide(x, np.abs(x), where=x != 0) + + @_scalar_output_to_0d_array def sign( x: np.ndarray, @@ -686,26 +673,41 @@ def sign( return np.sign(x, out=out) +sign.support_native_out = True + + @_scalar_output_to_0d_array def sin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sin(x, out=out) +sin.support_native_out = True + + @_scalar_output_to_0d_array def sinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sinh(x, out=out) +sinh.support_native_out = True + + @_scalar_output_to_0d_array def sqrt(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sqrt(x, out=out) +sqrt.support_native_out = True + + @_scalar_output_to_0d_array def square(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.square(x, out=out) +square.support_native_out = True + + @_scalar_output_to_0d_array def subtract( x1: Union[float, np.ndarray], @@ -723,16 +725,7 @@ def subtract( return np.subtract(x1, x2, out=out) -@_scalar_output_to_0d_array -def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.tan(x, out=out) - - -@_scalar_output_to_0d_array -def tanh( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tanh(x, out=out) +subtract.support_native_out = True @_scalar_output_to_0d_array @@ -748,6 +741,27 @@ def trapz( return np.trapz(y, x=x, dx=dx, axis=axis) +trapz.support_native_out = False + + +@_scalar_output_to_0d_array +def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.tan(x, out=out) + + +tan.support_native_out = True + + +@_scalar_output_to_0d_array +def tanh( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tanh(x, out=out) + + +tanh.support_native_out = True + + @_scalar_output_to_0d_array def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: if "int" in str(x.dtype): @@ -759,74 +773,192 @@ def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -abs.support_native_out = True -acos.support_native_out = True -acosh.support_native_out = True -add.support_native_out = True -asin.support_native_out = True -asinh.support_native_out = True -atan.support_native_out = True -atan2.support_native_out = True -atanh.support_native_out = True -bitwise_and.support_native_out = True -bitwise_invert.support_native_out = True -bitwise_left_shift.support_native_out = True -bitwise_or.support_native_out = True -bitwise_right_shift.support_native_out = True -bitwise_xor.support_native_out = True -ceil.support_native_out = True -cos.support_native_out = True -cosh.support_native_out = True -divide.support_native_out = True -equal.support_native_out = True -exp.support_native_out = True -exp2.support_native_out = True -expm1.support_native_out = True -floor.support_native_out = True -fmin.support_native_out = True -greater.support_native_out = True -greater_equal.support_native_out = True -isfinite.support_native_out = True -isnan.support_native_out = True -lcm.support_native_out = True -less.support_native_out = True -less_equal.support_native_out = True -log.support_native_out = True -log10.support_native_out = True -log1p.support_native_out = True -log2.support_native_out = True -logaddexp.support_native_out = True -logaddexp2.support_native_out = True -logical_and.support_native_out = True -logical_not.support_native_out = True -logical_or.support_native_out = True -logical_xor.support_native_out = True -multiply.support_native_out = True -negative.support_native_out = True -not_equal.support_native_out = True -positive.support_native_out = True -pow.support_native_out = True -remainder.support_native_out = True -round.support_native_out = True -sign.support_native_out = True -sin.support_native_out = True -sinh.support_native_out = True -sqrt.support_native_out = True -square.support_native_out = True -subtract.support_native_out = True -trapz.support_native_out = False -tan.support_native_out = True -tanh.support_native_out = True trunc.support_native_out = True + + +# Extra # +# ------# + + +@_scalar_output_to_0d_array +def erf(x, /, *, out: Optional[np.ndarray] = None): + a1 = 0.254829592 + a2 = -0.284496736 + a3 = 1.421413741 + a4 = -1.453152027 + a5 = 1.061405429 + p = 0.3275911 + + sign = np.sign(x) + x = np.abs(x) + + # A&S formula 7.1.26 + t = 1.0 / (1.0 + p * x) + y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) + ret = sign * y + if hasattr(x, "dtype"): + ret = np.asarray(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + erf.support_native_out = True + + +@_scalar_output_to_0d_array +def maximum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +): + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 >= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.maximum(x1, x2, out=out) + + maximum.support_native_out = True + + +@_scalar_output_to_0d_array +def minimum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 <= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.minimum(x1, x2, out=out) + + minimum.support_native_out = True + + +@_scalar_output_to_0d_array +def reciprocal( + x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + numerator = np.ones_like(x) + return np.true_divide(numerator, x, out=out) + + reciprocal.support_native_out = True + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.deg2rad(x, out=out) + + deg2rad.support_native_out = True + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.rad2deg(x, out=out) + + rad2deg.support_native_out = True + + +@_scalar_output_to_0d_array +def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.isreal(x) + + isreal.support_native_out = False + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def fmod( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmod( + x1, + x2, + out=None, + ) + + fmod.support_native_out = True + + +def angle( + z: np.ndarray, + /, + *, + deg: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.angle(z, deg=deg) + + angle.support_native_out = False + + +def gcd( + x1: Union[np.ndarray, int, list, tuple], + x2: Union[np.ndarray, float, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.gcd(x1, x2, out=out) + + gcd.support_native_out = True + + +def imag( + val: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.imag(val) + + imag.support_native_out = False + + +def nan_to_num( + x: np.ndarray, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) + + nan_to_num.support_native_out = False + + +def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.real(x) diff --git a/ivy/functional/backends/numpy/experimental/activations.py b/ivy/functional/backends/numpy/experimental/activations.py index 0ae99da3621d4..d3f8282a22de2 100644 --- a/ivy/functional/backends/numpy/experimental/activations.py +++ b/ivy/functional/backends/numpy/experimental/activations.py @@ -10,24 +10,6 @@ from . import backend_version -thresholded_relu.support_native_out = True -relu6.support_native_out = True -selu.support_native_out = True -silu.support_native_out = True -elu.support_native_out = True - - -@_scalar_output_to_0d_array -def elu( - x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None -) -> np.ndarray: - # exp = np.expm1(x) - ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret - - def logit( x: np.ndarray, /, @@ -46,10 +28,18 @@ def logit( return ret -@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) @_scalar_output_to_0d_array -def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return -(np.log1p(np.exp(-(input)))) +def thresholded_relu( + x: np.ndarray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.where(x > threshold, x, 0).astype(x.dtype) + + +thresholded_relu.support_native_out = True @_scalar_output_to_0d_array @@ -57,6 +47,15 @@ def relu6(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.minimum(np.maximum(x, 0, dtype=x.dtype), 6, out=out, dtype=x.dtype) +relu6.support_native_out = True + + +@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) +@_scalar_output_to_0d_array +def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return -(np.log1p(np.exp(-(input)))) + + @_scalar_output_to_0d_array def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: alpha = 1.6732632423543772848170429916717 @@ -67,6 +66,9 @@ def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret +selu.support_native_out = True + + @_scalar_output_to_0d_array def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: ret = np.asarray(x * (1 / (1 + np.exp(-x)))) @@ -78,12 +80,18 @@ def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.asarray(x * (1 / (1 + np.exp(-x)))).astype(x.dtype) +silu.support_native_out = True + + @_scalar_output_to_0d_array -def thresholded_relu( - x: np.ndarray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[np.ndarray] = None, +def elu( + x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None ) -> np.ndarray: - return np.where(x > threshold, x, 0).astype(x.dtype) + # exp = np.expm1(x) + ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret + + +elu.support_native_out = True diff --git a/ivy/functional/backends/numpy/experimental/creation.py b/ivy/functional/backends/numpy/experimental/creation.py index 150d49fce0060..c2e01f77c55e8 100644 --- a/ivy/functional/backends/numpy/experimental/creation.py +++ b/ivy/functional/backends/numpy/experimental/creation.py @@ -7,32 +7,37 @@ from ivy.functional.backends.numpy.device import _to_device import ivy - -vorbis_window.support_native_out = False -hann_window.support_native_out = False -kaiser_window.support_native_out = False -blackman_window.support_native_out = False +# Array API Standard # +# -------------------# -def blackman_window( - size: int, - /, +def vorbis_window( + window_length: np.ndarray, *, - periodic: bool = True, - dtype: Optional[np.dtype] = None, + dtype: np.dtype = np.float32, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if size < 2: - return np.ones([size], dtype=dtype) - if periodic: - count = np.arange(size) / size - else: - count = np.linspace(start=0, stop=size, num=size) + result = [] + for i in range(1, window_length * 2, 2): + temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) + result.append(round(temp, 8)) + return np.array(result, dtype=dtype) - return ( - (0.42 - 0.5 * np.cos(2 * np.pi * count)) - + (0.08 * np.cos(2 * np.pi * 2 * count)) - ).astype(dtype) + +vorbis_window.support_native_out = False + + +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[np.ndarray, ...]: + return tuple( + _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) + ) def hann_window( @@ -52,12 +57,7 @@ def hann_window( return (0.5 - 0.5 * np.cos(2 * np.pi * count)).astype(dtype) -def indices( - dimensions: Sequence, - dtype: np.dtype = np.int64, - sparse: bool = False, -) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: - return np.indices(dimensions, dtype=dtype, sparse=sparse) +hann_window.support_native_out = False def kaiser_window( @@ -76,30 +76,15 @@ def kaiser_window( return np.kaiser(M=window_length + 1, beta=beta)[:-1].astype(dtype) -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: str, -) -> Tuple[np.ndarray, ...]: - return tuple( - _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) - ) +kaiser_window.support_native_out = False -def trilu( - x: np.ndarray, - /, - *, - k: int = 0, - upper: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if upper: - return np.triu(x, k) - return np.tril(x, k) +def indices( + dimensions: Sequence, + dtype: np.dtype = np.int64, + sparse: bool = False, +) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: + return np.indices(dimensions, dtype=dtype, sparse=sparse) def unsorted_segment_min( @@ -128,6 +113,30 @@ def unsorted_segment_min( return res +def blackman_window( + size: int, + /, + *, + periodic: bool = True, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if size < 2: + return np.ones([size], dtype=dtype) + if periodic: + count = np.arange(size) / size + else: + count = np.linspace(start=0, stop=size, num=size) + + return ( + (0.42 - 0.5 * np.cos(2 * np.pi * count)) + + (0.08 * np.cos(2 * np.pi * 2 * count)) + ).astype(dtype) + + +blackman_window.support_native_out = False + + def unsorted_segment_sum( data: np.ndarray, segment_ids: np.ndarray, @@ -151,18 +160,14 @@ def unsorted_segment_sum( return res -# Array API Standard # -# -------------------# - - -def vorbis_window( - window_length: np.ndarray, +def trilu( + x: np.ndarray, + /, *, - dtype: np.dtype = np.float32, + k: int = 0, + upper: bool = True, out: Optional[np.ndarray] = None, ) -> np.ndarray: - result = [] - for i in range(1, window_length * 2, 2): - temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) - result.append(round(temp, 8)) - return np.array(result, dtype=dtype) + if upper: + return np.triu(x, k) + return np.tril(x, k) diff --git a/ivy/functional/backends/numpy/experimental/elementwise.py b/ivy/functional/backends/numpy/experimental/elementwise.py index 04479cf5508c6..cf91193091296 100644 --- a/ivy/functional/backends/numpy/experimental/elementwise.py +++ b/ivy/functional/backends/numpy/experimental/elementwise.py @@ -9,96 +9,49 @@ from . import backend_version -fmax.support_native_out = True -float_power.support_native_out = True -copysign.support_native_out = True -count_nonzero.support_native_out = False -nansum.support_native_out = True -isclose.support_native_out = False -signbit.support_native_out = True -diff.support_native_out = False -allclose.support_native_out = False -fix.support_native_out = True -nextafter.support_natvie_out = True -zeta.support_native_out = False -# --- LGAMMA --- # -LANCZOS_N = 13 -kBaseLanczosCoeff = 0.99999999999980993227684700473478 -kLanczosCoefficients = np.array( - [ - 676.520368121885098567009190444019, - -1259.13921672240287047156078755283, - 771.3234287776530788486528258894, - -176.61502916214059906584551354, - 12.507343278686904814458936853, - -0.13857109526572011689554707, - 9.984369578019570859563e-6, - 1.50563273514931155834e-7, - ] -) -# ---digamma---# -kLanczosGamma = 7 # aka g -lanczos_den_coeffs = np.array( - [ - 0.0, - 39916800.0, - 120543840.0, - 150917976.0, - 105258076.0, - 45995730.0, - 13339535.0, - 2637558.0, - 357423.0, - 32670.0, - 1925.0, - 66.0, - 1.0, - ] -) -lanczos_g = 6.024680040776729583740234375 -lanczos_num_coeffs = np.array( - [ - 23531376880.410759688572007674451636754734846804940, - 42919803642.649098768957899047001988850926355848959, - 35711959237.355668049440185451547166705960488635843, - 17921034426.037209699919755754458931112671403265390, - 6039542586.3520280050642916443072979210699388420708, - 1439720407.3117216736632230727949123939715485786772, - 248874557.86205415651146038641322942321632125127801, - 31426415.585400194380614231628318205362874684987640, - 2876370.6289353724412254090516208496135991145378768, - 186056.26539522349504029498971604569928220784236328, - 8071.6720023658162106380029022722506138218516325024, - 210.82427775157934587250973392071336271166969580291, - 2.5066282746310002701649081771338373386264310793408, - ] -) +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.sinc(x).astype(x.dtype) @_scalar_output_to_0d_array -def allclose( +def fmax( x1: np.ndarray, x2: np.ndarray, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[np.ndarray] = None, -) -> bool: - return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmax( + x1, + x2, + out=None, + where=True, + casting="same_kind", + order="K", + dtype=None, + subok=True, + ) -def conj( - x: np.ndarray, +fmax.support_native_out = True + + +@_scalar_output_to_0d_array +def float_power( + x1: Union[np.ndarray, float, list, tuple], + x2: Union[np.ndarray, float, list, tuple], /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ret = np.conj(x, out=out) - if x.dtype == bool: - return ret.astype("bool") - return ret + x1, x2 = promote_types_of_inputs(x1, x2) + return np.float_power(x1, x2, out=out) + + +float_power.support_native_out = True @_scalar_output_to_0d_array @@ -116,6 +69,9 @@ def copysign( return np.copysign(x1, x2, out=out) +copysign.support_native_out = True + + @_scalar_output_to_0d_array def count_nonzero( a: np.ndarray, @@ -134,6 +90,67 @@ def count_nonzero( return ret.astype(dtype) +count_nonzero.support_native_out = False + + +def nansum( + x: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[np.dtype] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +nansum.support_native_out = True + + +def isclose( + a: np.ndarray, + b: np.ndarray, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + if np.isscalar(ret): + return np.array(ret, dtype="bool") + return ret + + +isclose.support_native_out = False + + +def signbit( + x: Union[np.ndarray, float, int, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.signbit(x, out=out) + + +signbit.support_native_out = True + + +def hypot( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.hypot(x1, x2) + + def diff( x: Union[np.ndarray, list, tuple], /, @@ -149,6 +166,160 @@ def diff( return np.diff(x, n=n, axis=axis, prepend=prepend, append=append) +diff.support_native_out = False + + +@_scalar_output_to_0d_array +def allclose( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[np.ndarray] = None, +) -> bool: + return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) + + +allclose.support_native_out = False + + +def fix( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.fix(x, out=out) + + +fix.support_native_out = True + + +def nextafter( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nextafter(x1, x2) + + +nextafter.support_natvie_out = True + + +def zeta( + x: np.ndarray, + q: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) + inf_indices = np.logical_or(temp, np.equal(x, 1)) + temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + nan_indices = np.logical_or(temp, np.less(x, 1)) + n, res = 1, 1 / q**x + while n < 10000: + term = 1 / (q + n) ** x + n, res = n + 1, res + term + ret = np.round(res, decimals=4) + ret[nan_indices] = np.nan + ret[inf_indices] = np.inf + return ret + + +zeta.support_native_out = False + + +def gradient( + x: np.ndarray, + /, + *, + spacing: Union[int, list, tuple] = 1, + axis: Optional[Union[int, list, tuple]] = None, + edge_order: int = 1, +) -> Union[np.ndarray, List[np.ndarray]]: + if type(spacing) in (int, float): + return np.gradient(x, spacing, axis=axis, edge_order=edge_order) + return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) + + +def xlogy( + x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + x, y = promote_types_of_inputs(x, y) + if (x == 0).all(): + return 0.0 + else: + return x * np.log(y) + + +def conj( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ret = np.conj(x, out=out) + if x.dtype == bool: + return ret.astype("bool") + return ret + + +def ldexp( + x1: np.ndarray, + x2: Union[np.ndarray, int, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.ldexp(x1, x2, out=out) + + +def frexp( + x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None +) -> Tuple[np.ndarray, np.ndarray]: + if out is None: + return np.frexp(x, out=(None, None)) + else: + return np.frexp(x, out=out) + + +def modf( + x: np.ndarray, + /, + *, + out: Optional[Tuple[np.ndarray, np.ndarray]] = None, +) -> np.ndarray: + if out: + return np.modf(x, out=out) + return np.modf(x) + + +# ---digamma---# +kLanczosGamma = 7 # aka g +kBaseLanczosCoeff = 0.99999999999980993227684700473478 +kLanczosCoefficients = np.array( + [ + 676.520368121885098567009190444019, + -1259.13921672240287047156078755283, + 771.3234287776530788486528258894, + -176.61502916214059906584551354, + 12.507343278686904814458936853, + -0.13857109526572011689554707, + 9.984369578019570859563e-6, + 1.50563273514931155834e-7, + ] +) + + def digamma( x: np.ndarray, /, @@ -193,94 +364,64 @@ def digamma( ) -def fix( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.fix(x, out=out) - - -@_scalar_output_to_0d_array -def float_power( - x1: Union[np.ndarray, float, list, tuple], - x2: Union[np.ndarray, float, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.float_power(x1, x2, out=out) - +# --- LGAMMA --- # +LANCZOS_N = 13 +lanczos_g = 6.024680040776729583740234375 +lanczos_num_coeffs = np.array( + [ + 23531376880.410759688572007674451636754734846804940, + 42919803642.649098768957899047001988850926355848959, + 35711959237.355668049440185451547166705960488635843, + 17921034426.037209699919755754458931112671403265390, + 6039542586.3520280050642916443072979210699388420708, + 1439720407.3117216736632230727949123939715485786772, + 248874557.86205415651146038641322942321632125127801, + 31426415.585400194380614231628318205362874684987640, + 2876370.6289353724412254090516208496135991145378768, + 186056.26539522349504029498971604569928220784236328, + 8071.6720023658162106380029022722506138218516325024, + 210.82427775157934587250973392071336271166969580291, + 2.5066282746310002701649081771338373386264310793408, + ] +) +lanczos_den_coeffs = np.array( + [ + 0.0, + 39916800.0, + 120543840.0, + 150917976.0, + 105258076.0, + 45995730.0, + 13339535.0, + 2637558.0, + 357423.0, + 32670.0, + 1925.0, + 66.0, + 1.0, + ] +) -@_scalar_output_to_0d_array -def fmax( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmax( - x1, - x2, - out=None, - where=True, - casting="same_kind", - order="K", - dtype=None, - subok=True, - ) +def sinpi(x): + y = np.abs(x) % 2.0 + n = np.round(2.0 * y) + assert 0 <= n and n <= 4 -def frexp( - x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None -) -> Tuple[np.ndarray, np.ndarray]: - if out is None: - return np.frexp(x, out=(None, None)) + if n == 0: + r = np.sin(np.pi * y) + elif n == 1: + r = np.cos(np.pi * (y - 0.5)) + elif n == 2: + r = np.sin(np.pi * (1.0 - y)) + elif n == 3: + r = -np.cos(np.pi * (y - 1.5)) + elif n == 4: + r = np.sin(np.pi * (y - 2.0)) else: - return np.frexp(x, out=out) - - -def gradient( - x: np.ndarray, - /, - *, - spacing: Union[int, list, tuple] = 1, - axis: Optional[Union[int, list, tuple]] = None, - edge_order: int = 1, -) -> Union[np.ndarray, List[np.ndarray]]: - if type(spacing) in (int, float): - return np.gradient(x, spacing, axis=axis, edge_order=edge_order) - return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) - - -def hypot( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.hypot(x1, x2) - + raise Exception("Unreachable code") -def isclose( - a: np.ndarray, - b: np.ndarray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - if np.isscalar(ret): - return np.array(ret, dtype="bool") - return ret + return np.copysign(1.0, x) * r def lanczos_sum(x): @@ -299,16 +440,6 @@ def lanczos_sum(x): return num / den -def ldexp( - x1: np.ndarray, - x2: Union[np.ndarray, int, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.ldexp(x1, x2, out=out) - - # TODO: Replace with native lgamma implementation when available def lgamma( x: np.ndarray, @@ -349,108 +480,3 @@ def func(x): # Vectorize 'func' for element-wise operations on 'x', output matching 'x' dtype. vfunc = np.vectorize(func, otypes=[x.dtype]) return vfunc(x) - - -def modf( - x: np.ndarray, - /, - *, - out: Optional[Tuple[np.ndarray, np.ndarray]] = None, -) -> np.ndarray: - if out: - return np.modf(x, out=out) - return np.modf(x) - - -def nansum( - x: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[np.dtype] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -def nextafter( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nextafter(x1, x2) - - -def signbit( - x: Union[np.ndarray, float, int, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.signbit(x, out=out) - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.sinc(x).astype(x.dtype) - - -def sinpi(x): - y = np.abs(x) % 2.0 - n = np.round(2.0 * y) - assert 0 <= n and n <= 4 - - if n == 0: - r = np.sin(np.pi * y) - elif n == 1: - r = np.cos(np.pi * (y - 0.5)) - elif n == 2: - r = np.sin(np.pi * (1.0 - y)) - elif n == 3: - r = -np.cos(np.pi * (y - 1.5)) - elif n == 4: - r = np.sin(np.pi * (y - 2.0)) - else: - raise Exception("Unreachable code") - - return np.copysign(1.0, x) * r - - -def xlogy( - x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - x, y = promote_types_of_inputs(x, y) - if (x == 0).all(): - return 0.0 - else: - return x * np.log(y) - - -def zeta( - x: np.ndarray, - q: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) - inf_indices = np.logical_or(temp, np.equal(x, 1)) - temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - nan_indices = np.logical_or(temp, np.less(x, 1)) - n, res = 1, 1 / q**x - while n < 10000: - term = 1 / (q + n) ** x - n, res = n + 1, res + term - ret = np.round(res, decimals=4) - ret[nan_indices] = np.nan - ret[inf_indices] = np.inf - return ret diff --git a/ivy/functional/backends/numpy/experimental/layers.py b/ivy/functional/backends/numpy/experimental/layers.py index 4360cf4099908..757e434c7b834 100644 --- a/ivy/functional/backends/numpy/experimental/layers.py +++ b/ivy/functional/backends/numpy/experimental/layers.py @@ -19,10 +19,6 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): kernel, strides, depth_pooling = _depth_max_pooling_helper( x.shape, kernel, strides, dims=dims, data_format=data_format @@ -32,77 +28,82 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): - if isinstance(padding, str): - pad_specific = [ - _handle_padding(x_shape[i], strides[i], kernel[i], padding) - for i in range(dim) - ] - padding = [ - (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) - for i in range(dim) - ] - else: - pad_specific = [sum(padding[i]) for i in range(dim)] - - c = [] - if ceil_mode: - for i in range(dim): - padding[i], c_i = _padding_ceil_mode( - x_shape[i], kernel[i], padding[i], strides[i], True - ) - c.append(c_i) - pad_specific[i] = sum(padding[i]) - return padding, pad_specific, c - - -# --- Main --- # -# ------------ # - - -def avg_pool1d( +def max_pool1d( x: np.ndarray, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, data_format: str = "NWC", - count_include_pad: bool = False, + dilation: Union[int, Tuple[int]] = 1, ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] - elif len(kernel) == 1: - kernel = [kernel[0]] - - if isinstance(strides, int): - strides = [strides] - elif len(strides) == 1: - strides = [strides[0]] + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) - if data_format in ("NCW", "NCL"): + if data_format == "NCW": x = np.swapaxes(x, 1, 2) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) - x_shape = x.shape[1:-1] - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 1 + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) + x_shape = x.shape[1:2] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + if dilation[0] > 1: + filters = _add_dilations(filters, dilation[0], axis=0, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_list = [ + (pad_w // 2, pad_w - pad_w // 2), + ] + if ceil_mode: + pad_list[0] = _padding_ceil_mode( + x_shape[0], kernel[0], pad_list[0], strides[0] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) x_shape = x.shape new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] + new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] new_strides = ( x.strides[0], x.strides[1] * strides[0], @@ -110,80 +111,102 @@ def avg_pool1d( x.strides[2], ) + # B x OW x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - res = np.mean(sub_matrices, axis=2) - - if (not count_include_pad or ceil_mode) and any(pad_specific): - if not count_include_pad: - num_padded_values = np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[0], - "n": x.shape[1] - pad_specific[0], - "k": kernel[0], - "s": strides[0], - }, - unique={ - "i": np.arange(res.shape[1]), - }, - ), - dtype=res.dtype, - ) - else: - num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) - num_padded_values[-1] = c[0] - res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) + # B x OW x KW x I + sub_matrices = np.where( + filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf + ) - if data_format in ("NCW", "NCL"): - return res.swapaxes(1, 2) + res = sub_matrices.max(axis=(2)) + if depth_pooling: + res = np.swapaxes(res, 1, 2) + if data_format == "NCW": + res = np.swapaxes(res, 1, 2) return res -def avg_pool2d( +def max_pool2d( x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, data_format: str = "NHWC", - count_include_pad: bool = False, + dilation: Union[int, Tuple[int, ...]] = 1, ceil_mode: bool = False, - divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 2 - elif len(kernel) == 1: - kernel = [kernel[0]] * 2 - - if isinstance(strides, int): - strides = [strides] * 2 - elif len(strides) == 1: - strides = [strides[0]] * 2 + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) if data_format == "NCHW": x = np.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) - x_shape = list(x.shape[1:3]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 2 - ) - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) + x_shape = list(x.shape[1:3]) + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_list = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + x_shape = x.shape new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 @@ -196,101 +219,107 @@ def avg_pool2d( x.strides[2], x.strides[3], ) + # B x OH x OW x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) + # B x OH x OW x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf + ) + # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(3, 4)) - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): - if not count_include_pad: - num_padded_values = [ - np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[i], - "n": x.shape[i + 1] - pad_specific[i], - "k": kernel[i], - "s": strides[i], - }, - unique={ - "i": np.arange(res.shape[i + 1]), - }, - ), - dtype=res.dtype, - ) - for i in range(2) - ] - else: - num_padded_values = [] - for i in range(2): - num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) - num_pad[-1] = c[i] - num_padded_values.append(num_pad) - num_padded_values1 = num_padded_values[0][:, None] - num_padded_values2 = num_padded_values[1][None, :] - num_padded_values = ( - num_padded_values1 * kernel[1] - + num_padded_values2 * kernel[0] - - num_padded_values1 * num_padded_values2 - ) - kernel_mul = np.prod(kernel) - res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) + res = sub_matrices.max(axis=(3, 4)) + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 1)) if data_format == "NCHW": return np.transpose(res, (0, 3, 1, 2)) return res -def avg_pool3d( +def max_pool3d( x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, data_format: str = "NDHWC", - count_include_pad: bool = False, + dilation: Union[int, Tuple[int, ...]] = 1, ceil_mode: bool = False, - divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 3 - elif len(kernel) == 1: - kernel = [kernel[0]] * 3 - - if isinstance(strides, int): - strides = [strides] * 3 - elif len(strides) == 1: - strides = [strides[0]] * 3 + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) if data_format == "NCDHW": x = np.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) - x_shape = list(x.shape[1:4]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 3 + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) + x_shape = x.shape[1:4] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) + pad_list = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) x_shape = x.shape new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 @@ -307,262 +336,355 @@ def avg_pool3d( x.strides[3], x.strides[4], ) - # B x OH x OW x KH x KW x I + # B x OD x OH x OW x KD x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(4, 5, 6)) + # B x OD x OH x OW x KD x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf + ) - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): - if not count_include_pad: - num_padded_values = [ - np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[i], - "n": x.shape[i + 1] - pad_specific[i], - "k": kernel[i], - "s": strides[i], - }, - unique={ - "i": np.arange(res.shape[i + 1]), - }, - ), - dtype=res.dtype, - ) - for i in range(3) - ] - else: - num_padded_values = [] - for i in range(3): - num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) - num_pad[-1] = c[i] - num_padded_values.append(num_pad) - num_padded_values1 = num_padded_values[0].reshape((-1, 1, 1)) - num_padded_values2 = num_padded_values[1].reshape((1, -1, 1)) - num_padded_values3 = num_padded_values[2].reshape((1, 1, -1)) - num_padded_values = ( - num_padded_values1 * kernel[1] * kernel[2] - + num_padded_values2 * kernel[0] * kernel[2] - + num_padded_values3 * kernel[0] * kernel[1] - + num_padded_values1 * num_padded_values2 * num_padded_values3 - - num_padded_values1 * num_padded_values2 * kernel[2] - - num_padded_values1 * num_padded_values3 * kernel[1] - - num_padded_values2 * num_padded_values3 * kernel[0] - ) - kernel_mul = np.prod(kernel) - res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) + # B x OD x OH x OW x O + res = sub_matrices.max(axis=(4, 5, 6)) + + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 4, 1)) if data_format == "NCDHW": return np.transpose(res, (0, 4, 1, 2, 3)) return res -@with_supported_dtypes({"1.25.2 and below": ("float32", "float64")}, backend_version) -def dct( +def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): + if isinstance(padding, str): + pad_specific = [ + _handle_padding(x_shape[i], strides[i], kernel[i], padding) + for i in range(dim) + ] + padding = [ + (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) + for i in range(dim) + ] + else: + pad_specific = [sum(padding[i]) for i in range(dim)] + + c = [] + if ceil_mode: + for i in range(dim): + padding[i], c_i = _padding_ceil_mode( + x_shape[i], kernel[i], padding[i], strides[i], True + ) + c.append(c_i) + pad_specific[i] = sum(padding[i]) + return padding, pad_specific, c + + +def avg_pool1d( x: np.ndarray, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, /, *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if norm not in (None, "ortho"): - raise ValueError("Norm must be either None or 'ortho'") - if axis < 0: - axis = axis + len(x.shape) - if n is not None: - signal_len = x.shape[axis] - if n <= signal_len: - local_idx = [slice(None)] * len(x.shape) - local_idx[axis] = slice(None, n) - x = x[tuple(local_idx)] - else: - pad_idx = [[0, 0] for _ in range(len(x.shape))] - pad_idx[axis][1] = n - signal_len - x = np.pad(x, pad_idx) - real_zero = np.array(0.0, dtype=x.dtype) - axis_dim = x.shape[axis] - axis_dim_float = np.array(axis_dim, dtype=x.dtype) - cast_final = True if x.dtype != np.float64 else False + if isinstance(kernel, int): + kernel = [kernel] + elif len(kernel) == 1: + kernel = [kernel[0]] - if type == 1: - if norm: - raise ValueError("Normalization not supported for type-I DCT") - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(-2, 0, -1) - x = np.concatenate([x, x[tuple(axis_idx)]], axis=axis) - dct_out = np.real(np.fft.rfft(x, axis=axis)) + if isinstance(strides, int): + strides = [strides] + elif len(strides) == 1: + strides = [strides[0]] - elif type == 2: - cmplx = np.empty(axis_dim, dtype=np.complex64) - cmplx.real = real_zero - cmplx.imag = -np.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float + if data_format in ("NCW", "NCL"): + x = np.swapaxes(x, 1, 2) - scale_dims = [1] * len(x.shape) - scale_dims[axis] = axis_dim - scale = 2.0 * np.exp(cmplx).reshape(scale_dims) + x_shape = x.shape[1:-1] + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 1 + ) - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(None, axis_dim) - dct_out = np.real( - np.fft.rfft(x, n=2 * axis_dim, axis=axis)[tuple(axis_idx)] * scale - ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) - if norm == "ortho": - n1 = 0.5 * np.reciprocal(np.sqrt(axis_dim_float)) - n2 = n1 * math.sqrt(2.0) - sf = np.pad(np.expand_dims(n1, 0), (0, axis_dim - 1), constant_values=n2) - dct_out = sf.reshape(scale_dims) * dct_out + x_shape = x.shape + new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[1], + x.strides[2], + ) - elif type == 3: - cmplx = np.empty(axis_dim, dtype=np.complex64) - cmplx.real = real_zero - cmplx.imag = np.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) - scale_dims = [1] * len(x.shape) - scale_dims[axis] = axis_dim - scale = 2.0 * np.exp(cmplx).reshape(scale_dims) + res = np.mean(sub_matrices, axis=2) - if norm == "ortho": - n1 = np.sqrt(axis_dim_float) - n2 = n1 * np.sqrt(0.5) - sf = np.pad(np.expand_dims(n1, 0), (0, axis_dim - 1), constant_values=n2) - x = x * sf.reshape(scale_dims) + if (not count_include_pad or ceil_mode) and any(pad_specific): + if not count_include_pad: + num_padded_values = np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[0], + "n": x.shape[1] - pad_specific[0], + "k": kernel[0], + "s": strides[0], + }, + unique={ + "i": np.arange(res.shape[1]), + }, + ), + dtype=res.dtype, + ) else: - x = x * axis_dim_float - - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(None, axis_dim) - - x = x.astype(np.complex64) - x.imag = real_zero - dct_out = np.real(np.fft.irfft(scale * x, n=2 * axis_dim, axis=axis))[ - tuple(axis_idx) - ] - - elif type == 4: - dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(1, None, 2) - dct_out = dct_2[tuple(axis_idx)] - if norm == "ortho": - dct_out *= math.sqrt(0.5) * np.reciprocal(np.sqrt(axis_dim_float)) - - return dct_out.astype(np.float32) if cast_final else dct_out + num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) + num_padded_values[-1] = c[0] + res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) + if data_format in ("NCW", "NCL"): + return res.swapaxes(1, 2) -def dropout1d( - x: np.ndarray, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if training: - x_shape = x.shape - is_batched = len(x_shape) == 3 - if data_format == "NCW": - perm = (0, 2, 1) if is_batched else (1, 0) - x = np.transpose(x, perm) - x_shape = x.shape - mask = np.random.binomial(1, 1 - prob, x_shape) - res = np.where(mask, x / (1 - prob), 0) - if data_format == "NCW": - res = np.transpose(res, perm) - else: - res = x return res -def dropout2d( +def avg_pool2d( x: np.ndarray, - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, /, *, - training: bool = True, data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if training: - x_shape = x.shape - is_batched = len(x_shape) == 4 - if data_format == "NCHW": - perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) - x = np.transpose(x, perm) - x_shape = x.shape - mask = np.random.binomial(1, 1 - prob, x_shape) - res = np.where(mask, x / (1 - prob), 0) - if data_format == "NCHW": - perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) - res = np.transpose(res, perm) + if isinstance(kernel, int): + kernel = [kernel] * 2 + elif len(kernel) == 1: + kernel = [kernel[0]] * 2 + + if isinstance(strides, int): + strides = [strides] * 2 + elif len(strides) == 1: + strides = [strides[0]] * 2 + + if data_format == "NCHW": + x = np.transpose(x, (0, 2, 3, 1)) + + x_shape = list(x.shape[1:3]) + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 2 + ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) + + x_shape = x.shape + new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[1], + x.strides[2], + x.strides[3], + ) + # B x OH x OW x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OH x OW x O + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override else: - res = x + res = np.mean(sub_matrices, axis=(3, 4)) + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): + if not count_include_pad: + num_padded_values = [ + np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[i], + "n": x.shape[i + 1] - pad_specific[i], + "k": kernel[i], + "s": strides[i], + }, + unique={ + "i": np.arange(res.shape[i + 1]), + }, + ), + dtype=res.dtype, + ) + for i in range(2) + ] + else: + num_padded_values = [] + for i in range(2): + num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) + num_pad[-1] = c[i] + num_padded_values.append(num_pad) + num_padded_values1 = num_padded_values[0][:, None] + num_padded_values2 = num_padded_values[1][None, :] + num_padded_values = ( + num_padded_values1 * kernel[1] + + num_padded_values2 * kernel[0] + - num_padded_values1 * num_padded_values2 + ) + kernel_mul = np.prod(kernel) + res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) + + if data_format == "NCHW": + return np.transpose(res, (0, 3, 1, 2)) return res -def dropout3d( +def avg_pool3d( x: np.ndarray, - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, /, *, - training: bool = True, data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if training: - x_shape = x.shape - is_batched = len(x_shape) == 5 - if data_format == "NCDHW": - perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) - x = np.transpose(x, perm) - x_shape = x.shape - mask = np.random.binomial(1, 1 - prob, x_shape) - res = np.where(mask, x / (1 - prob), 0) - if data_format == "NCDHW": - perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) - res = np.transpose(res, perm) - else: - res = x - return res + if isinstance(kernel, int): + kernel = [kernel] * 3 + elif len(kernel) == 1: + kernel = [kernel[0]] * 3 + + if isinstance(strides, int): + strides = [strides] * 3 + elif len(strides) == 1: + strides = [strides[0]] * 3 + if data_format == "NCDHW": + x = np.transpose(x, (0, 2, 3, 4, 1)) -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def embedding( - weights: np.ndarray, - indices: np.ndarray, - /, - *, - max_norm: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False + x_shape = list(x.shape[1:4]) + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 3 ) - embeddings = np.take(weights, indices, axis=0) - if max_norm is not None: - norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = np.where( - norms > max_norm, embeddings * max_norm / norms, embeddings - ) - embeddings = np.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) + + x_shape = x.shape + new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 + new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[3] * strides[2], + x.strides[1], + x.strides[2], + x.strides[3], + x.strides[4], + ) + # B x OH x OW x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OH x OW x O + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override + else: + res = np.mean(sub_matrices, axis=(4, 5, 6)) + + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): + if not count_include_pad: + num_padded_values = [ + np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[i], + "n": x.shape[i + 1] - pad_specific[i], + "k": kernel[i], + "s": strides[i], + }, + unique={ + "i": np.arange(res.shape[i + 1]), + }, + ), + dtype=res.dtype, + ) + for i in range(3) + ] + else: + num_padded_values = [] + for i in range(3): + num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) + num_pad[-1] = c[i] + num_padded_values.append(num_pad) + num_padded_values1 = num_padded_values[0].reshape((-1, 1, 1)) + num_padded_values2 = num_padded_values[1].reshape((1, -1, 1)) + num_padded_values3 = num_padded_values[2].reshape((1, 1, -1)) + num_padded_values = ( + num_padded_values1 * kernel[1] * kernel[2] + + num_padded_values2 * kernel[0] * kernel[2] + + num_padded_values3 * kernel[0] * kernel[1] + + num_padded_values1 * num_padded_values2 * num_padded_values3 + - num_padded_values1 * num_padded_values2 * kernel[2] + - num_padded_values1 * num_padded_values3 * kernel[1] + - num_padded_values2 * num_padded_values3 * kernel[0] ) - return embeddings + kernel_mul = np.prod(kernel) + res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) + if data_format == "NCDHW": + return np.transpose(res, (0, 4, 1, 2, 3)) + return res def fft( @@ -602,39 +724,100 @@ def fft( return np.fft.fft(x, n, dim, norm).astype(out_dtype) -def fft2( +@with_supported_dtypes({"1.25.2 and below": ("float32", "float64")}, backend_version) +def dct( x: np.ndarray, + /, *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ivy.utils.assertions.check_elem_in_list( - norm, - ["backward", "ortho", "forward"], - message=f"Unrecognized normalization mode {norm}", - ) - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" + if norm not in (None, "ortho"): + raise ValueError("Norm must be either None or 'ortho'") + if axis < 0: + axis = axis + len(x.shape) + if n is not None: + signal_len = x.shape[axis] + if n <= signal_len: + local_idx = [slice(None)] * len(x.shape) + local_idx[axis] = slice(None, n) + x = x[tuple(local_idx)] + else: + pad_idx = [[0, 0] for _ in range(len(x.shape))] + pad_idx[axis][1] = n - signal_len + x = np.pad(x, pad_idx) + real_zero = np.array(0.0, dtype=x.dtype) + axis_dim = x.shape[axis] + axis_dim_float = np.array(axis_dim, dtype=x.dtype) + cast_final = True if x.dtype != np.float64 else False + + if type == 1: + if norm: + raise ValueError("Normalization not supported for type-I DCT") + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(-2, 0, -1) + x = np.concatenate([x, x[tuple(axis_idx)]], axis=axis) + dct_out = np.real(np.fft.rfft(x, axis=axis)) + + elif type == 2: + cmplx = np.empty(axis_dim, dtype=np.complex64) + cmplx.real = real_zero + cmplx.imag = -np.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float + + scale_dims = [1] * len(x.shape) + scale_dims[axis] = axis_dim + scale = 2.0 * np.exp(cmplx).reshape(scale_dims) + + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(None, axis_dim) + dct_out = np.real( + np.fft.rfft(x, n=2 * axis_dim, axis=axis)[tuple(axis_idx)] * scale ) - return np.fft.fft2(x, s, dim, norm).astype(np.complex128) + + if norm == "ortho": + n1 = 0.5 * np.reciprocal(np.sqrt(axis_dim_float)) + n2 = n1 * math.sqrt(2.0) + sf = np.pad(np.expand_dims(n1, 0), (0, axis_dim - 1), constant_values=n2) + dct_out = sf.reshape(scale_dims) * dct_out + + elif type == 3: + cmplx = np.empty(axis_dim, dtype=np.complex64) + cmplx.real = real_zero + cmplx.imag = np.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float + + scale_dims = [1] * len(x.shape) + scale_dims[axis] = axis_dim + scale = 2.0 * np.exp(cmplx).reshape(scale_dims) + + if norm == "ortho": + n1 = np.sqrt(axis_dim_float) + n2 = n1 * np.sqrt(0.5) + sf = np.pad(np.expand_dims(n1, 0), (0, axis_dim - 1), constant_values=n2) + x = x * sf.reshape(scale_dims) + else: + x = x * axis_dim_float + + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(None, axis_dim) + + x = x.astype(np.complex64) + x.imag = real_zero + dct_out = np.real(np.fft.irfft(scale * x, n=2 * axis_dim, axis=axis))[ + tuple(axis_idx) + ] + + elif type == 4: + dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(1, None, 2) + dct_out = dct_2[tuple(axis_idx)] + if norm == "ortho": + dct_out *= math.sqrt(0.5) * np.reciprocal(np.sqrt(axis_dim_float)) + + return dct_out.astype(np.float32) if cast_final else dct_out def idct( @@ -651,18 +834,95 @@ def idct( return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) -def ifft( +def dropout1d( x: np.ndarray, - dim: int, + prob: float, + /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + training: bool = True, + data_format: str = "NWC", out: Optional[np.ndarray] = None, ) -> np.ndarray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) + if training: + x_shape = x.shape + is_batched = len(x_shape) == 3 + if data_format == "NCW": + perm = (0, 2, 1) if is_batched else (1, 0) + x = np.transpose(x, perm) + x_shape = x.shape + mask = np.random.binomial(1, 1 - prob, x_shape) + res = np.where(mask, x / (1 - prob), 0) + if data_format == "NCW": + res = np.transpose(res, perm) + else: + res = x + return res + + +def dropout2d( + x: np.ndarray, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if training: + x_shape = x.shape + is_batched = len(x_shape) == 4 + if data_format == "NCHW": + perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) + x = np.transpose(x, perm) + x_shape = x.shape + mask = np.random.binomial(1, 1 - prob, x_shape) + res = np.where(mask, x / (1 - prob), 0) + if data_format == "NCHW": + perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) + res = np.transpose(res, perm) + else: + res = x + return res + + +def dropout3d( + x: np.ndarray, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if training: + x_shape = x.shape + is_batched = len(x_shape) == 5 + if data_format == "NCDHW": + perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) + x = np.transpose(x, perm) + x_shape = x.shape + mask = np.random.binomial(1, 1 - prob, x_shape) + res = np.where(mask, x / (1 - prob), 0) + if data_format == "NCDHW": + perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) + res = np.transpose(res, perm) + else: + res = x + return res + + +def ifft( + x: np.ndarray, + dim: int, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) if n is None: n = x.shape[dim] if n < -len(x.shape): @@ -683,343 +943,75 @@ def ifft( return np.asarray(np.fft.ifft(x, n, dim, norm), dtype=x.dtype) -def ifftn( +def fft2( x: np.ndarray, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), norm: str = "backward", out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.fft.ifftn(x, s, axes, norm).astype(x.dtype) - - -def max_pool1d( - x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims + ivy.utils.assertions.check_elem_in_list( + norm, + ["backward", "ortho", "forward"], + message=f"Unrecognized normalization mode {norm}", ) - - if data_format == "NCW": - x = np.swapaxes(x, 1, 2) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - x_shape = x.shape[1:2] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - if dilation[0] > 1: - filters = _add_dilations(filters, dilation[0], axis=0, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_list = [ - (pad_w // 2, pad_w - pad_w // 2), - ] - if ceil_mode: - pad_list[0] = _padding_ceil_mode( - x_shape[0], kernel[0], pad_list[0], strides[0] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - x_shape = x.shape - new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[1], - x.strides[2], - ) - - # B x OW x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OW x KW x I - sub_matrices = np.where( - filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf - ) - - res = sub_matrices.max(axis=(2)) - - if depth_pooling: - res = np.swapaxes(res, 1, 2) - if data_format == "NCW": - res = np.swapaxes(res, 1, 2) - return res + return np.fft.fft2(x, s, dim, norm).astype(np.complex128) -def max_pool2d( +def ifftn( x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, + norm: str = "backward", out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCHW": - x = np.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - x_shape = list(x.shape[1:3]) - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_list = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - x_shape = x.shape - new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[1], - x.strides[2], - x.strides[3], - ) - - # B x OH x OW x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OH x OW x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf - ) - - # B x OH x OW x O - res = sub_matrices.max(axis=(3, 4)) - - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 1)) - if data_format == "NCHW": - return np.transpose(res, (0, 3, 1, 2)) - return res + return np.fft.ifftn(x, s, axes, norm).astype(x.dtype) -def max_pool3d( - x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def embedding( + weights: np.ndarray, + indices: np.ndarray, /, *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, + max_norm: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False ) - if data_format == "NCDHW": - x = np.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + embeddings = np.take(weights, indices, axis=0) + if max_norm is not None: + norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = np.where( + norms > max_norm, embeddings * max_norm / norms, embeddings ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - x_shape = x.shape[1:4] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) - pad_list = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, + embeddings = np.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - x_shape = x.shape - new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 - new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[3] * strides[2], - x.strides[1], - x.strides[2], - x.strides[3], - x.strides[4], - ) - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf - ) - - # B x OD x OH x OW x O - res = sub_matrices.max(axis=(4, 5, 6)) - - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 4, 1)) - if data_format == "NCDHW": - return np.transpose(res, (0, 4, 1, 2, 3)) - return res + return embeddings def rfftn( diff --git a/ivy/functional/backends/numpy/experimental/linear_algebra.py b/ivy/functional/backends/numpy/experimental/linear_algebra.py index 54029e3dd0321..915b0b4ad9111 100644 --- a/ivy/functional/backends/numpy/experimental/linear_algebra.py +++ b/ivy/functional/backends/numpy/experimental/linear_algebra.py @@ -10,36 +10,6 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size -kron.support_native_out = False -eig.support_native_out = False -eigvals.support_native_out = False -multi_dot.support_native_out = True -cond.support_native_out = False -dot.support_native_out = True - - -def adjoint( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return np.conjugate(np.transpose(x, axes=axes)) - - -def cond( - x: np.ndarray, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[np.ndarray] = None, -) -> Any: - return np.linalg.cond(x, p=p) - - def diagflat( x: np.ndarray, /, @@ -114,14 +84,34 @@ def diagflat( return ret -def dot( +def kron( a: np.ndarray, b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.dot(a, b, out=out) + return np.kron(a, b) + + +kron.support_native_out = False + + +@with_supported_dtypes( + {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, + backend_version, +) +def matrix_exp( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + eig_vals, eig_vecs = np.linalg.eig(x) + exp_diag = np.exp(eig_vals) + exp_diag_mat = np.diag(exp_diag) + exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) + return exp_mat.astype(x.dtype) def eig( @@ -136,6 +126,9 @@ def eig( return e.astype(complex), v.astype(complex) +eig.support_native_out = False + + def eigvals(x: np.ndarray, /) -> np.ndarray: if ivy.dtype(x) == ivy.float16: x = x.astype(np.float32) @@ -143,14 +136,44 @@ def eigvals(x: np.ndarray, /) -> np.ndarray: return e.astype(complex) -def kron( - a: np.ndarray, - b: np.ndarray, +eigvals.support_native_out = False + + +def adjoint( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.kron(a, b) + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return np.conjugate(np.transpose(x, axes=axes)) + + +def multi_dot( + x: Sequence[np.ndarray], + /, + *, + out: Optional[np.array] = None, +) -> np.ndarray: + return np.linalg.multi_dot(x, out=out) + + +multi_dot.support_native_out = True + + +def cond( + x: np.ndarray, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[np.ndarray] = None, +) -> Any: + return np.linalg.cond(x, p=p) + + +cond.support_native_out = False def lu_factor( @@ -163,27 +186,14 @@ def lu_factor( raise IvyNotImplementedException() -@with_supported_dtypes( - {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, - backend_version, -) -def matrix_exp( - x: np.ndarray, +def dot( + a: np.ndarray, + b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - eig_vals, eig_vecs = np.linalg.eig(x) - exp_diag = np.exp(eig_vals) - exp_diag_mat = np.diag(exp_diag) - exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) - return exp_mat.astype(x.dtype) + return np.dot(a, b, out=out) -def multi_dot( - x: Sequence[np.ndarray], - /, - *, - out: Optional[np.array] = None, -) -> np.ndarray: - return np.linalg.multi_dot(x, out=out) +dot.support_native_out = True diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 0c03925ae8aa2..53436987e2187 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -19,207 +19,116 @@ from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array -# --- Helpers --- # -# --------------- # - - -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x - - -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = np.full(new_shape, padding_value, dtype=operand.dtype) - src_indices = np.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array - - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) - - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = np.pad(padded, pad_width, constant_values=padding_value) - return padded - - -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides - - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] - - -# --- Main --- # -# ------------ # - - -def atleast_1d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_1d(*arys) - - -def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: - return np.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_3d(*arys) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return np.broadcast_shapes(*shapes) - - -def concat_from_sequence( - input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], +def moveaxis( + a: np.ndarray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], /, *, - new_axis: int = 0, - axis: int = 0, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = np.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = np.stack(input_sequence, axis=axis) - return ret + return np.moveaxis(a, source, destination) -def dsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) +moveaxis.support_native_out = False -def dstack( - arrays: Sequence[np.ndarray], +def heaviside( + x1: np.ndarray, + x2: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.dstack(arrays) + return np.heaviside( + x1, + x2, + out=out, + ) -def expand( - x: np.ndarray, - shape: Union[List[int], List[Tuple]], +heaviside.support_native_out = True + + +def flipud( + m: np.ndarray, /, *, copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) - return np.broadcast_to(x, tuple(shape)) + return np.flipud(m) -def fill_diagonal( - a: np.ndarray, - v: Union[int, float], - /, - *, - wrap: bool = False, -) -> np.ndarray: - np.fill_diagonal(a, v, wrap=wrap) - return a +flipud.support_native_out = False -def fliplr( - m: np.ndarray, +def vstack( + arrays: Sequence[np.ndarray], /, *, - copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.fliplr(m) + return np.vstack(arrays) -def flipud( - m: np.ndarray, +def hstack( + arrays: Sequence[np.ndarray], /, *, - copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.flipud(m) + return np.hstack(arrays) -def heaviside( - x1: np.ndarray, - x2: np.ndarray, +def rot90( + m: np.ndarray, /, *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.heaviside( - x1, - x2, - out=out, - ) + return np.rot90(m, k, axes) -def hsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], +def top_k( + x: np.ndarray, + k: int, /, *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[np.ndarray, np.ndarray]] = None, +) -> Tuple[np.ndarray, np.ndarray]: + k = min(k, x.shape[axis]) + if not largest: + indices = np.argsort(x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + else: + indices = np.argsort(-x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + if not sorted: + indices = np.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) + val = np.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) -def hstack( - arrays: Sequence[np.ndarray], +def fliplr( + m: np.ndarray, /, *, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.hstack(arrays) + return np.fliplr(m) + + +fliplr.support_native_out = False def i0( @@ -231,16 +140,59 @@ def i0( return np.i0(x) -def moveaxis( - a: np.ndarray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.moveaxis(a, source, destination) +i0.support_native_out = False + + +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x + + +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides + + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] + + +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = np.full(new_shape, padding_value, dtype=operand.dtype) + src_indices = np.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) + + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = np.pad(padded, pad_width, constant_values=padding_value) + return padded def pad( @@ -316,16 +268,57 @@ def pad( ) -def rot90( - m: np.ndarray, +def vsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Sequence[int], np.ndarray], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), +) -> List[np.ndarray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def dsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def atleast_1d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_1d(*arys) + + +def dstack( + arrays: Sequence[np.ndarray], + /, + *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.rot90(m, k, axes) + return np.dstack(arrays) + + +def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: + return np.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_3d(*arys) @_scalar_output_to_0d_array @@ -373,28 +366,63 @@ def take_along_axis( return np.take_along_axis(arr, indices, axis) -def top_k( +def hsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +take_along_axis.support_native_out = False + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return np.broadcast_shapes(*shapes) + + +broadcast_shapes.support_native_out = False + + +def expand( x: np.ndarray, - k: int, + shape: Union[List[int], List[Tuple]], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[np.ndarray, np.ndarray]] = None, -) -> Tuple[np.ndarray, np.ndarray]: - k = min(k, x.shape[axis]) - if not largest: - indices = np.argsort(x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - else: - indices = np.argsort(-x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - if not sorted: - indices = np.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) - val = np.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + copy: Optional[bool] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) + return np.broadcast_to(x, tuple(shape)) + + +expand.support_native_out = False + + +def concat_from_sequence( + input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = np.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = np.stack(input_sequence, axis=axis) + return ret def unique_consecutive( @@ -440,34 +468,12 @@ def unique_consecutive( ) -def vsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Sequence[int], np.ndarray], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def vstack( - arrays: Sequence[np.ndarray], +def fill_diagonal( + a: np.ndarray, + v: Union[int, float], /, *, - out: Optional[np.ndarray] = None, + wrap: bool = False, ) -> np.ndarray: - return np.vstack(arrays) - - -moveaxis.support_native_out = False -heaviside.support_native_out = True -flipud.support_native_out = False -fliplr.support_native_out = False -i0.support_native_out = False -take_along_axis.support_native_out = False -broadcast_shapes.support_native_out = False -expand.support_native_out = False + np.fill_diagonal(a, v, wrap=wrap) + return a diff --git a/ivy/functional/backends/numpy/experimental/random.py b/ivy/functional/backends/numpy/experimental/random.py index e2af3988496df..16293d4eac407 100644 --- a/ivy/functional/backends/numpy/experimental/random.py +++ b/ivy/functional/backends/numpy/experimental/random.py @@ -10,26 +10,24 @@ ) -dirichlet.support_native_out = False - - -def bernoulli( - probs: Union[float, np.ndarray], +# dirichlet +def dirichlet( + alpha: Union[np.ndarray, float, Sequence[float]], + /, *, - logits: Optional[Union[float, np.ndarray]] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, dtype: Optional[np.dtype] = None, seed: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: + size = size if size is not None else len(alpha) + dtype = dtype if dtype is not None else np.float64 if seed is not None: np.random.seed(seed) - if logits is not None: - probs = np.asarray(ivy.softmax(logits), dtype=dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) + return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) + + +dirichlet.support_native_out = False def beta( @@ -49,23 +47,6 @@ def beta( return np.asarray(np.random.beta(alpha, beta, shape), dtype=dtype) -# dirichlet -def dirichlet( - alpha: Union[np.ndarray, float, Sequence[float]], - /, - *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: Optional[np.dtype] = None, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - size = size if size is not None else len(alpha) - dtype = dtype if dtype is not None else np.float64 - if seed is not None: - np.random.seed(seed) - return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) - - def gamma( alpha: Union[float, np.ndarray], beta: Union[float, np.ndarray], @@ -106,3 +87,22 @@ def poisson( else: ret = np.random.poisson(lam, shape) return np.asarray(ret, dtype=dtype) + + +def bernoulli( + probs: Union[float, np.ndarray], + *, + logits: Optional[Union[float, np.ndarray]] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, + dtype: Optional[np.dtype] = None, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if seed is not None: + np.random.seed(seed) + if logits is not None: + probs = np.asarray(ivy.softmax(logits), dtype=dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) diff --git a/ivy/functional/backends/numpy/experimental/searching.py b/ivy/functional/backends/numpy/experimental/searching.py index 297c9cdba2db2..23811e75b4077 100644 --- a/ivy/functional/backends/numpy/experimental/searching.py +++ b/ivy/functional/backends/numpy/experimental/searching.py @@ -7,9 +7,6 @@ from . import backend_version -unravel_index.support_native_out = False - - @with_supported_dtypes({"1.25.2 and below": ("int32", "int64")}, backend_version) def unravel_index( indices: np.ndarray, @@ -20,3 +17,6 @@ def unravel_index( ) -> Tuple[np.ndarray]: ret = np.asarray(np.unravel_index(indices, shape), dtype=np.int32) return tuple(ret) + + +unravel_index.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/sorting.py b/ivy/functional/backends/numpy/experimental/sorting.py index fd04f1f400214..cd734f8e69835 100644 --- a/ivy/functional/backends/numpy/experimental/sorting.py +++ b/ivy/functional/backends/numpy/experimental/sorting.py @@ -3,9 +3,6 @@ from typing import Optional, Union -lexsort.support_native_out = False - - # invert_permutation def invert_permutation( x: Union[np.ndarray, list, tuple], @@ -23,3 +20,6 @@ def lexsort( keys: np.ndarray, /, *, axis: int = -1, out: Optional[np.ndarray] = None ) -> np.ndarray: return np.asarray(np.lexsort(keys, axis=axis)) + + +lexsort.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/statistical.py b/ivy/functional/backends/numpy/experimental/statistical.py index 6efebf5a93ba6..12c3565148ec7 100644 --- a/ivy/functional/backends/numpy/experimental/statistical.py +++ b/ivy/functional/backends/numpy/experimental/statistical.py @@ -8,87 +8,194 @@ from copy import deepcopy -# --- Helpers --- # -# --------------- # +@with_unsupported_dtypes( + {"1.25.2 and below": ("bfloat16",)}, + backend_version, +) +def histogram( + a: np.ndarray, + /, + *, + bins: Optional[Union[int, np.ndarray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[np.dtype] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[np.ndarray] = None, + density: Optional[bool] = False, + out: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray]: + min_a = np.min(a) + max_a = np.max(a) + if isinstance(bins, np.ndarray) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + elif isinstance(bins, int): + range = (min_a, max_a) + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + if bins.size < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + bins_out = bins.copy() + if extend_lower_interval and min_a < bins[0]: + bins[0] = min_a + if extend_upper_interval and max_a > bins[-1]: + bins[-1] = max_a + if a.ndim > 0 and axis is not None: + inverted_shape_dims = list(np.flip(np.arange(a.ndim))) + if isinstance(axis, int): + axis = [axis] + shape_axes = 1 + for dimension in axis: + inverted_shape_dims.remove(dimension) + inverted_shape_dims.append(dimension) + shape_axes *= a.shape[dimension] + a_along_axis_1d = ( + a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) + ) + if weights is None: + ret = [] + for a_1d in a_along_axis_1d: + ret_1d = np.histogram( + a_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + else: + weights_along_axis_1d = ( + weights.transpose(inverted_shape_dims) + .flatten() + .reshape((-1, shape_axes)) + ) + ret = [] + for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): + ret_1d = np.histogram( + a_1d, + weights=weights_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + out_shape = list(a.shape) + for dimension in sorted(axis, reverse=True): + del out_shape[dimension] + out_shape.insert(0, len(bins) - 1) + ret = np.array(ret) + ret = ret.flatten() + index = np.zeros(len(out_shape), dtype=int) + ret_shaped = np.zeros(out_shape) + dim = 0 + i = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + while list(index) != list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + dim_full_flag = False + while index[dim] == out_shape[dim] - 1: + index[dim] = 0 + dim += 1 + dim_full_flag = True + index[dim] += 1 + i += 1 + if dim_full_flag: + dim = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + ret = ret_shaped + else: + ret = np.histogram( + a=a, bins=bins, range=range, weights=weights, density=density + )[0] + if dtype: + ret = ret.astype(dtype) + bins_out = np.array(bins_out).astype(dtype) + # TODO: weird error when returning bins: return ret, bins_out + return ret -def __find_cummax_indices( - x: np.ndarray, - axis: int = 0, +def median( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - indices = [] - if x[0] is np.ndarray: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - - else: - indice_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices.fill(0) - indice_list = sorted(indice_list, key=lambda i: i[1]) - for y, y_index in indice_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - else: - indices[y_index] = n1[tuple(multi_index[1:])] + if out is not None: + out = np.reshape(out, input.shape) + ret = np.median( + input, + axis=axis, + keepdims=keepdims, + out=out, + ) + if input.dtype in [np.uint64, np.int64, np.float64]: + return ret.astype(np.float64) + elif input.dtype in [np.float16]: + return ret.astype(input.dtype) else: - n = 0 - for index1, ret1 in enumerate(x): - if x[n] <= ret1 or index1 == 0: - n = index1 - indices.append(n) - return np.array(indices, dtype=np.int64) + return ret.astype(np.float32) -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] +median.support_native_out = True - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices +def nanmean( + a: np.ndarray, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - axis = tuple(axis) if isinstance(axis, list) else axis - return np.quantile( - x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out - ).astype(x.dtype) +nanmean.support_native_out = True + + +def _validate_quantile(q): + if isinstance(q, float): + q = np.asarray(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) + if not (np.all(0 <= q) and np.all(q <= 1)): + return False + return True + + +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis def _handle_axis(a, q, fn, keepdims=False, axis=None): @@ -154,39 +261,89 @@ def _quantile(a, q, axis=None): return out.astype(ret_dtype) -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + axis = tuple(axis) if isinstance(axis, list) else axis - if len(axis) == 0: - raise ValueError("Axis can't be empty!") + return np.quantile( + x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ).astype(x.dtype) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis +def quantile( + a: np.ndarray, + q: Union[float, np.ndarray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + interpolation: str = "linear", + out: Optional[np.ndarray] = None, +) -> np.ndarray: + # quantile method in numpy backend, always return an array with dtype=float64. + # in other backends, the output is the same dtype as the input. + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + out=out, + ) -def _validate_quantile(q): - if isinstance(q, float): - q = np.asarray(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False - else: - if not (np.all(0 <= q) and np.all(q <= 1)): - return False - return True +def corrcoef( + x: np.ndarray, + /, + *, + y: Optional[np.ndarray] = None, + rowvar: bool = True, + dtype: np.dtype = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dtype = dtype if dtype is not None else np.float64 + + return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) + + +@with_unsupported_dtypes( + {"1.25.0 and below": ("bfloat16",)}, + backend_version, +) +def nanmedian( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + ) -# --- Main --- # -# ------------ # +nanmedian.support_native_out = True def bincount( @@ -206,18 +363,7 @@ def bincount( return ret -def corrcoef( - x: np.ndarray, - /, - *, - y: Optional[np.ndarray] = None, - rowvar: bool = True, - dtype: np.dtype = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dtype = dtype if dtype is not None else np.float64 - - return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) +bincount.support_native_out = False def cov( @@ -244,6 +390,9 @@ def cov( ) +cov.support_native_out = False + + def cummax( x: np.ndarray, /, @@ -289,6 +438,58 @@ def cummax( return np.maximum.accumulate(x, axis=axis, dtype=x.dtype), indices +def __find_cummax_indices( + x: np.ndarray, + axis: int = 0, +) -> np.ndarray: + indices = [] + if x[0] is np.ndarray: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + + else: + indice_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices.fill(0) + indice_list = sorted(indice_list, key=lambda i: i[1]) + for y, y_index in indice_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + else: + indices[y_index] = n1[tuple(multi_index[1:])] + else: + n = 0 + for index1, ret1 in enumerate(x): + if x[n] <= ret1 or index1 == 0: + n = index1 + indices.append(n) + return np.array(indices, dtype=np.int64) + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + @with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) def cummin( x: np.ndarray, @@ -312,121 +513,6 @@ def cummin( return np.flip(x, axis=axis) -@with_unsupported_dtypes( - {"1.25.2 and below": ("bfloat16",)}, - backend_version, -) -def histogram( - a: np.ndarray, - /, - *, - bins: Optional[Union[int, np.ndarray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[np.dtype] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[np.ndarray] = None, - density: Optional[bool] = False, - out: Optional[np.ndarray] = None, -) -> Tuple[np.ndarray]: - min_a = np.min(a) - max_a = np.max(a) - if isinstance(bins, np.ndarray) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - elif isinstance(bins, int): - range = (min_a, max_a) - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - if bins.size < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - bins_out = bins.copy() - if extend_lower_interval and min_a < bins[0]: - bins[0] = min_a - if extend_upper_interval and max_a > bins[-1]: - bins[-1] = max_a - if a.ndim > 0 and axis is not None: - inverted_shape_dims = list(np.flip(np.arange(a.ndim))) - if isinstance(axis, int): - axis = [axis] - shape_axes = 1 - for dimension in axis: - inverted_shape_dims.remove(dimension) - inverted_shape_dims.append(dimension) - shape_axes *= a.shape[dimension] - a_along_axis_1d = ( - a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) - ) - if weights is None: - ret = [] - for a_1d in a_along_axis_1d: - ret_1d = np.histogram( - a_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - else: - weights_along_axis_1d = ( - weights.transpose(inverted_shape_dims) - .flatten() - .reshape((-1, shape_axes)) - ) - ret = [] - for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): - ret_1d = np.histogram( - a_1d, - weights=weights_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - out_shape = list(a.shape) - for dimension in sorted(axis, reverse=True): - del out_shape[dimension] - out_shape.insert(0, len(bins) - 1) - ret = np.array(ret) - ret = ret.flatten() - index = np.zeros(len(out_shape), dtype=int) - ret_shaped = np.zeros(out_shape) - dim = 0 - i = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - while list(index) != list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - dim_full_flag = False - while index[dim] == out_shape[dim] - 1: - index[dim] = 0 - dim += 1 - dim_full_flag = True - index[dim] += 1 - i += 1 - if dim_full_flag: - dim = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - ret = ret_shaped - else: - ret = np.histogram( - a=a, bins=bins, range=range, weights=weights, density=density - )[0] - if dtype: - ret = ret.astype(dtype) - bins_out = np.array(bins_out).astype(dtype) - # TODO: weird error when returning bins: return ret, bins_out - return ret - - def igamma( a: np.ndarray, /, @@ -442,89 +528,3 @@ def igamma_cal(a, x): igamma_vec = np.vectorize(igamma_cal) return igamma_vec(a, x).astype(a.dtype) - - -def median( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if out is not None: - out = np.reshape(out, input.shape) - ret = np.median( - input, - axis=axis, - keepdims=keepdims, - out=out, - ) - if input.dtype in [np.uint64, np.int64, np.float64]: - return ret.astype(np.float64) - elif input.dtype in [np.float16]: - return ret.astype(input.dtype) - else: - return ret.astype(np.float32) - - -def nanmean( - a: np.ndarray, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) - - -@with_unsupported_dtypes( - {"1.25.0 and below": ("bfloat16",)}, - backend_version, -) -def nanmedian( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out - ) - - -def quantile( - a: np.ndarray, - q: Union[float, np.ndarray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - interpolation: str = "linear", - out: Optional[np.ndarray] = None, -) -> np.ndarray: - # quantile method in numpy backend, always return an array with dtype=float64. - # in other backends, the output is the same dtype as the input. - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - out=out, - ) - - -median.support_native_out = True -nanmean.support_native_out = True -nanmedian.support_native_out = True -bincount.support_native_out = False -cov.support_native_out = False diff --git a/ivy/functional/backends/numpy/general.py b/ivy/functional/backends/numpy/general.py index f7d07be482b5d..e4c5a8d75ccac 100644 --- a/ivy/functional/backends/numpy/general.py +++ b/ivy/functional/backends/numpy/general.py @@ -17,11 +17,6 @@ from ...ivy.general import _broadcast_to -scatter_flat.support_native_out = True -scatter_nd.support_native_out = True -isin.support_native_out = True - - def array_equal(x0: np.ndarray, x1: np.ndarray, /) -> bool: return np.array_equal(x0, x1) @@ -34,51 +29,64 @@ def current_backend_str() -> str: return "numpy" -def gather( - params: np.ndarray, - indices: np.ndarray, +@_scalar_output_to_0d_array +def get_item( + x: np.ndarray, /, + query: Union[np.ndarray, Tuple], *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[np.ndarray] = None, + copy: bool = None, ) -> np.ndarray: - axis = axis % len(params.shape) - batch_dims = batch_dims % len(params.shape) - ivy.utils.assertions.check_gather_input_valid(params, indices, axis, batch_dims) - result = [] - if batch_dims == 0: - result = np.take(params, indices, axis) + return x.__getitem__(query) + + +@_scalar_output_to_0d_array +def set_item( + x: np.ndarray, + query: Union[np.ndarray, Tuple], + val: np.ndarray, + /, + *, + copy: Optional[bool] = False, +) -> np.ndarray: + if copy: + x = np.copy(x) + x.__setitem__(query, val) + return x + + +def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return x.copy() else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = np.take(p, i, axis - batch_dims) - result.append(r) - result = np.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return _to_device(result) + return x -def gather_nd( +def to_scalar(x: np.ndarray, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + +def to_list(x: np.ndarray, /) -> list: + return x.tolist() + + +def gather( params: np.ndarray, indices: np.ndarray, /, *, + axis: int = -1, batch_dims: int = 0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + axis = axis % len(params.shape) batch_dims = batch_dims % len(params.shape) + ivy.utils.assertions.check_gather_input_valid(params, indices, axis, batch_dims) result = [] if batch_dims == 0: - result = gather_nd_helper(params, indices) + result = np.take(params, indices, axis) else: for b in range(batch_dims): if b == 0: @@ -89,7 +97,7 @@ def gather_nd( ] for z in zip_list: p, i = z - r = gather_nd_helper(p, np.asarray(i, indices.dtype)) + r = np.take(p, i, axis - batch_dims) result.append(r) result = np.array(result) result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) @@ -127,15 +135,34 @@ def gather_nd_helper(params, indices): return res -@_scalar_output_to_0d_array -def get_item( - x: np.ndarray, +def gather_nd( + params: np.ndarray, + indices: np.ndarray, /, - query: Union[np.ndarray, Tuple], *, - copy: bool = None, + batch_dims: int = 0, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - return x.__getitem__(query) + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, np.asarray(i, indices.dtype)) + result.append(r) + result = np.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return _to_device(result) def get_num_dims(x, /, *, as_array=False): @@ -217,27 +244,6 @@ def is_native_array(x, /, *, exclusive=False): return False -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def isin( - elements: np.ndarray, - test_elements: np.ndarray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> np.ndarray: - return np.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - -def itemsize(x: np.ndarray) -> int: - return x.itemsize - - def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -282,6 +288,9 @@ def scatter_flat( return target +scatter_flat.support_native_out = True + + def scatter_nd( indices: np.ndarray, updates: np.ndarray, @@ -323,19 +332,7 @@ def scatter_nd( return _to_device(target) -@_scalar_output_to_0d_array -def set_item( - x: np.ndarray, - query: Union[np.ndarray, Tuple], - val: np.ndarray, - /, - *, - copy: Optional[bool] = False, -) -> np.ndarray: - if copy: - x = np.copy(x) - x.__setitem__(query, val) - return x +scatter_nd.support_native_out = True def shape( @@ -350,23 +347,6 @@ def shape( return ivy.Shape(x.shape) -def to_list(x: np.ndarray, /) -> list: - return x.tolist() - - -def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return x.copy() - else: - return x - - -def to_scalar(x: np.ndarray, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -451,3 +431,27 @@ def _vmap(*args): return res return _vmap + + +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def isin( + elements: np.ndarray, + test_elements: np.ndarray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> np.ndarray: + return np.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + +isin.support_native_out = True + + +def itemsize(x: np.ndarray) -> int: + return x.itemsize diff --git a/ivy/functional/backends/numpy/gradients.py b/ivy/functional/backends/numpy/gradients.py index 73621665e3d54..87be747e053b8 100644 --- a/ivy/functional/backends/numpy/gradients.py +++ b/ivy/functional/backends/numpy/gradients.py @@ -6,6 +6,24 @@ import ivy +def variable(x, /): + logging.warning( + "NumPy does not support autograd, declaring a 'variable' " + "is identical to declaring an 'array' when using numpy backend." + ) + return x + + +def is_variable(x, /, *, exclusive=False): + # NumPy does not support autograd, checking if x is a variable does have any meaning + # for NumPy. Return False. + return False + + +def variable_data(x, /): + return x + + def execute_with_gradients( func, xs, @@ -24,29 +42,23 @@ def execute_with_gradients( return func_ret, None -def grad(func, argnums=0): +def value_and_grad(func): logging.warning( - "NumPy does not support autograd, 'grad' " + "NumPy does not support autograd, 'value_and_grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grad = ivy.nested_map( + grads = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return grad + return y, grads return grad_fn -def is_variable(x, /, *, exclusive=False): - # NumPy does not support autograd, checking if x is a variable does have any meaning - # for NumPy. Return False. - return False - - def jac(func): logging.warning( "NumPy does not support autograd, 'jac' " @@ -62,38 +74,26 @@ def grad_fn(xs): return grad_fn -def stop_gradient(x, /, *, preserve_type=True, out=None): - logging.warning( - "NumPy does not support autograd, 'stop_gradient' " - "has no effect on the array, as gradients are not supported in the first place." - ) - return x - - -def value_and_grad(func): +def grad(func, argnums=0): logging.warning( - "NumPy does not support autograd, 'value_and_grad' " + "NumPy does not support autograd, 'grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grads = ivy.nested_map( + grad = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return y, grads + return grad return grad_fn -def variable(x, /): +def stop_gradient(x, /, *, preserve_type=True, out=None): logging.warning( - "NumPy does not support autograd, declaring a 'variable' " - "is identical to declaring an 'array' when using numpy backend." + "NumPy does not support autograd, 'stop_gradient' " + "has no effect on the array, as gradients are not supported in the first place." ) return x - - -def variable_data(x, /): - return x diff --git a/ivy/functional/backends/numpy/helpers.py b/ivy/functional/backends/numpy/helpers.py index 110cbebe2a53c..b5ae02d3a09f0 100644 --- a/ivy/functional/backends/numpy/helpers.py +++ b/ivy/functional/backends/numpy/helpers.py @@ -1,27 +1,23 @@ -import functools -from typing import Callable -import numpy as np - - -# --- Helpers --- # -# --------------- # - - -def _scalar_output_to_0d_array(function: Callable) -> Callable: - """ - Convert scalar outputs to 0d arrays. - - Sometimes NumPy functions return scalars e.g. `np.add` does when the - inputs are both 0 dimensional. - - We use this wrapper to handle such cases, and convert scalar outputs - to 0d arrays, since the array API standard dictates outputs must be - arrays. - """ - - @functools.wraps(function) - def new_function(*args, **kwargs): - ret = function(*args, **kwargs) - return np.asarray(ret) if np.isscalar(ret) else ret - - return new_function +import functools +from typing import Callable +import numpy as np + + +def _scalar_output_to_0d_array(function: Callable) -> Callable: + """ + Convert scalar outputs to 0d arrays. + + Sometimes NumPy functions return scalars e.g. `np.add` does when the + inputs are both 0 dimensional. + + We use this wrapper to handle such cases, and convert scalar outputs + to 0d arrays, since the array API standard dictates outputs must be + arrays. + """ + + @functools.wraps(function) + def new_function(*args, **kwargs): + ret = function(*args, **kwargs) + return np.asarray(ret) if np.isscalar(ret) else ret + + return new_function diff --git a/ivy/functional/backends/numpy/layers.py b/ivy/functional/backends/numpy/layers.py index ca278076265d0..52ad223304735 100644 --- a/ivy/functional/backends/numpy/layers.py +++ b/ivy/functional/backends/numpy/layers.py @@ -14,10 +14,6 @@ ) -# --- Helpers --- # -# --------------- # - - def _add_dilations(x, dilations, axis, values=0): return np.insert( x, @@ -117,10 +113,6 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters -# --- Main --- # -# ------------ # - - def conv1d( x: np.ndarray, filters: np.ndarray, @@ -288,6 +280,62 @@ def conv2d_transpose( return res +def depthwise_conv2d( + x: np.ndarray, + filters: np.ndarray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[np.ndarray] = None, +): + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + if data_format == "NHWC": + x = np.transpose(x, (3, 0, 1, 2)) + else: + x = np.transpose(x, (1, 0, 2, 3)) + filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters + filters = np.transpose(filters, (2, 0, 1)) + filters = np.expand_dims(filters, (-1, -2)) + filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) + filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) + if isinstance(padding, str): + if padding == "VALID": + out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) + out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) + else: + out_height = np.ceil(float(x.shape[2]) / float(strides[0])) + out_width = np.ceil(float(x.shape[3]) / float(strides[1])) + else: + out_height = np.ceil( + float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) + / float(strides[0]) + ) + out_width = np.ceil( + float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) + / float(strides[1]) + ) + if data_format == "NHWC": + outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) + else: + outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) + x = np.expand_dims(x, -1) + for i in range(x.shape[0]): + output = conv2d( + x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations + ) + if data_format == "NHWC": + outputs = np.append(outputs, output, axis=-1) + else: + outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) + return outputs + + def conv3d( x: np.ndarray, filters: np.ndarray, @@ -502,59 +550,3 @@ def conv_general_transpose( if data_format == "channel_first": return np.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res - - -def depthwise_conv2d( - x: np.ndarray, - filters: np.ndarray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[np.ndarray] = None, -): - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - if data_format == "NHWC": - x = np.transpose(x, (3, 0, 1, 2)) - else: - x = np.transpose(x, (1, 0, 2, 3)) - filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters - filters = np.transpose(filters, (2, 0, 1)) - filters = np.expand_dims(filters, (-1, -2)) - filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) - filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) - if isinstance(padding, str): - if padding == "VALID": - out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) - out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) - else: - out_height = np.ceil(float(x.shape[2]) / float(strides[0])) - out_width = np.ceil(float(x.shape[3]) / float(strides[1])) - else: - out_height = np.ceil( - float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) - / float(strides[0]) - ) - out_width = np.ceil( - float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) - / float(strides[1]) - ) - if data_format == "NHWC": - outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) - else: - outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) - x = np.expand_dims(x, -1) - for i in range(x.shape[0]): - output = conv2d( - x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations - ) - if data_format == "NHWC": - outputs = np.append(outputs, output, axis=-1) - else: - outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) - return outputs diff --git a/ivy/functional/backends/numpy/linear_algebra.py b/ivy/functional/backends/numpy/linear_algebra.py index 8f87923cbdbfc..9f9fcdbc17f1b 100644 --- a/ivy/functional/backends/numpy/linear_algebra.py +++ b/ivy/functional/backends/numpy/linear_algebra.py @@ -15,12 +15,6 @@ from . import backend_version -matmul.support_native_out = True -outer.support_native_out = True -trace.support_native_out = True -vector_to_skew_symmetric_matrix.support_native_out = True - - # Array API Standard # # -------------------# @@ -58,20 +52,6 @@ def det(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.linalg.det(x) -# Extra # -# ----- # - - -def diag( - x: np.ndarray, - /, - *, - k: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.diag(x, k=k) - - def diagonal( x: np.ndarray, /, @@ -84,15 +64,6 @@ def diagonal( return np.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) -@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) -def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] - ) - eigenvalues, eigenvectors = np.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def eigh( x: np.ndarray, /, *, UPLO: str = "L", out: Optional[np.ndarray] = None @@ -166,6 +137,9 @@ def matmul( return ret +matmul.support_native_out = True + + @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("float16", "bfloat16")}, backend_version) def matrix_norm( @@ -240,6 +214,9 @@ def outer( return np.outer(x1, x2, out=out) +outer.support_native_out = True + + @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def pinv( x: np.ndarray, @@ -326,27 +303,27 @@ def svdvals(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray return np.linalg.svd(x, compute_uv=False) -def tensordot( +def tensorsolve( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return np.tensordot(x1, x2, axes=axes) + return np.linalg.tensorsolve(x1, x2, axes=axes) -def tensorsolve( +def tensordot( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + axes: Union[int, Tuple[List[int], List[int]]] = 2, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.linalg.tensorsolve(x1, x2, axes=axes) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return np.tensordot(x1, x2, axes=axes) @_scalar_output_to_0d_array @@ -363,16 +340,7 @@ def trace( return np.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) -@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) -def vander( - x: np.ndarray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.vander(x, N=N, increasing=increasing).astype(x.dtype) +trace.support_native_out = True def vecdot( @@ -387,6 +355,15 @@ def vecdot( return np.tensordot(x1, x2, axes=(axis, axis)) +@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) +def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] + ) + eigenvalues, eigenvectors = np.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + def vector_norm( x: np.ndarray, /, @@ -427,6 +404,32 @@ def vector_norm( return res +# Extra # +# ----- # + + +def diag( + x: np.ndarray, + /, + *, + k: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.diag(x, k=k) + + +@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) +def vander( + x: np.ndarray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.vander(x, N=N, increasing=increasing).astype(x.dtype) + + @with_unsupported_dtypes( { "1.25.2 and below": ( @@ -454,3 +457,6 @@ def vector_to_skew_symmetric_matrix( row3 = np.concatenate((-a2s, a1s, zs), -1) # BS x 3 x 3 return np.concatenate((row1, row2, row3), -2, out=out) + + +vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/numpy/manipulation.py b/ivy/functional/backends/numpy/manipulation.py index 8e614f3d7335a..f258de2b5c59d 100644 --- a/ivy/functional/backends/numpy/manipulation.py +++ b/ivy/functional/backends/numpy/manipulation.py @@ -10,42 +10,10 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x -# --- Main --- # -# ------------ # - - -def as_strided( - x: np.ndarray, - shape: Union[ivy.NativeShape, Sequence[int]], - strides: Sequence[int], - /, -) -> np.ndarray: - return np.lib.stride_tricks.as_strided( - x, - shape=shape, - strides=strides, - ) - - -def clip( - x: np.ndarray, - x_min: Union[Number, np.ndarray], - x_max: Union[Number, np.ndarray], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) - - # Array API Standard # # -------------------# @@ -73,15 +41,7 @@ def concat( return ivy.astype(ret, highest_dtype, copy=False) -def constant_pad( - x: np.ndarray, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) +concat.support_native_out = True def expand_dims( @@ -128,18 +88,6 @@ def permute_dims( return np.transpose(x, axes) -@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) -def repeat( - x: np.ndarray, - /, - repeats: Union[int, List[int]], - *, - axis: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.repeat(x, repeats, axis) - - def reshape( x: np.ndarray, /, @@ -170,6 +118,38 @@ def roll( return np.roll(x, shift, axis) +def squeeze( + x: np.ndarray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + return np.squeeze(x, axis=axis) + + +def stack( + arrays: Union[Tuple[np.ndarray], List[np.ndarray]], + /, + *, + axis: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.stack(arrays, axis, out=out) + + +stack.support_native_out = True + + # Extra # # ------# @@ -211,33 +191,39 @@ def split( return np.split(x, num_or_size_splits, axis) -def squeeze( +@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) +def repeat( x: np.ndarray, /, + repeats: Union[int, List[int]], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, + axis: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - return np.squeeze(x, axis=axis) + return np.repeat(x, repeats, axis) -def stack( - arrays: Union[Tuple[np.ndarray], List[np.ndarray]], +def tile( + x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tile(x, repeats) + + +def constant_pad( + x: np.ndarray, /, + pad_width: List[List[int]], *, - axis: int = 0, + value: Number = 0.0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.stack(arrays, axis, out=out) + return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) + + +def zero_pad( + x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None +): + return np.pad(_flat_array_to_1_dim_array(x), pad_width) def swapaxes( @@ -252,12 +238,6 @@ def swapaxes( return np.swapaxes(x, axis0, axis1) -def tile( - x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tile(x, repeats) - - def unstack( x: np.ndarray, /, @@ -279,12 +259,28 @@ def unstack( return [np.squeeze(item, axis) for item in x_split] -def zero_pad( - x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None -): - return np.pad(_flat_array_to_1_dim_array(x), pad_width) +def clip( + x: np.ndarray, + x_min: Union[Number, np.ndarray], + x_max: Union[Number, np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) -concat.support_native_out = True -stack.support_native_out = True clip.support_native_out = True + + +def as_strided( + x: np.ndarray, + shape: Union[ivy.NativeShape, Sequence[int]], + strides: Sequence[int], + /, +) -> np.ndarray: + return np.lib.stride_tricks.as_strided( + x, + shape=shape, + strides=strides, + ) diff --git a/ivy/functional/backends/numpy/random.py b/ivy/functional/backends/numpy/random.py index 2f8db47cedd99..bba1591aa039a 100644 --- a/ivy/functional/backends/numpy/random.py +++ b/ivy/functional/backends/numpy/random.py @@ -14,6 +14,42 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, np.ndarray] = 0.0, + high: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, + seed: Optional[int] = None, +) -> np.ndarray: + if seed: + np.random.seed(seed) + shape = _check_bounds_and_get_shape(low, high, shape).shape + return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) + + +def random_normal( + *, + mean: Union[float, np.ndarray] = 0.0, + std: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: np.dtype, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + if seed: + np.random.seed(seed) + return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) + @with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) def multinomial( @@ -74,43 +110,6 @@ def randint( return np.random.randint(low, high, shape, dtype=dtype) -def random_normal( - *, - mean: Union[float, np.ndarray] = 0.0, - std: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: np.dtype, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - if seed: - np.random.seed(seed) - return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) - - -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, np.ndarray] = 0.0, - high: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, - seed: Optional[int] = None, -) -> np.ndarray: - if seed: - np.random.seed(seed) - shape = _check_bounds_and_get_shape(low, high, shape).shape - return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) - - def seed(*, seed_value: int = 0) -> None: np.random.seed(seed_value) return diff --git a/ivy/functional/backends/numpy/searching.py b/ivy/functional/backends/numpy/searching.py index 604c06c861b77..c5b2b7194025c 100644 --- a/ivy/functional/backends/numpy/searching.py +++ b/ivy/functional/backends/numpy/searching.py @@ -60,19 +60,6 @@ def argmin( return ret -# Extra # -# ----- # - - -def argwhere( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.argwhere(x) - - def nonzero( x: np.ndarray, /, @@ -108,3 +95,16 @@ def where( ) -> np.ndarray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(np.where(condition, x1, x2), x1.dtype, copy=False) + + +# Extra # +# ----- # + + +def argwhere( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.argwhere(x) diff --git a/ivy/functional/backends/numpy/sorting.py b/ivy/functional/backends/numpy/sorting.py index b3d949062b87a..aa3d84358220b 100644 --- a/ivy/functional/backends/numpy/sorting.py +++ b/ivy/functional/backends/numpy/sorting.py @@ -8,9 +8,6 @@ from . import backend_version -msort.support_native_out = False - - def argsort( x: np.ndarray, /, @@ -28,6 +25,22 @@ def argsort( ) +def sort( + x: np.ndarray, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + kind = "stable" if stable else "quicksort" + ret = np.asarray(np.sort(x, axis=axis, kind=kind)) + if descending: + ret = np.asarray((np.flip(ret, axis))) + return ret + + # msort @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def msort( @@ -36,6 +49,9 @@ def msort( return np.msort(a) +msort.support_native_out = False + + def searchsorted( x: np.ndarray, v: np.ndarray, @@ -73,19 +89,3 @@ def searchsorted( else: ret = np.searchsorted(x, v, side=side, sorter=sorter) return ret.astype(ret_dtype) - - -def sort( - x: np.ndarray, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - kind = "stable" if stable else "quicksort" - ret = np.asarray(np.sort(x, axis=axis, kind=kind)) - if descending: - ret = np.asarray((np.flip(ret, axis))) - return ret diff --git a/ivy/functional/backends/numpy/statistical.py b/ivy/functional/backends/numpy/statistical.py index 248318c8104ac..419ef9355d336 100644 --- a/ivy/functional/backends/numpy/statistical.py +++ b/ivy/functional/backends/numpy/statistical.py @@ -10,100 +10,23 @@ from ivy.utils.einsum_parser import legalise_einsum_expr -# --- Helpers --- # -# --------------- # - - -def _infer_dtype(dtype: np.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype - - -# --- Main --- # -# ------------ # - - -# Extra # -# ------# +# Array API Standard # +# -------------------# -@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) -def cumprod( +def min( x: np.ndarray, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return np.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - return np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumprod(x, -1, dtype=dtype) - return np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - return np.flip(x, axis=axis) - - -def cumsum( - x: np.ndarray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[np.dtype] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - - if exclusive or reverse: - if exclusive and reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - res = np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumsum(x, -1, dtype=dtype) - res = np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - res = np.flip(x, axis=axis) - return res - return np.cumsum(x, axis, dtype=dtype, out=out) + axis = tuple(axis) if isinstance(axis, list) else axis + return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) -@_scalar_output_to_0d_array -def einsum( - equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None -) -> np.ndarray: - equation = legalise_einsum_expr(*[equation, *operands]) - return np.einsum(equation, *operands, out=out) +min.support_native_out = True def max( @@ -118,6 +41,9 @@ def max( return np.asarray(np.amax(a=x, axis=axis, keepdims=keepdims, out=out)) +max.support_native_out = True + + @_scalar_output_to_0d_array def mean( x: np.ndarray, @@ -133,20 +59,14 @@ def mean( ) -# Array API Standard # -# -------------------# +mean.support_native_out = True -def min( - x: np.ndarray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - axis = tuple(axis) if isinstance(axis, list) else axis - return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) +def _infer_dtype(dtype: np.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype def prod( @@ -165,6 +85,9 @@ def prod( return np.asarray(np.prod(a=x, axis=axis, dtype=dtype, keepdims=keepdims, out=out)) +prod.support_native_out = True + + def std( x: np.ndarray, /, @@ -178,6 +101,9 @@ def std( return np.asarray(np.std(x, axis=axis, ddof=correction, keepdims=keepdims, out=out)) +std.support_native_out = True + + def sum( x: np.ndarray, /, @@ -201,6 +127,9 @@ def sum( ) +sum.support_native_out = True + + @_scalar_output_to_0d_array def var( x: np.ndarray, @@ -235,13 +164,94 @@ def var( ) -min.support_native_out = True -max.support_native_out = True -mean.support_native_out = True -prod.support_native_out = True -std.support_native_out = True -sum.support_native_out = True var.support_native_out = True + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) +def cumprod( + x: np.ndarray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return np.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + return np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumprod(x, -1, dtype=dtype) + return np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + return np.flip(x, axis=axis) + + cumprod.support_native_out = True + + +def cumsum( + x: np.ndarray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + + if exclusive or reverse: + if exclusive and reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + res = np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumsum(x, -1, dtype=dtype) + res = np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + res = np.flip(x, axis=axis) + return res + return np.cumsum(x, axis, dtype=dtype, out=out) + + cumsum.support_native_out = True + + +@_scalar_output_to_0d_array +def einsum( + equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None +) -> np.ndarray: + equation = legalise_einsum_expr(*[equation, *operands]) + return np.einsum(equation, *operands, out=out) + + einsum.support_native_out = True diff --git a/ivy/functional/backends/numpy/utility.py b/ivy/functional/backends/numpy/utility.py index 6904438177374..de5597f952940 100644 --- a/ivy/functional/backends/numpy/utility.py +++ b/ivy/functional/backends/numpy/utility.py @@ -6,10 +6,6 @@ import ivy -all.support_native_out = True -any.support_native_out = True - - def all( x: np.ndarray, /, @@ -24,6 +20,9 @@ def all( raise ivy.utils.exceptions.IvyIndexError(error) +all.support_native_out = True + + def any( x: np.ndarray, /, @@ -33,3 +32,6 @@ def any( out: Optional[np.ndarray] = None, ) -> np.ndarray: return np.asarray(np.any(x, axis=axis, keepdims=keepdims, out=out)) + + +any.support_native_out = True diff --git a/ivy/functional/backends/paddle/activations.py b/ivy/functional/backends/paddle/activations.py index 6979b59f13f8b..d090f69a66a4b 100644 --- a/ivy/functional/backends/paddle/activations.py +++ b/ivy/functional/backends/paddle/activations.py @@ -30,38 +30,14 @@ ] -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version -) -def gelu( - x: paddle.Tensor, - /, - *, - approximate: bool = False, - complex_mode="jax", - out: Optional[paddle.Tensor] = None, +def relu( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if paddle.is_complex(x): - sqrt_2_over_pi = 0.7978845608 - # the other magic number comes directly from the formula in - # https://doi.org/10.48550/arXiv.1606.08415 - return ( - 0.5 - * x - * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) - ) if x.dtype in unsupported_dtypes: - return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) - return F.gelu(x, approximate=approximate) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def hardswish( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return F.hardswish(x) + if paddle.is_complex(x): + return paddle.complex(F.relu(x.real()), F.relu(x.imag())) + return F.relu(x.cast("float32")).cast(x.dtype) + return F.relu(x) @with_unsupported_device_and_dtypes( @@ -86,47 +62,28 @@ def leaky_relu( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version + {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version ) -def log_softmax( +def gelu( x: paddle.Tensor, /, *, - axis: Optional[int] = None, + approximate: bool = False, + complex_mode="jax", out: Optional[paddle.Tensor] = None, -): - if axis is None: - axis = -1 - x_max = paddle_backend.max(x, axis=axis, keepdims=True) - x_max = paddle_backend.where( - paddle_backend.isfinite(x_max), - x_max, - paddle.zeros(shape=x_max.shape).astype(x_max.dtype), - ) - exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) - - s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) - ret = paddle_backend.log(s) - ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) - return ret - - -def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) - return F.mish(x.cast("float32")).cast(x.dtype) - return F.mish(x) - - -def relu( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: + if paddle.is_complex(x): + sqrt_2_over_pi = 0.7978845608 + # the other magic number comes directly from the formula in + # https://doi.org/10.48550/arXiv.1606.08415 + return ( + 0.5 + * x + * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) + ) if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return paddle.complex(F.relu(x.real()), F.relu(x.imag())) - return F.relu(x.cast("float32")).cast(x.dtype) - return F.relu(x) + return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) + return F.gelu(x, approximate=approximate) def sigmoid( @@ -185,3 +142,46 @@ def softplus( if threshold is not None: return ivy.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def log_softmax( + x: paddle.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +): + if axis is None: + axis = -1 + x_max = paddle_backend.max(x, axis=axis, keepdims=True) + x_max = paddle_backend.where( + paddle_backend.isfinite(x_max), + x_max, + paddle.zeros(shape=x_max.shape).astype(x_max.dtype), + ) + exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) + + s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) + ret = paddle_backend.log(s) + ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) + return ret + + +def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in unsupported_dtypes: + if paddle.is_complex(x): + return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) + return F.mish(x.cast("float32")).cast(x.dtype) + return F.mish(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def hardswish( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return F.hardswish(x) diff --git a/ivy/functional/backends/paddle/creation.py b/ivy/functional/backends/paddle/creation.py index c216ae2522833..7b01cae902d53 100644 --- a/ivy/functional/backends/paddle/creation.py +++ b/ivy/functional/backends/paddle/creation.py @@ -25,114 +25,6 @@ from . import backend_version from paddle.device import core - -# --- Helpers --- # -# --------------- # - - -def _differentiable_linspace(start, stop, num, *, dtype=None): - start = ivy.to_native(start) - num = paddle.to_tensor(num, stop_gradient=False) - if num == 1: - return paddle_backend.expand_dims(start, axis=0) - n_m_1 = paddle_backend.subtract(num, 1) - increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) - increment_tiled = paddle_backend.repeat(increment, n_m_1) - increments = paddle_backend.multiply( - increment_tiled, - paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), - ) - if isinstance(start, int) or start.ndim == 0: - start = paddle_backend.expand_dims(start, axis=0) - res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) - return res.cast(dtype) - - -def _linspace_helper(start, stop, num, axis=None, *, dtype=None): - num = num.detach().item() if isinstance(num, paddle.Tensor) else num - start_is_array = isinstance(start, paddle.Tensor) - stop_is_array = isinstance(stop, paddle.Tensor) - linspace_method = paddle.linspace - sos_shape = [] - if start_is_array: - start_shape = start.shape - sos_shape = start_shape - if num == 1: - if axis is not None: - return paddle_backend.expand_dims(start, axis=axis) - else: - return paddle_backend.expand_dims(start, axis=-1) - start = start.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not start.stop_gradient else paddle.linspace - ) - if stop_is_array: - stop_shape = stop.shape - sos_shape = stop_shape - if num == 1: - return ( - paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start - ) - stop = stop.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not stop.stop_gradient else paddle.linspace - ) - if start_is_array and stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=-1) - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = paddle_backend.subtract(stop, start) - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [ - linspace_method(strt, stp, num) - for strt, stp in zip( - paddle_backend.unstack(start, keepdims=True), - paddle_backend.unstack(stop, keepdims=True), - ) - ] - elif start_is_array and not stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=axis) - diff = stop - start - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(paddle.ones(start.shape).astype(start.dtype) * stop) - else: - res = [linspace_method(strt, stop, num) for strt in start] - elif not start_is_array and stop_is_array: - if num < stop.shape[0]: - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = stop - start - inc = diff / (num - 1) - res = [paddle.ones(stop.shape).astype(stop.dtype) * start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [linspace_method(start, stp, num) for stp in stop] - else: - return linspace_method(start, stop, num, dtype=dtype) - res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) - if axis is not None: - ndim = res.ndim - perm = list(range(ndim - 1)) - perm.insert(axis % (ndim + 1), ndim - 1) - res = paddle_backend.permute_dims(res, perm) - return res - - -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - -# --- Main --- # -# ------------ # - - # Array API Standard # # -------------------# @@ -216,21 +108,6 @@ def asarray( return paddle.to_tensor(obj, dtype=dtype, place=device) -def copy_array( - x: paddle.Tensor, - *, - to_ivy_array: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if 0 in x.shape: - new_arr = paddle.empty(x.shape, dtype=x.dtype) - else: - new_arr = x.clone() - if to_ivy_array: - return ivy.to_ivy(new_arr) - return new_arr - - def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -321,46 +198,6 @@ def from_dlpack(x, /, *, out: Optional[paddle.Tensor] = None): return paddle.utils.dlpack.from_dlpack(x_d) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def frombuffer( - buffer: bytes, - dtype: Optional[paddle.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> paddle.Tensor: - dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) - if str(dtype) == "bool": - dtype_bytes = 1 - dtype_str = str(dtype) - struct_format = { - "bool": "?", - "int8": "b", - "int16": "h", - "int32": "i", - "int64": "q", - "uint8": "B", - "float16": "e", - "float32": "f", - "float64": "d", - } - ret = [] - for i in range(0, len(buffer), dtype_bytes): - x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) - ret = ret + list(x) - if offset > 0: - offset = int(offset / dtype_bytes) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - ret = paddle.to_tensor(ret, dtype=dtype) - - return ret - - def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -400,6 +237,105 @@ def full_like( ) +def _linspace_helper(start, stop, num, axis=None, *, dtype=None): + num = num.detach().item() if isinstance(num, paddle.Tensor) else num + start_is_array = isinstance(start, paddle.Tensor) + stop_is_array = isinstance(stop, paddle.Tensor) + linspace_method = paddle.linspace + sos_shape = [] + if start_is_array: + start_shape = start.shape + sos_shape = start_shape + if num == 1: + if axis is not None: + return paddle_backend.expand_dims(start, axis=axis) + else: + return paddle_backend.expand_dims(start, axis=-1) + start = start.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not start.stop_gradient else paddle.linspace + ) + if stop_is_array: + stop_shape = stop.shape + sos_shape = stop_shape + if num == 1: + return ( + paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start + ) + stop = stop.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not stop.stop_gradient else paddle.linspace + ) + if start_is_array and stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=-1) + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = paddle_backend.subtract(stop, start) + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [ + linspace_method(strt, stp, num) + for strt, stp in zip( + paddle_backend.unstack(start, keepdims=True), + paddle_backend.unstack(stop, keepdims=True), + ) + ] + elif start_is_array and not stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=axis) + diff = stop - start + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(paddle.ones(start.shape).astype(start.dtype) * stop) + else: + res = [linspace_method(strt, stop, num) for strt in start] + elif not start_is_array and stop_is_array: + if num < stop.shape[0]: + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = stop - start + inc = diff / (num - 1) + res = [paddle.ones(stop.shape).astype(stop.dtype) * start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [linspace_method(start, stp, num) for stp in stop] + else: + return linspace_method(start, stop, num, dtype=dtype) + res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) + if axis is not None: + ndim = res.ndim + perm = list(range(ndim - 1)) + perm.insert(axis % (ndim + 1), ndim - 1) + res = paddle_backend.permute_dims(res, perm) + return res + + +def _differentiable_linspace(start, stop, num, *, dtype=None): + start = ivy.to_native(start) + num = paddle.to_tensor(num, stop_gradient=False) + if num == 1: + return paddle_backend.expand_dims(start, axis=0) + n_m_1 = paddle_backend.subtract(num, 1) + increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) + increment_tiled = paddle_backend.repeat(increment, n_m_1) + increments = paddle_backend.multiply( + increment_tiled, + paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), + ) + if isinstance(start, int) or start.ndim == 0: + start = paddle_backend.expand_dims(start, axis=0) + res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) + return res.cast(dtype) + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint16", "bfloat16", "float16")}}, backend_version ) @@ -505,58 +441,6 @@ def meshgrid( return res -def one_hot( - indices: paddle.Tensor, - depth: int, - /, - *, - on_value: Optional[paddle.Tensor] = None, - off_value: Optional[paddle.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[paddle.dtype] = None, - device: core.Place, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - on_none = on_value is None - off_none = off_value is None - expand_ret = False - if indices.ndim == 0: - expand_ret = True - indices = indices.cast("int64").unsqueeze(0) - if dtype is None: - if on_none and off_none: - dtype = paddle.float32 - else: - if not on_none: - dtype = paddle.to_tensor(on_value).dtype - elif not off_none: - dtype = paddle.to_tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = ( - paddle.to_tensor(1.0, dtype="float32") - if on_none - else paddle.to_tensor(on_value, dtype="float32") - ) - off_value = ( - paddle.to_tensor(0.0, dtype="float32") - if off_none - else paddle.to_tensor(off_value, dtype="float32") - ) - - res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) - - if not on_none or not off_none: - res = paddle.where(res == 1, on_value, off_value) - - if axis is not None: - res = paddle.moveaxis(res, -1, axis) - if expand_ret: - res = res.squeeze() - return res.cast(dtype) - - def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -616,22 +500,6 @@ def triu( return paddle.triu(x=x, diagonal=k) -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: Optional[int] = 0, - /, - *, - device: core.Place, -) -> Tuple[paddle.Tensor]: - # special case due to inconsistent behavior when n_cols=1 and n_rows=0 - if n_cols == 1 and n_rows == 0: - return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( - [], place=device, dtype="int64" - ) - return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) - - def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -658,3 +526,126 @@ def zeros_like( array = asarray + + +def copy_array( + x: paddle.Tensor, + *, + to_ivy_array: Optional[bool] = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if 0 in x.shape: + new_arr = paddle.empty(x.shape, dtype=x.dtype) + else: + new_arr = x.clone() + if to_ivy_array: + return ivy.to_ivy(new_arr) + return new_arr + + +def one_hot( + indices: paddle.Tensor, + depth: int, + /, + *, + on_value: Optional[paddle.Tensor] = None, + off_value: Optional[paddle.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[paddle.dtype] = None, + device: core.Place, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + on_none = on_value is None + off_none = off_value is None + expand_ret = False + if indices.ndim == 0: + expand_ret = True + indices = indices.cast("int64").unsqueeze(0) + if dtype is None: + if on_none and off_none: + dtype = paddle.float32 + else: + if not on_none: + dtype = paddle.to_tensor(on_value).dtype + elif not off_none: + dtype = paddle.to_tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = ( + paddle.to_tensor(1.0, dtype="float32") + if on_none + else paddle.to_tensor(on_value, dtype="float32") + ) + off_value = ( + paddle.to_tensor(0.0, dtype="float32") + if off_none + else paddle.to_tensor(off_value, dtype="float32") + ) + + res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) + + if not on_none or not off_none: + res = paddle.where(res == 1, on_value, off_value) + + if axis is not None: + res = paddle.moveaxis(res, -1, axis) + if expand_ret: + res = res.squeeze() + return res.cast(dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def frombuffer( + buffer: bytes, + dtype: Optional[paddle.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> paddle.Tensor: + dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) + if str(dtype) == "bool": + dtype_bytes = 1 + dtype_str = str(dtype) + struct_format = { + "bool": "?", + "int8": "b", + "int16": "h", + "int32": "i", + "int64": "q", + "uint8": "B", + "float16": "e", + "float32": "f", + "float64": "d", + } + ret = [] + for i in range(0, len(buffer), dtype_bytes): + x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) + ret = ret + list(x) + if offset > 0: + offset = int(offset / dtype_bytes) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + ret = paddle.to_tensor(ret, dtype=dtype) + + return ret + + +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: Optional[int] = 0, + /, + *, + device: core.Place, +) -> Tuple[paddle.Tensor]: + # special case due to inconsistent behavior when n_cols=1 and n_rows=0 + if n_cols == 1 and n_rows == 0: + return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( + [], place=device, dtype="int64" + ) + return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) diff --git a/ivy/functional/backends/paddle/data_type.py b/ivy/functional/backends/paddle/data_type.py index bb4550a92dcc8..1a67457c36079 100644 --- a/ivy/functional/backends/paddle/data_type.py +++ b/ivy/functional/backends/paddle/data_type.py @@ -22,6 +22,7 @@ paddle.complex128: "complex128", paddle.bool: "bool", } + native_dtype_dict = { "int8": paddle.int8, "int16": paddle.int16, @@ -101,51 +102,6 @@ def __repr__(self): ) -# Extra # -# ------# - - -def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: - if dtype_in is int: - return ivy.default_int_dtype() - if dtype_in is float: - return ivy.default_float_dtype() - if dtype_in is complex: - return ivy.default_complex_dtype() - if dtype_in is bool: - return ivy.Dtype("bool") - if isinstance(dtype_in, str): - if dtype_in in native_dtype_dict: - return ivy.Dtype(dtype_in) - else: - raise ivy.utils.exceptions.IvyException( - "Cannot convert to ivy dtype." - f" {dtype_in} is not supported by Paddle backend." - ) - return ivy.Dtype(ivy_dtype_dict[dtype_in]) - - -def as_native_dtype( - dtype_in: Union[paddle.dtype, str, bool, int, float] -) -> paddle.dtype: - if dtype_in is int: - return ivy.default_int_dtype(as_native=True) - if dtype_in is float: - return ivy.default_float_dtype(as_native=True) - if dtype_in is complex: - return ivy.default_complex_dtype(as_native=True) - if dtype_in is bool: - return paddle.bool - if not isinstance(dtype_in, str): - return dtype_in - if dtype_in in native_dtype_dict.keys(): - return native_dtype_dict[ivy.Dtype(dtype_in)] - else: - raise ivy.utils.exceptions.IvyException( - f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." - ) - - # Array API Standard # # -------------------# @@ -216,26 +172,6 @@ def broadcast_to( return paddle.broadcast_to(x, shape) -def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: - if as_native: - return ivy.to_native(x).dtype - return as_ivy_dtype(x.dtype) - - -def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: - dtype_str = as_ivy_dtype(dtype_in) - if "bool" in dtype_str: - return 1 - return int( - dtype_str.replace("paddle.", "") - .replace("uint", "") - .replace("int", "") - .replace("bfloat", "") - .replace("float", "") - .replace("complex", "") - ) - - @_handle_nestable_dtype_info def finfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Finfo: if isinstance(type, paddle.Tensor): @@ -259,6 +195,75 @@ def iinfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Iinfo: return Iinfo(np.iinfo(type)) +def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: + return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) + + +# Extra # +# ------# + + +def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: + if dtype_in is int: + return ivy.default_int_dtype() + if dtype_in is float: + return ivy.default_float_dtype() + if dtype_in is complex: + return ivy.default_complex_dtype() + if dtype_in is bool: + return ivy.Dtype("bool") + if isinstance(dtype_in, str): + if dtype_in in native_dtype_dict: + return ivy.Dtype(dtype_in) + else: + raise ivy.utils.exceptions.IvyException( + "Cannot convert to ivy dtype." + f" {dtype_in} is not supported by Paddle backend." + ) + return ivy.Dtype(ivy_dtype_dict[dtype_in]) + + +def as_native_dtype( + dtype_in: Union[paddle.dtype, str, bool, int, float] +) -> paddle.dtype: + if dtype_in is int: + return ivy.default_int_dtype(as_native=True) + if dtype_in is float: + return ivy.default_float_dtype(as_native=True) + if dtype_in is complex: + return ivy.default_complex_dtype(as_native=True) + if dtype_in is bool: + return paddle.bool + if not isinstance(dtype_in, str): + return dtype_in + if dtype_in in native_dtype_dict.keys(): + return native_dtype_dict[ivy.Dtype(dtype_in)] + else: + raise ivy.utils.exceptions.IvyException( + f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." + ) + + +def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: + if as_native: + return ivy.to_native(x).dtype + return as_ivy_dtype(x.dtype) + + +def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: + dtype_str = as_ivy_dtype(dtype_in) + if "bool" in dtype_str: + return 1 + return int( + dtype_str.replace("paddle.", "") + .replace("uint", "") + .replace("int", "") + .replace("bfloat", "") + .replace("float", "") + .replace("complex", "") + ) + + def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -266,7 +271,3 @@ def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: return True else: return False - - -def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: - return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) diff --git a/ivy/functional/backends/paddle/device.py b/ivy/functional/backends/paddle/device.py index f8295eaaf3d1d..c8e9d9ba385d9 100644 --- a/ivy/functional/backends/paddle/device.py +++ b/ivy/functional/backends/paddle/device.py @@ -13,26 +13,29 @@ from paddle.device import core -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper Paddle profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None +# API # +# ----# - def start(self): - self._start_time = time.perf_counter() - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) +def dev( + x: paddle.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, core.Place]: + return x.place if as_native else as_ivy_dev(x.place) - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def to_device( + x: paddle.Tensor, + device: core.Place, + /, + *, + stream: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + device = as_native_dev(device) + if device.is_cpu_place(): + return x.cpu() + elif device.is_gpu_place(): + return x.cuda(device.gpu_device_id()) def as_ivy_dev(device: core.Place, /): @@ -70,32 +73,31 @@ def as_native_dev( return native_dev -def clear_cached_mem_on_dev(device: str, /): +def clear_mem_on_dev(device: core.Place, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -def clear_mem_on_dev(device: core.Place, /): +def clear_cached_mem_on_dev(device: str, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -# API # -# ----# - - -def dev( - x: paddle.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, core.Place]: - return x.place if as_native else as_ivy_dev(x.place) +def num_gpus() -> int: + return paddle.device.cuda.device_count() def gpu_is_available() -> bool: return bool(paddle.device.cuda.device_count()) +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + return False + + def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( *args, device_shifting_dev=device_shifting_dev, **kwargs @@ -110,25 +112,23 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return ret -def num_gpus() -> int: - return paddle.device.cuda.device_count() +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper Paddle profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + def start(self): + self._start_time = time.perf_counter() -def to_device( - x: paddle.Tensor, - device: core.Place, - /, - *, - stream: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - device = as_native_dev(device) - if device.is_cpu_place(): - return x.cpu() - elif device.is_gpu_place(): - return x.cuda(device.gpu_device_id()) + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + def __enter__(self): + self.start() -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - return False + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/ivy/functional/backends/paddle/elementwise.py b/ivy/functional/backends/paddle/elementwise.py index cbba2fecde68e..f4664e17d527a 100644 --- a/ivy/functional/backends/paddle/elementwise.py +++ b/ivy/functional/backends/paddle/elementwise.py @@ -12,155 +12,187 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - -# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. -def _determine_sqrt_dtype_cast( - dtype: Type[paddle.Tensor], -) -> Tuple[Optional[str], Optional[str]]: - """ - Determine the appropriate casting dtype for sqrt operations. - - Returns: - (intermediate_dtype, output_dtype) - """ - - cast_and_return_float32_dtype = { - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.bool, - } - - if dtype in cast_and_return_float32_dtype: - return "float32", "float32" - elif dtype == paddle.int64: - return "float64", "float64" - elif dtype == paddle.float16: - return "float32", "float16" - elif dtype == paddle.bfloat16: - return "float32", "bfloat16" - else: - return None, None - - def _elementwise_helper(x1, x2): x1, x2 = ivy.promote_types_of_inputs(x1, x2) x1, x2 = paddle_backend.broadcast_arrays(x1, x2) return x1, x2, x1.dtype -# --- Main --- # -# ------------ # - - -def abs( - x: Union[float, paddle.Tensor], +def add( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() - if x.dtype in [ + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [ paddle.int8, - paddle.int16, paddle.uint8, paddle.float16, - paddle.bfloat16, paddle.bool, + paddle.bfloat16, ]: - return paddle.abs(x.astype("float32")).astype(x.dtype) - return paddle.abs(x) + x1, x2 = x1.astype("float32"), x2.astype("float32") + if alpha not in (1, None): + x2 = paddle_backend.multiply(x2, alpha) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.add(x1, x2).astype(ret_dtype) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.acos(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa - s1 = paddle_backend.sqrt(1 - x) - s2 = paddle_backend.sqrt(1 + x) - return paddle.complex( - 2.0 * paddle.atan2(s1.real(), s2.real()), - paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), - ) - return paddle.acos(x) +def bitwise_xor( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_xor(x1, x2) + + +def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: + return paddle.expm1(x) + return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) + + +def bitwise_invert( + x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.bitwise_not(x) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex64", + "complex128", + "bool", + ) + } + }, backend_version, ) -def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.acosh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa - s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) - s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) - return paddle.complex( - paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), - 2.0 * paddle.atan2(s1.imag(), s2.real()), - ) - return paddle.acosh(x) +def isfinite( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.isfinite(x) -def add( +def isinf( + x: paddle.Tensor, + /, + *, + detect_positive: bool = True, + detect_negative: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if detect_negative and detect_positive: + return paddle.isinf(x) + + if detect_negative: + return paddle_backend.equal(x, float("-inf")) + + if detect_positive: + return paddle_backend.equal(x, float("inf")) + + return paddle.zeros(shape=x.shape, dtype=bool) + + +def equal( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ + diff = paddle_backend.subtract(x1, x2) + ret = paddle_backend.logical_and( + paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) + ) + # ret result is sufficient for all cases except where the value is +/-INF of NaN + return paddle_backend.where( + paddle_backend.isnan(diff), + ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), + ret, + ) + + +def less_equal( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + if paddle.is_complex(x1): + real = paddle.less_equal(x1.real(), x2.real()) + imag = paddle.less_equal(x1.imag(), x2.imag()) + return paddle_backend.logical_and(real, imag) + return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) + + return paddle.less_equal(x1, x2) + + +def bitwise_and( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_and(x1, x2) + + +def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ paddle.int8, + paddle.int16, + paddle.int32, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, paddle.bool, - paddle.bfloat16, ]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if alpha not in (1, None): - x2 = paddle_backend.multiply(x2, alpha) - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.add(x1, x2).astype(ret_dtype) + if paddle.is_complex(x): + return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) + return paddle.ceil(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.ceil(x.astype("float64")).astype(x_dtype) + return paddle.ceil(x) -def angle( - input: paddle.Tensor, - /, - *, - deg: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - result = paddle.angle(input) - if deg: - result = paddle.rad2deg(result) - return result +def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) + return paddle.floor(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.floor(x.astype("float64")).astype(x_dtype) + return paddle.floor(x) @with_unsupported_device_and_dtypes( @@ -211,44 +243,86 @@ def asinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, +def sign( + x: paddle.Tensor, + /, + *, + np_variant: Optional[bool] = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, paddle.int32, paddle.int64, paddle.uint8, paddle.float16, + paddle.bfloat16, + paddle.bool, ]: - ret_dtype = x.dtype - return paddle.atan(x.astype("float32")).astype(ret_dtype) - if x.dtype in [paddle.complex64, paddle.complex128]: - atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) - return paddle.atan(x) + return paddle.sgn(x.astype("float32")).astype(x.dtype) + return paddle.sgn(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def atan2( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.atan2(x1, x2).astype(ret_dtype) +# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. +def _determine_sqrt_dtype_cast( + dtype: Type[paddle.Tensor], +) -> Tuple[Optional[str], Optional[str]]: + """ + Determine the appropriate casting dtype for sqrt operations. + + Returns: + (intermediate_dtype, output_dtype) + """ + + cast_and_return_float32_dtype = { + paddle.int8, + paddle.int16, + paddle.int32, + paddle.uint8, + paddle.bool, + } + + if dtype in cast_and_return_float32_dtype: + return "float32", "float32" + elif dtype == paddle.int64: + return "float64", "float64" + elif dtype == paddle.float16: + return "float32", "float16" + elif dtype == paddle.bfloat16: + return "float32", "bfloat16" + else: + return None, None + + +def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + """Calculate the square root with type handling.""" + + if paddle.is_complex(x): + angle = paddle.angle(x) + return paddle.complex( + paddle.cos(angle / 2), paddle.sin(angle / 2) + ) * paddle.sqrt(paddle.abs(x)) + + if x.dtype in {paddle.float32, paddle.float64}: + return paddle.sqrt(x) + + intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) + if intermediate_dtype: + result = paddle.sqrt(x.astype(intermediate_dtype)) + return result.astype(output_dtype) + + raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -258,95 +332,120 @@ def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.float16, ]: ret_dtype = x.dtype - return paddle.atanh(x.astype("float32")).astype(ret_dtype) + return paddle.cosh(x.astype("float32")).astype(ret_dtype) if paddle.is_complex(x): - return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) - return paddle.atanh(x) + re = x.real() + im = x.imag() + return paddle.complex( + paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) + ) + return paddle.cosh(x) -def bitwise_and( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_and(x1, x2) +def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + base = paddle.to_tensor(10.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log10(x.astype("float32")).astype(x.dtype) + return paddle.log10(x) -def bitwise_invert( - x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.bitwise_not(x) +def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + base = paddle.to_tensor(2.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log2(x.astype("float32")).astype(x.dtype) + return paddle.log2(x) -def bitwise_left_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( - ret_dtype - ) +def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle_backend.log(x + 1) + return paddle.log1p(x.astype("float32")).astype(x.dtype) + return paddle.log1p(x) -def bitwise_or( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_or(x1, x2) +def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) + return paddle.isnan(x.astype("float32")) + return paddle.isnan(x) -def bitwise_right_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], +def less( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( - ret_dtype - ) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + real = paddle.less_than(x1.real(), x2.real()) + imag = paddle.less_than(x1.imag(), x2.imag()) + return logical_and(real, imag) + return paddle.less_than(x1.astype("float32"), x2.astype("float32")) + + return paddle.less_than(x1, x2) -def bitwise_xor( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], +def multiply( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_xor(x1, x2) - - -def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) - return paddle.ceil(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.ceil(x.astype("float64")).astype(x_dtype) - return paddle.ceil(x) + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.multiply(x1, x2).astype(ret_dtype) @with_unsupported_device_and_dtypes( @@ -374,36 +473,16 @@ def cos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T return paddle.cos(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.cosh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) - ) - return paddle.cosh(x) - - -def deg2rad( +def logical_not( x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: - return paddle.deg2rad(x.astype("float32")).astype(x.dtype) - return paddle.deg2rad(x) + if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + return paddle.logical_and( + paddle.logical_not(x.real()), paddle.logical_not(x.imag()) + ) + return paddle.logical_not(x.astype("float32")) + return paddle.logical_not(x) def divide( @@ -421,101 +500,6 @@ def divide( return (x1 / x2).astype(ret_dtype) -def equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - diff = paddle_backend.subtract(x1, x2) - ret = paddle_backend.logical_and( - paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) - ) - # ret result is sufficient for all cases except where the value is +/-INF of NaN - return paddle_backend.where( - paddle_backend.isnan(diff), - ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), - ret, - ) - - -# Extra # -# ------# - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - # TODO: add support for complex x, supported in scipy only atm - if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: - return paddle.erf(x.astype("float32")).astype(x.dtype) - return paddle.erf(x) - - -def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.exp(x) - if paddle.is_complex(x): - return paddle.multiply( - paddle.exp(x.real()), - paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), - ) - return paddle_backend.pow(math.e, x).astype(x.dtype) - - -def exp2( - x: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.pow(2, x) - - -def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: - return paddle.expm1(x) - return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) - - -def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) - return paddle.floor(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.floor(x.astype("float64")).astype(x_dtype) - return paddle.floor(x) - - -def floor_divide( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int32, paddle.int64]: - return paddle.floor_divide(x1, x2) - return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) - - @with_supported_dtypes( {"2.5.1 and below": ("float64", "float32", "int64", "int64")}, backend_version, @@ -532,32 +516,6 @@ def fmin( return paddle.fmin(x1, x2) -def fmod( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) - return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version -) -def gcd( - x1: Union[paddle.Tensor, int, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.gcd(x1, x2) - - def greater( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], @@ -595,158 +553,90 @@ def greater_equal( @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, - backend_version, -) -def imag( - val: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.imag(val) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex64", - "complex128", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def isfinite( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.isfinite(x) - - -def isinf( - x: paddle.Tensor, - /, - *, - detect_positive: bool = True, - detect_negative: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if detect_negative and detect_positive: - return paddle.isinf(x) - - if detect_negative: - return paddle_backend.equal(x, float("-inf")) - - if detect_positive: - return paddle_backend.equal(x, float("inf")) - - return paddle.zeros(shape=x.shape, dtype=bool) - - -def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, + paddle.float16, ]: - if paddle.is_complex(x): - return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) - return paddle.isnan(x.astype("float32")) - return paddle.isnan(x) + return paddle.acos(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa + s1 = paddle_backend.sqrt(1 - x) + s2 = paddle_backend.sqrt(1 + x) + return paddle.complex( + 2.0 * paddle.atan2(s1.real(), s2.real()), + paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), + ) + return paddle.acos(x) -def isreal( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def logical_xor( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if paddle.is_complex(x): - return paddle.logical_not(x.imag().astype(bool)) - else: - return paddle.ones_like(x, dtype="bool") + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # with the XOR logic + # if paddle.is_complex(x1): + # return paddle.logical_xor( + # paddle.logical_xor(x1.real(), x2.real()), + # paddle.logical_xor(x1.imag(), x2.imag()), + # ) + return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_xor(x1, x2) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def lcm( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1_dtype = x1.dtype - x2_dtype = x2.dtype - if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): - return paddle.cast( - paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), - paddle.int16, - ) - elif x1_dtype != x2_dtype: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.lcm(x1, x2) - - -def less( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def logical_and( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - real = paddle.less_than(x1.real(), x2.real()) - imag = paddle.less_than(x1.imag(), x2.imag()) - return logical_and(real, imag) - return paddle.less_than(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_than(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # if paddle.is_complex(x1): + # return paddle.logical_and( + # paddle.logical_and(x1.real(), x2.real()), + # paddle.logical_and(x1.imag(), x2.imag()), + # ) + return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_and(x1, x2) -def less_equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def logical_or( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: if paddle.is_complex(x1): - if paddle.is_complex(x1): - real = paddle.less_equal(x1.real(), x2.real()) - imag = paddle.less_equal(x1.imag(), x2.imag()) - return paddle_backend.logical_and(real, imag) - return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_equal(x1, x2) + return paddle.logical_or( + paddle.logical_or(x1.real(), x2.real()), + paddle.logical_or(x1.imag(), x2.imag()), + ) + return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_or(x1, x2) -def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -754,17 +644,24 @@ def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) - return paddle.log(x.astype("float32")).astype(x.dtype) - return paddle.log(x) + return paddle.acosh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa + s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) + s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) + return paddle.complex( + paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), + 2.0 * paddle.atan2(s1.imag(), s2.real()), + ) + return paddle.acosh(x) -def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -772,38 +669,46 @@ def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - base = paddle.to_tensor(10.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log10(x.astype("float32")).astype(x.dtype) - return paddle.log10(x) + return paddle.sin(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) + ) + return paddle.sin(x) -def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle_backend.log(x + 1) - return paddle.log1p(x.astype("float32")).astype(x.dtype) - return paddle.log1p(x) +def negative( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + if x.dtype == paddle.bool: + return paddle.logical_not(x) + return paddle.neg(x) -def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def not_equal( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.logical_not(paddle_backend.equal(x1, x2)) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def tanh( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -811,255 +716,84 @@ def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - base = paddle.to_tensor(2.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log2(x.astype("float32")).astype(x.dtype) - return paddle.log2(x) + return paddle.tanh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + tanh_a = paddle.tanh(paddle.real(x)) + tan_b = paddle.tan(paddle.imag(x)) + return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) + return paddle.tanh(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def logaddexp( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def floor_divide( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - amax = paddle_backend.maximum(x1, x2) - return amax + paddle_backend.log( - paddle_backend.exp(x1 - amax) + paddle_backend.exp(x2 - amax) - ).astype(ret_dtype) + if x1.dtype in [paddle.int32, paddle.int64]: + return paddle.floor_divide(x1, x2) + return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def logaddexp2( - x1: Union[paddle.Tensor, float, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], +def bitwise_or( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.log2(ivy.exp2(x1) + ivy.exp2(x2)) + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_or(x1, x2) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def logical_and( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # if paddle.is_complex(x1): - # return paddle.logical_and( - # paddle.logical_and(x1.real(), x2.real()), - # paddle.logical_and(x1.imag(), x2.imag()), - # ) - return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_and(x1, x2) +def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + ret_dtype = x.dtype + return paddle.sinh(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) + ) + return paddle.sinh(x) -def logical_not( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def positive( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x): - return paddle.logical_and( - paddle.logical_not(x.real()), paddle.logical_not(x.imag()) - ) - return paddle.logical_not(x.astype("float32")) - return paddle.logical_not(x) + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + return x.clone() -def logical_or( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - return paddle.logical_or( - paddle.logical_or(x1.real(), x2.real()), - paddle.logical_or(x1.imag(), x2.imag()), - ) - return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_or(x1, x2) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def logical_xor( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # with the XOR logic - # if paddle.is_complex(x1): - # return paddle.logical_xor( - # paddle.logical_xor(x1.real(), x2.real()), - # paddle.logical_xor(x1.imag(), x2.imag()), - # ) - return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_xor(x1, x2) - - -def maximum( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - use_where: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if use_where: - return paddle_backend.where( - paddle_backend.greater_equal(x1, x2), x1, x2 - ).astype(ret_dtype) - return paddle.maximum(x1, x2).astype(ret_dtype) - - -def minimum( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - use_where: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") - - if use_where: - return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( - ret_dtype - ) - - return paddle.minimum(x1, x2).astype(ret_dtype) - - -def multiply( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.multiply(x1, x2).astype(ret_dtype) - - -def nan_to_num( - x: paddle.Tensor, - /, - *, - copy: Optional[bool] = True, - nan: Optional[Union[float, int]] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[paddle.Tensor] = None, +def square( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - with ivy.ArrayMode(False): - if ivy.is_int_dtype(x): - if posinf is None: - posinf = ivy.iinfo(x).max - if neginf is None: - neginf = ivy.iinfo(x).min - elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): - if posinf is None: - posinf = ivy.finfo(x).max - if neginf is None: - neginf = ivy.finfo(x).min - ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret > 0), - paddle.to_tensor(posinf, dtype=x.dtype), - ret, - ) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret < 0), - paddle.to_tensor(neginf, dtype=x.dtype), - ret, + if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + return paddle.square(x) + if paddle.is_complex(x): + return paddle.complex( + paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), + 2.0 * paddle.real(x) * paddle.imag(x), ) - if copy: - return ret.clone() - else: - x = ret - return x - - -def negative( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - if x.dtype == paddle.bool: - return paddle.logical_not(x) - return paddle.neg(x) - - -def not_equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.logical_not(paddle_backend.equal(x1, x2)) - - -def positive( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - return x.clone() + return paddle_backend.pow(x, 2).astype(x.dtype) @with_unsupported_device_and_dtypes( @@ -1092,72 +826,6 @@ def pow( return paddle.pow(x1, x2) -def rad2deg( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: - return paddle.rad2deg(x.astype("float32")).astype(x.dtype) - return paddle.rad2deg(x) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, - backend_version, -) -def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.real(x) - - -def reciprocal( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return paddle.reciprocal(x) - return paddle_backend.divide(1, x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def remainder( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - modulus: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if not modulus: - res = paddle_backend.divide(x1, x2) - res_floored = paddle_backend.where( - paddle_backend.greater_equal(res, 0.0), - paddle_backend.floor(res), - paddle_backend.ceil(res), - ) - diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) - return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) - - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.remainder(x1, x2).astype(ret_dtype) - - def round( x: paddle.Tensor, /, *, decimals: int = 0, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: @@ -1195,36 +863,146 @@ def _np_round(x, decimals): return _np_round(x, decimals).astype(x.dtype) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def sign( - x: paddle.Tensor, - /, - *, - np_variant: Optional[bool] = True, +def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) + return paddle.trunc(x.astype("float32")).astype(x.dtype) + return paddle.trunc(x) + + +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32")}, + backend_version, +) +def trapz( + y: paddle.Tensor, + /, + *, + x: Optional[paddle.Tensor] = None, + dx: Optional[float] = 1.0, + axis: Optional[int] = -1, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x is None: + d = dx + else: + if x.ndim == 1: + d = paddle.diff(x) + # reshape to correct shape + shape = [1] * y.ndim + shape[axis] = d.shape[0] + d = d.reshape(shape) + else: + d = paddle.diff(x, axis=axis) + + slice1 = [slice(None)] * y.ndim + slice2 = [slice(None)] * y.ndim + + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) + + with ivy.ArrayMode(False): + if y.shape[axis] < 2: + return ivy.zeros_like(ivy.squeeze(y, axis=axis)) + ret = ivy.sum( + ivy.divide( + ivy.multiply( + d, + ivy.add( + ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) + ), + ), + 2.0, + ), + axis=axis, + ) + + return ret + + +def abs( + x: Union[float, paddle.Tensor], + /, + *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() if x.dtype in [ paddle.int8, paddle.int16, - paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, paddle.bfloat16, paddle.bool, ]: - return paddle.sgn(x.astype("float32")).astype(x.dtype) - return paddle.sgn(x) + return paddle.abs(x.astype("float32")).astype(x.dtype) + return paddle.abs(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def logaddexp( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + amax = paddle_backend.maximum(x1, x2) + return amax + paddle_backend.log( + paddle_backend.exp(x1 - amax) + paddle_backend.exp(x2 - amax) + ).astype(ret_dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def logaddexp2( + x1: Union[paddle.Tensor, float, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.log2(ivy.exp2(x1) + ivy.exp2(x2)) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, + backend_version, +) +def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.real(x) @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -1233,21 +1011,19 @@ def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T paddle.uint8, paddle.float16, ]: - return paddle.sin(x.astype("float32")).astype(x.dtype) + ret_dtype = x.dtype + return paddle.tan(x.astype("float32")).astype(ret_dtype) if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) - ) - return paddle.sin(x) + tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) + return paddle.tan(x) @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -1257,47 +1033,63 @@ def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.float16, ]: ret_dtype = x.dtype - return paddle.sinh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) - ) - return paddle.sinh(x) - - -def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - """Calculate the square root with type handling.""" + return paddle.atan(x.astype("float32")).astype(ret_dtype) + if x.dtype in [paddle.complex64, paddle.complex128]: + atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) + return paddle.atan(x) - if paddle.is_complex(x): - angle = paddle.angle(x) - return paddle.complex( - paddle.cos(angle / 2), paddle.sin(angle / 2) - ) * paddle.sqrt(paddle.abs(x)) - if x.dtype in {paddle.float32, paddle.float64}: - return paddle.sqrt(x) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + backend_version, +) +def atan2( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.atan2(x1, x2).astype(ret_dtype) - intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) - if intermediate_dtype: - result = paddle.sqrt(x.astype(intermediate_dtype)) - return result.astype(output_dtype) - raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") +def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) + return paddle.log(x.astype("float32")).astype(x.dtype) + return paddle.log(x) -def square( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: +def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.square(x) + return paddle.exp(x) if paddle.is_complex(x): - return paddle.complex( - paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), - 2.0 * paddle.real(x) * paddle.imag(x), + return paddle.multiply( + paddle.exp(x.real()), + paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), ) - return paddle_backend.pow(x, 2).astype(x.dtype) + return paddle_backend.pow(math.e, x).astype(x.dtype) + + +def exp2( + x: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.pow(2, x) def subtract( @@ -1317,11 +1109,39 @@ def subtract( return paddle.subtract(x1, x2).astype(ret_dtype) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + backend_version, +) +def remainder( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + modulus: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if not modulus: + res = paddle_backend.divide(x1, x2) + res_floored = paddle_backend.where( + paddle_backend.greater_equal(res, 0.0), + paddle_backend.floor(res), + paddle_backend.ceil(res), + ) + diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) + return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) + + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.remainder(x1, x2).astype(ret_dtype) + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -1331,88 +1151,94 @@ def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T paddle.float16, ]: ret_dtype = x.dtype - return paddle.tan(x.astype("float32")).astype(ret_dtype) + return paddle.atanh(x.astype("float32")).astype(ret_dtype) if paddle.is_complex(x): - tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) - return paddle.tan(x) + return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) + return paddle.atanh(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def tanh( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +def bitwise_right_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.tanh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - tanh_a = paddle.tanh(paddle.real(x)) - tan_b = paddle.tan(paddle.imag(x)) - return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) - return paddle.tanh(x) + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( + ret_dtype + ) -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32")}, +def bitwise_left_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( + ret_dtype + ) + + +# Extra # +# ------# + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, backend_version, ) -def trapz( - y: paddle.Tensor, +def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + # TODO: add support for complex x, supported in scipy only atm + if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: + return paddle.erf(x.astype("float32")).astype(x.dtype) + return paddle.erf(x) + + +def minimum( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, - x: Optional[paddle.Tensor] = None, - dx: Optional[float] = 1.0, - axis: Optional[int] = -1, + use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if x is None: - d = dx - else: - if x.ndim == 1: - d = paddle.diff(x) - # reshape to correct shape - shape = [1] * y.ndim - shape[axis] = d.shape[0] - d = d.reshape(shape) + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x1): + use_where = True else: - d = paddle.diff(x, axis=axis) - - slice1 = [slice(None)] * y.ndim - slice2 = [slice(None)] * y.ndim - - slice1[axis] = slice(1, None) - slice2[axis] = slice(None, -1) + x1, x2 = x1.astype("float32"), x2.astype("float32") - with ivy.ArrayMode(False): - if y.shape[axis] < 2: - return ivy.zeros_like(ivy.squeeze(y, axis=axis)) - ret = ivy.sum( - ivy.divide( - ivy.multiply( - d, - ivy.add( - ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) - ), - ), - 2.0, - ), - axis=axis, + if use_where: + return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( + ret_dtype ) - return ret + return paddle.minimum(x1, x2).astype(ret_dtype) -def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def maximum( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + use_where: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [ paddle.int8, paddle.int16, paddle.uint8, @@ -1421,10 +1247,39 @@ def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) - return paddle.trunc(x.astype("float32")).astype(x.dtype) - return paddle.trunc(x) + if paddle.is_complex(x1): + use_where = True + else: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if use_where: + return paddle_backend.where( + paddle_backend.greater_equal(x1, x2), x1, x2 + ).astype(ret_dtype) + return paddle.maximum(x1, x2).astype(ret_dtype) + + +def reciprocal( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return paddle.reciprocal(x) + return paddle_backend.divide(1, x) + + +def deg2rad( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: + return paddle.deg2rad(x.astype("float32")).astype(x.dtype) + return paddle.deg2rad(x) + + +def rad2deg( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: + return paddle.rad2deg(x.astype("float32")).astype(x.dtype) + return paddle.rad2deg(x) def trunc_divide( @@ -1435,3 +1290,140 @@ def trunc_divide( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: return paddle_backend.trunc(paddle_backend.divide(x1, x2)) + + +def isreal( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if paddle.is_complex(x): + return paddle.logical_not(x.imag().astype(bool)) + else: + return paddle.ones_like(x, dtype="bool") + + +def fmod( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) + return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, + backend_version, +) +def lcm( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1_dtype = x1.dtype + x2_dtype = x2.dtype + if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): + return paddle.cast( + paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), + paddle.int16, + ) + elif x1_dtype != x2_dtype: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.lcm(x1, x2) + + +def angle( + input: paddle.Tensor, + /, + *, + deg: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + result = paddle.angle(input) + if deg: + result = paddle.rad2deg(result) + return result + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version +) +def gcd( + x1: Union[paddle.Tensor, int, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.gcd(x1, x2) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, + backend_version, +) +def imag( + val: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.imag(val) + + +def nan_to_num( + x: paddle.Tensor, + /, + *, + copy: Optional[bool] = True, + nan: Optional[Union[float, int]] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + if ivy.is_int_dtype(x): + if posinf is None: + posinf = ivy.iinfo(x).max + if neginf is None: + neginf = ivy.iinfo(x).min + elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): + if posinf is None: + posinf = ivy.finfo(x).max + if neginf is None: + neginf = ivy.finfo(x).min + ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret > 0), + paddle.to_tensor(posinf, dtype=x.dtype), + ret, + ) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret < 0), + paddle.to_tensor(neginf, dtype=x.dtype), + ret, + ) + if copy: + return ret.clone() + else: + x = ret + return x diff --git a/ivy/functional/backends/paddle/experimental/activations.py b/ivy/functional/backends/paddle/experimental/activations.py index 069bf40aeac5b..81bc6bdf25cb7 100644 --- a/ivy/functional/backends/paddle/experimental/activations.py +++ b/ivy/functional/backends/paddle/experimental/activations.py @@ -9,24 +9,6 @@ from . import backend_version -def elu( - x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.elu(x, alpha=alpha) - - if paddle.is_complex(x): - ret = ( - paddle_backend.where( - paddle_backend.greater(x, 0), - x, - paddle_backend.multiply(alpha, paddle_backend.expm1(x)), - ), - ) - return ret - return F.elu(x.cast("float32"), alpha).cast(x.dtype) - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -51,6 +33,28 @@ def logit(x: paddle.Tensor, /, *, eps: Optional[float] = None, out=None): ).cast(x.dtype) +def thresholded_relu( + x: paddle.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = 0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.thresholded_relu(x, threshold=threshold) + return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( + x.dtype + ) + + +def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.relu6(x) + if paddle.is_complex(x): + return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) + return F.relu6(x.cast("float32")).cast(x.dtype) + + def logsigmoid( input: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: @@ -65,14 +69,6 @@ def logsigmoid( return F.log_sigmoid(input.cast("float32")).cast(input.dtype) -def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.relu6(x) - if paddle.is_complex(x): - return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) - return F.relu6(x.cast("float32")).cast(x.dtype) - - def selu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: return F.selu(x) @@ -99,15 +95,19 @@ def silu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. return F.silu(x.cast("float32")).cast(x.dtype) -def thresholded_relu( - x: paddle.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = 0, - out: Optional[paddle.Tensor] = None, +def elu( + x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: - return F.thresholded_relu(x, threshold=threshold) - return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( - x.dtype - ) + return F.elu(x, alpha=alpha) + + if paddle.is_complex(x): + ret = ( + paddle_backend.where( + paddle_backend.greater(x, 0), + x, + paddle_backend.multiply(alpha, paddle_backend.expm1(x)), + ), + ) + return ret + return F.elu(x.cast("float32"), alpha).cast(x.dtype) diff --git a/ivy/functional/backends/paddle/experimental/creation.py b/ivy/functional/backends/paddle/experimental/creation.py index 09ab86a05324d..04d15e026a29d 100644 --- a/ivy/functional/backends/paddle/experimental/creation.py +++ b/ivy/functional/backends/paddle/experimental/creation.py @@ -15,11 +15,6 @@ import ivy from .. import backend_version - -# --- Helpers --- # -# --------------- # - - # noinspection PyProtectedMember # Helpers for calculating Window Functions # ---------------------------------------- @@ -34,28 +29,39 @@ def _kaiser_window(window_length, beta): ) / paddle_backend.i0(beta) -# --- Main --- # -# ------------ # +# Array API Standard # +# -------------------# -def blackman_window( - size: int, - /, +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, *, - periodic: Optional[bool] = True, dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if size < 2: - return paddle.ones([size], dtype=dtype) - if periodic: - count = paddle.arange(size) / size + if window_length < 2: + return paddle.ones([window_length], dtype=dtype) + if periodic is False: + return _kaiser_window(window_length, beta).cast(dtype) else: - count = paddle.linspace(start=0, stop=size, num=size) - return ( - (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) - + (0.08 * paddle.cos(2 * math.pi * 2 * count)) - ).cast(dtype) + return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) + + +def vorbis_window( + window_length: paddle.Tensor, + *, + dtype: Optional[paddle.dtype] = paddle.float32, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if window_length == 0: + return paddle.to_tensor([], dtype=dtype) + i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) + pi = paddle.full(shape=i.shape, fill_value=math.pi) + return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( + dtype + ) def hann_window( @@ -75,26 +81,6 @@ def hann_window( return (0.5 - 0.5 * paddle.cos(2 * math.pi * count)).cast(dtype) -# Array API Standard # -# -------------------# - - -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, - *, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if window_length < 2: - return paddle.ones([window_length], dtype=dtype) - if periodic is False: - return _kaiser_window(window_length, beta).cast(dtype) - else: - return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) - - def tril_indices( n_rows: int, n_cols: Optional[int] = None, @@ -115,32 +101,6 @@ def tril_indices( ) -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex", - ) - } - }, - backend_version, -) -def trilu( - x: paddle.Tensor, - /, - *, - k: int = 0, - upper: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if upper: - return paddle.triu(x=x, diagonal=k) - return paddle.tril(x=x, diagonal=k) - - @with_supported_dtypes( {"2.4.2 and below": ("float64", "float32", "int32", "int64")}, backend_version, @@ -174,6 +134,26 @@ def unsorted_segment_min( return res +def blackman_window( + size: int, + /, + *, + periodic: Optional[bool] = True, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if size < 2: + return paddle.ones([size], dtype=dtype) + if periodic: + count = paddle.arange(size) / size + else: + count = paddle.linspace(start=0, stop=size, num=size) + return ( + (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) + + (0.08 * paddle.cos(2 * math.pi * 2 * count)) + ).cast(dtype) + + def unsorted_segment_sum( data: paddle.Tensor, segment_ids: paddle.Tensor, @@ -208,16 +188,27 @@ def unsorted_segment_sum( return res -def vorbis_window( - window_length: paddle.Tensor, +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex", + ) + } + }, + backend_version, +) +def trilu( + x: paddle.Tensor, + /, *, - dtype: Optional[paddle.dtype] = paddle.float32, + k: int = 0, + upper: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if window_length == 0: - return paddle.to_tensor([], dtype=dtype) - i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) - pi = paddle.full(shape=i.shape, fill_value=math.pi) - return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( - dtype - ) + if upper: + return paddle.triu(x=x, diagonal=k) + return paddle.tril(x=x, diagonal=k) diff --git a/ivy/functional/backends/paddle/experimental/elementwise.py b/ivy/functional/backends/paddle/experimental/elementwise.py index 95b5be21b6aad..8fe58b9e1519a 100644 --- a/ivy/functional/backends/paddle/experimental/elementwise.py +++ b/ivy/functional/backends/paddle/experimental/elementwise.py @@ -17,102 +17,79 @@ from .. import backend_version -_BERNOULLI_COEFS = [ - 12, - -720, - 30240, - -1209600, - 47900160, - -1307674368000 / 691, - 74724249600, - -10670622842880000 / 3617, - 5109094217170944000 / 43867, - -802857662698291200000 / 174611, - 14101100039391805440000 / 77683, - -1693824136731743669452800000 / 236364091, - 186134520519971831808000000 / 657931, - -37893265687455865519472640000000 / 3392780147, - 759790291646040068357842010112000000 / 1723168255201, - -134196726836183700385281186201600000000 / 7709321041217, -] - - -# --- Helpers --- # -# --------------- # - - -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim +@with_supported_dtypes( + {"2.5.1 and below": ("float32", "float64")}, + backend_version, +) +def lgamma( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.lgamma(x) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, + backend_version, +) +def fmax( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x1.dtype != x2.dtype: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.fmax(x1, x2) -def _np_ndim(x): - return ivy.array(x).ndim +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) + return paddle.divide(paddle.sin(y), y) -# --- Main --- # -# ------------ # +def float_power( + x1: Union[paddle.Tensor, float, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1 = paddle.cast(x1, dtype="float64") + x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power + return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def allclose( - x1: paddle.Tensor, - x2: paddle.Tensor, +def frexp( + x: Union[paddle.Tensor, Number], /, *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, out: Optional[paddle.Tensor] = None, -) -> bool: - return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) +) -> paddle.Tensor: + raise IvyNotImplementedException -@with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - "float32", - "float64", - "int32", - "int64", - ) - }, - backend_version, -) -def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.conj(x) +def ldexp( + x1: Union[paddle.Tensor, Number], + x2: Union[paddle.Tensor, Number], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + out_dtype = x1.dtype + x1, x2 = promote_types_of_inputs(x1, x2) + with ivy.ArrayMode(False): + if ivy.any(ivy.less(x2, 0)): + pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 + neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 + ret = ivy.multiply(ivy.pow(2, pos_exp), x1) + ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) + else: + ret = ivy.multiply(ivy.pow(2, x2), x1) + return ivy.astype(ret, out_dtype, copy=False) def copysign( @@ -130,17 +107,38 @@ def copysign( return result -def count_nonzero( - a: paddle.Tensor, +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version +) +def nansum( + x: paddle.Tensor, /, *, - axis: Optional[Union[int, list, tuple]] = None, - keepdims: Optional[bool] = False, + axis: Optional[Union[Tuple[int, ...], int]] = None, dtype: Optional[paddle.dtype] = None, + keepdims: Optional[bool] = False, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) + result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) + if result.shape == [1]: + result = paddle.fluid.layers.squeeze(result, [0]) + return result + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def isclose( + a: paddle.Tensor, + b: paddle.Tensor, + /, + *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) def diff( @@ -169,22 +167,56 @@ def _tensor(val): ) -@with_supported_dtypes( +def signbit( + x: Union[paddle.Tensor, float, int, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.less( + paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 + ) + + +def hypot( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +@with_unsupported_device_and_dtypes( { - "2.5.0 and below": ( - "float32", - "float64", - ) + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } }, backend_version, ) -def digamma( - x: paddle.Tensor, +def allclose( + x1: paddle.Tensor, + x2: paddle.Tensor, /, *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.digamma(x) +) -> bool: + return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) def fix( @@ -197,41 +229,123 @@ def fix( return ivy.trunc(x) -def float_power( - x1: Union[paddle.Tensor, float, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def nextafter( + x1: paddle.Tensor, + x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1 = paddle.cast(x1, dtype="float64") - x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power - return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + with ivy.ArrayMode(False): + eps = ivy.finfo(x1.dtype).eps + return ivy.where( + ivy.equal(x1, x2), + x2, + ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), + ) -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, - backend_version, -) -def fmax( - x1: paddle.Tensor, - x2: paddle.Tensor, +_BERNOULLI_COEFS = [ + 12, + -720, + 30240, + -1209600, + 47900160, + -1307674368000 / 691, + 74724249600, + -10670622842880000 / 3617, + 5109094217170944000 / 43867, + -802857662698291200000 / 174611, + 14101100039391805440000 / 77683, + -1693824136731743669452800000 / 236364091, + 186134520519971831808000000 / 657931, + -37893265687455865519472640000000 / 3392780147, + 759790291646040068357842010112000000 / 1723168255201, + -134196726836183700385281186201600000000 / 7709321041217, +] + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "float16", + "bool", + ) + } + }, + backend_version, +) +def zeta( + x: paddle.Tensor, + q: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if x1.dtype != x2.dtype: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.fmax(x1, x2) + with ivy.ArrayMode(False): + s, a = ivy.promote_types_of_inputs(x, q) + s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) + N = M = ( + paddle.to_tensor(8.0, dtype="float32") + if q.dtype == paddle.float32 + else paddle.to_tensor(8.0, dtype="float64") + ) + assert M <= len(_BERNOULLI_COEFS) + k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) + S = paddle.sum((a_ + k) ** -s_, -1) + Q = ivy.divide((q + N) ** (1 - x), x - 1) + T0 = (q + N) ** -x + m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) + s_over_a = (s_ + m) / (a_ + N) + s_over_a = ivy.where( + s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a + ) + T1 = paddle.cumprod(s_over_a, -1)[..., ::2] + # t=np.array(T1) + T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) + coefs = paddle.unsqueeze( + paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), + tuple(range(a.ndim)), + ) + T1 = T1 / coefs + T = T0 * (0.5 + paddle.sum(T1, -1)) + ans = S + Q + T + mask = x < 1 + ans[mask] = ivy.nan + return ans -def frexp( - x: Union[paddle.Tensor, Number], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim + + +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis + + +def _np_ndim(x): + return ivy.array(x).ndim @with_supported_dtypes( @@ -464,60 +578,47 @@ def gradient( return outvals -def hypot( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, +def xlogy( + x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - raise IvyNotImplementedException() + x, y, ret_dtype = _elementwise_helper(x, y) + with ivy.ArrayMode(False): + x_ok = ivy.not_equal(x, 0.0) + safe_x = ivy.where(x_ok, x, 1.0) + safe_y = ivy.where(x_ok, y, 1.0) + return ivy.where( + x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) + ).cast(ret_dtype) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def isclose( +def count_nonzero( a: paddle.Tensor, - b: paddle.Tensor, - /, - *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -def ldexp( - x1: Union[paddle.Tensor, Number], - x2: Union[paddle.Tensor, Number], /, *, + axis: Optional[Union[int, list, tuple]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - out_dtype = x1.dtype - x1, x2 = promote_types_of_inputs(x1, x2) with ivy.ArrayMode(False): - if ivy.any(ivy.less(x2, 0)): - pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 - neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 - ret = ivy.multiply(ivy.pow(2, pos_exp), x1) - ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) - else: - ret = ivy.multiply(ivy.pow(2, x2), x1) - return ivy.astype(ret, out_dtype, copy=False) + return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) @with_supported_dtypes( - {"2.5.1 and below": ("float32", "float64")}, + { + "2.5.1 and below": ( + "complex64", + "complex128", + "float32", + "float64", + "int32", + "int64", + ) + }, backend_version, ) -def lgamma( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.lgamma(x) +def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.conj(x) def modf( @@ -527,128 +628,19 @@ def modf( return paddle.modf(x, out=out) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version -) -def nansum( - x: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[paddle.dtype] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) - if result.shape == [1]: - result = paddle.fluid.layers.squeeze(result, [0]) - return result - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def nextafter( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - with ivy.ArrayMode(False): - eps = ivy.finfo(x1.dtype).eps - return ivy.where( - ivy.equal(x1, x2), - x2, - ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), - ) - - -def signbit( - x: Union[paddle.Tensor, float, int, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.less( - paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 - ) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) - return paddle.divide(paddle.sin(y), y) - - -def xlogy( - x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x, y, ret_dtype = _elementwise_helper(x, y) - with ivy.ArrayMode(False): - x_ok = ivy.not_equal(x, 0.0) - safe_x = ivy.where(x_ok, x, 1.0) - safe_y = ivy.where(x_ok, y, 1.0) - return ivy.where( - x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) - ).cast(ret_dtype) - - -@with_unsupported_device_and_dtypes( +@with_supported_dtypes( { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "uint16", - "float16", - "bool", - ) - } + "2.5.0 and below": ( + "float32", + "float64", + ) }, backend_version, ) -def zeta( +def digamma( x: paddle.Tensor, - q: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - s, a = ivy.promote_types_of_inputs(x, q) - s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) - N = M = ( - paddle.to_tensor(8.0, dtype="float32") - if q.dtype == paddle.float32 - else paddle.to_tensor(8.0, dtype="float64") - ) - assert M <= len(_BERNOULLI_COEFS) - k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) - S = paddle.sum((a_ + k) ** -s_, -1) - Q = ivy.divide((q + N) ** (1 - x), x - 1) - T0 = (q + N) ** -x - m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) - s_over_a = (s_ + m) / (a_ + N) - s_over_a = ivy.where( - s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a - ) - T1 = paddle.cumprod(s_over_a, -1)[..., ::2] - # t=np.array(T1) - T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) - coefs = paddle.unsqueeze( - paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), - tuple(range(a.ndim)), - ) - T1 = T1 / coefs - T = T0 * (0.5 + paddle.sum(T1, -1)) - ans = S + Q + T - mask = x < 1 - ans[mask] = ivy.nan - return ans + return paddle.digamma(x) diff --git a/ivy/functional/backends/paddle/experimental/layers.py b/ivy/functional/backends/paddle/experimental/layers.py index 4fdba46985ad9..fafd86275beb3 100644 --- a/ivy/functional/backends/paddle/experimental/layers.py +++ b/ivy/functional/backends/paddle/experimental/layers.py @@ -13,11 +13,6 @@ ) from .. import backend_version - -# --- Helpers --- # -# --------------- # - - # local @@ -31,255 +26,6 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -# --- Main --- # -# ------------ # - - -def adaptive_max_pool2d( - input: paddle.Tensor, output_size: Union[Sequence[int], int] -) -> paddle.Tensor: - squeeze = input.ndim == 3 - x = paddle.unsqueeze(input, axis=0) if squeeze else input - ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) - return paddle.squeeze(ret, axis=0) if squeeze else ret - - -def avg_pool1d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool2d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool3d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def dct( - x: paddle.Tensor, - /, - *, - type: Optional[Literal[1, 2, 3, 4]] = 2, - n: Optional[int] = None, - axis: Optional[int] = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout1d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 3 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout2d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 4 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout3d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 5 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -def embedding( - weights: paddle.Tensor, - indices: paddle.Tensor, - /, - *, - max_norm: Optional[int] = None, - out=None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def fft( - x: paddle.Tensor, - dim: int, - /, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if not isinstance(dim, int): - raise IvyValueError(f"Expecting instead of {type(dim)}") - - if n is None: - n = x.shape[dim] - - if dim < -x.ndim or dim >= x.ndim: - raise IvyValueError( - f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" - ) - - if not isinstance(n, int): - raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") - - if n <= 1: - raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") - - valid_norm_modes = ["backward", "ortho", "forward"] - if norm not in valid_norm_modes: - raise IvyValueError( - f"Unrecognized normalization mode {norm}, expecting one of" - f" {valid_norm_modes}" - ) - - if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: - x = x.cast(paddle.complex128) - else: - x = x.cast(paddle.complex64) - - return paddle.fft.fft(x, n, dim, norm=norm) - - -@with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - ) - }, - backend_version, -) -def fft2( - x: paddle.Tensor, - *, - dim: Optional[Union[int, Tuple[int]]] = None, - norm: Optional[str] = "backward", - s: Optional[Union[int, Tuple[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - res = paddle.fft.fft2(x, s, dim, norm) - return res.astype("complex128") - - -def ifft( - x: paddle.Tensor, - dim: int, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def ifftn( - x: paddle.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.fft.ifftn(x, s, axes, norm) - - -def interpolate( - x: paddle.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -): - raise IvyNotImplementedException() - - @with_supported_device_and_dtypes( { "2.5.1 and below": { @@ -491,6 +237,230 @@ def max_pool3d( return res +def avg_pool1d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool2d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, + /, + *, + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool3d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, + /, + *, + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def dct( + x: paddle.Tensor, + /, + *, + type: Optional[Literal[1, 2, 3, 4]] = 2, + n: Optional[int] = None, + axis: Optional[int] = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def fft( + x: paddle.Tensor, + dim: int, + /, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if not isinstance(dim, int): + raise IvyValueError(f"Expecting instead of {type(dim)}") + + if n is None: + n = x.shape[dim] + + if dim < -x.ndim or dim >= x.ndim: + raise IvyValueError( + f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" + ) + + if not isinstance(n, int): + raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") + + if n <= 1: + raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") + + valid_norm_modes = ["backward", "ortho", "forward"] + if norm not in valid_norm_modes: + raise IvyValueError( + f"Unrecognized normalization mode {norm}, expecting one of" + f" {valid_norm_modes}" + ) + + if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: + x = x.cast(paddle.complex128) + else: + x = x.cast(paddle.complex64) + + return paddle.fft.fft(x, n, dim, norm=norm) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout1d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 3 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout2d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 4 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout3d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 5 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +def ifft( + x: paddle.Tensor, + dim: int, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def embedding( + weights: paddle.Tensor, + indices: paddle.Tensor, + /, + *, + max_norm: Optional[int] = None, + out=None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def interpolate( + x: paddle.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +): + raise IvyNotImplementedException() + + +def adaptive_max_pool2d( + input: paddle.Tensor, output_size: Union[Sequence[int], int] +) -> paddle.Tensor: + squeeze = input.ndim == 3 + x = paddle.unsqueeze(input, axis=0) if squeeze else input + ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) + return paddle.squeeze(ret, axis=0) if squeeze else ret + + +def ifftn( + x: paddle.Tensor, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: Optional[str] = "backward", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.fft.ifftn(x, s, axes, norm) + + @with_unsupported_dtypes( {"2.5.1 and below": ("bfloat16", "float16", "complex64", "complex128", "bool")}, backend_version, @@ -505,3 +475,24 @@ def rfftn( ) -> paddle.Tensor: result = paddle.fft.rfftn(x, s, axes, norm) return result.astype("complex128") + + +@with_supported_dtypes( + { + "2.5.1 and below": ( + "complex64", + "complex128", + ) + }, + backend_version, +) +def fft2( + x: paddle.Tensor, + *, + dim: Optional[Union[int, Tuple[int]]] = None, + norm: Optional[str] = "backward", + s: Optional[Union[int, Tuple[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + res = paddle.fft.fft2(x, s, dim, norm) + return res.astype("complex128") diff --git a/ivy/functional/backends/paddle/experimental/linear_algebra.py b/ivy/functional/backends/paddle/experimental/linear_algebra.py index e6f3e99f2bc4a..00f9e583ad2e7 100644 --- a/ivy/functional/backends/paddle/experimental/linear_algebra.py +++ b/ivy/functional/backends/paddle/experimental/linear_algebra.py @@ -12,29 +12,6 @@ from .. import backend_version -dot.support_native_out = True - - -def adjoint( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - _check_valid_dimension_size(x) - return paddle.moveaxis(x, -2, -1).conj() - - -def cond( - x: paddle.Tensor, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[paddle.Tensor] = None, -) -> Any: - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) @@ -68,14 +45,28 @@ def diagflat( )(diag) -def dot( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version +) +def kron( a: paddle.Tensor, b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.dot(a, b, out=out) + return paddle.kron(a, b) + + +def matrix_exp( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # TODO: this is elementwise exp, should be changed to matrix exp ASAP + # return paddle.exp(x) + raise IvyNotImplementedException() def eig( @@ -88,17 +79,24 @@ def eigvals(x: paddle.Tensor, /) -> paddle.Tensor: return paddle.linalg.eig(x)[0] -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version -) -def kron( - a: paddle.Tensor, - b: paddle.Tensor, +def adjoint( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.kron(a, b) + _check_valid_dimension_size(x) + return paddle.moveaxis(x, -2, -1).conj() + + +def cond( + x: paddle.Tensor, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[paddle.Tensor] = None, +) -> Any: + raise IvyNotImplementedException() def lu_factor( @@ -111,15 +109,17 @@ def lu_factor( raise IvyNotImplementedException() -def matrix_exp( - x: paddle.Tensor, +def dot( + a: paddle.Tensor, + b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - # TODO: this is elementwise exp, should be changed to matrix exp ASAP - # return paddle.exp(x) - raise IvyNotImplementedException() + return paddle.dot(a, b, out=out) + + +dot.support_native_out = True @with_supported_device_and_dtypes( diff --git a/ivy/functional/backends/paddle/experimental/losses.py b/ivy/functional/backends/paddle/experimental/losses.py index 54f155697c070..70441eb4655f3 100644 --- a/ivy/functional/backends/paddle/experimental/losses.py +++ b/ivy/functional/backends/paddle/experimental/losses.py @@ -26,21 +26,20 @@ }, backend_version, ) -def huber_loss( +def l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - delta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return paddle.fluid.layers.huber_loss(input, target, delta=delta) + return F.l1_loss(input, target, reduction=reduction) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( - "float16", "int8", "int16", "int32", @@ -54,20 +53,24 @@ def huber_loss( }, backend_version, ) -def l1_loss( +def smooth_l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, + beta: Optional[float] = 1.0, reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return F.l1_loss(input, target, reduction=reduction) + return paddle.nn.functional.smooth_l1_loss( + input, target, reduction=reduction, beta=beta + ) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( + "float16", "int8", "int16", "int32", @@ -81,14 +84,11 @@ def l1_loss( }, backend_version, ) -def smooth_l1_loss( +def huber_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - beta: Optional[float] = 1.0, - reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, ) -> paddle.Tensor: - return paddle.nn.functional.smooth_l1_loss( - input, target, reduction=reduction, beta=beta - ) + return paddle.fluid.layers.huber_loss(input, target, delta=delta) diff --git a/ivy/functional/backends/paddle/experimental/manipulation.py b/ivy/functional/backends/paddle/experimental/manipulation.py index c195b7894d40f..59e781c937591 100644 --- a/ivy/functional/backends/paddle/experimental/manipulation.py +++ b/ivy/functional/backends/paddle/experimental/manipulation.py @@ -44,6 +44,7 @@ -3.04682672343198398683e-1, 6.76795274409476084995e-1, ] + _i0B = [ -7.23318048787475395456e-18, -4.83050448594418207126e-18, @@ -73,174 +74,197 @@ ] -def atleast_1d( - *arys: paddle.Tensor, copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 1: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=0)) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res - - -def atleast_2d( - *arys: paddle.Tensor, copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 2: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res +def moveaxis( + a: paddle.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if isinstance(source, tuple): + source = list(source) + if isinstance(destination, tuple): + source = list(destination) + if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) + return paddle.moveaxis(a, source, destination) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, backend_version, ) -def atleast_3d( - *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim == 0: - result = ary.reshape((1, 1, 1)) - elif ary.ndim == 1: - result = ary[None, :, None] - elif ary.ndim == 2: - result = ary[:, :, None] - else: - result = ary - res.append(result) - if len(res) == 1: - return res[0] - else: - return res - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - def _broadcast_shape(s1, s2): - len_1 = len(s1) - len_2 = len(s2) - if len_1 == 0: - return () if len_2 == 0 else s2 - elif len_1 != 0 and len_2 == 0: - return s1 - else: - return paddle.broadcast_shape(s1, s2) - - if len(shapes) == 0: - raise ValueError("shapes=[] must be non-empty") - elif len(shapes) == 1: - return shapes[0] - result = _broadcast_shape(shapes[0], shapes[1]) - for i in range(2, len(shapes)): - result = _broadcast_shape(result, shapes[i]) - # paddle outputs -1 if the output dimension is 0 - result = [0 if dim == -1 else dim for dim in result] - return tuple(result) +def heaviside( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.heaviside(x1, x2) -def concat_from_sequence( - input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], +def flipud( + m: paddle.Tensor, /, *, - new_axis: int = 0, - axis: int = 0, + copy: Optional[bool] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - if new_axis == 0: - return ivy.concat(input_sequence, axis=axis) - elif new_axis == 1: - return ivy.stack(input_sequence, axis=axis) + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) + return paddle.flip(m, axis=0) -def dsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int16", "float16")}}, + backend_version, +) +def vstack( + arrays: Sequence[paddle.Tensor], /, *, - copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 3: - raise ivy.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=0) + else: + return ivy.stack(arrays, axis=0) -def dstack( +def hstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - arrays = ivy.atleast_2d(*arrays) - if not isinstance(arrays, list): - arrays = [arrays] - if arrays[0].ndim < 3: - return ivy.stack(arrays, axis=-1) + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=1) else: - return ivy.concat(arrays, axis=2) + return ivy.concat(arrays, axis=0) -def expand( - x: paddle.Tensor, - shape: Union[List[int], List[Tuple]], +@with_supported_device_and_dtypes( + { + "2.5.1 and above": { + "cpu": ( + "bool", + "int32", + "int64", + "float32", + "float64", + ), + "gpu": ("float16",), + }, + }, + backend_version, +) +def rot90( + m: paddle.Tensor, /, *, copy: Optional[bool] = None, + k: Optional[int] = 1, + axes: Optional[Tuple[int, int]] = (0, 1), out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle_backend.broadcast_to(x, shape) + if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) + return paddle.rot90(m, k=k, axes=axes) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, ) -def fill_diagonal( - a: paddle.Tensor, - v: Union[int, float], +def top_k( + x: paddle.Tensor, + k: int, /, *, - wrap: bool = False, + axis: int = -1, + largest: Optional[bool] = True, + sorted: bool = True, + out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] + ) + with ivy.ArrayMode(False): + indices = ivy.argsort(x, axis=axis, descending=largest) + indices = paddle.index_select(indices, paddle.arange(end=k), axis) + if not sorted: + indices = paddle.sort(indices, axis=axis) + val = ivy.take_along_axis(x, indices, axis) + return topk_res(val, indices) + + +def fliplr( + m: paddle.Tensor, + /, + *, + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - shape = a.shape - max_end = paddle.prod(paddle.to_tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() - end = max_end if end > max_end else end - a = paddle.reshape(a, (-1,)) - w = paddle.zeros(a.shape, dtype=bool) - ins = paddle.arange(0, max_end) - steps = paddle.arange(0, end, step) + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) + return paddle.flip(m, axis=1) - for i in steps: - i = ins == i - w = paddle.logical_or(w, i) - v = paddle.to_tensor(v, dtype=a.dtype) - a = paddle.where(w, v, a) - a = paddle.reshape(a, shape) - return a + +def i0( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + def _i0_1(x): + return paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), + ) + + def _i0_2(x): + return paddle_backend.divide( + paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl( + paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B + ), + ), + paddle_backend.sqrt(x), + ) + + def _chbevl(x, vals): + b0 = vals[0] + b1 = 0.0 + + for i in range(1, len(vals)): + b2 = b1 + b1 = b0 + b0 = paddle_backend.add( + paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] + ) + return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) + + x = paddle_backend.abs(x) + return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) def flatten( @@ -282,165 +306,105 @@ def _flatten(x, start_dim, end_dim): return _flatten(x, start_dim, end_dim) -def fliplr( - m: paddle.Tensor, +def vsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], /, *, copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) - return paddle.flip(m, axis=1) +) -> List[paddle.Tensor]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) -def flipud( - m: paddle.Tensor, +def dsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], /, *, copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) - return paddle.flip(m, axis=0) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def heaviside( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.heaviside(x1, x2) +) -> List[paddle.Tensor]: + if ary.ndim < 3: + raise ivy.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, +def atleast_1d( + *arys: paddle.Tensor, copy: Optional[bool] = None ) -> List[paddle.Tensor]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 1: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=0)) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res -def hstack( +def dstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=1) + arrays = ivy.atleast_2d(*arrays) + if not isinstance(arrays, list): + arrays = [arrays] + if arrays[0].ndim < 3: + return ivy.stack(arrays, axis=-1) else: - return ivy.concat(arrays, axis=0) - - -def i0( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - def _i0_1(x): - return paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), - ) - - def _i0_2(x): - return paddle_backend.divide( - paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl( - paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B - ), - ), - paddle_backend.sqrt(x), - ) - - def _chbevl(x, vals): - b0 = vals[0] - b1 = 0.0 - - for i in range(1, len(vals)): - b2 = b1 - b1 = b0 - b0 = paddle_backend.add( - paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] - ) - return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) - - x = paddle_backend.abs(x) - return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) + return ivy.concat(arrays, axis=2) -def moveaxis( - a: paddle.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if isinstance(source, tuple): - source = list(source) - if isinstance(destination, tuple): - source = list(destination) - if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) - return paddle.moveaxis(a, source, destination) +def atleast_2d( + *arys: paddle.Tensor, copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 2: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res -@with_supported_device_and_dtypes( - { - "2.5.1 and above": { - "cpu": ( - "bool", - "int32", - "int64", - "float32", - "float64", - ), - "gpu": ("float16",), - }, - }, +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, ) -def rot90( - m: paddle.Tensor, - /, - *, - copy: Optional[bool] = None, - k: Optional[int] = 1, - axes: Optional[Tuple[int, int]] = (0, 1), - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) - return paddle.rot90(m, k=k, axes=axes) +def atleast_3d( + *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim == 0: + result = ary.reshape((1, 1, 1)) + elif ary.ndim == 1: + result = ary[None, :, None] + elif ary.ndim == 2: + result = ary[:, :, None] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return res @with_unsupported_device_and_dtypes( @@ -520,31 +484,65 @@ def take_along_axis( return paddle.take_along_axis(arr, indices, axis) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def top_k( +def hsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[paddle.Tensor]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + def _broadcast_shape(s1, s2): + len_1 = len(s1) + len_2 = len(s2) + if len_1 == 0: + return () if len_2 == 0 else s2 + elif len_1 != 0 and len_2 == 0: + return s1 + else: + return paddle.broadcast_shape(s1, s2) + + if len(shapes) == 0: + raise ValueError("shapes=[] must be non-empty") + elif len(shapes) == 1: + return shapes[0] + result = _broadcast_shape(shapes[0], shapes[1]) + for i in range(2, len(shapes)): + result = _broadcast_shape(result, shapes[i]) + # paddle outputs -1 if the output dimension is 0 + result = [0 if dim == -1 else dim for dim in result] + return tuple(result) + + +def expand( x: paddle.Tensor, - k: int, + shape: Union[List[int], List[Tuple]], /, *, - axis: int = -1, - largest: Optional[bool] = True, - sorted: bool = True, - out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] - ) + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.broadcast_to(x, shape) + + +def concat_from_sequence( + input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: with ivy.ArrayMode(False): - indices = ivy.argsort(x, axis=axis, descending=largest) - indices = paddle.index_select(indices, paddle.arange(end=k), axis) - if not sorted: - indices = paddle.sort(indices, axis=axis) - val = ivy.take_along_axis(x, indices, axis) - return topk_res(val, indices) + if new_axis == 0: + return ivy.concat(input_sequence, axis=axis) + elif new_axis == 1: + return ivy.stack(input_sequence, axis=axis) @with_unsupported_device_and_dtypes( @@ -610,32 +608,35 @@ def unique_consecutive( ) -def vsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], - /, - *, - copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) - - @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int16", "float16")}}, - backend_version, + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) -def vstack( - arrays: Sequence[paddle.Tensor], +def fill_diagonal( + a: paddle.Tensor, + v: Union[int, float], /, *, - out: Optional[paddle.Tensor] = None, + wrap: bool = False, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=0) - else: - return ivy.stack(arrays, axis=0) + shape = a.shape + max_end = paddle.prod(paddle.to_tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() + end = max_end if end > max_end else end + a = paddle.reshape(a, (-1,)) + w = paddle.zeros(a.shape, dtype=bool) + ins = paddle.arange(0, max_end) + steps = paddle.arange(0, end, step) + + for i in steps: + i = ins == i + w = paddle.logical_or(w, i) + v = paddle.to_tensor(v, dtype=a.dtype) + a = paddle.where(w, v, a) + a = paddle.reshape(a, shape) + return a diff --git a/ivy/functional/backends/paddle/experimental/norms.py b/ivy/functional/backends/paddle/experimental/norms.py index 6694cd81aca28..c1fe25a638596 100644 --- a/ivy/functional/backends/paddle/experimental/norms.py +++ b/ivy/functional/backends/paddle/experimental/norms.py @@ -7,13 +7,6 @@ from . import backend_version -batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( - (x.ndim > 1 and x.ndim < 6) - and (scale is not None and scale.ndim == 1) - and (offset is not None and offset.ndim == 1) -) - - # TODO: add support for the rest of the dtypes # use numpy implementation with ivy functions @with_unsupported_device_and_dtypes( @@ -104,27 +97,11 @@ def batch_norm( return xnormalized, runningmean, runningvariance -def instance_norm( - x: paddle.Tensor, - mean: paddle.Tensor, - variance: paddle.Tensor, - /, - *, - scale: Optional[paddle.Tensor] = None, - offset: Optional[paddle.Tensor] = None, - training: Optional[bool] = False, - eps: Optional[float] = 1e-5, - momentum: Optional[float] = 1e-1, - data_format: Optional[str] = "NSC", - out: Optional[ - Tuple[ - paddle.Tensor, - paddle.Tensor, - paddle.Tensor, - ] - ] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: - raise IvyNotImplementedException() +batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( + (x.ndim > 1 and x.ndim < 6) + and (scale is not None and scale.ndim == 1) + and (offset is not None and offset.ndim == 1) +) def l1_normalize( @@ -155,6 +132,29 @@ def l2_normalize( raise IvyNotImplementedException() +def instance_norm( + x: paddle.Tensor, + mean: paddle.Tensor, + variance: paddle.Tensor, + /, + *, + scale: Optional[paddle.Tensor] = None, + offset: Optional[paddle.Tensor] = None, + training: Optional[bool] = False, + eps: Optional[float] = 1e-5, + momentum: Optional[float] = 1e-1, + data_format: Optional[str] = "NSC", + out: Optional[ + Tuple[ + paddle.Tensor, + paddle.Tensor, + paddle.Tensor, + ] + ] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: + raise IvyNotImplementedException() + + def lp_normalize( x: paddle.Tensor, /, *, p: float = 2, axis: int = None, out: paddle.Tensor = None ) -> paddle.Tensor: diff --git a/ivy/functional/backends/paddle/experimental/random.py b/ivy/functional/backends/paddle/experimental/random.py index 1c95f96b39179..542f665e4bb10 100644 --- a/ivy/functional/backends/paddle/experimental/random.py +++ b/ivy/functional/backends/paddle/experimental/random.py @@ -12,70 +12,6 @@ from paddle.device import core from ivy import with_supported_device_and_dtypes - -# bernoulli -@with_supported_device_and_dtypes( - { - "2.5.0 and above": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - }, - "2.4.2 and below": { - "cpu": ( - "float32", - "float64", - ), - "gpu": ("float16", "float32", "float64"), - }, - }, - backend_version, -) -def bernoulli( - probs: Union[float, paddle.Tensor], - *, - logits: Union[float, paddle.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: core.Place, - dtype: paddle.dtype, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - if probs is not None: - probs = probs - elif logits is not None: - probs = ivy.softmax(logits) - probs = paddle.cast(probs, dtype) - probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs - probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) - sample = paddle.bernoulli(probs) - return to_device(sample, device) - - -# beta -def beta( - alpha: Union[float, paddle.Tensor], - beta: Union[float, paddle.Tensor], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - device: core.Place = None, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - shape = _check_bounds_and_get_shape(alpha, beta, shape) - dtype = paddle.float32 if dtype is None else dtype - beta = paddle.cast(beta, alpha.dtype) - dist = paddle.distribution.Beta(alpha, beta) - sample = dist.sample(shape) - sample = paddle.cast(sample, dtype) - return to_device(sample, device) if device is not None else sample - - # dirichlet @@ -115,6 +51,29 @@ def dirichlet( return res +# beta +def beta( + alpha: Union[float, paddle.Tensor], + beta: Union[float, paddle.Tensor], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + device: core.Place = None, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + shape = _check_bounds_and_get_shape(alpha, beta, shape) + dtype = paddle.float32 if dtype is None else dtype + beta = paddle.cast(beta, alpha.dtype) + dist = paddle.distribution.Beta(alpha, beta) + sample = dist.sample(shape) + sample = paddle.cast(sample, dtype) + return to_device(sample, device) if device is not None else sample + + def gamma( alpha: Union[float, paddle.Tensor], beta: Union[float, paddle.Tensor], @@ -140,3 +99,43 @@ def poisson( out: Optional[paddle.Tensor] = None, ): raise IvyNotImplementedException() + + +# bernoulli +@with_supported_device_and_dtypes( + { + "2.5.0 and above": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + }, + "2.4.2 and below": { + "cpu": ( + "float32", + "float64", + ), + "gpu": ("float16", "float32", "float64"), + }, + }, + backend_version, +) +def bernoulli( + probs: Union[float, paddle.Tensor], + *, + logits: Union[float, paddle.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: core.Place, + dtype: paddle.dtype, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + if probs is not None: + probs = probs + elif logits is not None: + probs = ivy.softmax(logits) + probs = paddle.cast(probs, dtype) + probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs + probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) + sample = paddle.bernoulli(probs) + return to_device(sample, device) diff --git a/ivy/functional/backends/paddle/experimental/statistical.py b/ivy/functional/backends/paddle/experimental/statistical.py index 8047578332080..0cc4e40739733 100644 --- a/ivy/functional/backends/paddle/experimental/statistical.py +++ b/ivy/functional/backends/paddle/experimental/statistical.py @@ -11,120 +11,120 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def median( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # keepdims is set to True because in versions up to 2.5.1 + # there was a problem when the axis was defined and it was the + # only axis in the tensor so it needs to be handled manually -def __find_cummax( - x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None -) -> Tuple[paddle.Tensor, paddle.Tensor]: - indices = [] - values = [] - x_dtype = x.dtype if dtype is None else dtype - if ( - isinstance(x.tolist()[0], list) - and len(x[0].shape) >= 1 - and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) - ): - if axis >= 1: - if not isinstance(x, list): - x = x.tolist() - for ret1 in x: - value, indice = __find_cummax( - paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype - ) - indices.append(indice) - values.append(value) + ret_dtype = input.dtype + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(input): + ret = paddle.complex( + paddle.median(input.real(), axis=axis, keepdim=True), + paddle.median(input.imag(), axis=axis, keepdim=True), + ) else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] + ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) else: - if not isinstance(x, list): - x = x.tolist() - n = 0 - for idx, y in enumerate(x): - if x[n] > y: - values.append(x[n]) - elif x[n] <= y or idx == 0: - n = idx - values.append(y) - indices.append(n) + ret = paddle.median(input, axis=axis, keepdim=True) + if not keepdims: + ret = paddle_backend.squeeze(ret, axis=axis) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == input.ndim: + axis = None + if (input.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) - if isinstance(x, paddle.Tensor): - return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( - indices, dtype="int64" - ) + +def nanmean( + a: paddle.Tensor, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + ret_dtype = dtype if dtype is not None else a.dtype + a = a.cast( + ret_dtype + ) # this is necessary to match other FWs behaviour which cast before calculation + if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(a): + ret = paddle.complex( + paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), + paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), + ) + else: + ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) else: - return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") + ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == a.ndim: + axis = None + if (a.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) +def _validate_quantile(q): + if isinstance(q, float): + q = paddle.to_tensor(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False else: - indices.append((lst, tuple(prefix))) - return indices + if not (paddle.all(0 <= q) and paddle.all(q <= 1)): + return False + return True -def _compute_quantile_wrapper( - x, - q, - axis=None, - keepdims=False, - interpolation="linear", -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation not in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - return _handle_axis( - x, - q, - _quantile, - keepdims=keepdims, - axis=axis, - interpolation=interpolation, - ) +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis def _handle_axis(a, q, fn, keepdims=False, axis=None, interpolation="nearest"): @@ -216,39 +216,154 @@ def _quantile(a, q, axis=None, interpolation="nearest"): return out.astype(ret_dtype) -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis +def _compute_quantile_wrapper( + x, + q, + axis=None, + keepdims=False, + interpolation="linear", +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation not in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + return _handle_axis( + x, + q, + _quantile, + keepdims=keepdims, + axis=axis, + interpolation=interpolation, + ) -def _validate_quantile(q): - if isinstance(q, float): - q = paddle.to_tensor(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False - else: - if not (paddle.all(0 <= q) and paddle.all(q <= 1)): - return False - return True +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bfloat16", + "complex64", + "complex128", + ) + } + }, + backend_version, +) +def quantile( + a: paddle.Tensor, + q: Union[paddle.Tensor, float], + /, + *, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: Optional[bool] = False, + interpolation: Optional[str] = "linear", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + x=a, + q=q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) + + +def corrcoef( + x: paddle.Tensor, + /, + *, + y: Optional[paddle.Tensor] = None, + rowvar: Optional[bool] = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def histogram( + a: paddle.Tensor, + /, + *, + bins: Optional[Union[int, paddle.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[paddle.Tensor] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[paddle.Tensor] = None, + density: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> Tuple[paddle.Tensor]: + if range is None: + min_range = 0 + max_range = 0 + else: + min_range = range[0] + max_range = range[1] + return paddle.histogram(a, bins=bins, min=min_range, max=max_range) -# --- Main --- # -# ------------ # +def nanmedian( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + overwrite_input: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if dtype is None: + dtype = input.dtype + input = input.cast("float32") + paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bool", + ) + } + }, + backend_version, +) +def unravel_index( + indices: paddle.Tensor, + shape: Tuple[int], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if indices.ndim == 0: + indices = indices.unsqueeze(0) + coord = [] + indices = indices + for dim in reversed(shape): + coord.append((indices % dim).astype("int32")) + indices = paddle.floor(indices / dim) + + return tuple(reversed(coord)) @with_unsupported_device_and_dtypes( @@ -282,15 +397,34 @@ def bincount( ) -def corrcoef( - x: paddle.Tensor, +def igamma( + a: paddle.Tensor, /, *, - y: Optional[paddle.Tensor] = None, - rowvar: Optional[bool] = True, + x: paddle.Tensor, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - raise IvyNotImplementedException() + results = [] + ret_dtype = a.dtype if out is None else out.dtype + if paddle.float16 in [a.dtype, x.dtype]: + a = a.astype("float32") + x = x.astype("float32") + + for ai, xi in zip(a.flatten(), x.flatten()): + ai = ai.astype("float64") + xi = xi.astype("float64") + + def integrand(t): + return paddle.exp(-t) * paddle.pow(t, ai - 1) + + intervals = paddle.linspace(0, xi, 10001).astype("float64") + interval_width = xi / 10000 + values = integrand(intervals) + integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) + result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) + results.append(result) + + return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) def cov( @@ -402,6 +536,88 @@ def cummax( return ivy.flip(x, axis=axis), ivy.flip(indices, axis=axis) +def __find_cummax( + x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None +) -> Tuple[paddle.Tensor, paddle.Tensor]: + indices = [] + values = [] + x_dtype = x.dtype if dtype is None else dtype + if ( + isinstance(x.tolist()[0], list) + and len(x[0].shape) >= 1 + and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) + ): + if axis >= 1: + if not isinstance(x, list): + x = x.tolist() + for ret1 in x: + value, indice = __find_cummax( + paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype + ) + indices.append(indice) + values.append(value) + else: + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + else: + if not isinstance(x, list): + x = x.tolist() + n = 0 + for idx, y in enumerate(x): + if x[n] > y: + values.append(x[n]) + elif x[n] <= y or idx == 0: + n = idx + values.append(y) + indices.append(n) + + if isinstance(x, paddle.Tensor): + return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( + indices, dtype="int64" + ) + else: + return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16")}}, backend_version, @@ -429,227 +645,3 @@ def cummin( if reverse: cummin_x = paddle.flip(cummin_x, axis=[axis]) return cummin_x.cast(dtype) - - -def histogram( - a: paddle.Tensor, - /, - *, - bins: Optional[Union[int, paddle.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[paddle.Tensor] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[paddle.Tensor] = None, - density: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> Tuple[paddle.Tensor]: - if range is None: - min_range = 0 - max_range = 0 - else: - min_range = range[0] - max_range = range[1] - return paddle.histogram(a, bins=bins, min=min_range, max=max_range) - - -def igamma( - a: paddle.Tensor, - /, - *, - x: paddle.Tensor, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - results = [] - ret_dtype = a.dtype if out is None else out.dtype - if paddle.float16 in [a.dtype, x.dtype]: - a = a.astype("float32") - x = x.astype("float32") - - for ai, xi in zip(a.flatten(), x.flatten()): - ai = ai.astype("float64") - xi = xi.astype("float64") - - def integrand(t): - return paddle.exp(-t) * paddle.pow(t, ai - 1) - - intervals = paddle.linspace(0, xi, 10001).astype("float64") - interval_width = xi / 10000 - values = integrand(intervals) - integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) - result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) - results.append(result) - - return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def median( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # keepdims is set to True because in versions up to 2.5.1 - # there was a problem when the axis was defined and it was the - # only axis in the tensor so it needs to be handled manually - - ret_dtype = input.dtype - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(input): - ret = paddle.complex( - paddle.median(input.real(), axis=axis, keepdim=True), - paddle.median(input.imag(), axis=axis, keepdim=True), - ) - else: - ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) - else: - ret = paddle.median(input, axis=axis, keepdim=True) - if not keepdims: - ret = paddle_backend.squeeze(ret, axis=axis) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == input.ndim: - axis = None - if (input.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) - - -def nanmean( - a: paddle.Tensor, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - ret_dtype = dtype if dtype is not None else a.dtype - a = a.cast( - ret_dtype - ) # this is necessary to match other FWs behaviour which cast before calculation - if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(a): - ret = paddle.complex( - paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), - paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), - ) - else: - ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) - else: - ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) - - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == a.ndim: - axis = None - if (a.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) - - -def nanmedian( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - overwrite_input: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if dtype is None: - dtype = input.dtype - input = input.cast("float32") - paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bfloat16", - "complex64", - "complex128", - ) - } - }, - backend_version, -) -def quantile( - a: paddle.Tensor, - q: Union[paddle.Tensor, float], - /, - *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: Optional[bool] = False, - interpolation: Optional[str] = "linear", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - x=a, - q=q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bool", - ) - } - }, - backend_version, -) -def unravel_index( - indices: paddle.Tensor, - shape: Tuple[int], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if indices.ndim == 0: - indices = indices.unsqueeze(0) - coord = [] - indices = indices - for dim in reversed(shape): - coord.append((indices % dim).astype("int32")) - indices = paddle.floor(indices / dim) - - return tuple(reversed(coord)) diff --git a/ivy/functional/backends/paddle/general.py b/ivy/functional/backends/paddle/general.py index 31ad5024f28f7..f4258c4f452b9 100644 --- a/ivy/functional/backends/paddle/general.py +++ b/ivy/functional/backends/paddle/general.py @@ -13,8 +13,24 @@ from ivy.utils.exceptions import _check_inplace_update_support -# --- Helpers --- # -# --------------- # +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, paddle.Tensor): + if exclusive and not x.stop_gradient: + return False + return True + return False + + +def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: + return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) + + +def container_types(): + return [] + + +def current_backend_str() -> str: + return "paddle" def _check_query(query): @@ -29,20 +45,59 @@ def _check_query(query): ) -# --- Main --- # -# ------------ # +def get_item( + x: paddle.Tensor, + /, + query: Union[paddle.Tensor, Tuple], + *, + copy: bool = None, +) -> paddle.Tensor: + dtype = x.dtype + if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: + ret = x.cast("float32").__getitem__(query).cast(dtype) + elif dtype in [paddle.complex64, paddle.complex128]: + ret = paddle.complex( + x.real().__getitem__(query), + x.imag().__getitem__(query), + ) + else: + ret = x.__getitem__(query) + return ret -def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: - return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) +get_item.partial_mixed_handler = ( + lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape +) -def container_types(): - return [] +def to_numpy( + x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif paddle.is_tensor(x): + if copy: + return np.array(x) + else: + return np.asarray(x) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") -def current_backend_str() -> str: - return "paddle" +def to_scalar(x: paddle.Tensor, /) -> Number: + if isinstance(x, (Number, complex)): + return x + return x.item() + + +def to_list(x: paddle.Tensor, /) -> list: + return x.tolist() def gather( @@ -232,26 +287,6 @@ def gather_nd( return out -def get_item( - x: paddle.Tensor, - /, - query: Union[paddle.Tensor, Tuple], - *, - copy: bool = None, -) -> paddle.Tensor: - dtype = x.dtype - if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: - ret = x.cast("float32").__getitem__(query).cast(dtype) - elif dtype in [paddle.complex64, paddle.complex128]: - ret = paddle.complex( - x.real().__getitem__(query), - x.imag().__getitem__(query), - ) - else: - ret = x.__getitem__(query) - return ret - - def get_num_dims( x: paddle.Tensor, /, *, as_array: bool = False ) -> Union[paddle.Tensor, int]: @@ -321,48 +356,6 @@ def inplace_variables_supported(): return False -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, paddle.Tensor): - if exclusive and not x.stop_gradient: - return False - return True - return False - - -def isin( - elements: paddle.Tensor, - test_elements: paddle.Tensor, - /, - *, - assume_unique: Optional[bool] = False, - invert: Optional[bool] = False, -) -> paddle.Tensor: - input_shape = elements.shape - if elements.ndim == 0: - elements = paddle_backend.expand_dims(elements, axis=0) - if test_elements.ndim == 0: - test_elements = paddle_backend.expand_dims(test_elements, axis=0) - if not assume_unique: - test_elements = paddle_backend.unique_values(test_elements) - - elements = elements.reshape([-1]) - test_elements = test_elements.reshape([-1]) - - output = paddle_backend.any( - paddle_backend.equal( - paddle_backend.expand_dims(elements, axis=-1), test_elements - ), - axis=-1, - ) - return paddle_backend.logical_xor( - paddle_backend.reshape(output, input_shape), invert - ) - - -def itemsize(x: paddle.Tensor) -> int: - return x.element_size() - - def multiprocessing(context=None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -519,36 +512,6 @@ def shape( return ivy.Shape(x.shape) -def to_list(x: paddle.Tensor, /) -> list: - return x.tolist() - - -def to_numpy( - x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif paddle.is_tensor(x): - if copy: - return np.array(x) - else: - return np.asarray(x) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") - - -def to_scalar(x: paddle.Tensor, /) -> Number: - if isinstance(x, (Number, complex)): - return x - return x.item() - - def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -632,6 +595,35 @@ def _vmap(*args, **kwargs): return _vmap -get_item.partial_mixed_handler = ( - lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape -) +def isin( + elements: paddle.Tensor, + test_elements: paddle.Tensor, + /, + *, + assume_unique: Optional[bool] = False, + invert: Optional[bool] = False, +) -> paddle.Tensor: + input_shape = elements.shape + if elements.ndim == 0: + elements = paddle_backend.expand_dims(elements, axis=0) + if test_elements.ndim == 0: + test_elements = paddle_backend.expand_dims(test_elements, axis=0) + if not assume_unique: + test_elements = paddle_backend.unique_values(test_elements) + + elements = elements.reshape([-1]) + test_elements = test_elements.reshape([-1]) + + output = paddle_backend.any( + paddle_backend.equal( + paddle_backend.expand_dims(elements, axis=-1), test_elements + ), + axis=-1, + ) + return paddle_backend.logical_xor( + paddle_backend.reshape(output, input_shape), invert + ) + + +def itemsize(x: paddle.Tensor) -> int: + return x.element_size() diff --git a/ivy/functional/backends/paddle/gradients.py b/ivy/functional/backends/paddle/gradients.py index f7a7fee303be4..9faf1f8d2dca0 100644 --- a/ivy/functional/backends/paddle/gradients.py +++ b/ivy/functional/backends/paddle/gradients.py @@ -20,45 +20,24 @@ ) -# --- Helpers --- # -# --------------- # - - -def _get_jac_one_arg_fn(grad_fn, xs, out_idx): - nested_indices = iter(ivy.all_nested_indices(xs)) - - def one_arg_fn(x): - idx = next(nested_indices) - new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x - ret = grad_fn(new_xs) - for i in out_idx: - ret = ret[i] +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = x.astype(ivy.default_float_dtype()) + if not x.is_leaf: + ret = x.detach() + ret.stop_gradient = False return ret + ret = paddle_backend.copy_array(x) + ret.stop_gradient = False + return ret - return one_arg_fn +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, paddle.Tensor) and not x.stop_gradient -def _get_one_out_fn(grad_fn, xs, fn_ret): - out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) - def one_out_fn(o): - out_idx = next(out_nested_indices) - out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape - one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) - jacobian = ivy.nested_map( - xs, - lambda x: jacobian_to_ivy( - paddle.incubate.autograd.Jacobian( - one_arg_fn, ivy.to_native(x.expand_dims()) - ), - x.shape, - out_shape, - ), - shallow=False, - ) - return jacobian - - return one_out_fn +def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: + return x.value() def _grad_func(y, xs, retain_grads): @@ -124,10 +103,6 @@ def grad_(x): return grads -# --- Main --- # -# ------------ # - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -184,6 +159,100 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = paddle.grad(y, x, allow_unused=True)[0] + grad = grad if grad is not None else paddle.zeros_like(x) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def stop_gradient( + x: Optional[paddle.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[paddle.Tensor] = None, +): + is_var = is_variable(x) + x.stop_gradient = True + if is_var and preserve_type: + return variable(x) + return x + + +def _get_jac_one_arg_fn(grad_fn, xs, out_idx): + nested_indices = iter(ivy.all_nested_indices(xs)) + + def one_arg_fn(x): + idx = next(nested_indices) + new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x + ret = grad_fn(new_xs) + for i in out_idx: + ret = ret[i] + return ret + + return one_arg_fn + + +def _get_one_out_fn(grad_fn, xs, fn_ret): + out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) + + def one_out_fn(o): + out_idx = next(out_nested_indices) + out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape + one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) + jacobian = ivy.nested_map( + xs, + lambda x: jacobian_to_ivy( + paddle.incubate.autograd.Jacobian( + one_arg_fn, ivy.to_native(x.expand_dims()) + ), + x.shape, + out_shape, + ), + shallow=False, + ) + return jacobian + + return one_out_fn + + +def jacobian_to_ivy(jacobian, in_shape, out_shape): + jac_ivy = ivy.to_ivy(jacobian[:]) + jac_shape = out_shape + in_shape + jac_reshaped = jac_ivy.reshape(jac_shape) + return jac_reshaped + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(xs): + fn_ret = grad_fn(xs) + one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) + jacobian = ivy.nested_map(fn_ret, one_out_fn) + return jacobian + + return callback_fn + + def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -221,82 +290,5 @@ def _inner(x): return _nth_derivative(grad.nth) -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, paddle.Tensor) and not x.stop_gradient - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(xs): - fn_ret = grad_fn(xs) - one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) - jacobian = ivy.nested_map(fn_ret, one_out_fn) - return jacobian - - return callback_fn - - -def jacobian_to_ivy(jacobian, in_shape, out_shape): - jac_ivy = ivy.to_ivy(jacobian[:]) - jac_shape = out_shape + in_shape - jac_reshaped = jac_ivy.reshape(jac_shape) - return jac_reshaped - - -def stop_gradient( - x: Optional[paddle.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[paddle.Tensor] = None, -): - is_var = is_variable(x) - x.stop_gradient = True - if is_var and preserve_type: - return variable(x) - return x - - -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = paddle.grad(y, x, allow_unused=True)[0] - grad = grad if grad is not None else paddle.zeros_like(x) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = x.astype(ivy.default_float_dtype()) - if not x.is_leaf: - ret = x.detach() - ret.stop_gradient = False - return ret - ret = paddle_backend.copy_array(x) - ret.stop_gradient = False - return ret - - -def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: - return x.value() - - grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/paddle/layers.py b/ivy/functional/backends/paddle/layers.py index 462d2b96f23e8..f70c374bc70de 100644 --- a/ivy/functional/backends/paddle/layers.py +++ b/ivy/functional/backends/paddle/layers.py @@ -19,8 +19,8 @@ from . import backend_version -# --- Helpers --- # -# --------------- # +def _is_list_or_tuple(inp): + return isinstance(inp, (list, tuple)) def _convert_to_list(value, n, name="padding", _type=int): @@ -37,27 +37,6 @@ def _convert_to_list(value, n, name="padding", _type=int): return value_list -def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): - if filter_format == "channel_first": - filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) - - # adding dilation in input - x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations - for i in range(dims): - if x_dilations[i] > 1: - h = x.shape[1 + i] - new_height = h + (h - 1) * (x_dilations[i] - 1) - h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] - x = paddle_backend.swapaxes(x, 1 + i, -1) - x = paddle.matmul(x, h) - x = paddle_backend.swapaxes(x, -1, 1 + i) - return x, filters - - -def _is_list_or_tuple(inp): - return isinstance(inp, (list, tuple)) - - def _pad_before_conv(x, filters, strides, padding, dims, dilations, data_format): dilations = _convert_to_list(dilations, dims, "dilations") strides = _convert_to_list(strides, dims, "strides") @@ -143,8 +122,21 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding -# --- Main --- # -# ------------ # +def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): + if filter_format == "channel_first": + filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) + + # adding dilation in input + x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations + for i in range(dims): + if x_dilations[i] > 1: + h = x.shape[1 + i] + new_height = h + (h - 1) * (x_dilations[i] - 1) + h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] + x = paddle_backend.swapaxes(x, 1 + i, -1) + x = paddle.matmul(x, h) + x = paddle_backend.swapaxes(x, -1, 1 + i) + return x, filters def conv1d( @@ -267,6 +259,21 @@ def conv2d_transpose( return res +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: paddle.Tensor, + filters: paddle.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: Optional[str] = "NHWC", + dilations: Optional[Union[int, Tuple[int, int]]] = 1, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, @@ -487,18 +494,3 @@ def conv_general_transpose( if data_format == "channel_last": res = res.transpose(0, *range(2, dims + 2), 1) return res - - -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: paddle.Tensor, - filters: paddle.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: Optional[str] = "NHWC", - dilations: Optional[Union[int, Tuple[int, int]]] = 1, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/paddle/linear_algebra.py b/ivy/functional/backends/paddle/linear_algebra.py index 61e2f5744c1dc..8a99ab77110fc 100644 --- a/ivy/functional/backends/paddle/linear_algebra.py +++ b/ivy/functional/backends/paddle/linear_algebra.py @@ -108,33 +108,6 @@ def det(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T return ret -# Extra # -# ----- # - - -def diag( - x: paddle.Tensor, - /, - *, - k: int = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) - ) - return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) - return paddle.diag(x, offset=k) - - def diagonal( x: paddle.Tensor, /, @@ -163,16 +136,6 @@ def diagonal( return paddle.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) -def eig( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> Tuple[paddle.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] - ) - eigenvalues, eigenvectors = paddle.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - def eigh( x: paddle.Tensor, /, @@ -355,6 +318,16 @@ def matrix_norm( return ret +def eig( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> Tuple[paddle.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] + ) + eigenvalues, eigenvectors = paddle.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -450,6 +423,18 @@ def pinv( return paddle.linalg.pinv(x, rcond=rtol) +def tensorsolve( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # Implemented as a composite function in ivy.functional.ivy.linear_algebra + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -575,18 +560,6 @@ def tensordot( return ret.squeeze().cast(ret_dtype) if x1.ndim == axes else ret.cast(ret_dtype) -def tensorsolve( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # Implemented as a composite function in ivy.functional.ivy.linear_algebra - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( { "2.5.1 and below": { @@ -615,28 +588,6 @@ def trace( return ret.squeeze() if x.ndim <= 2 else ret -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, - backend_version, -) -def vander( - x: paddle.Tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - return paddle.pow( - paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), - paddle.arange(start, stop, step, dtype=x.dtype), - ) - - def vecdot( x1: paddle.Tensor, x2: paddle.Tensor, @@ -682,6 +633,55 @@ def vector_norm( ) +# Extra # +# ----- # + + +def diag( + x: paddle.Tensor, + /, + *, + k: int = 0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) + ) + return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) + return paddle.diag(x, offset=k) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, + backend_version, +) +def vander( + x: paddle.Tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + return paddle.pow( + paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), + paddle.arange(start, stop, step, dtype=x.dtype), + ) + + @with_unsupported_dtypes( {"2.5.1 and below": ("unsigned", "int8", "int16", "float16")}, backend_version, diff --git a/ivy/functional/backends/paddle/manipulation.py b/ivy/functional/backends/paddle/manipulation.py index 0902034d2362a..a52336b3c46a0 100644 --- a/ivy/functional/backends/paddle/manipulation.py +++ b/ivy/functional/backends/paddle/manipulation.py @@ -14,33 +14,6 @@ from ...ivy.manipulation import _calculate_out_shape -# --- Helpers --- # -# --------------- # - - -def _reshape_fortran_paddle(x, shape): - if len(x.shape) > 0: - x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) - return paddle_backend.permute_dims( - paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] - ) - - -# --- Main --- # -# ------------ # - - -def clip( - x: paddle.Tensor, - x_min: Union[Number, paddle.Tensor], - x_max: Union[Number, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) - - # Array API Standard # # -------------------# @@ -81,35 +54,6 @@ def concat( return ret -def constant_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - paddings = [] - pad_width = list(pad_width) - for item in pad_width: - if len(item) != 2: - raise ivy.utils.exceptions.IvyException("Length of each item should be 2") - else: - paddings.append(item[0]) - paddings.append(item[1]) - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.nn.functional.pad( - x.cast("float32"), pad=paddings, value=value - ).cast(x.dtype) - return paddle.nn.functional.pad(x=x, pad=paddings, value=value) - - def expand_dims( x: paddle.Tensor, /, @@ -153,50 +97,12 @@ def permute_dims( return paddle.transpose(x, axes) -def repeat( - x: paddle.Tensor, - /, - repeats: Union[int, Iterable[int]], - *, - axis: int = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # handle the case when repeats contains 0 as paddle doesn't support it - if (isinstance(repeats, Number) and repeats == 0) or ( - isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 - ): - if axis is None: - return paddle.to_tensor([], dtype=x.dtype) - else: - shape = x.shape - shape[axis] = 0 - return paddle.zeros(shape=shape).cast(x.dtype) - - if isinstance(repeats, paddle.Tensor) and repeats.size == 1: - repeats = repeats.item() - - if axis is not None: - axis = axis % x.ndim - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), - paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), - ) - - return paddle.repeat_interleave( - x.cast("float32"), repeats=repeats, axis=axis - ).cast(x.dtype) - - return paddle.repeat_interleave(x, repeats=repeats, axis=axis) +def _reshape_fortran_paddle(x, shape): + if len(x.shape) > 0: + x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) + return paddle_backend.permute_dims( + paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] + ) def reshape( @@ -264,62 +170,6 @@ def roll( return paddle.roll(x, shift, axis) -# Extra # -# ------# - - -def split( - x: paddle.Tensor, - /, - *, - copy: Optional[bool] = None, - num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, - axis: Optional[int] = 0, - with_remainder: Optional[bool] = False, -) -> List[paddle.Tensor]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - num_or_size_splits = x.shape[axis] - elif isinstance(num_or_size_splits, paddle.Tensor): - num_or_size_splits = num_or_size_splits.cast("int32") - num_or_size_splits = num_or_size_splits.tolist() - elif isinstance(num_or_size_splits, int): - num_chunks = x.shape[axis] // num_or_size_splits - remainder = x.shape[axis] % num_or_size_splits - if remainder != 0: - if with_remainder: - num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] - else: - raise ivy.utils.exceptions.IvyException( - "Split size is not compatible with input shape" - ) - - if isinstance(num_or_size_splits, (list, tuple)): - if sum(num_or_size_splits) < x.shape[axis]: - num_or_size_splits + type(num_or_size_splits)([-1]) - elif sum(num_or_size_splits) > x.shape[axis]: - raise ivy.utils.exceptions.IvyException( - "total split size is not compatible with input shape," - f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" - ) - - if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x): - imag_list = paddle.split(x.imag(), num_or_size_splits, axis) - real_list = paddle.split(x.real(), num_or_size_splits, axis) - return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] - ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) - return [tensor.cast(x.dtype) for tensor in ret] - return paddle.split(x, num_or_size_splits, axis) - - def squeeze( x: paddle.Tensor, /, @@ -391,18 +241,106 @@ def stack( return paddle.stack(arrays, axis=axis) -def swapaxes( +# Extra # +# ------# + + +def split( x: paddle.Tensor, - axis0: int, - axis1: int, /, *, copy: Optional[bool] = None, + num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, + axis: Optional[int] = 0, + with_remainder: Optional[bool] = False, +) -> List[paddle.Tensor]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + num_or_size_splits = x.shape[axis] + elif isinstance(num_or_size_splits, paddle.Tensor): + num_or_size_splits = num_or_size_splits.cast("int32") + num_or_size_splits = num_or_size_splits.tolist() + elif isinstance(num_or_size_splits, int): + num_chunks = x.shape[axis] // num_or_size_splits + remainder = x.shape[axis] % num_or_size_splits + if remainder != 0: + if with_remainder: + num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] + else: + raise ivy.utils.exceptions.IvyException( + "Split size is not compatible with input shape" + ) + + if isinstance(num_or_size_splits, (list, tuple)): + if sum(num_or_size_splits) < x.shape[axis]: + num_or_size_splits + type(num_or_size_splits)([-1]) + elif sum(num_or_size_splits) > x.shape[axis]: + raise ivy.utils.exceptions.IvyException( + "total split size is not compatible with input shape," + f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" + ) + + if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + imag_list = paddle.split(x.imag(), num_or_size_splits, axis) + real_list = paddle.split(x.real(), num_or_size_splits, axis) + return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] + ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) + return [tensor.cast(x.dtype) for tensor in ret] + return paddle.split(x, num_or_size_splits, axis) + + +def repeat( + x: paddle.Tensor, + /, + repeats: Union[int, Iterable[int]], + *, + axis: int = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - axes = [x for x in range(x.ndim)] - axes[axis0], axes[axis1] = axes[axis1], axes[axis0] - return paddle_backend.permute_dims(x, axes) + # handle the case when repeats contains 0 as paddle doesn't support it + if (isinstance(repeats, Number) and repeats == 0) or ( + isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 + ): + if axis is None: + return paddle.to_tensor([], dtype=x.dtype) + else: + shape = x.shape + shape[axis] = 0 + return paddle.zeros(shape=shape).cast(x.dtype) + + if isinstance(repeats, paddle.Tensor) and repeats.size == 1: + repeats = repeats.item() + + if axis is not None: + axis = axis % x.ndim + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), + paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), + ) + + return paddle.repeat_interleave( + x.cast("float32"), repeats=repeats, axis=axis + ).cast(x.dtype) + + return paddle.repeat_interleave(x, repeats=repeats, axis=axis) def tile( @@ -445,6 +383,70 @@ def tile( return paddle.tile(x, repeats) +def constant_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + paddings = [] + pad_width = list(pad_width) + for item in pad_width: + if len(item) != 2: + raise ivy.utils.exceptions.IvyException("Length of each item should be 2") + else: + paddings.append(item[0]) + paddings.append(item[1]) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.nn.functional.pad( + x.cast("float32"), pad=paddings, value=value + ).cast(x.dtype) + return paddle.nn.functional.pad(x=x, pad=paddings, value=value) + + +def zero_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[paddle.Tensor] = None, +): + return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) + + +def swapaxes( + x: paddle.Tensor, + axis0: int, + axis1: int, + /, + *, + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axes = [x for x in range(x.ndim)] + axes[axis0], axes[axis1] = axes[axis1], axes[axis0] + return paddle_backend.permute_dims(x, axes) + + +def clip( + x: paddle.Tensor, + x_min: Union[Number, paddle.Tensor], + x_max: Union[Number, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) + + def unstack( x: paddle.Tensor, /, @@ -480,13 +482,3 @@ def unstack( if keepdims: return [paddle_backend.expand_dims(r, axis=axis) for r in ret] return ret - - -def zero_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[paddle.Tensor] = None, -): - return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) diff --git a/ivy/functional/backends/paddle/random.py b/ivy/functional/backends/paddle/random.py index 85d02e5b69146..3ca640ccf531b 100644 --- a/ivy/functional/backends/paddle/random.py +++ b/ivy/functional/backends/paddle/random.py @@ -19,69 +19,39 @@ ) from . import backend_version - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "float32", - "float64", - ) - } - }, - backend_version, -) -def multinomial( - population_size: int, - num_samples: int, - /, - *, - batch_size: int = 1, - probs: Optional[paddle.Tensor] = None, - replace: bool = True, - device: Place, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if probs is None: - probs = paddle.ones((batch_size, num_samples)) / population_size - probs = paddle.cast(probs, paddle.float32) - if seed: - paddle.seed(seed) - x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) - return x +# Extra # +# ------# @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def randint( - low: Union[int, paddle.Tensor], - high: Union[int, paddle.Tensor], - /, +def random_uniform( *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + low: Union[float, paddle.Tensor] = 0.0, + high: Union[float, paddle.Tensor] = 1.0, + shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: paddle.dtype, device: Place, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - seed: Optional[int] = None, + seed=None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) - _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - range = high - low + # Set range and seed + rng = high - low if seed: _ = paddle.seed(seed) + random_base = paddle.uniform(shape, min=0.0, max=1.0) - _retval = paddle.cast( - paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype + return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( + dtype ) - return _retval if shape else _retval.squeeze(axis=0) @with_unsupported_device_and_dtypes( @@ -110,39 +80,68 @@ def random_normal( return paddle.normal(mean, std).cast(dtype) -# Extra # -# ------# +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "float32", + "float64", + ) + } + }, + backend_version, +) +def multinomial( + population_size: int, + num_samples: int, + /, + *, + batch_size: int = 1, + probs: Optional[paddle.Tensor] = None, + replace: bool = True, + device: Place, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if probs is None: + probs = paddle.ones((batch_size, num_samples)) / population_size + probs = paddle.cast(probs, paddle.float32) + if seed: + paddle.seed(seed) + x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) + return x @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def random_uniform( +def randint( + low: Union[int, paddle.Tensor], + high: Union[int, paddle.Tensor], + /, *, - low: Union[float, paddle.Tensor] = 0.0, - high: Union[float, paddle.Tensor] = 1.0, - shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: paddle.dtype, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, device: Place, - seed=None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + seed: Optional[int] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) + _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - # Set range and seed - rng = high - low + range = high - low if seed: _ = paddle.seed(seed) - random_base = paddle.uniform(shape, min=0.0, max=1.0) - return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( - dtype + _retval = paddle.cast( + paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype ) + return _retval if shape else _retval.squeeze(axis=0) def seed(*, seed_value: int = 0) -> None: diff --git a/ivy/functional/backends/paddle/searching.py b/ivy/functional/backends/paddle/searching.py index 36fd90f771b40..9519e03c23916 100644 --- a/ivy/functional/backends/paddle/searching.py +++ b/ivy/functional/backends/paddle/searching.py @@ -94,31 +94,6 @@ def argmin( return ret.astype(dtype) -# Extra # -# ----- # - - -def argwhere( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.ndim == 0: - return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") - if x.dtype in [ - paddle.int8, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - ]: - if paddle.is_complex(x): - real_idx = paddle.nonzero(x.real()) - imag_idx = paddle.nonzero(x.imag()) - idx = paddle.concat([real_idx, imag_idx], axis=0) - return paddle.unique(idx, axis=0) - return paddle.nonzero(x.cast("float32")) - return paddle.nonzero(x) - - def nonzero( x: paddle.Tensor, /, @@ -200,3 +175,28 @@ def where( result = paddle.where(condition, x1, x2) return result.squeeze().cast(ret_dtype) if scalar_out else result.cast(ret_dtype) + + +# Extra # +# ----- # + + +def argwhere( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.ndim == 0: + return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") + if x.dtype in [ + paddle.int8, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + ]: + if paddle.is_complex(x): + real_idx = paddle.nonzero(x.real()) + imag_idx = paddle.nonzero(x.imag()) + idx = paddle.concat([real_idx, imag_idx], axis=0) + return paddle.unique(idx, axis=0) + return paddle.nonzero(x.cast("float32")) + return paddle.nonzero(x) diff --git a/ivy/functional/backends/paddle/sorting.py b/ivy/functional/backends/paddle/sorting.py index b12e326d19c8c..6a1d8f38f10d8 100644 --- a/ivy/functional/backends/paddle/sorting.py +++ b/ivy/functional/backends/paddle/sorting.py @@ -33,13 +33,29 @@ def argsort( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def msort( - a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None +def sort( + x: paddle.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.sort(a, axis=0) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( + x.dtype + ) + return paddle.sort(x, axis=axis, descending=descending) @with_unsupported_device_and_dtypes( @@ -99,26 +115,10 @@ def searchsorted( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, backend_version, ) -def sort( - x: paddle.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[paddle.Tensor] = None, +def msort( + a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( - x.dtype - ) - return paddle.sort(x, axis=axis, descending=descending) + return paddle.sort(a, axis=0) diff --git a/ivy/functional/backends/paddle/statistical.py b/ivy/functional/backends/paddle/statistical.py index b03ff9d943fcd..4d0297b1b7aac 100644 --- a/ivy/functional/backends/paddle/statistical.py +++ b/ivy/functional/backends/paddle/statistical.py @@ -1,3 +1,7 @@ +# global + +torch_scatter = None + from typing import Union, Optional, Sequence import paddle @@ -12,163 +16,45 @@ # local from . import backend_version -# global - -torch_scatter = None - - -# --- Helpers --- # -# --------------- # - - -def _std(x, axis, correction, keepdim): - u = paddle_backend.mean(x, axis=axis, keepdims=True) - out = paddle_backend.sum( - paddle_backend.pow(paddle_backend.subtract(x, u), 2), - axis=axis, - keepdims=keepdim, - ) - num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() - num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() - n = num_elm_out / num_elm_in - out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) - if correction: - n = paddle_backend.sqrt( - paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) - ) - out = paddle_backend.multiply(out, n) - return out - - -# --- Main --- # -# ------------ # +# Array API Standard # +# -------------------# -# Extra # -# ----- # -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") - } - }, - backend_version, -) -def cumprod( +def min( x: paddle.Tensor, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[paddle.dtype] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ - paddle.uint8, + ret_dtype = x.dtype + if x.dtype in [ paddle.int8, paddle.int16, - paddle.float16, - ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumprod(x, dim=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumprod(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) - else: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def cumsum( - x: paddle.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ paddle.uint8, - paddle.int8, paddle.float16, + paddle.bfloat16, + paddle.complex64, + paddle.complex128, paddle.bool, ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumsum(x, axis=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumsum(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) + if paddle.is_complex(x): + real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) + imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) + ret = paddle.complex(real, imag) + else: + ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) else: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -def einsum( - equation: str, - *operands: paddle.Tensor, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() + ret = paddle.amin(x, axis=axis, keepdim=keepdims) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == x.ndim: + axis = None + if (x.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) def max( @@ -242,47 +128,6 @@ def mean( return ret.astype(ret_dtype) -# Array API Standard # -# -------------------# - - -def min( - x: paddle.Tensor, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - ret_dtype = x.dtype - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bfloat16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) - imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) - ret = paddle.complex(real, imag) - else: - ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) - else: - ret = paddle.amin(x, axis=axis, keepdim=keepdims) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == x.ndim: - axis = None - if (x.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) - - def prod( x: paddle.Tensor, /, @@ -305,6 +150,25 @@ def prod( return ret +def _std(x, axis, correction, keepdim): + u = paddle_backend.mean(x, axis=axis, keepdims=True) + out = paddle_backend.sum( + paddle_backend.pow(paddle_backend.subtract(x, u), 2), + axis=axis, + keepdims=keepdim, + ) + num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() + num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() + n = num_elm_out / num_elm_in + out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) + if correction: + n = paddle_backend.sqrt( + paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) + ) + out = paddle_backend.multiply(out, n) + return out + + def std( x: paddle.Tensor, /, @@ -353,3 +217,130 @@ def var( ) -> paddle.Tensor: ret = paddle_backend.pow(_std(x, axis, correction, keepdims), 2).cast(x.dtype) return ret + + +# Extra # +# ----- # +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") + } + }, + backend_version, +) +def cumprod( + x: paddle.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ + paddle.uint8, + paddle.int8, + paddle.int16, + paddle.float16, + ]: + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumprod(x, dim=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumprod(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) + else: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def cumsum( + x: paddle.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ + paddle.uint8, + paddle.int8, + paddle.float16, + paddle.bool, + ]: + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumsum(x, axis=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumsum(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) + else: + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +def einsum( + equation: str, + *operands: paddle.Tensor, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/tensorflow/activations.py b/ivy/functional/backends/tensorflow/activations.py index 719d6b8d6cccf..de43f95f9a8f7 100644 --- a/ivy/functional/backends/tensorflow/activations.py +++ b/ivy/functional/backends/tensorflow/activations.py @@ -30,11 +30,6 @@ def gelu( return tf.nn.gelu(x, approximate) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return x * tf.nn.relu6(x + 3) / 6 - - def leaky_relu( x: Tensor, /, @@ -46,23 +41,6 @@ def leaky_relu( return tf.nn.leaky_relu(x, alpha) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def log_softmax( - x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None -): - return tf.nn.log_softmax(x, axis) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def mish( - x: Tensor, - /, - *, - out: Optional[Tensor] = None, -) -> Tensor: - return x * tf.math.tanh(tf.math.softplus(x)) - - def relu(x: Tensor, /, *, complex_mode="jax", out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu(x) @@ -111,3 +89,25 @@ def softplus( if threshold is not None: return tf.where(x_beta > threshold, x, res) return res + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def log_softmax( + x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None +): + return tf.nn.log_softmax(x, axis) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def mish( + x: Tensor, + /, + *, + out: Optional[Tensor] = None, +) -> Tensor: + return x * tf.math.tanh(tf.math.softplus(x)) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return x * tf.nn.relu6(x + 3) / 6 diff --git a/ivy/functional/backends/tensorflow/control_flow_ops.py b/ivy/functional/backends/tensorflow/control_flow_ops.py index 1f600eebbe861..9b6f42c456d75 100644 --- a/ivy/functional/backends/tensorflow/control_flow_ops.py +++ b/ivy/functional/backends/tensorflow/control_flow_ops.py @@ -1,20 +1,42 @@ import tensorflow as tf +# def if_exp(cond, if_true, if_false, expr_repr): +# def true_fn(): +# return if_true() +# +# def false_fn(): +# return if_false() +# +# return tf.cond(cond, true_fn, false_fn) -# --- Helpers --- # -# --------------- # +def if_else(cond, body_fn, orelse_fn, vars): + # back-compatibility + if isinstance(cond, bool): + v = cond + cond = lambda *_: v + cond = bool(cond(*vars)) + # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) + # use pythonic placeholder until the graph compiler supports callable arguments + if cond: + return body_fn(*vars) + else: + return orelse_fn(*vars) -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} +def while_loop(test_fn, body_fn, vars): + def body_fn_wrapper(*loop_vars): + return body_fn(*loop_vars) + + def test_fn_wrapper(*loop_vars): + return test_fn(*loop_vars) -# --- Main --- # -# ------------ # + if not vars: + vars = (0,) + + return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) def for_loop( @@ -48,40 +70,9 @@ def empty_function(*args): return _dict_to_tuple(vars_dict) -# def if_exp(cond, if_true, if_false, expr_repr): -# def true_fn(): -# return if_true() -# -# def false_fn(): -# return if_false() -# -# return tf.cond(cond, true_fn, false_fn) - - -def if_else(cond, body_fn, orelse_fn, vars): - # back-compatibility - if isinstance(cond, bool): - v = cond - cond = lambda *_: v - cond = bool(cond(*vars)) - # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) - - # use pythonic placeholder until the graph compiler supports callable arguments - - if cond: - return body_fn(*vars) - else: - return orelse_fn(*vars) - - -def while_loop(test_fn, body_fn, vars): - def body_fn_wrapper(*loop_vars): - return body_fn(*loop_vars) - - def test_fn_wrapper(*loop_vars): - return test_fn(*loop_vars) +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} - if not vars: - vars = (0,) - return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) diff --git a/ivy/functional/backends/tensorflow/creation.py b/ivy/functional/backends/tensorflow/creation.py index 9db7dc926d4a9..25ce085eb766d 100644 --- a/ivy/functional/backends/tensorflow/creation.py +++ b/ivy/functional/backends/tensorflow/creation.py @@ -20,18 +20,6 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - -# --- Main --- # -# ------------ # - - # Array API Standard # # -------------------# @@ -109,17 +97,6 @@ def asarray( return tf.identity(ret) if copy else ret -def copy_array( - x: Union[tf.Tensor, tf.Variable], - *, - to_ivy_array: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if to_ivy_array: - return ivy.to_ivy(tf.identity(x)) - return tf.identity(x) - - def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -204,27 +181,6 @@ def from_dlpack( return tf.experimental.dlpack.from_dlpack(dlcapsule) -@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) -def frombuffer( - buffer: bytes, - dtype: Optional[tf.DType] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(buffer, bytearray): - buffer = bytes(buffer) - ret = tf.io.decode_raw(buffer, dtype) - dtype = tf.dtypes.as_dtype(dtype) - if offset > 0: - offset = int(offset / dtype.size) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - - return ret - - def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -254,6 +210,10 @@ def full_like( return tf.experimental.numpy.full_like(x, fill_value, dtype=dtype) +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + def linspace( start: Union[tf.Tensor, tf.Variable, float], stop: Union[tf.Tensor, tf.Variable, float], @@ -305,23 +265,6 @@ def meshgrid( return res -def one_hot( - indices: Union[tf.Tensor, tf.Variable], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[tf.DType] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.one_hot( - indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype - ) - - def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -367,29 +310,6 @@ def triu( return tf.experimental.numpy.triu(x, k) -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: str, -) -> Tuple[Union[tf.Tensor, tf.Variable]]: - n_cols = n_rows if n_cols is None else n_cols - - if n_rows < 0 or n_cols < 0: - n_rows, n_cols = 0, 0 - - ret = [[], []] - - for i in range(0, min(n_rows, n_cols - k), 1): - for j in range(max(0, k + i), n_cols, 1): - ret[0].append(i) - ret[1].append(j) - - return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) - - def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -416,3 +336,75 @@ def zeros_like( array = asarray + + +def copy_array( + x: Union[tf.Tensor, tf.Variable], + *, + to_ivy_array: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if to_ivy_array: + return ivy.to_ivy(tf.identity(x)) + return tf.identity(x) + + +def one_hot( + indices: Union[tf.Tensor, tf.Variable], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[tf.DType] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.one_hot( + indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype + ) + + +@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) +def frombuffer( + buffer: bytes, + dtype: Optional[tf.DType] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(buffer, bytearray): + buffer = bytes(buffer) + ret = tf.io.decode_raw(buffer, dtype) + dtype = tf.dtypes.as_dtype(dtype) + if offset > 0: + offset = int(offset / dtype.size) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + + return ret + + +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[Union[tf.Tensor, tf.Variable]]: + n_cols = n_rows if n_cols is None else n_cols + + if n_rows < 0 or n_cols < 0: + n_rows, n_cols = 0, 0 + + ret = [[], []] + + for i in range(0, min(n_rows, n_cols - k), 1): + for j in range(max(0, k + i), n_cols, 1): + ret[0].append(i) + ret[1].append(j) + + return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) diff --git a/ivy/functional/backends/tensorflow/data_type.py b/ivy/functional/backends/tensorflow/data_type.py index d6100609a38d8..819a8fc381d75 100644 --- a/ivy/functional/backends/tensorflow/data_type.py +++ b/ivy/functional/backends/tensorflow/data_type.py @@ -28,6 +28,7 @@ tf.complex128: "complex128", tf.bool: "bool", } + native_dtype_dict = { "int8": tf.int8, "int16": tf.int16, @@ -90,6 +91,93 @@ def __repr__(self): ) +# Array API Standard # +# -------------------# + + +def astype( + x: Union[tf.Tensor, tf.Variable], + dtype: Union[DType, str], + /, + *, + copy: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return tf.experimental.numpy.copy(x) if copy else x + return tf.cast(x, dtype) + + +def broadcast_arrays( + *arrays: Union[tf.Tensor, tf.Variable], +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(arrays) > 1: + try: + desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + if len(arrays) > 2: + for i in range(2, len(arrays)): + try: + desired_shape = tf.broadcast_dynamic_shape( + desired_shape, arrays[i].shape + ) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + else: + return [arrays[0]] + result = [] + for tensor in arrays: + result.append(tf.broadcast_to(tensor, desired_shape)) + + return result + + +def broadcast_to( + x: Union[tf.Tensor, tf.Variable], + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if tf.rank(x) > len(shape): + return tf.broadcast_to(tf.reshape(x, -1), shape) + return tf.broadcast_to(x, shape) + + +@_handle_nestable_dtype_info +def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + if ivy.as_native_dtype(type) == tf.bfloat16: + return Finfo(Bfloat16Finfo()) + return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def result_type( + *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], +) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return tf.experimental.numpy.result_type(arrays_and_dtypes) + + result = tf.experimental.numpy.result_type( + arrays_and_dtypes[0], arrays_and_dtypes[1] + ) + for i in range(2, len(arrays_and_dtypes)): + result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) + + # Extra # # ------# @@ -159,62 +247,6 @@ def as_native_dtype( ) -# Array API Standard # -# -------------------# - - -def astype( - x: Union[tf.Tensor, tf.Variable], - dtype: Union[DType, str], - /, - *, - copy: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return tf.experimental.numpy.copy(x) if copy else x - return tf.cast(x, dtype) - - -def broadcast_arrays( - *arrays: Union[tf.Tensor, tf.Variable], -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(arrays) > 1: - try: - desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - if len(arrays) > 2: - for i in range(2, len(arrays)): - try: - desired_shape = tf.broadcast_dynamic_shape( - desired_shape, arrays[i].shape - ) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - else: - return [arrays[0]] - result = [] - for tensor in arrays: - result.append(tf.broadcast_to(tensor, desired_shape)) - - return result - - -def broadcast_to( - x: Union[tf.Tensor, tf.Variable], - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if tf.rank(x) > len(shape): - return tf.broadcast_to(tf.reshape(x, -1), shape) - return tf.broadcast_to(x, shape) - - def dtype( x: Union[tf.Tensor, tf.Variable, np.ndarray], *, as_native: bool = False ) -> ivy.Dtype: @@ -237,22 +269,6 @@ def dtype_bits(dtype_in: Union[tf.DType, str, np.dtype], /) -> int: ) -@_handle_nestable_dtype_info -def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - if ivy.as_native_dtype(type) == tf.bfloat16: - return Finfo(Bfloat16Finfo()) - return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) - - def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -262,16 +278,6 @@ def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: return False -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def result_type( - *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], -) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return tf.experimental.numpy.result_type(arrays_and_dtypes) - - result = tf.experimental.numpy.result_type( - arrays_and_dtypes[0], arrays_and_dtypes[1] - ) - for i in range(2, len(arrays_and_dtypes)): - result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) +# ToDo: +# 1. result_type: Add support for bfloat16 with int16 +# 2. can_cast : Add support for complex64, complex128 diff --git a/ivy/functional/backends/tensorflow/device.py b/ivy/functional/backends/tensorflow/device.py index ca7c545ecba6e..d0872e5e2da9d 100644 --- a/ivy/functional/backends/tensorflow/device.py +++ b/ivy/functional/backends/tensorflow/device.py @@ -4,6 +4,9 @@ Collection of TensorFlow general functions, wrapped to fit Ivy syntax and signature. """ + +# global +_round = round import tensorflow as tf from typing import Union, Optional @@ -11,33 +14,6 @@ import ivy from ivy.functional.ivy.device import Profiler as BaseProfiler -# global -_round = round - - -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._options = tf.profiler.experimental.ProfilerOptions( - host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 - ) - - def start(self): - tf.profiler.experimental.start(self._save_dir, options=self._options) - - def stop(self): - tf.profiler.experimental.stop() - - def __enter__(self): - self.start() - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - -# --- Helpers --- # -# --------------- # - def _same_device(dev_a, dev_b): if dev_a is None or dev_b is None: @@ -47,8 +23,34 @@ def _same_device(dev_a, dev_b): ) -# --- Main --- # -# ------------ # +def dev( + x: Union[tf.Tensor, tf.Variable], + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, str]: + dv = x.device + if as_native: + return dv + return as_ivy_dev(dv) + + +def to_device( + x: Union[tf.Tensor, tf.Variable], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if device is None: + return x + device = as_native_dev(device) + current_dev = dev(x) + if not _same_device(current_dev, device): + with tf.device("/" + device.upper()): + return tf.identity(x) + return x def as_ivy_dev(device: str, /): @@ -77,48 +79,12 @@ def clear_cached_mem_on_dev(device: str, /): return None -def dev( - x: Union[tf.Tensor, tf.Variable], - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, str]: - dv = x.device - if as_native: - return dv - return as_ivy_dev(dv) - - -def gpu_is_available() -> bool: - return len(tf.config.list_physical_devices("GPU")) > 0 - - -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - default_device = ivy.default_device(device_shifting_dev, as_native=True) - with tf.device(default_device): - return fn(*args, **kwargs) - - def num_gpus() -> int: return len(tf.config.list_physical_devices("GPU")) -def to_device( - x: Union[tf.Tensor, tf.Variable], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if device is None: - return x - device = as_native_dev(device) - current_dev = dev(x) - if not _same_device(current_dev, device): - with tf.device("/" + device.upper()): - return tf.identity(x) - return x +def gpu_is_available() -> bool: + return len(tf.config.list_physical_devices("GPU")) > 0 def tpu_is_available() -> bool: @@ -131,3 +97,29 @@ def tpu_is_available() -> bool: return True except ValueError: return False + + +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + default_device = ivy.default_device(device_shifting_dev, as_native=True) + with tf.device(default_device): + return fn(*args, **kwargs) + + +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._options = tf.profiler.experimental.ProfilerOptions( + host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 + ) + + def start(self): + tf.profiler.experimental.start(self._save_dir, options=self._options) + + def stop(self): + tf.profiler.experimental.stop() + + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/ivy/functional/backends/tensorflow/elementwise.py b/ivy/functional/backends/tensorflow/elementwise.py index 38c717e0b4ce2..1e060e8e8cae3 100644 --- a/ivy/functional/backends/tensorflow/elementwise.py +++ b/ivy/functional/backends/tensorflow/elementwise.py @@ -10,9 +10,6 @@ from . import backend_version -gcd.support_native_out = False - - def abs( x: Union[float, tf.Tensor, tf.Variable], /, @@ -59,32 +56,6 @@ def add( return tf.add(x1, x2) -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def angle( - input: Union[tf.Tensor, tf.Variable], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if deg: - return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) - else: - return tf.math.angle(input, name=None) - - def asin( x: Union[tf.Tensor, tf.Variable], /, @@ -247,16 +218,6 @@ def cosh( return tf.cosh(x) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def deg2rad( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.deg2rad(x) - - def divide( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -284,20 +245,6 @@ def equal( return tf.math.equal(x1, x2) -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def erf( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.erf(x) - - def exp( x: Union[tf.Tensor, tf.Variable], /, @@ -366,37 +313,6 @@ def fmin( return ret -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, - backend_version, -) -def fmod( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - # tf.math.floormod returns wrong results - res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) - return tf.where(x1 < 0, -res, res) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def gcd( - x1: Union[tf.Tensor, tf.Variable, int, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - return tf.experimental.numpy.gcd(x1, x2) - - @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def greater( x1: Union[float, tf.Tensor, tf.Variable], @@ -421,28 +337,6 @@ def greater_equal( return tf.math.greater_equal(x1, x2) -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def imag( - val: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.imag(val, name=None) - - def isfinite( x: Union[tf.Tensor, tf.Variable], /, @@ -493,15 +387,6 @@ def isnan( return tf.math.is_nan(x) -def isreal( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isreal(x) - - @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) def lcm( x1: Union[tf.Tensor, tf.Variable], @@ -586,6 +471,16 @@ def logaddexp( return tf.experimental.numpy.logaddexp(x1, x2) +@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) +def real( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.real(x) + + @with_unsupported_dtypes( { "2.13.0 and below": ( @@ -660,64 +555,6 @@ def logical_xor( return tf.math.logical_xor(tf.cast(x1, tf.bool), tf.cast(x2, tf.bool)) -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def maximum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.maximum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def minimum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.minimum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) - - def multiply( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -729,31 +566,6 @@ def multiply( return tf.math.multiply(x1, x2) -def nan_to_num( - x: Union[tf.Tensor, tf.Variable], - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - posinf = posinf if posinf is not None else x.dtype.max - neginf = neginf if neginf is not None else x.dtype.min - posinf = tf.constant(posinf, x.dtype) - neginf = tf.constant(neginf, x.dtype) - nan = tf.constant(nan, x.dtype) - ret = tf.where(tf.math.is_nan(x), nan, x) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) - if copy: - return ret - else: - x = ret - return x - - @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) @@ -811,49 +623,6 @@ def pow( return tf.experimental.numpy.power(x1, x2) -def rad2deg( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rad2deg(x) - - -@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) -def real( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.real(x) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "int8", - "int16", - "int32", - "int64", - ) - }, - backend_version, -) -def reciprocal( - x: Union[float, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.reciprocal(x) - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def remainder( x1: Union[float, tf.Tensor, tf.Variable], @@ -1014,3 +783,234 @@ def trunc( else: ret = (tf.math.floor if ret >= 0 else tf.math.ceil)(ret) return ret + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def erf( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.erf(x) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def maximum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.maximum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def minimum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.minimum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "int8", + "int16", + "int32", + "int64", + ) + }, + backend_version, +) +def reciprocal( + x: Union[float, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.reciprocal(x) + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def deg2rad( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.deg2rad(x) + + +def rad2deg( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.rad2deg(x) + + +def isreal( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.isreal(x) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, + backend_version, +) +def fmod( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + # tf.math.floormod returns wrong results + res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) + return tf.where(x1 < 0, -res, res) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def gcd( + x1: Union[tf.Tensor, tf.Variable, int, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + return tf.experimental.numpy.gcd(x1, x2) + + +gcd.support_native_out = False + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def angle( + input: Union[tf.Tensor, tf.Variable], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if deg: + return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) + else: + return tf.math.angle(input, name=None) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def imag( + val: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.imag(val, name=None) + + +def nan_to_num( + x: Union[tf.Tensor, tf.Variable], + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + posinf = posinf if posinf is not None else x.dtype.max + neginf = neginf if neginf is not None else x.dtype.min + posinf = tf.constant(posinf, x.dtype) + neginf = tf.constant(neginf, x.dtype) + nan = tf.constant(nan, x.dtype) + ret = tf.where(tf.math.is_nan(x), nan, x) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) + if copy: + return ret + else: + x = ret + return x diff --git a/ivy/functional/backends/tensorflow/experimental/activations.py b/ivy/functional/backends/tensorflow/experimental/activations.py index 0f5126a6678ff..f798d7b907d98 100644 --- a/ivy/functional/backends/tensorflow/experimental/activations.py +++ b/ivy/functional/backends/tensorflow/experimental/activations.py @@ -10,15 +10,6 @@ from . import backend_version -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: - alpha = tf.cast(alpha, x.dtype) - ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) - - def logit( x: Union[tf.Tensor, tf.Variable], /, @@ -34,9 +25,16 @@ def logit( return tf.cast(tf.math.log(x / (1 - x)), x_dtype) -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return tf.math.log_sigmoid(input) +@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) +def thresholded_relu( + x: Tensor, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[Tensor] = None, +) -> Tensor: + threshold = tf.cast(threshold, x.dtype) + return tf.cast(tf.where(x > threshold, x, 0), x.dtype) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -44,6 +42,11 @@ def relu6(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu6(x) +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return tf.math.log_sigmoid(input) + + @with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) def selu(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: ret = tf.nn.selu(x) @@ -65,13 +68,10 @@ def silu( return ivy.astype(ret, x.dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) -def thresholded_relu( - x: Tensor, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[Tensor] = None, -) -> Tensor: - threshold = tf.cast(threshold, x.dtype) - return tf.cast(tf.where(x > threshold, x, 0), x.dtype) +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: + alpha = tf.cast(alpha, x.dtype) + ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) diff --git a/ivy/functional/backends/tensorflow/experimental/creation.py b/ivy/functional/backends/tensorflow/experimental/creation.py index 4462e9d61977e..96cd0c15b84be 100644 --- a/ivy/functional/backends/tensorflow/experimental/creation.py +++ b/ivy/functional/backends/tensorflow/experimental/creation.py @@ -8,40 +8,28 @@ from .. import backend_version -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=tf.result_type(size, 0.0)) - if periodic: - count = tf.arange(size) / size - else: - count = tf.linspace(start=0, stop=size, num=size) - - return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( - 0.08 * tf.cos(2 * tf.pi * 2 * count) - ) +# Array API Standard # +# -------------------# -def hann_window( - size: int, - /, - *, +@with_unsupported_device_and_dtypes( + {"2.13.0 and below": {"cpu": ("bfloat16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=dtype) - if periodic: - return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] + if window_length < 2: + return tf.ones([window_length], dtype=dtype) + if periodic is False: + return tf.signal.kaiser_window(window_length, beta, dtype=dtype) else: - return tf.signal.hann_window(size, periodic=False, dtype=dtype) + return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] def kaiser_bessel_derived_window( @@ -54,28 +42,29 @@ def kaiser_bessel_derived_window( return tf.signal.kaiser_bessel_derived_window(window_length, beta, dtype) -# Array API Standard # -# -------------------# +def vorbis_window( + window_length: Union[tf.Tensor, tf.Variable], + *, + dtype: tf.DType = tf.dtypes.float32, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) -@with_unsupported_device_and_dtypes( - {"2.13.0 and below": {"cpu": ("bfloat16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if window_length < 2: - return tf.ones([window_length], dtype=dtype) - if periodic is False: - return tf.signal.kaiser_window(window_length, beta, dtype=dtype) + if size < 2: + return tf.ones([size], dtype=dtype) + if periodic: + return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] else: - return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] + return tf.signal.hann_window(size, periodic=False, dtype=dtype) def tril_indices( @@ -101,20 +90,6 @@ def tril_indices( return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) -@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) -def trilu( - x: Union[tf.Tensor, tf.Variable], - /, - *, - k: int = 0, - upper: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if upper: - return tf.experimental.numpy.triu(x, k) - return tf.experimental.numpy.tril(x, k) - - def unsorted_segment_min( data: tf.Tensor, segment_ids: tf.Tensor, @@ -123,6 +98,26 @@ def unsorted_segment_min( return tf.math.unsorted_segment_min(data, segment_ids, num_segments) +def blackman_window( + size: int, + /, + *, + periodic: bool = True, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if size < 2: + return tf.ones([size], dtype=tf.result_type(size, 0.0)) + if periodic: + count = tf.arange(size) / size + else: + count = tf.linspace(start=0, stop=size, num=size) + + return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( + 0.08 * tf.cos(2 * tf.pi * 2 * count) + ) + + def unsorted_segment_sum( data: tf.Tensor, segment_ids: tf.Tensor, @@ -131,10 +126,15 @@ def unsorted_segment_sum( return tf.math.unsorted_segment_sum(data, segment_ids, num_segments) -def vorbis_window( - window_length: Union[tf.Tensor, tf.Variable], +@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) +def trilu( + x: Union[tf.Tensor, tf.Variable], + /, *, - dtype: tf.DType = tf.dtypes.float32, + k: int = 0, + upper: bool = True, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) + if upper: + return tf.experimental.numpy.triu(x, k) + return tf.experimental.numpy.tril(x, k) diff --git a/ivy/functional/backends/tensorflow/experimental/elementwise.py b/ivy/functional/backends/tensorflow/experimental/elementwise.py index 487ed04c624b4..96dba1c88a4dd 100644 --- a/ivy/functional/backends/tensorflow/experimental/elementwise.py +++ b/ivy/functional/backends/tensorflow/experimental/elementwise.py @@ -11,54 +11,61 @@ from .. import backend_version -# --- Helpers --- # -# --------------- # - - -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim - - -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +@with_supported_dtypes( + {"2.13.0 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.lgamma(x) -# --- Main --- # -# ------------ # +def sinc( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x = ivy.pi * x + return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) -def allclose( +@with_supported_dtypes( + {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version +) +def fmax( x1: Union[tf.Tensor, tf.Variable], x2: Union[tf.Tensor, tf.Variable], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> bool: - return tf.experimental.numpy.allclose( - x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan - ) +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + x1 = tf.where(tf.math.is_nan(x1), x2, x1) + x2 = tf.where(tf.math.is_nan(x2), x1, x2) + return tf.experimental.numpy.maximum(x1, x2) -def conj( - x: Union[tf.Tensor, tf.Variable], +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def float_power( + x1: Union[tf.Tensor, tf.Variable, float, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.conj(x) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): + out_dtype = tf.complex128 + else: + out_dtype = tf.float64 + return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) def copysign( @@ -96,35 +103,67 @@ def count_nonzero( ) -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def diff( - x: Union[tf.Tensor, tf.Variable, list, tuple], +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def nansum( + x: Union[tf.Tensor, tf.Variable], /, *, - n: int = 1, - axis: int = -1, - prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, - append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[tf.DType] = None, + keepdims: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if n == 0: - return x - if prepend is not None: - x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) - if append is not None: - x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) - return tf.experimental.numpy.diff(x, n=n, axis=axis) + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) -def digamma( - x: Union[tf.Tensor, tf.Variable], +def isclose( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.digamma(x) + return tf.experimental.numpy.isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan + ) + + +def signbit( + x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.signbit(x) + + +def hypot( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) + + +def allclose( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> bool: + return tf.experimental.numpy.allclose( + x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan + ) @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) @@ -137,56 +176,74 @@ def fix( return tf.cast(tf.where(x > 0, tf.math.floor(x), tf.math.ceil(x)), x.dtype) +@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) +def nextafter( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.nextafter(x1, x2) + + @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) -def float_power( - x1: Union[tf.Tensor, tf.Variable, float, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], +def diff( + x: Union[tf.Tensor, tf.Variable, list, tuple], /, *, + n: int = 1, + axis: int = -1, + prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, + append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): - out_dtype = tf.complex128 - else: - out_dtype = tf.float64 - return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) + if n == 0: + return x + if prepend is not None: + x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) + if append is not None: + x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) + return tf.experimental.numpy.diff(x, n=n, axis=axis) @with_supported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version + { + "2.13.0 and below": ( + "float32", + "float64", + ) + }, + backend_version, ) -def fmax( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], +def zeta( + x: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - x1 = tf.where(tf.math.is_nan(x1), x2, x1) - x2 = tf.where(tf.math.is_nan(x2), x1, x2) - return tf.experimental.numpy.maximum(x1, x2) + return tf.math.zeta(x, q) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) -def frexp( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[ - Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] - ] = None, -) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: - e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) - e = tf.cast(e, x.dtype) - while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): - e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) - m = x / tf.math.pow(2, e) - e = tf.cast(e, tf.int32) - return m, e +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim + + +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis def gradient( @@ -371,29 +428,36 @@ def gradient( return outvals -def hypot( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float16", + "float32", + "float64", + "complex64", + "complex128", + ) + }, + backend_version, +) +def xlogy( + x: Union[tf.Tensor, tf.Variable], + y: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) + x, y = promote_types_of_inputs(x, y) + return tf.math.xlogy(x, y) -def isclose( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], +def conj( + x: Union[tf.Tensor, tf.Variable], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan - ) + return tf.math.conj(x) @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) @@ -417,17 +481,22 @@ def ldexp( return tf.cast(ret, out_dtype) -@with_supported_dtypes( - {"2.13.0 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) +def frexp( x: Union[tf.Tensor, tf.Variable], /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.lgamma(x) + out: Optional[ + Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] + ] = None, +) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: + e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) + e = tf.cast(e, x.dtype) + while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): + e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) + m = x / tf.math.pow(2, e) + e = tf.cast(e, tf.int32) + return m, e def modf( @@ -439,87 +508,10 @@ def modf( return tf.math.modf(x) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def nansum( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[tf.DType] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) - - -@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) -def nextafter( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.nextafter(x1, x2) - - -def signbit( - x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.signbit(x) - - -def sinc( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x = ivy.pi * x - return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) - - -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float16", - "float32", - "float64", - "complex64", - "complex128", - ) - }, - backend_version, -) -def xlogy( - x: Union[tf.Tensor, tf.Variable], - y: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x, y = promote_types_of_inputs(x, y) - return tf.math.xlogy(x, y) - - -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float32", - "float64", - ) - }, - backend_version, -) -def zeta( +def digamma( x: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.zeta(x, q) + return tf.math.digamma(x) diff --git a/ivy/functional/backends/tensorflow/experimental/layers.py b/ivy/functional/backends/tensorflow/experimental/layers.py index 2a528a74cdf3f..94c60963871b5 100644 --- a/ivy/functional/backends/tensorflow/experimental/layers.py +++ b/ivy/functional/backends/tensorflow/experimental/layers.py @@ -20,10 +20,6 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode, _get_size -# --- Helpers --- # -# --------------- # - - def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # Determine depth pooling kernel, strides, depth_pooling = _depth_max_pooling_helper( @@ -34,71 +30,229 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def _fft2_helper(x, shape, axes): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) +def max_pool1d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + if data_format == "NCW": + x = tf.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - shape = shape_initialization(shape, axes, x) + if not depth_pooling: + new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + padding = [(pad_w // 2, pad_w - pad_w // 2)] - rank = rank_initialization(axes) + if ceil_mode: + padding[0] = _padding_ceil_mode( + x.shape[1], new_kernel[0], padding[0], strides[0] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - perm = get_perm(input_rank_tensor, axes) + if depth_pooling: + res = tf.transpose(res, (0, 2, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCW": + return tf.transpose(res, (0, 2, 1)) + return res - x = transpose_x(x, perm, perform_transpose) - x = fft2_operations(x, rank) +def max_pool2d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) - x = transpose_x(x, tf.argsort(perm), perform_transpose) + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) - return x + if not depth_pooling: + new_kernel = [ + kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) + ] + if isinstance(padding, str): + pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) + padding = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + x_shape = x.shape[1:-1] -def _fft2_norm( - x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", -): - n = tf.constant(s[0] * s[1], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.sqrt(n) - elif norm == "forward": - return x / n + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) -def _fft_norm( + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def max_pool3d( x: Union[tf.Tensor, tf.Variable], - dim: int, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - norm: str = "backward", -): - n = tf.constant(x.shape[dim], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) - elif norm == "forward": - return x / tf.cast(n, x.dtype) + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCDHW": + x = tf.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + if not depth_pooling: + x_shape = x.shape[1:-1] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + padding = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) + + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 4, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCDHW": + return tf.transpose(res, (0, 4, 1, 2, 3)) + return res def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): @@ -126,123 +280,13 @@ def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): return padding, pad_specific, c -def _ifft_norm( +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) +def avg_pool1d( x: Union[tf.Tensor, tf.Variable], - dim: int, - *, - norm: str = "backward", -): - n = x.shape[dim] - if norm == "backward": - return x - elif norm == "ortho": - return x * math.sqrt(n) - elif norm == "forward": - return x * n - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -def _ifftn_helper(x, shape, axes, norm): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor - ) - - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = ifft_operations(x, rank, norm_factor) - - x = transpose_x(x, tf.argsort(perm), perform_transpose) - - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - - return x - - -def _rfftn_helper(x, shape, axes, norm): - x = rfft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor - ) - - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = rfft_operations(x, rank, norm_factor) - - x = transpose_x(x, tf.argsort(perm), perform_transpose) - - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - - return x - - -def _right_pad_or_crop(tensor, shape): - input_shape = tf.shape(tensor) - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - with tf.control_dependencies( - [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] - ): - shape = tf.identity(shape) - shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) - - pad_sizes = tf.math.maximum(shape - input_shape, 0) - pad_sizes = tf.expand_dims(pad_sizes, -1) - pad_sizes = tf.concat( - [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 - ) - tensor = tf.pad(tensor, pad_sizes, constant_values=0) - - crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) - tensor = tf.slice(tensor, crop_tensor, shape) - return tensor - - -# --- Main --- # -# ------------ # - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) -def avg_pool1d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, *, data_format: str = "NWC", count_include_pad: bool = False, @@ -503,15 +547,31 @@ def avg_pool3d( return res -def axes_initialization(shape, axes, input_shape, input_rank_tensor): - if axes is None: - axes = ( - tf.range(-tf.size(input_shape), 0) - if shape is None - else tf.range(-tf.size(shape), 0) - ) - axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) - return axes +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def pool( + x: Union[tf.Tensor, tf.Variable], + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, + *, + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.nn.pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ) @with_supported_dtypes({"2.13.0 and below": ("float32", "float64")}, backend_version) @@ -539,6 +599,112 @@ def dct( return dct_out +def idct( + x: Union[tf.Tensor, tf.Variable], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> tf.Tensor: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def _fft_norm( + x: Union[tf.Tensor, tf.Variable], + dim: int, + /, + *, + norm: str = "backward", +): + n = tf.constant(x.shape[dim], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) + elif norm == "forward": + return x / tf.cast(n, x.dtype) + else: + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + + +def _ifft_norm( + x: Union[tf.Tensor, tf.Variable], + dim: int, + *, + norm: str = "backward", +): + n = x.shape[dim] + if norm == "backward": + return x + elif norm == "ortho": + return x * math.sqrt(n) + elif norm == "forward": + return x * n + else: + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft( + x: Union[tf.Tensor, tf.Variable], + dim: int, + /, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if x.shape[dim] != n: + s = list(x.shape) + if s[dim] > n: + index = [slice(None)] * len(s) + index[dim] = slice(0, n) + x = x[tuple(index)] + del index + else: + s[dim] = n - s[dim] + z = tf.zeros(s, x.dtype) + x = tf.concat([x, z], dim) + del s + operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" + if dim != -1 or dim != len(x.shape) - 1: + permute = [i for i in range(len(x.shape))] + permute[dim], permute[-1] = permute[-1], permute[dim] + x = tf.transpose(x, permute) + ret = tf.signal.fft(x, operation_name) + ret = tf.transpose(ret, permute) + del permute + else: + ret = tf.signal.fft(x, operation_name) + ret = _fft_norm(ret, dim, norm=norm) + return ret + + def dropout( x: Union[tf.Tensor, tf.Variable], prob: float, @@ -625,177 +791,6 @@ def dropout3d( return res -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def embedding( - weights: Union[tf.Tensor, tf.Variable], - indices: Union[tf.Tensor, tf.Variable], - /, - *, - max_norm: Optional[float] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) - - -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft( - x: Union[tf.Tensor, tf.Variable], - dim: int, - /, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.shape[dim] != n: - s = list(x.shape) - if s[dim] > n: - index = [slice(None)] * len(s) - index[dim] = slice(0, n) - x = x[tuple(index)] - del index - else: - s[dim] = n - s[dim] - z = tf.zeros(s, x.dtype) - x = tf.concat([x, z], dim) - del s - operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" - if dim != -1 or dim != len(x.shape) - 1: - permute = [i for i in range(len(x.shape))] - permute[dim], permute[-1] = permute[-1], permute[dim] - x = tf.transpose(x, permute) - ret = tf.signal.fft(x, operation_name) - ret = tf.transpose(ret, permute) - del permute - else: - ret = tf.signal.fft(x, operation_name) - ret = _fft_norm(ret, dim, norm=norm) - return ret - - -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft2( - x: Union[tf.Tensor, tf.Variable], - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if len(x.shape) > 2: - result = _fft2_helper(x, s, dim) - else: - x_new = trans_x_to_s(x, s, dim) - x_complex = tf.cast(x_new, tf.complex128) - result = tf.signal.fft2d(x_complex) - - result = _fft2_norm(result, s, dim, norm) - if x.dtype == tf.complex64: - result = tf.cast(result, dtype=tf.complex128) - return result - - -def fft2_operations(x, rank): - if x.shape.rank == 1: - x = tf.signal.fft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} - ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.fft(x), - 1: lambda: tf.signal.fft2d(x), - 2: lambda: tf.signal.fft3d(x), - }, - ) - return x - - -# --- IFFTN --- # -def fft_input_validation(x): - if not x.dtype.is_complex: - raise TypeError( - "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( - x.dtype - ) - ) - return x - - -def get_perm(input_rank_tensor, axes): - all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) - perm = tf.concat( - [ - tf.boolean_mask( - all_dims, - tf.foldl( - lambda acc, elem: tf.math.logical_and( - acc, tf.math.not_equal(all_dims, elem) - ), - axes, - initializer=tf.fill(all_dims.shape, True), - ), - ), - axes, - ], - 0, - ) - return perm - - -def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): - if perform_padding: - pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) - pad_shape = tf.tensor_scatter_nd_update( - pad_shape, tf.expand_dims(axes, -1), shape - ) - x = _right_pad_or_crop(x, pad_shape) - return x - - -def idct( - x: Union[tf.Tensor, tf.Variable], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - def ifft( x: Union[tf.Tensor, tf.Variable], dim: int, @@ -851,41 +846,19 @@ def ifft( return ret -def ifft_operations(x, rank, norm_factor): - if x.shape.rank == 1: - x = tf.signal.ifft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} - ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.ifft(x), - 1: lambda: tf.signal.ifft2d(x), - 2: lambda: tf.signal.ifft3d(x), - }, - ) - x = x * norm_factor - return x - - -def ifftn( - x: Union[tf.Tensor, tf.Variable], - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def embedding( + weights: Union[tf.Tensor, tf.Variable], + indices: Union[tf.Tensor, tf.Variable], + /, *, - norm: Optional[str] = "backward", + max_norm: Optional[float] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - result = _ifftn_helper(x, s, axes, norm) - - if out is not None: - out = result - return out - else: - return result + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) def interpolate( @@ -946,292 +919,417 @@ def interpolate( return ret -def max_pool1d( +interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (not align_corners and (len(x.shape) - 2) < 2) + and mode not in ["nearest", "area", "bicubic", "nd"] +) + + +def _fft2_norm( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", +): + n = tf.constant(s[0] * s[1], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.sqrt(n) + elif norm == "forward": + return x / n + else: + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if data_format == "NCW": - x = tf.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) +def trans_x_to_s( + x: Union[tf.Tensor, tf.Variable], + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), +) -> Union[tf.Tensor, tf.Variable]: + """Change the shape of the input array x to the desired output shape s.""" + if x.dtype != tf.complex128 and x.dtype != tf.complex64: + x = tf.cast(x, tf.float32) + x_shape = x.shape + if dim == (-1, -2) or dim == (1, 0): + s = (s[1], s[0]) + if s[0] >= x_shape[0] and s[1] >= x_shape[1]: + paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) + x_new = tf.pad(x, paddings=paddings) + elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): + x_new = x[: s[0], : s[1]] + if s[0] != x_new.shape[0]: + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif s[1] != x_new.shape[1]: + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], 1) + elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], axis=1) + else: + x_new = x[: s[0], : s[1]] + return x_new - if not depth_pooling: - new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - padding = [(pad_w // 2, pad_w - pad_w // 2)] - if ceil_mode: - padding[0] = _padding_ceil_mode( - x.shape[1], new_kernel[0], padding[0], strides[0] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) +def fft2_operations(x, rank): + if x.shape.rank == 1: + x = tf.signal.fft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} + ) else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.fft(x), + 1: lambda: tf.signal.fft2d(x), + 2: lambda: tf.signal.fft3d(x), + }, + ) + return x - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - if depth_pooling: - res = tf.transpose(res, (0, 2, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCW": - return tf.transpose(res, (0, 2, 1)) - return res +def _fft2_helper(x, shape, axes): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) -def max_pool2d( + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = fft2_operations(x, rank) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft2( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if len(x.shape) > 2: + result = _fft2_helper(x, s, dim) + else: + x_new = trans_x_to_s(x, s, dim) + x_complex = tf.cast(x_new, tf.complex128) + result = tf.signal.fft2d(x_complex) - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + result = _fft2_norm(result, s, dim, norm) + if x.dtype == tf.complex64: + result = tf.cast(result, dtype=tf.complex128) + return result + + +# --- IFFTN --- # +def fft_input_validation(x): + if not x.dtype.is_complex: + raise TypeError( + "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( + x.dtype + ) ) + return x - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - if not depth_pooling: - new_kernel = [ - kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) +def shape_and_axes_validation(shape, axes, input_rank_tensor): + if shape is not None: + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + checks_shape = [ + tf.debugging.assert_less_equal( + tf.size(shape), + input_rank_tensor, + message=( + "Argument `shape` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(shape), + ) ] - if isinstance(padding, str): - pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) - padding = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] + with tf.control_dependencies(checks_shape): + shape = tf.identity(shape) - x_shape = x.shape[1:-1] + if axes is not None: + axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) + checks_axes = [ + tf.debugging.assert_less_equal( + tf.size(axes), + input_rank_tensor, + message=( + "Argument `axes` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(axes), + ), + tf.debugging.assert_less( + axes, + input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + tf.debugging.assert_greater_equal( + axes, + -input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + ] + with tf.control_dependencies(checks_axes): + axes = tf.identity(axes) - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" + if shape is not None and axes is not None: + checks_shape_axes = [ + tf.debugging.assert_equal( + tf.size(shape), + tf.size(axes), + message=( + "Arguments `shape` and `axes` must have equal length. " + "Received: {}, {}" + ).format(shape, axes), ) + ] + with tf.control_dependencies(checks_shape_axes): + shape, axes = tf.identity_n([shape, axes]) - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) + return shape, axes - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res + +def axes_initialization(shape, axes, input_shape, input_rank_tensor): + if axes is None: + axes = ( + tf.range(-tf.size(input_shape), 0) + if shape is None + else tf.range(-tf.size(shape), 0) + ) + axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) + return axes -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def max_pool3d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims +def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): + perform_padding = shape is not None + perform_transpose = tf.math.logical_not( + tf.math.reduce_all( + tf.math.equal( + axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) + ) + ) ) + return perform_padding, perform_transpose - if data_format == "NCDHW": - x = tf.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel + +def shape_initialization(shape, axes, x): + if shape is None: + shape = tf.gather(tf.shape(x), axes, axis=0) + return shape + + +def rank_initialization(axes): + rank = tf.size(axes) + with tf.control_dependencies( + [ + tf.debugging.assert_less_equal( + rank, 3, message="N-D FFT supported only up to 3-D." + ) + ] + ): + rank = tf.identity(rank) + + return rank + + +def norm_initialization(norm, shape, x): + if norm == "backward": + norm_factor = tf.constant(1, x.dtype) + elif norm == "forward" or norm == "ortho": + norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) + if norm == "ortho": + norm_factor = tf.math.sqrt(norm_factor) + return norm_factor + + +def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): + if perform_padding: + pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) + pad_shape = tf.tensor_scatter_nd_update( + pad_shape, tf.expand_dims(axes, -1), shape ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides + x = _right_pad_or_crop(x, pad_shape) + return x + + +def get_perm(input_rank_tensor, axes): + all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) + perm = tf.concat( + [ + tf.boolean_mask( + all_dims, + tf.foldl( + lambda acc, elem: tf.math.logical_and( + acc, tf.math.not_equal(all_dims, elem) + ), + axes, + initializer=tf.fill(all_dims.shape, True), + ), + ), + axes, + ], + 0, + ) + return perm + + +def ifft_operations(x, rank, norm_factor): + if x.shape.rank == 1: + x = tf.signal.ifft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + else: + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.ifft(x), + 1: lambda: tf.signal.ifft2d(x), + 2: lambda: tf.signal.ifft3d(x), + }, ) + x = x * norm_factor + return x - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" + +def transpose_x(x, perm, perform_transpose): + x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) + return x + + +def static_output_shape(input_shape, shape, axes): + output_shape = input_shape.as_list() + if shape is not None: + if axes is None: + axes = list(range(-len(shape), 0)) + if isinstance(shape, tf.Tensor): + if isinstance(axes, tf.Tensor): + output_shape = [None] * len(output_shape) + else: + for ax in axes: + output_shape[ax] = None + else: + for idx, ax in enumerate(axes): + output_shape[ax] = shape[idx] + return tf.TensorShape(output_shape) + + +def _right_pad_or_crop(tensor, shape): + input_shape = tf.shape(tensor) + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + with tf.control_dependencies( + [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] + ): + shape = tf.identity(shape) + shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) + + pad_sizes = tf.math.maximum(shape - input_shape, 0) + pad_sizes = tf.expand_dims(pad_sizes, -1) + pad_sizes = tf.concat( + [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 ) + tensor = tf.pad(tensor, pad_sizes, constant_values=0) - if not depth_pooling: - x_shape = x.shape[1:-1] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - padding = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] + crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) + tensor = tf.slice(tensor, crop_tensor, shape) + return tensor - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) +def _ifftn_helper(x, shape, axes, norm): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 4, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCDHW": - return tf.transpose(res, (0, 4, 1, 2, 3)) - return res + shape = shape_initialization(shape, axes, x) + rank = rank_initialization(axes) -def norm_initialization(norm, shape, x): - if norm == "backward": - norm_factor = tf.constant(1, x.dtype) - elif norm == "forward" or norm == "ortho": - norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) - if norm == "ortho": - norm_factor = tf.math.sqrt(norm_factor) - return norm_factor + norm_factor = norm_initialization(norm, shape, x) + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) -def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): - perform_padding = shape is not None - perform_transpose = tf.math.logical_not( - tf.math.reduce_all( - tf.math.equal( - axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) - ) - ) - ) - return perform_padding, perform_transpose + perm = get_perm(input_rank_tensor, axes) + x = transpose_x(x, perm, perform_transpose) -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def pool( + x = ifft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +def ifftn( x: Union[tf.Tensor, tf.Variable], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, + norm: Optional[str] = "backward", out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.nn.pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ) + result = _ifftn_helper(x, s, axes, norm) + if out is not None: + out = result + return out + else: + return result -def rank_initialization(axes): - rank = tf.size(axes) - with tf.control_dependencies( - [ - tf.debugging.assert_less_equal( - rank, 3, message="N-D FFT supported only up to 3-D." - ) - ] - ): - rank = tf.identity(rank) - return rank +""" +RFFTN Function +""" def rfft_input_validation(x): @@ -1266,6 +1364,40 @@ def rfft_operations(x, rank, norm_factor): return x +def _rfftn_helper(x, shape, axes, norm): + x = rfft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + norm_factor = norm_initialization(norm, shape, x) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = rfft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + @with_supported_device_and_dtypes( { "2.5.0 and above": { @@ -1295,139 +1427,3 @@ def rfftn( else: # return result return tf.cast(result, tf.complex128) - - -def shape_and_axes_validation(shape, axes, input_rank_tensor): - if shape is not None: - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - checks_shape = [ - tf.debugging.assert_less_equal( - tf.size(shape), - input_rank_tensor, - message=( - "Argument `shape` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(shape), - ) - ] - with tf.control_dependencies(checks_shape): - shape = tf.identity(shape) - - if axes is not None: - axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) - checks_axes = [ - tf.debugging.assert_less_equal( - tf.size(axes), - input_rank_tensor, - message=( - "Argument `axes` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(axes), - ), - tf.debugging.assert_less( - axes, - input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - tf.debugging.assert_greater_equal( - axes, - -input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - ] - with tf.control_dependencies(checks_axes): - axes = tf.identity(axes) - - if shape is not None and axes is not None: - checks_shape_axes = [ - tf.debugging.assert_equal( - tf.size(shape), - tf.size(axes), - message=( - "Arguments `shape` and `axes` must have equal length. " - "Received: {}, {}" - ).format(shape, axes), - ) - ] - with tf.control_dependencies(checks_shape_axes): - shape, axes = tf.identity_n([shape, axes]) - - return shape, axes - - -def shape_initialization(shape, axes, x): - if shape is None: - shape = tf.gather(tf.shape(x), axes, axis=0) - return shape - - -def static_output_shape(input_shape, shape, axes): - output_shape = input_shape.as_list() - if shape is not None: - if axes is None: - axes = list(range(-len(shape), 0)) - if isinstance(shape, tf.Tensor): - if isinstance(axes, tf.Tensor): - output_shape = [None] * len(output_shape) - else: - for ax in axes: - output_shape[ax] = None - else: - for idx, ax in enumerate(axes): - output_shape[ax] = shape[idx] - return tf.TensorShape(output_shape) - - -def trans_x_to_s( - x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), -) -> Union[tf.Tensor, tf.Variable]: - """Change the shape of the input array x to the desired output shape s.""" - if x.dtype != tf.complex128 and x.dtype != tf.complex64: - x = tf.cast(x, tf.float32) - x_shape = x.shape - if dim == (-1, -2) or dim == (1, 0): - s = (s[1], s[0]) - if s[0] >= x_shape[0] and s[1] >= x_shape[1]: - paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) - x_new = tf.pad(x, paddings=paddings) - elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): - x_new = x[: s[0], : s[1]] - if s[0] != x_new.shape[0]: - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif s[1] != x_new.shape[1]: - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], 1) - elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], axis=1) - else: - x_new = x[: s[0], : s[1]] - return x_new - - -def transpose_x(x, perm, perform_transpose): - x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) - return x - - -interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (not align_corners and (len(x.shape) - 2) < 2) - and mode not in ["nearest", "area", "bicubic", "nd"] -) -"""RFFTN Function.""" diff --git a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py index e59d5c46a79f9..c457e4ceb0ac0 100644 --- a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py @@ -11,59 +11,33 @@ from .. import backend_version -dot.support_native_out = True - - -def adjoint( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - _check_valid_dimension_size(x) - return tf.linalg.adjoint(x) - - -@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( - x: Union[tf.Tensor, tf.Variable], +@with_unsupported_dtypes( + {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version +) +def eigh_tridiagonal( + alpha: Union[tf.Tensor, tf.Variable], + beta: Union[tf.Tensor, tf.Variable], /, *, - p: Optional[Union[None, int, str]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - svd = tf.linalg.svd(x, compute_uv=False) - if len(x.shape) >= 3: - ax = len(x.shape) // 2 - elif len(x.shape) >= 3 and p == -1: - ax = [-1, -2] - else: - ax = None - if p is None or p == 2: - k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) - elif p == "nuc": - svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) - k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) - elif p == "fro": - k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] - ) - elif p < 0: - if p == -1: - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=0), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) - elif p == -2: - k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) - elif p == -float("inf"): - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=1), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) - else: - k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord=p, axis=[-2, -1] - ) - return k + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] + ] = None, + tol: Optional[float] = None, +) -> Union[ + tf.Tensor, + tf.Variable, + Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], +]: + return tf.linalg.eigh_tridiagonal( + alpha, + beta, + eigvals_only=eigvals_only, + select=select, + select_range=select_range, + tol=tol, + ) def diagflat( @@ -101,14 +75,23 @@ def diagflat( return ret -def dot( - a: tf.Tensor, - b: tf.Tensor, +def kron( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], /, *, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - return tf.tensordot(a, b, axes=1) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.kron(a, b) + + +def matrix_exp( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.linalg.expm(x) def eig( @@ -122,35 +105,6 @@ def eig( return tf.linalg.eig(x) -@with_unsupported_dtypes( - {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version -) -def eigh_tridiagonal( - alpha: Union[tf.Tensor, tf.Variable], - beta: Union[tf.Tensor, tf.Variable], - /, - *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] - ] = None, - tol: Optional[float] = None, -) -> Union[ - tf.Tensor, - tf.Variable, - Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], -]: - return tf.linalg.eigh_tridiagonal( - alpha, - beta, - eigvals_only=eigvals_only, - select=select, - select_range=select_range, - tol=tol, - ) - - def eigvals( x: Union[tf.Tensor, tf.Variable], /, @@ -160,33 +114,14 @@ def eigvals( return tf.linalg.eigvals(x) -def kron( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.kron(a, b) - - -def lu_factor( - x: Union[tf.Tensor, tf.Variable], - /, - *, - pivot: Optional[bool] = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - raise IvyNotImplementedException() - - -def matrix_exp( +def adjoint( x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.linalg.expm(x) + _check_valid_dimension_size(x) + return tf.linalg.adjoint(x) @with_supported_dtypes( @@ -214,3 +149,68 @@ def multi_dot( raise ValueError("Expecting at least two tensors.") dot_out = _reduce(tf.matmul, x) return dot_out + + +@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( + x: Union[tf.Tensor, tf.Variable], + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + svd = tf.linalg.svd(x, compute_uv=False) + if len(x.shape) >= 3: + ax = len(x.shape) // 2 + elif len(x.shape) >= 3 and p == -1: + ax = [-1, -2] + else: + ax = None + if p is None or p == 2: + k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) + elif p == "nuc": + svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) + k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) + elif p == "fro": + k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] + ) + elif p < 0: + if p == -1: + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=0), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) + elif p == -2: + k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) + elif p == -float("inf"): + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=1), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) + else: + k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord=p, axis=[-2, -1] + ) + return k + + +def lu_factor( + x: Union[tf.Tensor, tf.Variable], + /, + *, + pivot: Optional[bool] = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + raise IvyNotImplementedException() + + +def dot( + a: tf.Tensor, + b: tf.Tensor, + /, + *, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + return tf.tensordot(a, b, axes=1) + + +dot.support_native_out = True diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index bf0be94c43b24..7e0a806ee1e19 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -10,126 +10,99 @@ import ivy -def atleast_1d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_1d(*arys) - - -def atleast_2d( - *arys: Union[tf.Tensor, tf.Variable], - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], +def moveaxis( + a: Union[tf.Tensor, tf.Variable], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_3d(*arys) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.moveaxis(a, source, destination) -def broadcast_shapes( - *shapes: Union[List[int], List[Tuple]], -) -> Tuple[int, ...]: - if len(shapes) > 1: - desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) - if len(shapes) > 2: - for i in range(2, len(shapes)): - desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) - else: - return [shapes[0]] - return tuple(desired_shape.numpy().tolist()) +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def heaviside( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) -def concat_from_sequence( - input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], +def flipud( + m: Union[tf.Tensor, tf.Variable], /, *, - new_axis: int = 0, - axis: int = 0, + copy: Optional[bool] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - highest_dtype = input_sequence[0].dtype - for i in input_sequence: - highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) - - if new_axis == 0: - ret = tf.concat(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = tf.stack(input_sequence, axis=axis) - return ret + return tf.experimental.numpy.flipud(m) -def dsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], +def vstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], /, *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vstack(arrays) -def dstack( +def hstack( arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.dstack(arrays) + return tf.experimental.numpy.hstack(arrays) -def expand( - x: Union[tf.Tensor, tf.Variable], - shape: Union[List[int], List[Tuple]], +def rot90( + m: Union[tf.Tensor, tf.Variable], /, *, copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Union[tf.Tensor, tf.Variable] = None, ) -> Union[tf.Tensor, tf.Variable]: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape.num_elements() / tf.reduce_prod( - [s for s in shape if s > 0] - ) - return tf.broadcast_to(x, shape) + return tf.experimental.numpy.rot90(m, k, axes) -def fill_diagonal( - a: tf.Tensor, - v: Union[int, float], +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) +def top_k( + x: tf.Tensor, + k: int, /, *, - wrap: bool = False, -): - shape = tf.shape(a) - max_end = tf.math.reduce_prod(shape) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, +) -> Tuple[tf.Tensor, tf.Tensor]: + k = min(k, x.shape[axis]) + if not largest: + indices = tf.experimental.numpy.argsort(x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) else: - step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) - a = tf.reshape(a, (-1,)) - end = min(end, max_end) - indices = [[i] for i in range(0, end, step)] - ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) - a = tf.tensor_scatter_nd_update(a, indices, ups) - a = tf.reshape(a, shape) - return a + indices = tf.experimental.numpy.argsort(-x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) + if not sorted: + indices = tf.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) + val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) + indices = tf.dtypes.cast(indices, tf.int64) + return topk_res(val, indices) def fliplr( @@ -142,80 +115,72 @@ def fliplr( return tf.experimental.numpy.fliplr(m) -def flipud( - m: Union[tf.Tensor, tf.Variable], +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def i0( + x: Union[tf.Tensor, tf.Variable], /, *, - copy: Optional[bool] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.flipud(m) + return tf.math.bessel_i0(x, name=None) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def heaviside( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], +def vsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) -def hsplit( +def dsplit( ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Tuple[int, ...]], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, copy: Optional[bool] = None, ) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + if len(ary.shape) < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.hstack(arrays) +def atleast_1d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_1d(*arys) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def i0( - x: Union[tf.Tensor, tf.Variable], +def dstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.bessel_i0(x, name=None) + return tf.experimental.numpy.dstack(arrays) -def moveaxis( - a: Union[tf.Tensor, tf.Variable], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, +def atleast_2d( + *arys: Union[tf.Tensor, tf.Variable], copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.moveaxis(a, source, destination) +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_2d(*arys) -def rot90( - m: Union[tf.Tensor, tf.Variable], - /, - *, +def atleast_3d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Union[tf.Tensor, tf.Variable] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rot90(m, k, axes) +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_3d(*arys) def take_along_axis( @@ -263,36 +228,69 @@ def take_along_axis( return tf.experimental.numpy.take_along_axis(arr, indices, axis) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) -def top_k( - x: tf.Tensor, - k: int, +def hsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Tuple[int, ...]], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, -) -> Tuple[tf.Tensor, tf.Tensor]: - k = min(k, x.shape[axis]) - if not largest: - indices = tf.experimental.numpy.argsort(x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +def broadcast_shapes( + *shapes: Union[List[int], List[Tuple]], +) -> Tuple[int, ...]: + if len(shapes) > 1: + desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) + if len(shapes) > 2: + for i in range(2, len(shapes)): + desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) else: - indices = tf.experimental.numpy.argsort(-x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) - if not sorted: - indices = tf.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) - val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) - indices = tf.dtypes.cast(indices, tf.int64) - return topk_res(val, indices) + return [shapes[0]] + return tuple(desired_shape.numpy().tolist()) + + +def expand( + x: Union[tf.Tensor, tf.Variable], + shape: Union[List[int], List[Tuple]], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape.num_elements() / tf.reduce_prod( + [s for s in shape if s > 0] + ) + return tf.broadcast_to(x, shape) + + +def concat_from_sequence( + input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + highest_dtype = input_sequence[0].dtype + for i in input_sequence: + highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) + + if new_axis == 0: + ret = tf.concat(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = tf.stack(input_sequence, axis=axis) + return ret def unique_consecutive( @@ -348,24 +346,26 @@ def unique_consecutive( ) -def vsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], - /, - *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def vstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], +def fill_diagonal( + a: tf.Tensor, + v: Union[int, float], /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vstack(arrays) + wrap: bool = False, +): + shape = tf.shape(a) + max_end = tf.math.reduce_prod(shape) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) + a = tf.reshape(a, (-1,)) + end = min(end, max_end) + indices = [[i] for i in range(0, end, step)] + ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) + a = tf.tensor_scatter_nd_update(a, indices, ups) + a = tf.reshape(a, shape) + return a diff --git a/ivy/functional/backends/tensorflow/experimental/norms.py b/ivy/functional/backends/tensorflow/experimental/norms.py index 22f00921bcf96..5fc38b60283af 100644 --- a/ivy/functional/backends/tensorflow/experimental/norms.py +++ b/ivy/functional/backends/tensorflow/experimental/norms.py @@ -4,6 +4,30 @@ from . import backend_version +def l1_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + +def l2_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) def batch_norm( x: Union[tf.Tensor, tf.Variable], @@ -129,30 +153,6 @@ def instance_norm( ) -def l1_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - -def l2_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - def lp_normalize( x: Union[tf.Tensor, tf.Variable], /, diff --git a/ivy/functional/backends/tensorflow/experimental/random.py b/ivy/functional/backends/tensorflow/experimental/random.py index fc31d6a1c8d3f..a5bbf97481be0 100644 --- a/ivy/functional/backends/tensorflow/experimental/random.py +++ b/ivy/functional/backends/tensorflow/experimental/random.py @@ -15,51 +15,6 @@ ) -def bernoulli( - probs: Union[float, tf.Tensor, tf.Variable], - *, - logits: Union[float, tf.Tensor, tf.Variable] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: DType, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if seed is not None: - tf.random.set_seed(seed) - if logits is not None: - logits = tf.cast(logits, dtype) - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - probs = tf.cast(probs, dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return tfp.distributions.Bernoulli( - logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True - ).sample(shape, seed) - - -def beta( - alpha: Union[float, tf.Tensor, tf.Variable], - beta: Union[float, tf.Tensor, tf.Variable], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, - dtype: Optional[Union[DType, ivy.Dtype]] = None, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not dtype: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - shape = _check_bounds_and_get_shape(alpha, beta, shape).shape - alpha = tf.cast(alpha, dtype) - beta = tf.cast(beta, dtype) - return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) - - # dirichlet @with_unsupported_dtypes( { @@ -99,6 +54,26 @@ def dirichlet( ) +def beta( + alpha: Union[float, tf.Tensor, tf.Variable], + beta: Union[float, tf.Tensor, tf.Variable], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, + dtype: Optional[Union[DType, ivy.Dtype]] = None, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not dtype: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + shape = _check_bounds_and_get_shape(alpha, beta, shape).shape + alpha = tf.cast(alpha, dtype) + beta = tf.cast(beta, dtype) + return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) + + def gamma( alpha: Union[float, tf.Tensor, tf.Variable], beta: Union[float, tf.Tensor, tf.Variable], @@ -142,3 +117,28 @@ def poisson( if tf.reduce_any(lam < 0): return tf.where(lam < 0, fill_value, ret) return ret + + +def bernoulli( + probs: Union[float, tf.Tensor, tf.Variable], + *, + logits: Union[float, tf.Tensor, tf.Variable] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: DType, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if seed is not None: + tf.random.set_seed(seed) + if logits is not None: + logits = tf.cast(logits, dtype) + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + probs = tf.cast(probs, dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return tfp.distributions.Bernoulli( + logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True + ).sample(shape, seed) diff --git a/ivy/functional/backends/tensorflow/experimental/statistical.py b/ivy/functional/backends/tensorflow/experimental/statistical.py index 2a5feef312527..9cd91437a5f4a 100644 --- a/ivy/functional/backends/tensorflow/experimental/statistical.py +++ b/ivy/functional/backends/tensorflow/experimental/statistical.py @@ -14,79 +14,202 @@ from copy import deepcopy -# --- Helpers --- # -# --------------- # +def histogram( + a: tf.Tensor, + /, + *, + bins: Optional[Union[int, tf.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[tf.DType] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[tf.Tensor] = None, + density: Optional[bool] = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + min_a = tf.reduce_min(a) + max_a = tf.reduce_max(a) + if isinstance(bins, tf.Tensor) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + if isinstance(bins, int): + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + elif isinstance(bins, int): + range = (min_a, max_a) + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + if tf.shape(bins)[0] < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + if min_a < bins[0] and not extend_lower_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_lower_interval to deal with this." + ) + if max_a > bins[-1] and not extend_upper_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_upper_interval to deal with this." + ) + ret = tfp.stats.histogram( + x=a, + edges=bins, + axis=axis, + weights=weights, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + name="histogram", + ) + if density: + pass + # TODO: Tensorflow native dtype argument is not working + if dtype: + ret = tf.cast(ret, dtype) + bins = tf.cast(bins, dtype) + # TODO: weird error when returning bins: return ret, bins + return ret -def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: - values, indices = [], [] - if ( - isinstance(x[0], tf.Tensor) - and isinstance(x[0].numpy().tolist(), list) - and len(x[0].numpy().tolist()) >= 1 - ): - if axis >= 1: - for ret1 in x: - value, indice = __find_cummax(ret1, axis=axis - 1) - indices.append(indice) - values.append(value) - else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - else: - x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) - values, indices = tf.scan( - lambda a, b: ( - a - if a > b - or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 - else b - ), - (x, x_indices), +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float", + "complex", ) - - return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( - tf.convert_to_tensor(indices), dtype=tf.int64 + }, + backend_version, +) +def median( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tfp.stats.percentile( + input, + 50.0, + axis=axis, + interpolation="midpoint", + keepdims=keepdims, ) -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] +def nanmean( + a: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) + +def _validate_quantile(q): + if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: + for i in range(tf.size(q)): + if not (0.0 <= q[i] <= 1.0): + return False else: - indices.append((lst, tuple(prefix))) - return indices + if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): + return False + return True + + +def to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis + + +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = tf.experimental.numpy.ndim(a) + axis_arg = deepcopy(axis) + if axis is not None: + axis = to_positive_axis(axis, nd) + + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) + + for i, s in enumerate(sorted(keep)): + a = tf.experimental.numpy.moveaxis(a, s, i) + a = tf.reshape( + a, + [ + *a.shape[:nkeep], + -1, + ], + ) + axis_arg = -1 + + ret = fn(a, q, axis=axis_arg) + + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret + + +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if tf.experimental.numpy.ndim(q) > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = tf.reshape(a, [-1]) + elif axis != 0: + a = tf.experimental.numpy.moveaxis(a, axis, 0) + axis = 0 + + n = a.shape[axis] + + indices = q * (n - 1) + + a = tf.sort(a, axis) + + indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) + indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) + + weights = indices - tf.cast(indices_below, dtype=ret_dtype) + + indices_below = tf.clip_by_value(indices_below, 0, n - 1) + indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) + tensor_upper = tf.gather(a, indices_upper, axis=axis) + tensor_below = tf.gather(a, indices_below, axis=axis) + + pred = weights <= 0.5 + out = tf.where(pred, tensor_below, tensor_upper) + + return tf.cast(out, ret_dtype) def _compute_quantile_wrapper( @@ -124,39 +247,50 @@ def _compute_quantile_wrapper( ) -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = tf.experimental.numpy.ndim(a) - axis_arg = deepcopy(axis) - if axis is not None: - axis = to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) +def quantile( + a: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, float], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) - for i, s in enumerate(sorted(keep)): - a = tf.experimental.numpy.moveaxis(a, s, i) - a = tf.reshape( - a, - [ - *a.shape[:nkeep], - -1, - ], - ) - axis_arg = -1 - ret = fn(a, q, axis=axis_arg) +def corrcoef( + x: tf.Tensor, + /, + *, + y: tf.Tensor, + rowvar: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> tf.Tensor: + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = tf.concat([x, y], axis=axis) - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] + if rowvar: + mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) + cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) + else: + mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) + cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) - return ret + cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) + cor = cov2_t @ cov_t @ cov2_t + return cor def _nanmedian_helper(input, axis=None, keepdims=False): @@ -318,52 +452,22 @@ def _nanmedian_helper(input, axis=None, keepdims=False): return result -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if tf.experimental.numpy.ndim(q) > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = tf.reshape(a, [-1]) - elif axis != 0: - a = tf.experimental.numpy.moveaxis(a, axis, 0) - axis = 0 - - n = a.shape[axis] - - indices = q * (n - 1) - - a = tf.sort(a, axis) - - indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) - indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) - - weights = indices - tf.cast(indices_below, dtype=ret_dtype) - - indices_below = tf.clip_by_value(indices_below, 0, n - 1) - indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) - tensor_upper = tf.gather(a, indices_upper, axis=axis) - tensor_below = tf.gather(a, indices_below, axis=axis) - - pred = weights <= 0.5 - out = tf.where(pred, tensor_below, tensor_upper) - - return tf.cast(out, ret_dtype) - +def nanmedian( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if overwrite_input: + copied_input = tf.identity(input) + return _nanmedian_helper(copied_input, axis, keepdims) -def _validate_quantile(q): - if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: - for i in range(tf.size(q)): - if not (0.0 <= q[i] <= 1.0): - return False else: - if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): - return False - return True - - -# --- Main --- # -# ------------ # + result = _nanmedian_helper(input, axis, keepdims) + return result @with_supported_device_and_dtypes( @@ -401,30 +505,19 @@ def bincount( ) -def corrcoef( - x: tf.Tensor, - /, - *, - y: tf.Tensor, - rowvar: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = tf.concat([x, y], axis=axis) - - if rowvar: - mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) - cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) - else: - mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) - cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) - - cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) - cor = cov2_t @ cov_t @ cov2_t - return cor +@with_supported_device_and_dtypes( + { + "2.13.0 and below": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def igamma( + a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None +) -> tf.Tensor: + return tf.math.igamma(a, x) @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) @@ -587,6 +680,77 @@ def cummax( return __find_cummax(x, axis=axis) +def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: + values, indices = [], [] + if ( + isinstance(x[0], tf.Tensor) + and isinstance(x[0].numpy().tolist(), list) + and len(x[0].numpy().tolist()) >= 1 + ): + if axis >= 1: + for ret1 in x: + value, indice = __find_cummax(ret1, axis=axis - 1) + indices.append(indice) + values.append(value) + else: + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + else: + x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) + values, indices = tf.scan( + lambda a, b: ( + a + if a > b + or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 + else b + ), + (x, x_indices), + ) + + return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( + tf.convert_to_tensor(indices), dtype=tf.int64 + ) + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "complex")}, backend_version, @@ -617,175 +781,3 @@ def cummin( return cummin_x else: return tf.cast(cummin_x, dtype) - - -def histogram( - a: tf.Tensor, - /, - *, - bins: Optional[Union[int, tf.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[tf.DType] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[tf.Tensor] = None, - density: Optional[bool] = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - min_a = tf.reduce_min(a) - max_a = tf.reduce_max(a) - if isinstance(bins, tf.Tensor) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - if isinstance(bins, int): - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - elif isinstance(bins, int): - range = (min_a, max_a) - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - if tf.shape(bins)[0] < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - if min_a < bins[0] and not extend_lower_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_lower_interval to deal with this." - ) - if max_a > bins[-1] and not extend_upper_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_upper_interval to deal with this." - ) - ret = tfp.stats.histogram( - x=a, - edges=bins, - axis=axis, - weights=weights, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - name="histogram", - ) - if density: - pass - # TODO: Tensorflow native dtype argument is not working - if dtype: - ret = tf.cast(ret, dtype) - bins = tf.cast(bins, dtype) - # TODO: weird error when returning bins: return ret, bins - return ret - - -@with_supported_device_and_dtypes( - { - "2.13.0 and below": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def igamma( - a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None -) -> tf.Tensor: - return tf.math.igamma(a, x) - - -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float", - "complex", - ) - }, - backend_version, -) -def median( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tfp.stats.percentile( - input, - 50.0, - axis=axis, - interpolation="midpoint", - keepdims=keepdims, - ) - - -def nanmean( - a: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) - - -def nanmedian( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if overwrite_input: - copied_input = tf.identity(input) - return _nanmedian_helper(copied_input, axis, keepdims) - - else: - result = _nanmedian_helper(input, axis, keepdims) - return result - - -def quantile( - a: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, float], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) - - -def to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis diff --git a/ivy/functional/backends/tensorflow/general.py b/ivy/functional/backends/tensorflow/general.py index e6052c0b5e49e..ba239549e3fdd 100644 --- a/ivy/functional/backends/tensorflow/general.py +++ b/ivy/functional/backends/tensorflow/general.py @@ -23,26 +23,12 @@ _round = round -# --- Helpers --- # -# --------------- # - - -def _check_query(query): - return not isinstance(query, list) and ( - not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) - ) - - -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view - - -# --- Main --- # -# ------------ # +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, (tf.Tensor, tf.Variable)): + if exclusive and isinstance(x, tf.Variable): + return False + return True + return False def array_equal( @@ -62,6 +48,58 @@ def current_backend_str() -> str: return "tensorflow" +def _check_query(query): + return not isinstance(query, list) and ( + not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) + ) + + +def get_item( + x: Union[tf.Tensor, tf.Variable], + /, + query: Union[tf.Tensor, tf.Variable, Tuple], + *, + copy: bool = None, +) -> Union[tf.Tensor, tf.Variable]: + return x.__getitem__(query) + + +get_item.partial_mixed_handler = lambda x, query, **kwargs: ( + all(_check_query(i) for i in query) + if isinstance(query, tuple) + else _check_query(query) +) + + +def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: + # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions + if ( + ivy.is_array(x) + and get_num_dims(x) == 0 + and ivy.as_native_dtype(x.dtype) is tf.bfloat16 + ): + x = tf.expand_dims(x, 0) + if copy: + return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) + else: + return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) + if copy: + return np.array(tf.convert_to_tensor(x)) + else: + return np.asarray(tf.convert_to_tensor(x)) + + +def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: + ret = to_numpy(x).item() + if x.dtype == tf.bfloat16: + return float(ret) + return ret + + +def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: + return x.numpy().tolist() + + def gather( params: Union[tf.Tensor, tf.Variable], indices: Union[tf.Tensor, tf.Variable], @@ -77,43 +115,6 @@ def gather( return tf.gather(params, indices, axis=axis, batch_dims=batch_dims) -def gather_nd( - params: Union[tf.Tensor, tf.Variable], - indices: Union[tf.Tensor, tf.Variable], - /, - *, - batch_dims: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - try: - return tf.gather_nd(params, indices, batch_dims=batch_dims) - except Exception: # fall back to compositional implementation - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = list(zip(params, indices)) - else: - zip_list = [ - (p, i) - for z in [zip(p1, i1) for p1, i1 in zip_list] - for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = tf.stack(result) - result = tf.reshape( - result, tf.concat([params.shape[0:batch_dims], result.shape[1:]], 0) - ) - return result - - def gather_nd_helper(params, indices): indices_shape = tf.shape(indices) params_shape = tf.shape(params) @@ -146,14 +147,41 @@ def gather_nd_helper(params, indices): return res -def get_item( - x: Union[tf.Tensor, tf.Variable], +def gather_nd( + params: Union[tf.Tensor, tf.Variable], + indices: Union[tf.Tensor, tf.Variable], /, - query: Union[tf.Tensor, tf.Variable, Tuple], *, - copy: bool = None, + batch_dims: int = 0, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return x.__getitem__(query) + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + try: + return tf.gather_nd(params, indices, batch_dims=batch_dims) + except Exception: # fall back to compositional implementation + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = list(zip(params, indices)) + else: + zip_list = [ + (p, i) + for z in [zip(p1, i1) for p1, i1 in zip_list] + for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = tf.stack(result) + result = tf.reshape( + result, tf.concat([params.shape[0:batch_dims], result.shape[1:]], 0) + ) + return result def get_num_dims(x, /, *, as_array=False): @@ -258,47 +286,16 @@ def inplace_update( return val -def inplace_variables_supported(): - return True - - -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, (tf.Tensor, tf.Variable)): - if exclusive and isinstance(x, tf.Variable): - return False - return True - return False - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def isin( - elements: tf.Tensor, - test_elements: tf.Tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> tf.Tensor: - input_shape = elements.shape - - if tf.rank(elements) == 0: - elements = tf.reshape(elements, [1]) - if tf.rank(test_elements) == 0: - test_elements = tf.reshape(test_elements, [1]) - if not assume_unique: - test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] - - elements = tf.reshape(elements, [-1]) - test_elements = tf.reshape(test_elements, [-1]) - - output = tf.reduce_any( - tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 - ) - return tf.reshape(output, input_shape) ^ invert +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view -def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: - return x.dtype.size +def inplace_variables_supported(): + return True def multiprocessing(context: Optional[str] = None): @@ -350,6 +347,9 @@ def scatter_flat( return res +scatter_flat.support_native_out = True + + @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def scatter_nd( indices: Union[tf.Tensor, tf.Variable], @@ -406,6 +406,9 @@ def scatter_nd( return res +scatter_nd.support_native_out = True + + def shape( x: Union[tf.Tensor, tf.Variable], /, @@ -418,35 +421,6 @@ def shape( return ivy.Shape(x.shape) -def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: - return x.numpy().tolist() - - -def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: - # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions - if ( - ivy.is_array(x) - and get_num_dims(x) == 0 - and ivy.as_native_dtype(x.dtype) is tf.bfloat16 - ): - x = tf.expand_dims(x, 0) - if copy: - return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) - else: - return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) - if copy: - return np.array(tf.convert_to_tensor(x)) - else: - return np.asarray(tf.convert_to_tensor(x)) - - -def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: - ret = to_numpy(x).item() - if x.dtype == tf.bfloat16: - return float(ret) - return ret - - def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -533,10 +507,32 @@ def _vmap(*args, **kwargs): return _vmap -get_item.partial_mixed_handler = lambda x, query, **kwargs: ( - all(_check_query(i) for i in query) - if isinstance(query, tuple) - else _check_query(query) -) -scatter_flat.support_native_out = True -scatter_nd.support_native_out = True +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def isin( + elements: tf.Tensor, + test_elements: tf.Tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> tf.Tensor: + input_shape = elements.shape + + if tf.rank(elements) == 0: + elements = tf.reshape(elements, [1]) + if tf.rank(test_elements) == 0: + test_elements = tf.reshape(test_elements, [1]) + if not assume_unique: + test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] + + elements = tf.reshape(elements, [-1]) + test_elements = tf.reshape(test_elements, [-1]) + + output = tf.reduce_any( + tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 + ) + return tf.reshape(output, input_shape) ^ invert + + +def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: + return x.dtype.size diff --git a/ivy/functional/backends/tensorflow/gradients.py b/ivy/functional/backends/tensorflow/gradients.py index 6968f2f245ac8..3eb4207d2df36 100644 --- a/ivy/functional/backends/tensorflow/gradients.py +++ b/ivy/functional/backends/tensorflow/gradients.py @@ -21,8 +21,17 @@ ) -# --- Helpers --- # -# --------------- # +def variable(x, /): + with tf.device(ivy.dev(x, as_native=True)): + return tf.Variable(x, trainable=True) + + +def is_variable(x, /, *, exclusive=False): + return isinstance(x, tf.Variable) + + +def variable_data(x: tf.Variable, /) -> tf.Variable: + return x.value() def _grad_func(y, xs, xs_required, tape): @@ -54,10 +63,6 @@ def _grad_func(y, xs, xs_required, tape): return grads -# --- Main --- # -# ------------ # - - def execute_with_gradients( func, xs: Union[tf.Tensor, tf.Variable], @@ -112,6 +117,81 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) +def value_and_grad(func): + def grad_fn(xs): + grads = ivy.nested_map( + xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False + ) + with tf.GradientTape(watch_accessed_variables=False) as tape: + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + tape.watch(xs) + y = func(xs) + y = y.to_native(y) + grads_ = tape.gradient(y, xs) + grads_ = ivy.nested_map( + grads_, + lambda x: ivy.to_ivy(x), + include_derived=True, + ) + grads_ = ivy.to_ivy(grads_) + grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) + grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) + xs = ivy.to_ivy(xs) + if isinstance(xs, ivy.Array): + grads = grads_ + else: + ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) + y = ivy.to_ivy(y) + return y, grads + + return grad_fn + + +def stop_gradient( + x: Union[tf.Tensor, tf.Variable], + /, + *, + preserve_type: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + is_var = is_variable(x) + x = tf.stop_gradient(x) + if is_var and preserve_type: + return variable(x) + return x + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(x_in): + with tf.GradientTape(persistent=True) as tape: + ivy.nested_map(x_in, ivy.copy_array) + x_in = ivy.to_native(x_in, nested=True) + tape.watch(x_in) + y = grad_fn(x_in) + + # Deal with multiple outputs + if not isinstance(y, ivy.NativeArray): + jacobian = ivy.nested_map( + y, + lambda yi: ivy.to_ivy( + tape.jacobian(yi, x_in, unconnected_gradients="zero"), + nested=True, + ), + include_derived=True, + ) + else: + jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) + return jacobian + + return callback_fn + + def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -160,93 +240,5 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) -def is_variable(x, /, *, exclusive=False): - return isinstance(x, tf.Variable) - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(x_in): - with tf.GradientTape(persistent=True) as tape: - ivy.nested_map(x_in, ivy.copy_array) - x_in = ivy.to_native(x_in, nested=True) - tape.watch(x_in) - y = grad_fn(x_in) - - # Deal with multiple outputs - if not isinstance(y, ivy.NativeArray): - jacobian = ivy.nested_map( - y, - lambda yi: ivy.to_ivy( - tape.jacobian(yi, x_in, unconnected_gradients="zero"), - nested=True, - ), - include_derived=True, - ) - else: - jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) - return jacobian - - return callback_fn - - -def stop_gradient( - x: Union[tf.Tensor, tf.Variable], - /, - *, - preserve_type: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - is_var = is_variable(x) - x = tf.stop_gradient(x) - if is_var and preserve_type: - return variable(x) - return x - - -def value_and_grad(func): - def grad_fn(xs): - grads = ivy.nested_map( - xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False - ) - with tf.GradientTape(watch_accessed_variables=False) as tape: - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - tape.watch(xs) - y = func(xs) - y = y.to_native(y) - grads_ = tape.gradient(y, xs) - grads_ = ivy.nested_map( - grads_, - lambda x: ivy.to_ivy(x), - include_derived=True, - ) - grads_ = ivy.to_ivy(grads_) - grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) - grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) - xs = ivy.to_ivy(xs) - if isinstance(xs, ivy.Array): - grads = grads_ - else: - ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) - y = ivy.to_ivy(y) - return y, grads - - return grad_fn - - -def variable(x, /): - with tf.device(ivy.dev(x, as_native=True)): - return tf.Variable(x, trainable=True) - - -def variable_data(x: tf.Variable, /) -> tf.Variable: - return x.value() - - grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/tensorflow/layers.py b/ivy/functional/backends/tensorflow/layers.py index d361fdd9827dc..c3070846e21cb 100644 --- a/ivy/functional/backends/tensorflow/layers.py +++ b/ivy/functional/backends/tensorflow/layers.py @@ -17,10 +17,6 @@ ) -# --- Helpers --- # -# --------------- # - - def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): if filter_format == "channel_first": filters = tf.transpose(filters, (*range(2, dims + 2), 1, 0)) @@ -37,24 +33,6 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters -def _output_shape( - x_shape, filter_shape, output_shape, strides, padding, dims, dilations -): - dilations = [dilations] * dims if isinstance(dilations, int) else dilations - strides = [strides] * dims if isinstance(strides, int) else strides - if output_shape is None: - out_shape = [ - _deconv_length( - x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] - ) - for i in range(dims) - ] - output_shape = [x_shape[0], *out_shape, filter_shape[-2]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] - return output_shape - - def _pad_before_conv(x, filters, strides, padding, dims, dilations): dilations = [dilations] * dims if isinstance(dilations, int) else dilations strides = [strides] * dims if isinstance(strides, int) else strides @@ -86,8 +64,22 @@ def _pad_before_conv(x, filters, strides, padding, dims, dilations): ) -# --- Main --- # -# ------------ # +def _output_shape( + x_shape, filter_shape, output_shape, strides, padding, dims, dilations +): + dilations = [dilations] * dims if isinstance(dilations, int) else dilations + strides = [strides] * dims if isinstance(strides, int) else strides + if output_shape is None: + out_shape = [ + _deconv_length( + x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] + ) + for i in range(dims) + ] + output_shape = [x_shape[0], *out_shape, filter_shape[-2]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] + return output_shape @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) @@ -213,6 +205,32 @@ def conv2d_transpose( return res +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def depthwise_conv2d( + x: Union[tf.Tensor, tf.Variable], + filters: Union[tf.Tensor, tf.Variable], + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + if tf.rank(filters) == 3: + filters = tf.expand_dims(filters, -1) + x = _pad_before_conv(x, filters, strides, padding, 2, dilations) + strides = [1, strides[0], strides[1], 1] + res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res + + @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def conv3d( x: Union[tf.Tensor, tf.Variable], @@ -473,29 +491,3 @@ def conv_general_transpose( if data_format == "channel_first": res = tf.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def depthwise_conv2d( - x: Union[tf.Tensor, tf.Variable], - filters: Union[tf.Tensor, tf.Variable], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - if tf.rank(filters) == 3: - filters = tf.expand_dims(filters, -1) - x = _pad_before_conv(x, filters, strides, padding, 2, dilations) - strides = [1, strides[0], strides[1], 1] - res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res diff --git a/ivy/functional/backends/tensorflow/linear_algebra.py b/ivy/functional/backends/tensorflow/linear_algebra.py index 28c91cb2bcac4..5281feade12b5 100644 --- a/ivy/functional/backends/tensorflow/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/linear_algebra.py @@ -79,21 +79,6 @@ def det( return tf.linalg.det(x) -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def diag( - x: Union[tf.Tensor, tf.Variable], - /, - *, - k: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.diag(x, k=k) - - @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def diagonal( x: Union[tf.Tensor, tf.Variable], @@ -629,21 +614,6 @@ def trace( return tf.experimental.numpy.trace(x, offset=offset, axis1=axis1, axis2=axis2) -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, - backend_version, -) -def vander( - x: Union[tf.Tensor, tf.Variable], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -697,6 +667,36 @@ def vector_norm( return tf.reduce_sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def diag( + x: Union[tf.Tensor, tf.Variable], + /, + *, + k: int = 0, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.diag(x, k=k) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, + backend_version, +) +def vander( + x: Union[tf.Tensor, tf.Variable], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes( { "2.13.0 and below": ( diff --git a/ivy/functional/backends/tensorflow/manipulation.py b/ivy/functional/backends/tensorflow/manipulation.py index a26c08e2c8abe..41800bca81d4f 100644 --- a/ivy/functional/backends/tensorflow/manipulation.py +++ b/ivy/functional/backends/tensorflow/manipulation.py @@ -15,47 +15,12 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _reshape_fortran_tf(x, shape): if len(x.shape) > 0: x = tf.transpose(x) return tf.transpose(tf.reshape(x, shape[::-1])) -# --- Main --- # -# ------------ # - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def clip( - x: Union[tf.Tensor, tf.Variable], - x_min: Union[Number, tf.Tensor, tf.Variable], - x_max: Union[Number, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): - promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) - promoted_type = ivy.as_native_dtype( - ivy.promote_types(promoted_type, x_max.dtype) - ) - x = tf.cast(x, promoted_type) - x_min = tf.cast(x_min, promoted_type) - x_max = tf.cast(x_max, promoted_type) - if tf.size(x) == 0: - ret = x - elif x.dtype == tf.bool: - ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) - ret = tf.cast(ret, x.dtype) - else: - ret = tf.clip_by_value(x, x_min, x_max) - return ret - - # Array API Standard # # -------------------# @@ -89,14 +54,6 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None -): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width, constant_values=value) - - def expand_dims( x: Union[tf.Tensor, tf.Variable], /, @@ -149,18 +106,6 @@ def permute_dims( return tf.transpose(x, perm=axes) -@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) -def repeat( - x: Union[tf.Tensor, tf.Variable], - /, - repeats: Union[int, List[int]], - *, - axis: int = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.repeat(x, repeats, axis) - - def reshape( x: Union[tf.Tensor, tf.Variable], /, @@ -203,47 +148,6 @@ def roll( return ret -# Extra # -# ------# - - -def split( - x: Union[tf.Tensor, tf.Variable], - /, - *, - copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> Union[tf.Tensor, tf.Variable]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - dim_size = tf.shape(x)[axis] - num_or_size_splits = int(dim_size) - if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): - num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) - num_or_size_splits = num_or_size_splits.numpy().tolist() - elif isinstance(num_or_size_splits, int) and with_remainder: - num_chunks = x.shape[axis] / num_or_size_splits - num_chunks_int = math.floor(num_chunks) - remainder = num_chunks - num_chunks_int - if remainder != 0: - num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ - int(remainder * num_or_size_splits) - ] - - return tf.split(x, num_or_size_splits, axis) - - def squeeze( x: Union[tf.Tensor, tf.Variable], /, @@ -298,25 +202,57 @@ def stack( raise ivy.utils.exceptions.IvyIndexError(e) -def swapaxes( - x, - axis0, - axis1, +# Extra # +# ------# + + +def split( + x: Union[tf.Tensor, tf.Variable], /, *, copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> Union[tf.Tensor, tf.Variable]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + dim_size = tf.shape(x)[axis] + num_or_size_splits = int(dim_size) + if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): + num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) + num_or_size_splits = num_or_size_splits.numpy().tolist() + elif isinstance(num_or_size_splits, int) and with_remainder: + num_chunks = x.shape[axis] / num_or_size_splits + num_chunks_int = math.floor(num_chunks) + remainder = num_chunks - num_chunks_int + if remainder != 0: + num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ + int(remainder * num_or_size_splits) + ] + + return tf.split(x, num_or_size_splits, axis) + + +@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) +def repeat( + x: Union[tf.Tensor, tf.Variable], + /, + repeats: Union[int, List[int]], + *, + axis: int = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -): - x_shape = x.shape - num_dims = len(x_shape) - axis0 %= num_dims - axis1 %= num_dims - config = list(range(num_dims)) - config.pop(axis0) - config.insert(axis0, axis1) - config.pop(axis1) - config.insert(axis1, axis0) - return tf.transpose(x, config) +) -> Union[tf.Tensor, tf.Variable]: + return tf.repeat(x, repeats, axis) @with_unsupported_dtypes( @@ -357,6 +293,68 @@ def tile( return tf.tile(x, repeats) +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None +): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width, constant_values=value) + + +def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width) + + +def swapaxes( + x, + axis0, + axis1, + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +): + x_shape = x.shape + num_dims = len(x_shape) + axis0 %= num_dims + axis1 %= num_dims + config = list(range(num_dims)) + config.pop(axis0) + config.insert(axis0, axis1) + config.pop(axis1) + config.insert(axis1, axis0) + return tf.transpose(x, config) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def clip( + x: Union[tf.Tensor, tf.Variable], + x_min: Union[Number, tf.Tensor, tf.Variable], + x_max: Union[Number, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): + promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) + promoted_type = ivy.as_native_dtype( + ivy.promote_types(promoted_type, x_max.dtype) + ) + x = tf.cast(x, promoted_type) + x_min = tf.cast(x_min, promoted_type) + x_max = tf.cast(x_max, promoted_type) + if tf.size(x) == 0: + ret = x + elif x.dtype == tf.bool: + ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) + ret = tf.cast(ret, x.dtype) + else: + ret = tf.clip_by_value(x, x_min, x_max) + return ret + + def unstack( x: Union[tf.Tensor, tf.Variable], /, @@ -371,9 +369,3 @@ def unstack( if keepdims: return [tf.expand_dims(r, axis) for r in ret] return ret - - -def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width) diff --git a/ivy/functional/backends/tensorflow/random.py b/ivy/functional/backends/tensorflow/random.py index ae0f6b4a26c31..b0f3b97a3d877 100644 --- a/ivy/functional/backends/tensorflow/random.py +++ b/ivy/functional/backends/tensorflow/random.py @@ -22,6 +22,47 @@ from . import backend_version +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, tf.Tensor, tf.Variable] = 0.0, + high: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, + dtype: DType, + device: str, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + shape = _check_bounds_and_get_shape(low, high, shape).shape + low = tf.cast(low, dtype) + high = tf.cast(high, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) + + +def random_normal( + *, + mean: Union[float, tf.Tensor, tf.Variable] = 0.0, + std: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: DType, + seed: Optional[int] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + mean = tf.cast(mean, dtype) + std = tf.cast(std, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) + + @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) def multinomial( population_size: int, @@ -101,47 +142,6 @@ def randint( return tf.cast(tf.random.uniform(shape, low, high, "float32", seed=seed), dtype) -def random_normal( - *, - mean: Union[float, tf.Tensor, tf.Variable] = 0.0, - std: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: DType, - seed: Optional[int] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - mean = tf.cast(mean, dtype) - std = tf.cast(std, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) - - -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, tf.Tensor, tf.Variable] = 0.0, - high: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, - dtype: DType, - device: str, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - shape = _check_bounds_and_get_shape(low, high, shape).shape - low = tf.cast(low, dtype) - high = tf.cast(high, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) - - def seed(*, seed_value: int = 0) -> None: tf.random.set_seed(seed_value) return diff --git a/ivy/functional/backends/tensorflow/searching.py b/ivy/functional/backends/tensorflow/searching.py index 13d461c8c54c5..047441cff9b08 100644 --- a/ivy/functional/backends/tensorflow/searching.py +++ b/ivy/functional/backends/tensorflow/searching.py @@ -77,29 +77,6 @@ def argmin( return tf.cast(ret, dtype) if dtype is not None else ret -# Extra # -# ----- # - - -def argwhere( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(x, tf.Variable): - x_ndim = x.shape.rank - else: - x_ndim = x.ndim - if x_ndim == 0: - return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") - where_x = tf.experimental.numpy.nonzero(x) - res = tf.experimental.numpy.concatenate( - [tf.expand_dims(item, -1) for item in where_x], -1 - ) - return res - - def nonzero( x: Union[tf.Tensor, tf.Variable], /, @@ -137,3 +114,26 @@ def where( ) -> Union[tf.Tensor, tf.Variable]: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return tf.cast(tf.experimental.numpy.where(condition, x1, x2), x1.dtype) + + +# Extra # +# ----- # + + +def argwhere( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(x, tf.Variable): + x_ndim = x.shape.rank + else: + x_ndim = x.ndim + if x_ndim == 0: + return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") + where_x = tf.experimental.numpy.nonzero(x) + res = tf.experimental.numpy.concatenate( + [tf.expand_dims(item, -1) for item in where_x], -1 + ) + return res diff --git a/ivy/functional/backends/tensorflow/sorting.py b/ivy/functional/backends/tensorflow/sorting.py index 2e1202a943ee0..06403a1d358ca 100644 --- a/ivy/functional/backends/tensorflow/sorting.py +++ b/ivy/functional/backends/tensorflow/sorting.py @@ -27,6 +27,29 @@ def argsort( return tf.cast(ret, dtype=tf.int64) +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def sort( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # TODO: handle stable sort when it's supported in tensorflow + # currently it supports only quicksort (unstable) + direction = "DESCENDING" if descending else "ASCENDING" + x = tf.convert_to_tensor(x) + is_bool = x.dtype.is_bool + if is_bool: + x = tf.cast(x, tf.int32) + ret = tf.sort(x, axis=axis, direction=direction) + if is_bool: + ret = tf.cast(ret, dtype=tf.bool) + return ret + + # msort @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def msort( @@ -80,26 +103,3 @@ def searchsorted( if is_supported_int_ret_dtype: return tf.searchsorted(x, v, side=side, out_type=ret_dtype) return tf.cast(tf.searchsorted(x, v, side=side), ret_dtype) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def sort( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # TODO: handle stable sort when it's supported in tensorflow - # currently it supports only quicksort (unstable) - direction = "DESCENDING" if descending else "ASCENDING" - x = tf.convert_to_tensor(x) - is_bool = x.dtype.is_bool - if is_bool: - x = tf.cast(x, tf.int32) - ret = tf.sort(x, axis=axis, direction=direction) - if is_bool: - ret = tf.cast(ret, dtype=tf.bool) - return ret diff --git a/ivy/functional/backends/tensorflow/statistical.py b/ivy/functional/backends/tensorflow/statistical.py index 4ab04d803cd90..53fc8b5bb57e6 100644 --- a/ivy/functional/backends/tensorflow/statistical.py +++ b/ivy/functional/backends/tensorflow/statistical.py @@ -9,82 +9,21 @@ from . import backend_version from ivy.utils.einsum_parser import legalise_einsum_expr - -# --- Helpers --- # -# --------------- # - - -def _infer_dtype(dtype: tf.DType): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype - - -# --- Main --- # -# ------------ # - - -# Extra # -# ------# +# Array API Standard # +# -------------------# -@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) -def cumprod( +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def min( x: Union[tf.Tensor, tf.Variable], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumprod(x, axis, exclusive, reverse) - - -def cumsum( - x: Union[tf.Tensor, tf.Variable], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumsum(x, axis, exclusive, reverse) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("unsigned", "int8", "int16")}, - backend_version, -) -def einsum( - equation: str, - *operands: Union[tf.Tensor, tf.Variable], + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - dtype = _get_promoted_type_of_operands(operands) - equation = legalise_einsum_expr(*[equation, *operands]) - return tf.cast(tf.einsum(equation, *operands), dtype) + axis = tuple(axis) if isinstance(axis, list) else axis + return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -112,21 +51,11 @@ def mean( return tf.math.reduce_mean(x, axis=axis, keepdims=keepdims) -# Array API Standard # -# -------------------# - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def min( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - axis = tuple(axis) if isinstance(axis, list) else axis - return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) +def _infer_dtype(dtype: tf.DType): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype def prod( @@ -215,3 +144,65 @@ def var( * size / (size - correction) ) + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) +def cumprod( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumprod(x, axis, exclusive, reverse) + + +def cumsum( + x: Union[tf.Tensor, tf.Variable], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumsum(x, axis, exclusive, reverse) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("unsigned", "int8", "int16")}, + backend_version, +) +def einsum( + equation: str, + *operands: Union[tf.Tensor, tf.Variable], + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = _get_promoted_type_of_operands(operands) + equation = legalise_einsum_expr(*[equation, *operands]) + return tf.cast(tf.einsum(equation, *operands), dtype) diff --git a/ivy/functional/backends/torch/activations.py b/ivy/functional/backends/torch/activations.py index f40e405bf579b..afb85ebe96f5d 100644 --- a/ivy/functional/backends/torch/activations.py +++ b/ivy/functional/backends/torch/activations.py @@ -17,7 +17,23 @@ from . import backend_version -sigmoid.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def relu( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None +) -> torch.Tensor: + return torch.relu(x) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def leaky_relu( + x: torch.Tensor, + /, + *, + alpha: float = 0.2, + complex_mode="jax", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.leaky_relu(x, alpha) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -36,31 +52,43 @@ def gelu( return torch.nn.functional.gelu(x) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "complex", - "float16", - ) - }, - backend_version, -) -def hardswish( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + if not ivy.is_array(x): + x = torch.tensor(x) + return torch.sigmoid(x, out=out) + + +sigmoid.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def softmax( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nn.functional.hardswish(x) + if axis is None: + axis = -1 + return torch.nn.functional.softmax(x, axis) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def leaky_relu( +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def softplus( x: torch.Tensor, /, *, - alpha: float = 0.2, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, complex_mode="jax", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nn.functional.leaky_relu(x, alpha) + kwargs = { + k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None + } + return torch.nn.functional.softplus(x, **kwargs) @with_unsupported_dtypes( @@ -97,44 +125,16 @@ def mish(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.mish(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def relu( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.relu(x) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - if not ivy.is_array(x): - x = torch.tensor(x) - return torch.sigmoid(x, out=out) - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def softmax( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if axis is None: - axis = -1 - return torch.nn.functional.softmax(x, axis) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def softplus( - x: torch.Tensor, - /, - *, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, - complex_mode="jax", - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "complex", + "float16", + ) + }, + backend_version, +) +def hardswish( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - kwargs = { - k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None - } - return torch.nn.functional.softplus(x, **kwargs) + return torch.nn.functional.hardswish(x) diff --git a/ivy/functional/backends/torch/creation.py b/ivy/functional/backends/torch/creation.py index 8f1603b9c8f79..7c5a41482e44c 100644 --- a/ivy/functional/backends/torch/creation.py +++ b/ivy/functional/backends/torch/creation.py @@ -25,10 +25,6 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - # noinspection PyProtectedMember @@ -51,28 +47,6 @@ def _differentiable_linspace(start, stop, num, *, device, dtype=None): return res -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - -def _stack_tensors(x, dtype): - if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): - for i, item in enumerate(x): - x[i] = _stack_tensors(item, dtype) - x = torch.stack(x) - else: - if isinstance(x, (list, tuple)): - if isinstance(x[0], torch.Tensor): - x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) - else: - x = torch.as_tensor(x, dtype=dtype) - return x - - -# --- Main --- # -# ------------ # - - @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def arange( start: float, @@ -104,6 +78,23 @@ def arange( return torch.arange(start, stop, step, dtype=dtype, device=device) +arange.support_native_out = True + + +def _stack_tensors(x, dtype): + if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): + for i, item in enumerate(x): + x[i] = _stack_tensors(item, dtype) + x = torch.stack(x) + else: + if isinstance(x, (list, tuple)): + if isinstance(x[0], torch.Tensor): + x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) + else: + x = torch.as_tensor(x, dtype=dtype) + return x + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) @asarray_to_native_arrays_and_back @asarray_infer_device @@ -146,17 +137,6 @@ def asarray( return ret.clone().detach() if copy else ret -def copy_array( - x: torch.Tensor, - *, - to_ivy_array: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if to_ivy_array: - return ivy.to_ivy(x.clone()) - return x.clone() - - def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -172,6 +152,9 @@ def empty( ) +empty.support_native_out = True + + def empty_like( x: torch.Tensor, /, @@ -242,23 +225,14 @@ def eye( return ret +eye.support_native_out = True + + def from_dlpack(x, /, *, out: Optional[torch.Tensor] = None): x = x.detach() if x.requires_grad else x return torch.utils.dlpack.from_dlpack(x) -def frombuffer( - buffer: bytes, - dtype: Optional[torch.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> torch.Tensor: - buffer_copy = copy.deepcopy(buffer) - dtype = ivy.as_native_dtype(dtype) - - return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) - - def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -280,6 +254,9 @@ def full( ) +full.support_native_out = True + + def full_like( x: torch.Tensor, /, @@ -293,6 +270,10 @@ def full_like( return torch.full_like(x, fill_value, dtype=dtype, device=device) +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + @with_unsupported_device_and_dtypes( {"2.0.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -344,6 +325,9 @@ def linspace( return ans.to(dtype) +linspace.support_native_out = True + + def linspace_helper(start, stop, num, axis=None, *, dtype=None, device): num = num.detach().numpy().item() if isinstance(num, torch.Tensor) else num start_is_array = isinstance(start, torch.Tensor) @@ -444,46 +428,6 @@ def meshgrid( return res -def one_hot( - indices: torch.Tensor, - depth: int, - /, - *, - on_value: Optional[torch.Tensor] = None, - off_value: Optional[torch.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = torch.float32 - else: - if not on_none: - dtype = torch.tensor(on_value).dtype - elif not off_none: - dtype = torch.tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) - off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) - - res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) - - if not on_none or not off_none: - res = torch.where(res == 1, on_value, off_value) - - if axis is not None: - res = torch.moveaxis(res, -1, axis) - - return res.to(device, dtype) - - def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -494,21 +438,18 @@ def ones( return torch.ones(shape, dtype=dtype, device=device, out=out) -def ones_like_v_0p1p12_to_0p2p0( +ones.support_native_out = True + + +def ones_like_v_0p4p0_and_above( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -): - if len(x.shape) == 1: - for i in range(x.shape[0]): - x[i] = 1 - return x - for i in range(x.shape[0]): - x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) - return x +) -> torch.Tensor: + return torch.ones_like(x, dtype=dtype, device=device) def ones_like_v_0p3p0_to_0p3p1( @@ -522,15 +463,21 @@ def ones_like_v_0p3p0_to_0p3p1( return torch.ones_like(x, out=out) -def ones_like_v_0p4p0_and_above( +def ones_like_v_0p1p12_to_0p2p0( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.ones_like(x, dtype=dtype, device=device) +): + if len(x.shape) == 1: + for i in range(x.shape[0]): + x[i] = 1 + return x + for i in range(x.shape[0]): + x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) + return x def tril( @@ -539,26 +486,16 @@ def tril( return torch.tril(x, diagonal=k, out=out) +tril.support_native_out = True + + def triu( x: torch.Tensor, /, *, k: int = 0, out: Optional[torch.Tensor] = None ) -> torch.Tensor: return torch.triu(x, diagonal=k, out=out) -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: torch.device, -) -> Tuple[torch.Tensor]: - n_cols = n_rows if n_cols is None else n_cols - return tuple( - torch.triu_indices( - row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device - ) - ) +triu.support_native_out = True def zeros( @@ -571,6 +508,9 @@ def zeros( return torch.zeros(shape, dtype=dtype, device=device, out=out) +zeros.support_native_out = True + + def zeros_like( x: torch.Tensor, /, @@ -587,12 +527,82 @@ def zeros_like( array = asarray -arange.support_native_out = True -empty.support_native_out = True -eye.support_native_out = True -full.support_native_out = True -linspace.support_native_out = True -ones.support_native_out = True -tril.support_native_out = True -triu.support_native_out = True -zeros.support_native_out = True + + +def copy_array( + x: torch.Tensor, + *, + to_ivy_array: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if to_ivy_array: + return ivy.to_ivy(x.clone()) + return x.clone() + + +def one_hot( + indices: torch.Tensor, + depth: int, + /, + *, + on_value: Optional[torch.Tensor] = None, + off_value: Optional[torch.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = torch.float32 + else: + if not on_none: + dtype = torch.tensor(on_value).dtype + elif not off_none: + dtype = torch.tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) + off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) + + res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) + + if not on_none or not off_none: + res = torch.where(res == 1, on_value, off_value) + + if axis is not None: + res = torch.moveaxis(res, -1, axis) + + return res.to(device, dtype) + + +def frombuffer( + buffer: bytes, + dtype: Optional[torch.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> torch.Tensor: + buffer_copy = copy.deepcopy(buffer) + dtype = ivy.as_native_dtype(dtype) + + return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) + + +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: torch.device, +) -> Tuple[torch.Tensor]: + n_cols = n_rows if n_cols is None else n_cols + return tuple( + torch.triu_indices( + row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device + ) + ) diff --git a/ivy/functional/backends/torch/data_type.py b/ivy/functional/backends/torch/data_type.py index 67f5aeabdaf84..05b4ccb37a4c4 100644 --- a/ivy/functional/backends/torch/data_type.py +++ b/ivy/functional/backends/torch/data_type.py @@ -23,6 +23,7 @@ torch.complex128: "complex128", torch.bool: "bool", } + native_dtype_dict = { "int8": torch.int8, "int16": torch.int16, @@ -67,6 +68,73 @@ def smallest_normal(self): return self._torch_finfo.tiny +# Array API Standard # +# -------------------# + + +def astype( + x: torch.Tensor, + dtype: torch.dtype, + /, + *, + copy: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return x.clone() if copy else x + return x.to(dtype) + + +def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: + try: + return list(torch.broadcast_tensors(*arrays)) + except RuntimeError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: torch.Tensor, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return torch.broadcast_to(x.reshape(-1), shape) + return torch.broadcast_to(x, shape) + + +@_handle_nestable_dtype_info +def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return Finfo(torch.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return torch.iinfo(ivy.as_native_dtype(type)) + + +def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: + input = [] + for val in arrays_and_dtypes: + torch_val = as_native_dtype(val) + if isinstance(torch_val, torch.dtype): + torch_val = torch.tensor(1, dtype=torch_val) + input.append(torch_val) + + result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) + + for i in range(2, len(input)): + result = torch.tensor(1, dtype=torch.result_type(result, input[i])) + return as_ivy_dtype(result.dtype) + + # Extra # # ------# @@ -136,44 +204,6 @@ def as_native_dtype( ) -# Array API Standard # -# -------------------# - - -def astype( - x: torch.Tensor, - dtype: torch.dtype, - /, - *, - copy: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return x.clone() if copy else x - return x.to(dtype) - - -def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: - try: - return list(torch.broadcast_tensors(*arrays)) - except RuntimeError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: torch.Tensor, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return torch.broadcast_to(x.reshape(-1), shape) - return torch.broadcast_to(x, shape) - - def dtype(x: Union[torch.tensor, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -194,20 +224,6 @@ def dtype_bits(dtype_in: Union[torch.dtype, str, np.dtype], /) -> int: ) -@_handle_nestable_dtype_info -def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return Finfo(torch.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return torch.iinfo(ivy.as_native_dtype(type)) - - def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -215,18 +231,3 @@ def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: return True else: return False - - -def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: - input = [] - for val in arrays_and_dtypes: - torch_val = as_native_dtype(val) - if isinstance(torch_val, torch.dtype): - torch_val = torch.tensor(1, dtype=torch_val) - input.append(torch_val) - - result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) - - for i in range(2, len(input)): - result = torch.tensor(1, dtype=torch.result_type(result, input[i])) - return as_ivy_dtype(result.dtype) diff --git a/ivy/functional/backends/torch/device.py b/ivy/functional/backends/torch/device.py index 4a86dcf02f129..e7f768a7f5fe1 100644 --- a/ivy/functional/backends/torch/device.py +++ b/ivy/functional/backends/torch/device.py @@ -18,26 +18,37 @@ torch_scatter = None +# API # +# ----# -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._prof = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True - ) - - def start(self): - self._prof.__enter__() - def stop(self): - self._prof.__exit__(None, None, None) - self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) +def dev( + x: torch.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, torch.device]: + dv = x.device + if as_native: + if isinstance(dv, torch.device): + dv = dv.type + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return torch.device(dv.replace("gpu", "mps")) + return torch.device(dv.replace("gpu", "cuda")) + return as_ivy_dev(dv) - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def to_device( + x: torch.Tensor, + device: torch.device, + /, + *, + stream: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if device is None: + return x + ret = x.to(as_native_dev(device)) + if isinstance(x, torch.nn.Parameter): + return torch.nn.Parameter(ret) + return ret def as_ivy_dev(device: torch.device, /): @@ -78,21 +89,10 @@ def clear_cached_mem_on_dev(device: Union[ivy.Device, torch.device], /) -> None: mps.empty_cache() -# API # -# ----# - - -def dev( - x: torch.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, torch.device]: - dv = x.device - if as_native: - if isinstance(dv, torch.device): - dv = dv.type - elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return torch.device(dv.replace("gpu", "mps")) - return torch.device(dv.replace("gpu", "cuda")) - return as_ivy_dev(dv) +def num_gpus() -> int: + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return 1 + return torch.cuda.device_count() def gpu_is_available() -> bool: @@ -103,6 +103,13 @@ def gpu_is_available() -> bool: return False +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + if importlib.util.find_spec("torch_xla") is not None: + return True + return False + + def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( *args, device_shifting_dev=device_shifting_dev, **kwargs @@ -114,30 +121,22 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return fn(*args, **kwargs) -def num_gpus() -> int: - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return 1 - return torch.cuda.device_count() +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._prof = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True + ) + def start(self): + self._prof.__enter__() -def to_device( - x: torch.Tensor, - device: torch.device, - /, - *, - stream: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if device is None: - return x - ret = x.to(as_native_dev(device)) - if isinstance(x, torch.nn.Parameter): - return torch.nn.Parameter(ret) - return ret + def stop(self): + self._prof.__exit__(None, None, None) + self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) + def __enter__(self): + self.start() -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - if importlib.util.find_spec("torch_xla") is not None: - return True - return False + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/ivy/functional/backends/torch/elementwise.py b/ivy/functional/backends/torch/elementwise.py index 1d03e1b29cfcc..a0a3f5c11e053 100644 --- a/ivy/functional/backends/torch/elementwise.py +++ b/ivy/functional/backends/torch/elementwise.py @@ -13,49 +13,12 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _cast_for_unary_op(x): if not isinstance(x, torch.Tensor): x = torch.tensor(x) return x -# --- Main --- # -# ------------ # - - -@handle_numpy_arrays_in_specific_backend -def abs( - x: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x = _cast_for_unary_op(x) - if x.dtype is torch.bool: - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.abs(x, out=out) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.acos(x, out=out) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.acosh(x, out=out) - - @handle_numpy_arrays_in_specific_backend def add( x1: Union[float, torch.Tensor], @@ -71,69 +34,48 @@ def add( return torch.add(x1, x2, out=out) -def angle( - input: torch.Tensor, +add.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def bitwise_xor( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, - deg: Optional[bool] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if deg: - return torch.angle(input, out=out) * (180 / pi) - else: - return torch.angle(input, out=out) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.asin(x, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_xor(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.asinh(x, out=out) +bitwise_xor.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.atan(x, out=out) +def imag( + val: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if val.dtype not in (torch.complex64, torch.complex128): + ret = torch.imag(val.to(torch.complex64)) + return ret.to(val.dtype) + return torch.imag(val) -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version -) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) -@handle_numpy_arrays_in_specific_backend -def atan2( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.atan2(x1, x2, out=out) +imag.support_native_out = False -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.atanh(x, out=out) + return torch.expm1(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_and( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_and(x1, x2, out=out) +expm1.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -145,49 +87,68 @@ def bitwise_invert( return torch.bitwise_not(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +bitwise_invert.support_native_out = True + + @handle_numpy_arrays_in_specific_backend -def bitwise_left_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.isfinite(x) + + +@handle_numpy_arrays_in_specific_backend +def isinf( + x: torch.Tensor, /, *, + detect_positive: bool = True, + detect_negative: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_left_shift(x1, x2, out=out) + x = _cast_for_unary_op(x) + if detect_negative and detect_positive: + return torch.isinf(x) + elif detect_negative: + return torch.isneginf(x) + elif detect_positive: + return torch.isposinf(x) + return torch.full_like(x, False, dtype=torch.bool) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_or( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def equal( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_or(x1, x2, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.eq(x1, x2, out=out) + + +equal.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_right_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def less_equal( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) - return torch.bitwise_right_shift(x1, x2, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.less_equal(x1, x2, out=out) + + +less_equal.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_xor( +def bitwise_and( x1: Union[int, bool, torch.Tensor], x2: Union[int, bool, torch.Tensor], /, @@ -195,7 +156,10 @@ def bitwise_xor( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_xor(x1, x2, out=out) + return torch.bitwise_and(x1, x2, out=out) + + +bitwise_and.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @@ -209,162 +173,135 @@ def ceil(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.ceil(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.cos(x, out=out) +ceil.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.cosh(x, out=out) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.deg2rad(x, out=out) + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.floor(x, out=out) -@handle_numpy_arrays_in_specific_backend -def divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2) - if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): - ret = ivy.astype(ret, x1.dtype, copy=False) - else: - ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) - return ret +floor.support_native_out = True -@handle_numpy_arrays_in_specific_backend -def equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmin( + x1: torch.Tensor, + x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.eq(x1, x2, out=out) + return torch.fmin(x1, x2, out=None) -# Extra # -# ------# +fmin.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.erf(x, out=out) + return torch.asin(x, out=out) + + +asin.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.exp(x, out=out) + return torch.asinh(x, out=out) + + +asinh.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def exp2( - x: Union[torch.Tensor, float, list, tuple], +def sign( + x: torch.Tensor, /, *, + np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.exp2(x, out=out) + x = _cast_for_unary_op(x) + if "complex" in str(x.dtype): + if np_variant: + return torch.where( + x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j + ) + return torch.sgn(x, out=out) + return torch.sign(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +sign.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.expm1(x, out=out) + return torch.sqrt(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +sqrt.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.floor(x, out=out) + return torch.cosh(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +cosh.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor_divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.exists(out): - if not ivy.is_float_dtype(out): - return ivy.inplace_update( - out, torch.floor(torch.div(x1, x2)).type(out.dtype) - ) - return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) +def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log10(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmin( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.fmin(x1, x2, out=None) +log10.support_native_out = True -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "complex")}, - backend_version, -) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def fmod( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmod(x1, x2, out=None) +def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log2(x, out=out) -def gcd( - x1: Union[torch.Tensor, int, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.gcd(x1, x2, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log1p(x, out=out) + + +log1p.support_native_out = True + + +@handle_numpy_arrays_in_specific_backend +def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.isnan(x) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater( +def less( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -372,12 +309,14 @@ def greater( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater(x1, x2, out=out) + return torch.lt(x1, x2, out=out) + + +less.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater_equal( +def multiply( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -385,73 +324,56 @@ def greater_equal( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater_equal(x1, x2, out=out) + return torch.multiply(x1, x2, out=out) -def imag( - val: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if val.dtype not in (torch.complex64, torch.complex128): - ret = torch.imag(val.to(torch.complex64)) - return ret.to(val.dtype) - return torch.imag(val) +multiply.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.isfinite(x) + return torch.cos(x, out=out) -@handle_numpy_arrays_in_specific_backend -def isinf( - x: torch.Tensor, - /, - *, - detect_positive: bool = True, - detect_negative: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x = _cast_for_unary_op(x) - if detect_negative and detect_positive: - return torch.isinf(x) - elif detect_negative: - return torch.isneginf(x) - elif detect_positive: - return torch.isposinf(x) - return torch.full_like(x, False, dtype=torch.bool) +cos.support_native_out = True @handle_numpy_arrays_in_specific_backend -def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def logical_not( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.isnan(x) + return torch.logical_not(x.type(torch.bool), out=out) -@handle_numpy_arrays_in_specific_backend -def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.isreal(x) +logical_not.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def lcm( - x1: torch.Tensor, - x2: torch.Tensor, +def divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.lcm(x1, x2, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + ret = torch.div(x1, x2) + if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): + ret = ivy.astype(ret, x1.dtype, copy=False) + else: + ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) + return ret + + +divide.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less( +def greater( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -459,12 +381,15 @@ def less( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.lt(x1, x2, out=out) + return torch.greater(x1, x2, out=out) + + +greater.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less_equal( +def greater_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -472,60 +397,46 @@ def less_equal( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.less_equal(x1, x2, out=out) + return torch.greater_equal(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log(x, out=out) +greater_equal.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.log10(x, out=out) + return torch.acos(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log1p(x, out=out) +acos.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log2(x, out=out) +def lcm( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.lcm(x1, x2, out=out) + + +lcm.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def logaddexp( +def logical_xor( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.logaddexp(x1, x2, out=out) + return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def logaddexp2( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - if not ivy.is_float_dtype(x1): - x1 = x1.type(ivy.default_float_dtype(as_native=True)) - x2 = x2.type(ivy.default_float_dtype(as_native=True)) - return torch.logaddexp2(x1, x2, out=out) +logical_xor.support_native_out = True @handle_numpy_arrays_in_specific_backend @@ -535,12 +446,7 @@ def logical_and( return torch.logical_and(x1.type(torch.bool), x2.type(torch.bool), out=out) -@handle_numpy_arrays_in_specific_backend -def logical_not( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.logical_not(x.type(torch.bool), out=out) +logical_and.support_native_out = True @handle_numpy_arrays_in_specific_backend @@ -550,72 +456,27 @@ def logical_or( return torch.logical_or(x1.type(torch.bool), x2.type(torch.bool), out=out) +logical_or.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def logical_xor( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) +def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.acosh(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +acosh.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def maximum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 >= x2, x1, x2, out=out) - return torch.maximum(x1, x2, out=out) - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def minimum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 <= x2, x1, x2, out=out) - return torch.minimum(x1, x2, out=out) - - -@handle_numpy_arrays_in_specific_backend -def multiply( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.multiply(x1, x2, out=out) +def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sin(x, out=out) -def nan_to_num( - x: torch.Tensor, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if copy: - return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) - else: - x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) - return x +sin.support_native_out = True @handle_numpy_arrays_in_specific_backend @@ -626,28 +487,11 @@ def negative( return torch.neg(x, out=out) -@handle_numpy_arrays_in_specific_backend -def not_equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.not_equal(x1, x2, out=out) - - -@handle_numpy_arrays_in_specific_backend -def positive( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.positive(x) +negative.support_native_out = True @handle_numpy_arrays_in_specific_backend -def pow( +def not_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -655,91 +499,59 @@ def pow( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.pow(x1, x2, out=out) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.rad2deg(x, out=out) + return torch.not_equal(x1, x2, out=out) -def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.real(x) +not_equal.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def reciprocal( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +def tanh( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None ) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.reciprocal(x, out=out) + return torch.tanh(x, out=out) + + +tanh.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def remainder( +def floor_divide( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, - modulus: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if not modulus: - res = x1 / x2 - res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) - diff = res - res_floored - diff, x2 = ivy.promote_types_of_inputs(diff, x2) - if ivy.exists(out): - if out.dtype != x2.dtype: - return ivy.inplace_update( - out, torch.round(torch.mul(diff, x2)).to(out.dtype) - ) - return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) - return torch.remainder(x1, x2, out=out).to(x1.dtype) + if ivy.exists(out): + if not ivy.is_float_dtype(out): + return ivy.inplace_update( + out, torch.floor(torch.div(x1, x2)).type(out.dtype) + ) + return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def round( - x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.round(x, decimals=decimals, out=out) +floor_divide.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def sign( - x: torch.Tensor, +def bitwise_or( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, - np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - if "complex" in str(x.dtype): - if np_variant: - return torch.where( - x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j - ) - return torch.sgn(x, out=out) - return torch.sign(x, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_or(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sin(x, out=out) +bitwise_or.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -749,11 +561,15 @@ def sinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.sinh(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +sinh.support_native_out = True + + @handle_numpy_arrays_in_specific_backend -def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def positive( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.sqrt(x, out=out) + return torch.positive(x) @handle_numpy_arrays_in_specific_backend @@ -762,35 +578,37 @@ def square(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.T return torch.square(x, out=out) +square.support_native_out = True + + @handle_numpy_arrays_in_specific_backend -def subtract( +def pow( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if alpha not in (1, None): - return torch.subtract(x1, x2, alpha=alpha, out=out) - return torch.subtract(x1, x2, out=out) + return torch.pow(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tan(x, out=out) +pow.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def tanh( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None +def round( + x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tanh(x, out=out) + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.round(x, decimals=decimals, out=out) + + +round.support_native_out = True def trapz( @@ -816,6 +634,9 @@ def trapz( return torch.trapezoid(y, x=x, dim=axis) +trapz.support_native_out = False + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: @@ -828,87 +649,388 @@ def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return ret -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +trunc.support_native_out = True + + @handle_numpy_arrays_in_specific_backend -def trunc_divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def abs( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = _cast_for_unary_op(x) + if x.dtype is torch.bool: + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.abs(x, out=out) + + +abs.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def logaddexp( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2, rounding_mode="trunc") - if ivy.is_float_dtype(x1.dtype): - ret = ret.to(x1.dtype) - else: - ret = ret.to(ivy.default_float_dtype(as_native=True)) - return ret + return torch.logaddexp(x1, x2, out=out) -add.support_native_out = True -bitwise_xor.support_native_out = True -imag.support_native_out = False -expm1.support_native_out = True -bitwise_invert.support_native_out = True -equal.support_native_out = True -less_equal.support_native_out = True -bitwise_and.support_native_out = True -ceil.support_native_out = True -floor.support_native_out = True -fmin.support_native_out = True -asin.support_native_out = True -asinh.support_native_out = True -sign.support_native_out = True -sqrt.support_native_out = True -cosh.support_native_out = True -log10.support_native_out = True -log1p.support_native_out = True -less.support_native_out = True -multiply.support_native_out = True -cos.support_native_out = True -logical_not.support_native_out = True -divide.support_native_out = True -greater.support_native_out = True -greater_equal.support_native_out = True -acos.support_native_out = True -lcm.support_native_out = True -logical_xor.support_native_out = True -logical_and.support_native_out = True -logical_or.support_native_out = True -acosh.support_native_out = True -sin.support_native_out = True -negative.support_native_out = True -not_equal.support_native_out = True -tanh.support_native_out = True -floor_divide.support_native_out = True -bitwise_or.support_native_out = True -sinh.support_native_out = True -square.support_native_out = True -pow.support_native_out = True -round.support_native_out = True -trapz.support_native_out = False -trunc.support_native_out = True -abs.support_native_out = True logaddexp.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def logaddexp2( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + if not ivy.is_float_dtype(x1): + x1 = x1.type(ivy.default_float_dtype(as_native=True)) + x2 = x2.type(ivy.default_float_dtype(as_native=True)) + return torch.logaddexp2(x1, x2, out=out) + + logaddexp2.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.tan(x, out=out) + + tan.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.atan(x, out=out) + + atan.support_native_out = True + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version +) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) +@handle_numpy_arrays_in_specific_backend +def atan2( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.atan2(x1, x2, out=out) + + atan2.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log(x, out=out) + + log.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.exp(x, out=out) + + exp.support_native_out = True + + +@handle_numpy_arrays_in_specific_backend +def exp2( + x: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.exp2(x, out=out) + + exp2.support_native_out = True + + +@handle_numpy_arrays_in_specific_backend +def subtract( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + alpha: Optional[Union[int, float]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if alpha not in (1, None): + return torch.subtract(x1, x2, alpha=alpha, out=out) + return torch.subtract(x1, x2, out=out) + + subtract.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def remainder( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + modulus: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if not modulus: + res = x1 / x2 + res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) + diff = res - res_floored + diff, x2 = ivy.promote_types_of_inputs(diff, x2) + if ivy.exists(out): + if out.dtype != x2.dtype: + return ivy.inplace_update( + out, torch.round(torch.mul(diff, x2)).to(out.dtype) + ) + return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) + return torch.remainder(x1, x2, out=out).to(x1.dtype) + + remainder.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.atanh(x, out=out) + + atanh.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def bitwise_right_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) + return torch.bitwise_right_shift(x1, x2, out=out) + + bitwise_right_shift.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def bitwise_left_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_left_shift(x1, x2, out=out) + + bitwise_left_shift.support_native_out = True + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.erf(x, out=out) + + erf.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def minimum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 <= x2, x1, x2, out=out) + return torch.minimum(x1, x2, out=out) + + minimum.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def maximum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 >= x2, x1, x2, out=out) + return torch.maximum(x1, x2, out=out) + + maximum.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def reciprocal( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.reciprocal(x, out=out) + + reciprocal.support_native_out = True + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.deg2rad(x, out=out) + + deg2rad.support_native_out = True + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.rad2deg(x, out=out) + + rad2deg.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def trunc_divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + ret = torch.div(x1, x2, rounding_mode="trunc") + if ivy.is_float_dtype(x1.dtype): + ret = ret.to(x1.dtype) + else: + ret = ret.to(ivy.default_float_dtype(as_native=True)) + return ret + + +@handle_numpy_arrays_in_specific_backend +def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.isreal(x) + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "complex")}, + backend_version, +) +@handle_numpy_arrays_in_specific_backend +def fmod( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmod(x1, x2, out=None) + + fmod.support_native_out = True + + +def gcd( + x1: Union[torch.Tensor, int, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.gcd(x1, x2, out=out) + + gcd.support_native_out = True + + +def angle( + input: torch.Tensor, + /, + *, + deg: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if deg: + return torch.angle(input, out=out) * (180 / pi) + else: + return torch.angle(input, out=out) + + angle.support_native_out = True + + +def nan_to_num( + x: torch.Tensor, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if copy: + return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) + else: + x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) + return x + + +def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.real(x) diff --git a/ivy/functional/backends/torch/experimental/activations.py b/ivy/functional/backends/torch/experimental/activations.py index 008312f394795..98d72ac526c45 100644 --- a/ivy/functional/backends/torch/experimental/activations.py +++ b/ivy/functional/backends/torch/experimental/activations.py @@ -10,16 +10,6 @@ from . import backend_version -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def elu( - x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - ret = torch.nn.functional.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def logit( x: torch.Tensor, @@ -31,11 +21,15 @@ def logit( return torch.logit(x, eps=eps, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def logsigmoid( - input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def thresholded_relu( + x: torch.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nn.functional.logsigmoid(input) + return torch.threshold(x, threshold=threshold, value=0) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -43,6 +37,13 @@ def relu6(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return torch.nn.functional.relu6(x) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def logsigmoid( + input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + return torch.nn.functional.logsigmoid(input) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def selu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: ret = torch.nn.functional.selu(x) @@ -56,12 +57,11 @@ def silu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.silu(x) -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def thresholded_relu( - x: torch.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = None, - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def elu( + x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.threshold(x, threshold=threshold, value=0) + ret = torch.nn.functional.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) diff --git a/ivy/functional/backends/torch/experimental/creation.py b/ivy/functional/backends/torch/experimental/creation.py index 14dd6d3c12b86..a1b94a8a894cf 100644 --- a/ivy/functional/backends/torch/experimental/creation.py +++ b/ivy/functional/backends/torch/experimental/creation.py @@ -12,26 +12,33 @@ ) from .. import backend_version +# noinspection PyProtectedMember -vorbis_window.support_native_out = False -hann_window.support_native_out = False -blackman_window.support_native_out = False -trilu.support_native_out = True +# Array API Standard # +# -------------------# -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def blackman_window( - size: int, - /, - *, + +@with_unsupported_device_and_dtypes( + {"2.0.1 and below": {"cpu": ("float16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[torch.dtype] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.blackman_window( - size, - periodic=periodic, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.kaiser_window( + window_length, + periodic, + beta, dtype=dtype, + layout=torch.strided, + device=None, + requires_grad=False, ) @@ -57,6 +64,29 @@ def hamming_window( ) +def vorbis_window( + window_length: torch.tensor, + *, + dtype: torch.dtype = torch.float32, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.tensor( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, + ) + + +vorbis_window.support_native_out = False + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def hann_window( size: int, @@ -73,34 +103,7 @@ def hann_window( ) -# noinspection PyProtectedMember - - -# Array API Standard # -# -------------------# - - -@with_unsupported_device_and_dtypes( - {"2.0.1 and below": {"cpu": ("float16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, - *, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.kaiser_window( - window_length, - periodic, - beta, - dtype=dtype, - layout=torch.strided, - device=None, - requires_grad=False, - ) +hann_window.support_native_out = False def tril_indices( @@ -123,19 +126,6 @@ def tril_indices( ) -def trilu( - x: torch.Tensor, - /, - *, - k: int = 0, - upper: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if upper: - return torch.triu(x, diagonal=k, out=out) - return torch.tril(x, diagonal=k, out=out) - - def unsorted_segment_min( data: torch.Tensor, segment_ids: torch.Tensor, @@ -162,6 +152,25 @@ def unsorted_segment_min( return res +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def blackman_window( + size: int, + /, + *, + periodic: bool = True, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.blackman_window( + size, + periodic=periodic, + dtype=dtype, + ) + + +blackman_window.support_native_out = False + + def unsorted_segment_sum( data: torch.Tensor, segment_ids: torch.Tensor, @@ -187,21 +196,17 @@ def unsorted_segment_sum( return res -def vorbis_window( - window_length: torch.tensor, +def trilu( + x: torch.Tensor, + /, *, - dtype: torch.dtype = torch.float32, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.tensor( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, - ) + k: int = 0, + upper: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if upper: + return torch.triu(x, diagonal=k, out=out) + return torch.tril(x, diagonal=k, out=out) + + +trilu.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/elementwise.py b/ivy/functional/backends/torch/experimental/elementwise.py index ab2d5a0c2125d..dcb4e30dcf1dc 100644 --- a/ivy/functional/backends/torch/experimental/elementwise.py +++ b/ivy/functional/backends/torch/experimental/elementwise.py @@ -14,69 +14,49 @@ from .. import backend_version -# --- Helpers --- # -# --------------- # +@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) +def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.lgamma(x, out=out) -def _are_suitable_types_for_torch_lerp(input, end, weight): - suitable_types = [ - torch.int8, - torch.int16, - torch.int32, - torch.int64, - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ] - - if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): - return False - else: - if input.dtype not in suitable_types or end.dtype not in suitable_types: - return False +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmax( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmax(x1, x2, out=None) - if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): - return False - else: - if isinstance(weight, torch.Tensor): - if weight.dtype not in [ - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ]: - return False - return True +fmax.support_native_out = True -# --- Main --- # -# ------------ # +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sinc(x, out=out) -def allclose( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[torch.Tensor] = None, -) -> bool: - ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - return torch.tensor(ret) +sinc.support_native_out = True -def conj( - x: torch.Tensor, +def float_power( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - conj_x = torch.conj(x) - return torch.resolve_conj(input=conj_x) + # Native out is supported but with restrictions leading + # to failures hence letting ivy handle it. + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.float_power(x1, x2, out=out) + + +float_power.support_native_out = True def copysign( @@ -93,6 +73,9 @@ def copysign( return torch.copysign(torch.as_tensor(x1), x2, out=out) +copysign.support_native_out = True + + def count_nonzero( a: torch.Tensor, /, @@ -124,6 +107,42 @@ def count_nonzero( return x +count_nonzero.support_native_out = False + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def nansum( + x: torch.Tensor, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[torch.dtype] = None, + keepdims: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) + + +nansum.support_native_out = False + + +def isclose( + a: torch.Tensor, + b: torch.Tensor, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + +isclose.support_native_out = False + + def diff( x: Union[torch.Tensor, list, tuple], /, @@ -148,59 +167,86 @@ def diff( return torch.diff(x, n=n, dim=axis, prepend=prepend, append=append) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def digamma( - x: torch.Tensor, +def signbit( + x: Union[torch.Tensor, float, int, list, tuple], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.digamma(x, out=out) + return torch.signbit(x, out=out) + + +signbit.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def fix( - x: torch.Tensor, +def hypot( + x1: torch.Tensor, + x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.fix(x, out=out) + return torch.hypot(x1, x2) -def float_power( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], +def allclose( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[torch.Tensor] = None, +) -> bool: + ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) + return torch.tensor(ret) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def fix( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # Native out is supported but with restrictions leading - # to failures hence letting ivy handle it. - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.float_power(x1, x2, out=out) + return torch.fix(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmax( +fix.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def nextafter( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmax(x1, x2, out=None) + return torch.nextafter(x1, x2) -def frexp( +nextafter.support_native_out = True + + +def zeta( x: torch.Tensor, + q: torch.Tensor, /, *, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - mantissa, exponent = torch.frexp(x, out=out) - return mantissa, exponent + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) + temp = torch.logical_and(temp, torch.le(q, 0)) + nan_indices = torch.logical_or(temp, torch.lt(x, 1)) + result = torch.special.zeta(x, q) + result.masked_fill_(nan_indices, float("nan")) + return result + + +zeta.support_native_out = False def gradient( @@ -224,28 +270,25 @@ def gradient( return grad -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def hypot( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.hypot(x1, x2) +@with_supported_dtypes( + {"2.0.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def xlogy( + x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None +) -> torch.tensor: + x, y = promote_types_of_inputs(x, y) + return torch.xlogy(x, y, out=out) -def isclose( - a: torch.Tensor, - b: torch.Tensor, +def conj( + x: torch.Tensor, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + conj_x = torch.conj(x) + return torch.resolve_conj(input=conj_x) def ldexp( @@ -258,6 +301,39 @@ def ldexp( return torch.ldexp(x1, x2, out=out) +def _are_suitable_types_for_torch_lerp(input, end, weight): + suitable_types = [ + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ] + + if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): + return False + else: + if input.dtype not in suitable_types or end.dtype not in suitable_types: + return False + + if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): + return False + else: + if isinstance(weight, torch.Tensor): + if weight.dtype not in [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ]: + return False + + return True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def lerp( input: torch.Tensor, @@ -270,100 +346,40 @@ def lerp( return torch.lerp(input, end, weight, out=out) -@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) -def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.lgamma(x, out=out) +lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( + _are_suitable_types_for_torch_lerp(input, end, weight) +) +lerp.support_native_out = True -def modf( +def frexp( x: torch.Tensor, /, *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - modf_x = torch.modf(x) - return torch.resolve_modf(input=modf_x) + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + mantissa, exponent = torch.frexp(x, out=out) + return mantissa, exponent -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def nansum( +def modf( x: torch.Tensor, /, *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[torch.dtype] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def nextafter( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nextafter(x1, x2) - - -def signbit( - x: Union[torch.Tensor, float, int, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.signbit(x, out=out) + modf_x = torch.modf(x) + return torch.resolve_modf(input=modf_x) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sinc(x, out=out) - - -@with_supported_dtypes( - {"2.0.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def xlogy( - x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None -) -> torch.tensor: - x, y = promote_types_of_inputs(x, y) - return torch.xlogy(x, y, out=out) - - -def zeta( +def digamma( x: torch.Tensor, - q: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) - temp = torch.logical_and(temp, torch.le(q, 0)) - nan_indices = torch.logical_or(temp, torch.lt(x, 1)) - result = torch.special.zeta(x, q) - result.masked_fill_(nan_indices, float("nan")) - return result + return torch.special.digamma(x, out=out) -fmax.support_native_out = True -sinc.support_native_out = True -float_power.support_native_out = True -copysign.support_native_out = True -count_nonzero.support_native_out = False -nansum.support_native_out = False -isclose.support_native_out = False -signbit.support_native_out = True -fix.support_native_out = True -nextafter.support_native_out = True -zeta.support_native_out = False -lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( - _are_suitable_types_for_torch_lerp(input, end, weight) -) -lerp.support_native_out = True digamma.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/layers.py b/ivy/functional/backends/torch/experimental/layers.py index 0c26508af9aca..9d35ad4f31232 100644 --- a/ivy/functional/backends/torch/experimental/layers.py +++ b/ivy/functional/backends/torch/experimental/layers.py @@ -16,26 +16,6 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode -# --- Helpers --- # -# --------------- # - - -def _add_ceil_pad_to_pad_list(num_pad, k, c): - return num_pad + (num_pad - ((k * num_pad) / (k - c))) - - -def _adjust_num_padded_values_to_ceil( - pad_specific, num_padded_values, x_shape, kernel, strides, dims -): - for i in range(dims): - pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] - _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) - num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( - num_padded_values[i][-1], kernel[i], c - ) - return num_padded_values - - def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_first"): # Determine depth pooling kernel, strides, depth_pooling = _depth_max_pooling_helper( @@ -46,49 +26,271 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def _get_specific_pad(x_shape, kernel, strides, padding, dims): - if isinstance(padding, str): - if padding == "SAME": - pad_specific = [ - _handle_padding(x_shape[i], strides[i], kernel[i], padding) - for i in range(dims - 1, -1, -1) - ] - pad_list_top = [pad_specific[i] // 2 for i in range(dims)] - pad_list_bot = [pad_specific[i] - pad_specific[i] // 2 for i in range(dims)] - padding = [None] * len(pad_list_top) * 2 - padding[::2] = pad_list_top - padding[1::2] = pad_list_bot - pad_specific = pad_specific[::-1] +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def max_pool1d( + x: torch.Tensor, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NWC": + x = x.permute((0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + if not depth_pooling: + new_kernel = [dilation[0] * (kernel[0] - 1) + 1] + + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2] else: - pad_specific = [0] * dims - padding = [0] * dims * 2 + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) else: - if isinstance(padding, int): - padding = [(padding, padding)] * dims - pad_specific = [sum(padding[i]) for i in range(dims)] - padding = [item for sublist in padding for item in sublist[::-1]][::-1] - return padding, pad_specific + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) -# --- Main --- # -# ------------ # + if depth_pooling: + res = torch.permute(res, (0, 2, 1)) + if data_format == "NWC": + res = res.permute((0, 2, 1)) + return res -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool1d(input, output_size): - return torch.nn.functional.adaptive_avg_pool1d(input, output_size) +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def max_pool2d( + x: torch.Tensor, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NHWC": + x = x.permute(0, 3, 1, 2) + kernel = ( + [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 3, 1, 2]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 3, 1, 2]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool2d(input, output_size): - return torch.nn.functional.adaptive_avg_pool2d(input, output_size) + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + x_shape = list(x.shape[2:]) + if not depth_pooling: + new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_max_pool2d( - input: torch.Tensor, output_size: Union[Sequence[int], int] + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] + else: + # torch pad takes width padding first, then height padding + padding = (padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) + if depth_pooling: + res = torch.permute(res, (0, 2, 1, 3)) + if data_format == "NHWC": + return res.permute(0, 2, 3, 1) + return res + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def max_pool3d( + x: torch.Tensor, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nn.functional.adaptive_max_pool2d(input, output_size) + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NDHWC": + x = x.permute(0, 4, 1, 2, 3) + kernel = ( + [kernel[i] for i in [0, 4, 1, 2, 3]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 4, 1, 2, 3]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 4, 1, 2, 3]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + if not depth_pooling: + x_shape = x.shape[2:] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] + + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + pad_list = [ + pad_w // 2, + pad_w - pad_w // 2, + pad_h // 2, + pad_h - pad_h // 2, + pad_d // 2, + pad_d - pad_d // 2, + ] + else: + # torch pad takes width padding first, then height, then depth + padding = (padding[2], padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) + + if depth_pooling: + res = res.permute(0, 2, 1, 3, 4) + if data_format == "NDHWC": + res = res.permute(0, 2, 3, 4, 1) + return res + + +def _add_ceil_pad_to_pad_list(num_pad, k, c): + return num_pad + (num_pad - ((k * num_pad) / (k - c))) + + +def _get_specific_pad(x_shape, kernel, strides, padding, dims): + if isinstance(padding, str): + if padding == "SAME": + pad_specific = [ + _handle_padding(x_shape[i], strides[i], kernel[i], padding) + for i in range(dims - 1, -1, -1) + ] + pad_list_top = [pad_specific[i] // 2 for i in range(dims)] + pad_list_bot = [pad_specific[i] - pad_specific[i] // 2 for i in range(dims)] + padding = [None] * len(pad_list_top) * 2 + padding[::2] = pad_list_top + padding[1::2] = pad_list_bot + pad_specific = pad_specific[::-1] + else: + pad_specific = [0] * dims + padding = [0] * dims * 2 + else: + if isinstance(padding, int): + padding = [(padding, padding)] * dims + pad_specific = [sum(padding[i]) for i in range(dims)] + padding = [item for sublist in padding for item in sublist[::-1]][::-1] + return padding, pad_specific @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -161,6 +363,18 @@ def avg_pool1d( return res +def _adjust_num_padded_values_to_ceil( + pad_specific, num_padded_values, x_shape, kernel, strides, dims +): + for i in range(dims): + pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] + _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) + num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( + num_padded_values[i][-1], kernel[i], c + ) + return num_padded_values + + @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -400,238 +614,39 @@ def dct( return dct_out elif type == 3: - scale_dims = [1] * len(x.shape) - scale_dims[axis] = axis_dim - complex_part = torch.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float - scale = 2.0 * torch.exp( - torch.complex(real_zero, complex_part.type(real_zero.type())) - ).view(scale_dims) - if norm == "ortho": - n1 = torch.sqrt(axis_dim_float) - n2 = n1 * math.sqrt(0.5) - scale_dims = [1] * len(x.shape) - scale_dims[axis] = axis_dim - sf = torch.nn.functional.pad(n1.unsqueeze(0), (0, axis_dim - 1), value=n2) - x = x * sf.view(scale_dims) - else: - x = x * axis_dim_float - - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(None, axis_dim) - dct_out = torch.real( - torch.fft.irfft( - scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis - ) - )[axis_idx] - return dct_out - - elif type == 4: - dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(1, None, 2) - dct_out = dct_2[axis_idx] - if norm == "ortho": - dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) - return dct_out - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def dropout( - x: torch.Tensor, - prob: float, - /, - *, - scale: bool = True, - dtype: torch.dtype = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x = ivy.astype(x, dtype) if dtype else x - res = torch.nn.functional.dropout(x, prob, training=training) - res = torch.multiply(res, (1.0 - prob)) if not scale else res - return res - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def dropout1d( - x: torch.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - is_batched = len(x.shape) == 3 - if data_format == "NWC": - perm = (0, 2, 1) if is_batched else (1, 0) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout1d(x, prob, training=training) - if data_format == "NWC": - res = torch.permute(res, perm) - return res - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def dropout2d( - x: torch.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NHWC", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - is_batched = len(x.shape) == 4 - if data_format == "NHWC": - perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout2d(x, prob, training=training) - if data_format == "NHWC": - perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) - res = torch.permute(res, perm) - return res - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def dropout3d( - x: torch.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - is_batched = len(x.shape) == 5 - if data_format == "NDHWC": - perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout3d(x, prob, training=training) - if data_format == "NDHWC": - perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) - res = torch.permute(res, perm) - return res - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def embedding( - weights: torch.Tensor, - indices: torch.Tensor, - /, - *, - max_norm: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def fft( - x: torch.Tensor, - dim: int, - /, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.dtype in [torch.int64, torch.float64, torch.complex128]: - out_dtype = torch.complex128 - else: - out_dtype = torch.complex64 - return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def fft2( - x: torch.Tensor, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.tensor( - torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 - ) + scale_dims = [1] * len(x.shape) + scale_dims[axis] = axis_dim + complex_part = torch.arange(axis_dim_float) * math.pi * 0.5 / axis_dim_float + scale = 2.0 * torch.exp( + torch.complex(real_zero, complex_part.type(real_zero.type())) + ).view(scale_dims) + if norm == "ortho": + n1 = torch.sqrt(axis_dim_float) + n2 = n1 * math.sqrt(0.5) + scale_dims = [1] * len(x.shape) + scale_dims[axis] = axis_dim + sf = torch.nn.functional.pad(n1.unsqueeze(0), (0, axis_dim - 1), value=n2) + x = x * sf.view(scale_dims) + else: + x = x * axis_dim_float + + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(None, axis_dim) + dct_out = torch.real( + torch.fft.irfft( + scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis + ) + )[axis_idx] + return dct_out + + elif type == 4: + dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(1, None, 2) + dct_out = dct_2[axis_idx] + if norm == "ortho": + dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) + return dct_out def idct( @@ -648,9 +663,19 @@ def idct( return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) -def ifft( +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def fft( x: torch.Tensor, dim: int, + /, *, norm: str = "backward", n: Union[int, Tuple[int]] = None, @@ -677,120 +702,11 @@ def ifft( ) if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() - - -def ifftn( - x: torch.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) - - -def interpolate( - x: torch.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Literal[ - "linear", - "bilinear", - "trilinear", - "nearest", - "area", - "nearest_exact", - "tf_area", - "bicubic", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", - ] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: bool = False, - out: Optional[torch.Tensor] = None, -): - return torch.nn.functional.interpolate( - x, - size=size, - mode=mode, - align_corners=align_corners, - antialias=antialias, - scale_factor=scale_factor, - recompute_scale_factor=recompute_scale_factor, - ) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def max_pool1d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NWC": - x = x.permute((0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) - - if not depth_pooling: - new_kernel = [dilation[0] * (kernel[0] - 1) + 1] - - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2] - else: - pad_list = [item for sublist in padding for item in sublist] - - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) + if x.dtype in [torch.int64, torch.float64, torch.complex128]: + out_dtype = torch.complex128 else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) - - if depth_pooling: - res = torch.permute(res, (0, 2, 1)) - if data_format == "NWC": - res = res.permute((0, 2, 1)) - return res + out_dtype = torch.complex64 + return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) @with_unsupported_dtypes( @@ -802,75 +718,73 @@ def max_pool1d( }, backend_version, ) -def max_pool2d( +def dropout( x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + prob: float, /, *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, + scale: bool = True, + dtype: torch.dtype = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + x = ivy.astype(x, dtype) if dtype else x + res = torch.nn.functional.dropout(x, prob, training=training) + res = torch.multiply(res, (1.0 - prob)) if not scale else res + return res - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - kernel = ( - [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 3, 1, 2]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 3, 1, 2]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) +dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( + kwargs.get("noise_shape") is None and kwargs.get("seed") is None +) - x_shape = list(x.shape[2:]) - if not depth_pooling: - new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] - else: - # torch pad takes width padding first, then height padding - padding = (padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout1d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 3 + if data_format == "NWC": + perm = (0, 2, 1) if is_batched else (1, 0) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout1d(x, prob, training=training) + if data_format == "NWC": + res = torch.permute(res, perm) + return res - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = torch.permute(res, (0, 2, 1, 3)) +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout2d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 4 if data_format == "NHWC": - return res.permute(0, 2, 3, 1) + perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout2d(x, prob, training=training) + if data_format == "NHWC": + perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) + res = torch.permute(res, perm) return res @@ -883,87 +797,184 @@ def max_pool2d( }, backend_version, ) -def max_pool3d( +def dropout3d( x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + prob: float, /, *, + training: bool = True, data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - + is_batched = len(x.shape) == 5 if data_format == "NDHWC": - x = x.permute(0, 4, 1, 2, 3) - kernel = ( - [kernel[i] for i in [0, 4, 1, 2, 3]] - if len(kernel) == (dims + 2) - else kernel + perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout3d(x, prob, training=training) + if data_format == "NDHWC": + perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) + res = torch.permute(res, perm) + return res + + +def ifft( + x: torch.Tensor, + dim: int, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" ) - strides = ( - [strides[i] for i in [0, 4, 1, 2, 3]] - if len(strides) == (dims + 2) - else strides + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " ) - padding = ( - [padding[i] for i in [0, 4, 1, 2, 3]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def embedding( + weights: torch.Tensor, + indices: torch.Tensor, + /, + *, + max_norm: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False ) + return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) - if not depth_pooling: - x_shape = x.shape[2:] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - pad_list = [ - pad_w // 2, - pad_w - pad_w // 2, - pad_h // 2, - pad_h - pad_h // 2, - pad_d // 2, - pad_d - pad_d // 2, - ] - else: - # torch pad takes width padding first, then height, then depth - padding = (padding[2], padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] +embedding.support_native_out = False - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), + +def interpolate( + x: torch.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Literal[ + "linear", + "bilinear", + "trilinear", + "nearest", + "area", + "nearest_exact", + "tf_area", + "bicubic", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", + ] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: bool = False, + out: Optional[torch.Tensor] = None, +): + return torch.nn.functional.interpolate( + x, + size=size, + mode=mode, + align_corners=align_corners, + antialias=antialias, + scale_factor=scale_factor, + recompute_scale_factor=recompute_scale_factor, + ) + + +interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ + "tf_area", + "nd", + "bicubic_tensorflow", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", +] + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_max_pool2d( + input: torch.Tensor, output_size: Union[Sequence[int], int] +) -> torch.Tensor: + return torch.nn.functional.adaptive_max_pool2d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool1d(input, output_size): + return torch.nn.functional.adaptive_avg_pool1d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool2d(input, output_size): + return torch.nn.functional.adaptive_avg_pool2d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def fft2( + x: torch.Tensor, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return torch.tensor( + torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 + ) - res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = res.permute(0, 2, 1, 3, 4) - if data_format == "NDHWC": - res = res.permute(0, 2, 3, 4, 1) - return res +def ifftn( + x: torch.Tensor, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: Optional[str] = "backward", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -999,18 +1010,3 @@ def rfftn( return torch.tensor( torch.fft.rfftn(x, s, axes, norm=norm, out=out), dtype=torch.complex128 ) - - -dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( - kwargs.get("noise_shape") is None and kwargs.get("seed") is None -) -embedding.support_native_out = False -interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ - "tf_area", - "nd", - "bicubic_tensorflow", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", -] diff --git a/ivy/functional/backends/torch/experimental/linear_algebra.py b/ivy/functional/backends/torch/experimental/linear_algebra.py index 1c3a668fe152c..68e98bff30a60 100644 --- a/ivy/functional/backends/torch/experimental/linear_algebra.py +++ b/ivy/functional/backends/torch/experimental/linear_algebra.py @@ -12,37 +12,6 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size -diagflat.support_native_out = False -kron.support_native_out = True -matrix_exp.support_native_out = True -eig.support_native_out = False -eigvals.support_native_out = False -multi_dot.support_native_out = True -cond.support_native_out = False -dot.support_native_out = True - - -def adjoint( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_dimension_size(x) - return torch.adjoint(x).resolve_conj() - - -@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( - x: torch.Tensor, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.cond(x, p=p, out=out) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def diagflat( x: torch.Tensor, @@ -128,14 +97,32 @@ def diagflat( return ret -def dot( +diagflat.support_native_out = False + + +def kron( a: torch.Tensor, b: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, +) -> torch.tensor: + return torch.kron(a, b, out=out) + + +kron.support_native_out = True + + +def matrix_exp( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.matmul(a, b) + return torch.linalg.matrix_exp(x) + + +matrix_exp.support_native_out = True def eig( @@ -146,46 +133,73 @@ def eig( return torch.linalg.eig(x) +eig.support_native_out = False + + def eigvals(x: torch.Tensor, /) -> torch.Tensor: if not torch.is_complex(x): x = x.to(torch.complex128) return torch.linalg.eigvals(x) -def kron( - a: torch.Tensor, - b: torch.Tensor, +eigvals.support_native_out = False + + +def adjoint( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.kron(a, b, out=out) +) -> torch.Tensor: + _check_valid_dimension_size(x) + return torch.adjoint(x).resolve_conj() -def lu_factor( - x: torch.Tensor, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def multi_dot( + x: Sequence[torch.Tensor], /, *, - pivot: Optional[bool] = True, out: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor]: - raise IvyNotImplementedException() +) -> torch.Tensor: + return torch.linalg.multi_dot(x, out=out) -def matrix_exp( +multi_dot.support_native_out = True + + +@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( x: torch.Tensor, /, *, + p: Optional[Union[None, int, str]] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.linalg.matrix_exp(x) + return torch.linalg.cond(x, p=p, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def multi_dot( - x: Sequence[torch.Tensor], +cond.support_native_out = False + + +def lu_factor( + x: torch.Tensor, + /, + *, + pivot: Optional[bool] = True, + out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + raise IvyNotImplementedException() + + +def dot( + a: torch.Tensor, + b: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.linalg.multi_dot(x, out=out) + return torch.matmul(a, b) + + +dot.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/losses.py b/ivy/functional/backends/torch/experimental/losses.py index e3b6e9485e543..db3442685a678 100644 --- a/ivy/functional/backends/torch/experimental/losses.py +++ b/ivy/functional/backends/torch/experimental/losses.py @@ -3,24 +3,6 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, - backend_version, -) -def huber_loss( - input: torch.Tensor, - target: torch.Tensor, - /, - *, - reduction: Optional[str] = "mean", - delta: Optional[float] = 1.0, -) -> torch.Tensor: - return torch.nn.functional.huber_loss( - input, target, reduction=reduction, delta=delta - ) - - # Assuming ivy and backend_version are imported and defined properly @@ -70,3 +52,20 @@ def smooth_l1_loss( beta=beta, reduction=reduction, ) + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, + backend_version, +) +def huber_loss( + input: torch.Tensor, + target: torch.Tensor, + /, + *, + reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, +) -> torch.Tensor: + return torch.nn.functional.huber_loss( + input, target, reduction=reduction, delta=delta + ) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index 42793dacfbbf2..8683957a72c97 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -10,80 +10,63 @@ import ivy +def moveaxis( + a: torch.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.moveaxis(a, source, destination) + + moveaxis.support_native_out = False -heaviside.support_native_out = True -flipud.support_native_out = False -fliplr.support_native_out = False -i0.support_native_out = True -flatten.partial_mixed_handler = ( - lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" -) -take_along_axis.support_native_out = True -broadcast_shapes.support_native_out = False -expand.support_native_out = False -def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_1d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +def heaviside( + x1: torch.tensor, + x2: torch.tensor, + /, + *, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.heaviside( + x1, + x2, + out=out, + ) -def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_2d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +heaviside.support_native_out = True -def atleast_3d( - *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[torch.Tensor]: - transformed = torch.atleast_3d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +def flipud( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.flipud(m) -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return tuple(torch.broadcast_shapes(*shapes)) +flipud.support_native_out = False -def concat_from_sequence( - input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], +def vstack( + arrays: Sequence[torch.Tensor], /, *, - new_axis: int = 0, - axis: int = 0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = torch.cat(input_sequence, dim=axis) - return ret - elif new_axis == 1: - ret = torch.stack(input_sequence, dim=axis) - return ret - - -def dsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], - /, - *, - copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + if not isinstance(arrays, tuple): + arrays = tuple(arrays) + return torch.vstack(arrays, out=None) -def dstack( +def hstack( arrays: Sequence[torch.Tensor], /, *, @@ -91,49 +74,71 @@ def dstack( ) -> torch.Tensor: if not isinstance(arrays, tuple): arrays = tuple(arrays) - return torch.dstack(arrays, out=out) + return torch.hstack(arrays, out=None) -def expand( - x: torch.Tensor, - shape: Union[List[int], List[Tuple]], +def rot90( + m: torch.Tensor, /, *, copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return x.expand(shape) + return torch.rot90(m, k, axes) -def fill_diagonal( - a: torch.Tensor, - v: Union[int, float], +def top_k( + x: torch.Tensor, + k: int, /, *, - wrap: bool = False, -) -> torch.Tensor: - shape = a.shape - max_end = torch.prod(torch.tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] + ) + if not largest: + indices = torch.argsort(x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) else: - step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + indices = torch.argsort(-x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) + if not sorted: + indices = torch.sort(indices, dim=axis)[0] + val = torch.gather(x, axis, indices) + return topk_res(val, indices) - end = max_end if end > max_end else end - a = torch.reshape(a, (-1,)) - w = torch.zeros(a.shape, dtype=bool).to(a.device) - ins = torch.arange(0, max_end).to(a.device) - steps = torch.arange(0, end, step).to(a.device) - for i in steps: - i = ins == i - w = torch.logical_or(w, i) - a = torch.where(w, v, a) - a = torch.reshape(a, shape) - return a +def fliplr( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.fliplr(m) + + +fliplr.support_native_out = False + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def i0( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.i0(x, out=out) + + +i0.support_native_out = True def flatten( @@ -149,53 +154,47 @@ def flatten( return torch.flatten(x, start_dim=start_dim, end_dim=end_dim) -def fliplr( - m: torch.Tensor, - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.fliplr(m) +flatten.partial_mixed_handler = ( + lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" +) -def flipud( - m: torch.Tensor, +def vsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.flipud(m) - - -def heaviside( - x1: torch.tensor, - x2: torch.tensor, - /, - *, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.heaviside( - x1, - x2, - out=out, - ) +) -> List[torch.Tensor]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) -def hsplit( +def dsplit( ary: torch.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, copy: Optional[bool] = None, ) -> List[torch.Tensor]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_1d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed + + +def dstack( arrays: Sequence[torch.Tensor], /, *, @@ -203,41 +202,23 @@ def hstack( ) -> torch.Tensor: if not isinstance(arrays, tuple): arrays = tuple(arrays) - return torch.hstack(arrays, out=None) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def i0( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.i0(x, out=out) + return torch.dstack(arrays, out=out) -def moveaxis( - a: torch.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.moveaxis(a, source, destination) +def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_2d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -def rot90( - m: torch.Tensor, - /, - *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.rot90(m, k, axes) +def atleast_3d( + *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[torch.Tensor]: + transformed = torch.atleast_3d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -282,30 +263,59 @@ def take_along_axis( return torch.take_along_dim(arr, indices, axis, out=out) -def top_k( +def hsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[torch.Tensor]: + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +take_along_axis.support_native_out = True + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return tuple(torch.broadcast_shapes(*shapes)) + + +broadcast_shapes.support_native_out = False + + +def expand( x: torch.Tensor, - k: int, + shape: Union[List[int], List[Tuple]], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] - ) - if not largest: - indices = torch.argsort(x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - else: - indices = torch.argsort(-x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - if not sorted: - indices = torch.sort(indices, dim=axis)[0] - val = torch.gather(x, axis, indices) - return topk_res(val, indices) + copy: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return x.expand(shape) + + +expand.support_native_out = False + + +def concat_from_sequence( + input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = torch.cat(input_sequence, dim=axis) + return ret + elif new_axis == 1: + ret = torch.stack(input_sequence, dim=axis) + return ret @with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) @@ -332,26 +342,32 @@ def unique_consecutive( ) -def vsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], +def fill_diagonal( + a: torch.Tensor, + v: Union[int, float], /, *, - copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + wrap: bool = False, +) -> torch.Tensor: + shape = a.shape + max_end = torch.prod(torch.tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + end = max_end if end > max_end else end + a = torch.reshape(a, (-1,)) + w = torch.zeros(a.shape, dtype=bool).to(a.device) + ins = torch.arange(0, max_end).to(a.device) + steps = torch.arange(0, end, step).to(a.device) -def vstack( - arrays: Sequence[torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(arrays, tuple): - arrays = tuple(arrays) - return torch.vstack(arrays, out=None) + for i in steps: + i = ins == i + w = torch.logical_or(w, i) + a = torch.where(w, v, a) + a = torch.reshape(a, shape) + return a diff --git a/ivy/functional/backends/torch/experimental/norms.py b/ivy/functional/backends/torch/experimental/norms.py index 9524918277dc9..5a3e0c31f764b 100644 --- a/ivy/functional/backends/torch/experimental/norms.py +++ b/ivy/functional/backends/torch/experimental/norms.py @@ -5,27 +5,31 @@ from .. import backend_version +def l1_normalize( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) + + l1_normalize.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def l2_normalize( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) + + l2_normalize.support_native_out = True -batch_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) - ) -) -instance_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) - ) -) -lp_normalize.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -69,29 +73,15 @@ def batch_norm( return xnormalized, runningmean, runningvariance -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def group_norm( - x: torch.Tensor, - num_groups: int = 1, - /, - *, - offset: Optional[torch.Tensor] = None, - scale: Optional[torch.Tensor] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", - out: Optional[torch.Tensor] = None, -): - xdims = x.ndim - if data_format == "NSC": - x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) - xnormalized = torch.nn.functional.group_norm( - x, num_groups, weight=scale, bias=offset, eps=eps +batch_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) ) - - if data_format == "NSC": - xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) - - return xnormalized +) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -137,25 +127,40 @@ def instance_norm( return xnormalized, runningmean, runningvariance -def l1_normalize( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) +instance_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) + ) +) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def l2_normalize( +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def group_norm( x: torch.Tensor, + num_groups: int = 1, /, *, - axis: Optional[int] = None, + offset: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) +): + xdims = x.ndim + if data_format == "NSC": + x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) + xnormalized = torch.nn.functional.group_norm( + x, num_groups, weight=scale, bias=offset, eps=eps + ) + + if data_format == "NSC": + xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) + + return xnormalized @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -168,3 +173,6 @@ def lp_normalize( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: return torch.nn.functional.normalize(x, p=p, dim=axis, out=out) + + +lp_normalize.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/random.py b/ivy/functional/backends/torch/experimental/random.py index b8e58426ed24b..a65cee4f57431 100644 --- a/ivy/functional/backends/torch/experimental/random.py +++ b/ivy/functional/backends/torch/experimental/random.py @@ -12,46 +12,23 @@ ) -# --- Helpers --- # -# --------------- # - - -def _poisson_with_neg_lam(lam, fill_value, device, dtype): - if torch.any(lam < 0): - pos_lam = torch.where(lam < 0, 0, lam) - ret = torch.poisson(pos_lam).type(dtype).to(device) - ret = torch.where(lam < 0, fill_value, ret) - else: - ret = torch.poisson(lam).type(dtype).to(device) - return ret - - -# --- Main --- # -# ------------ # - - -def bernoulli( - probs: Union[float, torch.Tensor], +# dirichlet +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def dirichlet( + alpha: Union[torch.tensor, float, Sequence[float]], + /, *, - logits: Union[float, torch.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: torch.device, - dtype: torch.dtype, - seed: Optional[int] = None, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, out: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: - if seed: + size = size if size is not None else len(alpha) + if seed is not None: torch.manual_seed(seed) - if logits is not None: - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return ( - torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) - .sample(shape) - .to(device, dtype) + return torch.tensor( + torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), + dtype=dtype, ) @@ -76,26 +53,6 @@ def beta( return ret -# dirichlet -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def dirichlet( - alpha: Union[torch.tensor, float, Sequence[float]], - /, - *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - out: Optional[torch.Tensor] = None, - seed: Optional[int] = None, - dtype: Optional[torch.dtype] = None, -) -> torch.Tensor: - size = size if size is not None else len(alpha) - if seed is not None: - torch.manual_seed(seed) - return torch.tensor( - torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), - dtype=dtype, - ) - - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, torch.Tensor], @@ -117,6 +74,16 @@ def gamma( return ret +def _poisson_with_neg_lam(lam, fill_value, device, dtype): + if torch.any(lam < 0): + pos_lam = torch.where(lam < 0, 0, lam) + ret = torch.poisson(pos_lam).type(dtype).to(device) + ret = torch.where(lam < 0, fill_value, ret) + else: + ret = torch.poisson(lam).type(dtype).to(device) + return ret + + def poisson( lam: Union[float, torch.Tensor], *, @@ -137,3 +104,28 @@ def poisson( _check_shapes_broadcastable(lam.shape, list_shape) lam = torch.broadcast_to(lam, list_shape) return _poisson_with_neg_lam(lam, fill_value, device, dtype) + + +def bernoulli( + probs: Union[float, torch.Tensor], + *, + logits: Union[float, torch.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: torch.device, + dtype: torch.dtype, + seed: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if seed: + torch.manual_seed(seed) + if logits is not None: + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return ( + torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) + .sample(shape) + .to(device, dtype) + ) diff --git a/ivy/functional/backends/torch/experimental/searching.py b/ivy/functional/backends/torch/experimental/searching.py index 2c2a8040bdf21..456d6204f8056 100644 --- a/ivy/functional/backends/torch/experimental/searching.py +++ b/ivy/functional/backends/torch/experimental/searching.py @@ -3,9 +3,6 @@ import torch -unravel_index.support_native_out = False - - def unravel_index( indices: torch.Tensor, shape: Tuple[int], @@ -19,3 +16,6 @@ def unravel_index( output.append(temp % dim) temp = temp // dim return tuple(reversed(output)) + + +unravel_index.support_native_out = False diff --git a/ivy/functional/backends/torch/experimental/sorting.py b/ivy/functional/backends/torch/experimental/sorting.py index db91a013dde4e..c6b7ee0f5bb06 100644 --- a/ivy/functional/backends/torch/experimental/sorting.py +++ b/ivy/functional/backends/torch/experimental/sorting.py @@ -6,9 +6,6 @@ import ivy -lexsort.support_native_out = False - - # invert_permutation def invert_permutation( x: Union[torch.Tensor, list, tuple], @@ -47,3 +44,6 @@ def lexsort( # only valid for torch > 2.0.1 result = result[temp] return result + + +lexsort.support_native_out = False diff --git a/ivy/functional/backends/torch/experimental/statistical.py b/ivy/functional/backends/torch/experimental/statistical.py index 78b4e43e4db7e..cc1d4d5fa05ce 100644 --- a/ivy/functional/backends/torch/experimental/statistical.py +++ b/ivy/functional/backends/torch/experimental/statistical.py @@ -10,333 +10,6 @@ from copy import deepcopy -# --- Helpers --- # -# --------------- # - - -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - return torch.quantile( - x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out - ) - else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - - -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = a.ndim - axis_arg = deepcopy(axis) - if axis is not None: - axis = _to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) - - for i, s in enumerate(sorted(keep)): - a = torch.moveaxis(a, s, i) - a = a.view( - [ - *a.shape[:nkeep], - -1, - ] - ) - axis_arg = -1 - - ret = fn(a, q, axis=axis_arg) - - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] - - return ret - - -def _nanmedian(input, axis, keepdims): - dtype = input.dtype - temp = input.to(torch.float64) - num_dim = len(temp.size()) - keepdim_shape = list(temp.size()) - q = 0.5 - - axis = [axis] if isinstance(axis, int) else list(axis) - - for i in axis: - keepdim_shape[i] = 1 - axis = [num_dim + x if x < 0 else x for x in axis] - axis.sort() - dimension = len(temp.size()) - while len(axis) > 0: - axis1 = axis[0] - for axis2 in range(axis1 + 1, dimension): - temp = torch.transpose(temp, axis1, axis2) - axis1 = axis2 - axis = [x - 1 for x in axis] - axis.pop(0) - dimension = dimension - 1 - temp = torch.flatten(temp, start_dim=dimension - len(axis)) - ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") - if keepdims: - keepdim_shape = tuple(keepdim_shape) - ret = ret.reshape(keepdim_shape) - - if dtype in [torch.int32, torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=torch.float16) - else: - ret = torch.asarray(ret, dtype=torch.float32) - - return ret - - -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if isinstance(q, float): - q = torch.as_tensor(q) - if isinstance(q, torch.Tensor) and q.ndim > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = a.flatten() - - n = a.shape[axis] - - indices = q * (n - 1) - - a = torch.sort(a, axis)[axis] - indices_below = torch.floor(indices).to(torch.int64) - indices_upper = torch.ceil(indices).to(torch.int64) - - weights = indices - indices_below.to(torch.float64) - - indices_below = torch.clip(indices_below, 0, n - 1) - indices_upper = torch.clip(indices_upper, 0, n - 1) - tensor_upper = torch.index_select(a, 0, indices_upper) - tensor_below = torch.index_select(a, 0, indices_below) - - pred = weights <= 0.5 - out = torch.where(pred, tensor_below, tensor_upper) - return out.to(ret_dtype) - - -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis - - -def _validate_quantile(q): - if isinstance(q, float): - q = torch.as_tensor(q) - if q.ndim == 1 and torch.numel(q) < 10: - for i in range(torch.numel(q)): - if not (0.0 <= q[i] <= 1.0): - return False - else: - if not (torch.all(0 <= q) and torch.all(q <= 1)): - return False - return True - - -# --- Main --- # -# ------------ # - - -def bincount( - x: torch.Tensor, - /, - *, - weights: Optional[torch.Tensor] = None, - minlength: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if weights is None: - ret = torch.bincount(x, minlength=minlength) - ret = ret.to(x.dtype) - else: - ret = torch.bincount(x, weights=weights, minlength=minlength) - ret = ret.to(weights.dtype) - return ret - - -def corrcoef( - x: torch.Tensor, - /, - *, - y: Optional[torch.Tensor] = None, - rowvar: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = torch.concat([x, y], dim=axis) - xarr = xarr.T if not rowvar else xarr - - return torch.corrcoef(xarr) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def cov( - x1: torch.Tensor, - x2: torch.Tensor = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[torch.Tensor] = None, - aweights: Optional[torch.Tensor] = None, - dtype: Optional[torch.dtype] = None, -) -> torch.Tensor: - # dtype casts separately - if fweights is not None: - fweights = fweights.type(torch.int64) - if aweights is not None: - aweights = aweights.type(torch.float64) - - if x1.dim() > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if x2.dim() > 2: - raise ValueError("x2 has more than 2 dimensions") - - if ddof is None: - if bias == 0: - ddof = 1 - else: - ddof = 0 - - if dtype is None: - x1 = x1.type(torch.float64) - if x2 is not None: - x2 = x2.type(torch.float64) - else: - x1 = x1.type(dtype) - if x2 is not None: - x2 = x2.type(dtype) - - X = x1 - if not rowVar and len(x1.shape) != 1: - X = torch.t(x1) - - if x2 is not None: - if not rowVar and len(x2.shape) != 1: - x2 = torch.t(x2) - X = torch.vstack((X, x2)) - - return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, - backend_version, -) -def cummax( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - if x.dtype in (torch.bool, torch.float16): - x = x.to(dtype=torch.float64) - elif x.dtype in (torch.int16, torch.int8, torch.uint8): - x = x.to(dtype=torch.int64) - elif x.dtype in (torch.complex64, torch.complex128): - x = x.real.to(dtype=torch.float64) - - if exclusive or reverse: - if exclusive and reverse: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - x1, x2 = torch.concat( - (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 - ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x1, x2 = torch.cummax(x, -1) - res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - else: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - return res1, res2 - - return torch.cummax(x, axis, out=out) - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cummin( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - if not (reverse): - ret = torch.cummin(x, axis)[0] - else: - ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] - ret = torch.flip(ret, (axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret.to(dtype)) - return ret.to(dtype) - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -463,54 +136,254 @@ def histogram( return ret -def igamma( +histogram.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) +def median( + input: torch.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if isinstance(axis, tuple): + if len(axis) == 1: + axis = axis[0] + ret = quantile( + input, + 0.5, + axis=axis, + keepdims=keepdims, + interpolation="midpoint", + ) + if input.dtype in [torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif input.dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=input.dtype) + else: + ret = torch.asarray(ret, dtype=torch.float32) + return ret + + +median.support_native_out = False + + +def nanmean( a: torch.Tensor, /, *, - x: torch.Tensor, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.gammainc(a, x, out=out) + return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) -def median( - input: torch.Tensor, +nanmean.support_native_out = True + + +def _validate_quantile(q): + if isinstance(q, float): + q = torch.as_tensor(q) + if q.ndim == 1 and torch.numel(q) < 10: + for i in range(torch.numel(q)): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (torch.all(0 <= q) and torch.all(q <= 1)): + return False + return True + + +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis + + +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = a.ndim + axis_arg = deepcopy(axis) + if axis is not None: + axis = _to_positive_axis(axis, nd) + + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) + + for i, s in enumerate(sorted(keep)): + a = torch.moveaxis(a, s, i) + a = a.view( + [ + *a.shape[:nkeep], + -1, + ] + ) + axis_arg = -1 + + ret = fn(a, q, axis=axis_arg) + + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret + + +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if isinstance(q, float): + q = torch.as_tensor(q) + if isinstance(q, torch.Tensor) and q.ndim > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = a.flatten() + + n = a.shape[axis] + + indices = q * (n - 1) + + a = torch.sort(a, axis)[axis] + indices_below = torch.floor(indices).to(torch.int64) + indices_upper = torch.ceil(indices).to(torch.int64) + + weights = indices - indices_below.to(torch.float64) + + indices_below = torch.clip(indices_below, 0, n - 1) + indices_upper = torch.clip(indices_upper, 0, n - 1) + tensor_upper = torch.index_select(a, 0, indices_upper) + tensor_below = torch.index_select(a, 0, indices_below) + + pred = weights <= 0.5 + out = torch.where(pred, tensor_below, tensor_upper) + return out.to(ret_dtype) + + +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + return torch.quantile( + x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out + ) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def quantile( + a: torch.Tensor, + q: Union[torch.Tensor, float], /, *, - axis: Optional[Union[Tuple[int], int]] = None, + axis: Optional[Union[Sequence[int], int]] = None, keepdims: bool = False, + interpolation: str = "linear", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if isinstance(axis, tuple): - if len(axis) == 1: - axis = axis[0] - ret = quantile( - input, - 0.5, + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, axis=axis, keepdims=keepdims, - interpolation="midpoint", + interpolation=interpolation, + out=out, ) - if input.dtype in [torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif input.dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=input.dtype) - else: - ret = torch.asarray(ret, dtype=torch.float32) - return ret -def nanmean( - a: torch.Tensor, +quantile.support_native_out = True + + +def corrcoef( + x: torch.Tensor, /, *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[torch.dtype] = None, + y: Optional[torch.Tensor] = None, + rowvar: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = torch.concat([x, y], dim=axis) + xarr = xarr.T if not rowvar else xarr + + return torch.corrcoef(xarr) + + +def _nanmedian(input, axis, keepdims): + dtype = input.dtype + temp = input.to(torch.float64) + num_dim = len(temp.size()) + keepdim_shape = list(temp.size()) + q = 0.5 + + axis = [axis] if isinstance(axis, int) else list(axis) + + for i in axis: + keepdim_shape[i] = 1 + axis = [num_dim + x if x < 0 else x for x in axis] + axis.sort() + dimension = len(temp.size()) + while len(axis) > 0: + axis1 = axis[0] + for axis2 in range(axis1 + 1, dimension): + temp = torch.transpose(temp, axis1, axis2) + axis1 = axis2 + axis = [x - 1 for x in axis] + axis.pop(0) + dimension = dimension - 1 + temp = torch.flatten(temp, start_dim=dimension - len(axis)) + ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") + if keepdims: + keepdim_shape = tuple(keepdim_shape) + ret = ret.reshape(keepdim_shape) + + if dtype in [torch.int32, torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=torch.float16) + else: + ret = torch.asarray(ret, dtype=torch.float32) + + return ret @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -565,33 +438,166 @@ def nanmedian( return _nanmedian(input, axis, keepdims) -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def quantile( - a: torch.Tensor, - q: Union[torch.Tensor, float], +nanmedian.support_native_out = True + + +def bincount( + x: torch.Tensor, /, *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: bool = False, - interpolation: str = "linear", + weights: Optional[torch.Tensor] = None, + minlength: int = 0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - out=out, - ) + if weights is None: + ret = torch.bincount(x, minlength=minlength) + ret = ret.to(x.dtype) + else: + ret = torch.bincount(x, weights=weights, minlength=minlength) + ret = ret.to(weights.dtype) + return ret -histogram.support_native_out = True -median.support_native_out = False -nanmean.support_native_out = True -quantile.support_native_out = True -nanmedian.support_native_out = True bincount.support_native_out = False + + +def igamma( + a: torch.Tensor, + /, + *, + x: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.special.gammainc(a, x, out=out) + + igamma.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def cov( + x1: torch.Tensor, + x2: torch.Tensor = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[torch.Tensor] = None, + aweights: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + # dtype casts separately + if fweights is not None: + fweights = fweights.type(torch.int64) + if aweights is not None: + aweights = aweights.type(torch.float64) + + if x1.dim() > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if x2.dim() > 2: + raise ValueError("x2 has more than 2 dimensions") + + if ddof is None: + if bias == 0: + ddof = 1 + else: + ddof = 0 + + if dtype is None: + x1 = x1.type(torch.float64) + if x2 is not None: + x2 = x2.type(torch.float64) + else: + x1 = x1.type(dtype) + if x2 is not None: + x2 = x2.type(dtype) + + X = x1 + if not rowVar and len(x1.shape) != 1: + X = torch.t(x1) + + if x2 is not None: + if not rowVar and len(x2.shape) != 1: + x2 = torch.t(x2) + X = torch.vstack((X, x2)) + + return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) + + cov.support_native_out = False + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, + backend_version, +) +def cummax( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if x.dtype in (torch.bool, torch.float16): + x = x.to(dtype=torch.float64) + elif x.dtype in (torch.int16, torch.int8, torch.uint8): + x = x.to(dtype=torch.int64) + elif x.dtype in (torch.complex64, torch.complex128): + x = x.real.to(dtype=torch.float64) + + if exclusive or reverse: + if exclusive and reverse: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + x1, x2 = torch.concat( + (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 + ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x1, x2 = torch.cummax(x, -1) + res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + else: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + return res1, res2 + + return torch.cummax(x, axis, out=out) + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cummin( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + if not (reverse): + ret = torch.cummin(x, axis)[0] + else: + ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] + ret = torch.flip(ret, (axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret.to(dtype)) + return ret.to(dtype) diff --git a/ivy/functional/backends/torch/general.py b/ivy/functional/backends/torch/general.py index c75cc289b5bc7..f8cd14960cddb 100644 --- a/ivy/functional/backends/torch/general.py +++ b/ivy/functional/backends/torch/general.py @@ -4,6 +4,11 @@ from numbers import Number from operator import mul from typing import Optional, Union, Sequence, Callable, List, Tuple + +try: + import functorch +except ImportError: + functorch = () # for torch 1.10.1 import numpy as np import torch @@ -13,18 +18,9 @@ from . import backend_version, is_variable from ...ivy.general import _broadcast_to -try: - import functorch -except ImportError: - functorch = () # for torch 1.10.1 - torch_scatter = None -# --- Helpers --- # -# --------------- # - - def _parse_index(indices, ndims): ind = list() for so in indices: @@ -48,8 +44,12 @@ def _parse_index(indices, ndims): return ind -# --- Main --- # -# ------------ # +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, torch.Tensor): + if exclusive and x.requires_grad: + return False + return True + return False @with_unsupported_dtypes({"2.0.1 and below": ("complex", "bfloat16")}, backend_version) @@ -66,6 +66,111 @@ def current_backend_str() -> str: return "torch" +def neg_step(query): + return ( + not isinstance(query, (int, bool)) + and not ivy.is_array(query) + and query is not None + and query is not Ellipsis + and ( + (isinstance(query, slice) and query.step is not None and query.step < 0) + or ( + not isinstance(query, slice) + and any( + isinstance(q, slice) and q.step is not None and q.step < 0 + for q in query + ) + ) + ) + ) + + +def get_item( + x: torch.Tensor, + /, + query: Union[torch.Tensor, Tuple], + *, + copy: bool = None, +) -> torch.Tensor: + return x.__getitem__(query) + + +get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) + + +def set_item( + x: torch.Tensor, + query: Union[torch.Tensor, Tuple], + val: torch.Tensor, + /, + *, + copy: Optional[bool] = False, +) -> torch.Tensor: + if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: + val = val.to(x.dtype) + if copy: + x = x.clone() + x.__setitem__(query, val) + return x + + +set_item.partial_mixed_handler = ( + lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad +) + + +def to_numpy( + x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif torch.is_tensor(x): + x = x.resolve_neg().resolve_conj() + if copy: + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16") + return x.detach().cpu().numpy() + else: + raise ivy.utils.exceptions.IvyException( + "Overwriting the same address is not supported for torch." + ) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + +def to_scalar(x: torch.Tensor, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + +def to_list(x: torch.Tensor, /) -> list: + if isinstance(x, np.ndarray): + return x.tolist() + elif torch.is_tensor(x): + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16").tolist() + else: + return x.detach().cpu().numpy().tolist() + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + def gather( params: torch.Tensor, indices: torch.Tensor, @@ -102,36 +207,6 @@ def gather( return result -def gather_nd( - params: torch.Tensor, - indices: torch.Tensor, - /, - *, - batch_dims: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = torch.stack(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result - - def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -162,14 +237,34 @@ def gather_nd_helper(params, indices): return res -def get_item( - x: torch.Tensor, +def gather_nd( + params: torch.Tensor, + indices: torch.Tensor, /, - query: Union[torch.Tensor, Tuple], *, - copy: bool = None, + batch_dims: int = 0, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return x.__getitem__(query) + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = torch.stack(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result def get_num_dims( @@ -243,37 +338,6 @@ def inplace_variables_supported(): return True -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, torch.Tensor): - if exclusive and x.requires_grad: - return False - return True - return False - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version -) -def isin( - elements: torch.tensor, - test_elements: torch.tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> torch.tensor: - return torch.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - -def itemsize(x: torch.tensor) -> int: - return x.element_size() - - def multiprocessing(context: Optional[str] = None): import torch.multiprocessing @@ -282,25 +346,6 @@ def multiprocessing(context: Optional[str] = None): return torch.multiprocessing.get_context(context) -def neg_step(query): - return ( - not isinstance(query, (int, bool)) - and not ivy.is_array(query) - and query is not None - and query is not Ellipsis - and ( - (isinstance(query, slice) and query.step is not None and query.step < 0) - or ( - not isinstance(query, slice) - and any( - isinstance(q, slice) and q.step is not None and q.step < 0 - for q in query - ) - ) - ) - ) - - @with_unsupported_dtypes( { "2.0.1 and below": ("bfloat16",), @@ -350,6 +395,9 @@ def scatter_flat( return res +scatter_flat.support_native_out = True + + @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -444,20 +492,7 @@ def scatter_nd( return res -def set_item( - x: torch.Tensor, - query: Union[torch.Tensor, Tuple], - val: torch.Tensor, - /, - *, - copy: Optional[bool] = False, -) -> torch.Tensor: - if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: - val = val.to(x.dtype) - if copy: - x = x.clone() - x.__setitem__(query, val) - return x +scatter_nd.support_native_out = True def shape( @@ -472,58 +507,6 @@ def shape( return ivy.Shape(x.shape) -def to_list(x: torch.Tensor, /) -> list: - if isinstance(x, np.ndarray): - return x.tolist() - elif torch.is_tensor(x): - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16").tolist() - else: - return x.detach().cpu().numpy().tolist() - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - -def to_numpy( - x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif torch.is_tensor(x): - x = x.resolve_neg().resolve_conj() - if copy: - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16") - return x.detach().cpu().numpy() - else: - raise ivy.utils.exceptions.IvyException( - "Overwriting the same address is not supported for torch." - ) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - -def to_scalar(x: torch.Tensor, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def vmap( func: Callable, @@ -540,10 +523,27 @@ def _vmap(*args): return _vmap -get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) -set_item.partial_mixed_handler = ( - lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version ) -scatter_flat.support_native_out = True -scatter_nd.support_native_out = True +def isin( + elements: torch.tensor, + test_elements: torch.tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> torch.tensor: + return torch.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + isin.support_native_out = True + + +def itemsize(x: torch.tensor) -> int: + return x.element_size() diff --git a/ivy/functional/backends/torch/gradients.py b/ivy/functional/backends/torch/gradients.py index ad88bc84f0d3d..012e8aafe4a17 100644 --- a/ivy/functional/backends/torch/gradients.py +++ b/ivy/functional/backends/torch/gradients.py @@ -19,8 +19,20 @@ ) -# --- Helpers --- # -# --------------- # +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = ivy.astype(x, ivy.default_float_dtype()).to_native() + if not x.is_leaf: + return x.detach().requires_grad_() + return x.clone().requires_grad_() + + +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, torch.Tensor) and x.requires_grad + + +def variable_data(x: torch.Tensor, /) -> torch.Tensor: + return x.data def _grad_func(y, xs, retain_grads): @@ -80,10 +92,6 @@ def grad_(x): return grads -# --- Main --- # -# ------------ # - - def execute_with_gradients( func, xs: torch.Tensor, @@ -132,6 +140,61 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = torch.autograd.grad(y, x, allow_unused=True)[0] + grad = ( + grad + if grad is not None + else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) + ) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def stop_gradient( + x: Optional[torch.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[torch.Tensor] = None, +): + if is_variable(x) and preserve_type: + if x.grad_fn: + x = x.detach() + x.requires_grad = True + elif x.grad: + x.grad.data.zero_() + return x + return x.detach() + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + callback_fn = lambda x_in: ivy.to_ivy( + torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), + nested=True, + include_derived=True, + ) + return callback_fn + + def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -194,76 +257,5 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, torch.Tensor) and x.requires_grad - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - callback_fn = lambda x_in: ivy.to_ivy( - torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), - nested=True, - include_derived=True, - ) - return callback_fn - - -def stop_gradient( - x: Optional[torch.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[torch.Tensor] = None, -): - if is_variable(x) and preserve_type: - if x.grad_fn: - x = x.detach() - x.requires_grad = True - elif x.grad: - x.grad.data.zero_() - return x - return x.detach() - - -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = torch.autograd.grad(y, x, allow_unused=True)[0] - grad = ( - grad - if grad is not None - else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) - ) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = ivy.astype(x, ivy.default_float_dtype()).to_native() - if not x.is_leaf: - return x.detach().requires_grad_() - return x.clone().requires_grad_() - - -def variable_data(x: torch.Tensor, /) -> torch.Tensor: - return x.data - - grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/torch/layers.py b/ivy/functional/backends/torch/layers.py index ecb3f3f29d9bd..0ae11f08e2e9b 100644 --- a/ivy/functional/backends/torch/layers.py +++ b/ivy/functional/backends/torch/layers.py @@ -11,8 +11,22 @@ from ivy.functional.ivy.layers import _handle_padding, _deconv_length -# --- Helpers --- # -# --------------- # +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, + backend_version, +) +def linear( + x: torch.Tensor, + weight: torch.Tensor, + /, + *, + bias: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.linear(x, weight, bias) + + +linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): @@ -112,10 +126,6 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding -# --- Main --- # -# ------------ # - - @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version, @@ -278,6 +288,48 @@ def conv2d_transpose( return res +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + "complex", + ) + }, + backend_version, +) +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: torch.Tensor, + filters: torch.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NHWC": + x = x.permute(0, 3, 1, 2) + filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters + filters = torch.unsqueeze(filters, -1) + dims_in = filters.shape[-2] + filters = filters.permute(2, 3, 0, 1) + x, padding = _pad_before_conv( + x, filters, strides, padding, 2, dilations, "channel_first" + ) + # noinspection PyArgumentEqualDefault + res = torch.nn.functional.conv2d( + x, filters, None, strides, padding, dilations, dims_in + ) + if data_format == "NHWC": + return res.permute(0, 2, 3, 1) + return res + + @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version ) @@ -492,63 +544,3 @@ def conv_general_transpose( if data_format == "channel_last": res = res.permute(0, *range(2, dims + 2), 1) return res - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - "complex", - ) - }, - backend_version, -) -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: torch.Tensor, - filters: torch.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters - filters = torch.unsqueeze(filters, -1) - dims_in = filters.shape[-2] - filters = filters.permute(2, 3, 0, 1) - x, padding = _pad_before_conv( - x, filters, strides, padding, 2, dilations, "channel_first" - ) - # noinspection PyArgumentEqualDefault - res = torch.nn.functional.conv2d( - x, filters, None, strides, padding, dilations, dims_in - ) - if data_format == "NHWC": - return res.permute(0, 2, 3, 1) - return res - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, - backend_version, -) -def linear( - x: torch.Tensor, - weight: torch.Tensor, - /, - *, - bias: Optional[torch.Tensor] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.linear(x, weight, bias) - - -linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 diff --git a/ivy/functional/backends/torch/linear_algebra.py b/ivy/functional/backends/torch/linear_algebra.py index 341eef4f46c7f..950c2ff6fd295 100644 --- a/ivy/functional/backends/torch/linear_algebra.py +++ b/ivy/functional/backends/torch/linear_algebra.py @@ -14,27 +14,6 @@ from .elementwise import _cast_for_unary_op -cholesky.support_native_out = True -cross.support_native_out = True -det.support_native_out = True -eigh.support_native_out = True -eigvalsh.support_native_out = True -inner.support_native_out = True -inv.support_native_out = True -matmul.support_native_out = True -matrix_norm.support_native_out = True -eig.support_native_out = True -matrix_power.support_native_out = True -matrix_rank.support_native_out = True -outer.support_native_out = True -pinv.support_native_out = True -slogdet.support_native_out = True -svdvals.support_native_out = True -vecdot.support_native_out = True -vector_norm.support_native_out = True -vector_to_skew_symmetric_matrix.support_native_out = True - - # Array API Standard # # -------------------# @@ -61,6 +40,9 @@ def cholesky( return ret +cholesky.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) def cross( x1: torch.Tensor, @@ -86,24 +68,15 @@ def cross( ) +cross.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def det(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: return torch.linalg.det(x, out=out) -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def diag( - x: torch.Tensor, - /, - *, - k: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.diag(x, diagonal=k) +det.support_native_out = True def diagonal( @@ -118,17 +91,6 @@ def diagonal( return torch.diagonal(x, offset=offset, dim1=axis1, dim2=axis2) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def eig( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] - ) - eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) - return result_tuple(eigenvalues, eigenvectors) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -140,6 +102,9 @@ def eigh( return result_tuple(eigenvalues, eigenvectors) +eigh.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigvalsh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -147,6 +112,9 @@ def eigvalsh( return torch.linalg.eigvalsh(x, UPLO=UPLO, out=out) +eigvalsh.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def inner( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None @@ -164,6 +132,9 @@ def inner( return torch.inner(x1, x2, out=out) +inner.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def inv( x: torch.Tensor, @@ -186,6 +157,9 @@ def inv( return ret +inv.support_native_out = True + + @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "bool")}, backend_version ) @@ -217,6 +191,9 @@ def matmul( return torch.matmul(x1, x2, out=out) +matmul.support_native_out = True + + @with_supported_dtypes({"2.0.1 and below": ("float", "complex")}, backend_version) def matrix_norm( x: torch.Tensor, @@ -230,6 +207,23 @@ def matrix_norm( return torch.linalg.matrix_norm(x, ord=ord, dim=axis, keepdim=keepdims, out=out) +matrix_norm.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def eig( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] + ) + eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) + return result_tuple(eigenvalues, eigenvectors) + + +eig.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_power( x: torch.Tensor, n: int, /, *, out: Optional[torch.Tensor] = None @@ -237,6 +231,9 @@ def matrix_power( return torch.linalg.matrix_power(x, n, out=out) +matrix_power.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_rank( x: torch.Tensor, @@ -274,6 +271,9 @@ def matrix_rank( return ret +matrix_rank.support_native_out = True + + def matrix_transpose( x: torch.Tensor, /, *, conjugate: bool = False, out: Optional[torch.Tensor] = None ) -> torch.Tensor: @@ -294,6 +294,9 @@ def outer( return torch.outer(x1, x2, out=out) +outer.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def pinv( x: torch.Tensor, @@ -307,6 +310,21 @@ def pinv( return torch.linalg.pinv(x, rtol, out=out) +pinv.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def tensorsolve( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.tensorsolve(x1, x2, dims=axes) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def qr( x: torch.Tensor, @@ -341,6 +359,9 @@ def slogdet( return results(sign, logabsdet) +slogdet.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def solve( x1: torch.Tensor, @@ -399,6 +420,9 @@ def svdvals(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch. return torch.linalg.svdvals(x, out=out) +svdvals.support_native_out = True + + # ToDo: re-add int32 support once # (https://github.com/pytorch/pytorch/issues/84530) is fixed @with_unsupported_dtypes({"2.0.1 and below": ("int32",)}, backend_version) @@ -424,18 +448,6 @@ def tensordot( return ret -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def tensorsolve( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.tensorsolve(x1, x2, dims=axes) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def trace( x: torch.Tensor, @@ -453,30 +465,6 @@ def trace( return ret -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def vander( - x: torch.tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - # torch.vander hasn't been used as it produces 0 gradients - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - ret = torch.pow( - torch.transpose(torch.unsqueeze(x, 0), 0, 1), - torch.arange(start, stop, step), - out=out, - ) - if ret.dtype != x.dtype: - return ret.to(x.dtype) - return ret - - def vecdot( x1: torch.Tensor, x2: torch.Tensor, @@ -497,6 +485,9 @@ def vecdot( return torch.tensordot(x1, x2, dims=([axis], [axis])).to(dtype) +vecdot.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("integer",)}, backend_version) def vector_norm( x: torch.Tensor, @@ -516,6 +507,48 @@ def vector_norm( return torch.linalg.vector_norm(x, ord, axis, keepdims, out=out) +vector_norm.support_native_out = True + + +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def diag( + x: torch.Tensor, + /, + *, + k: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.tensor: + return torch.diag(x, diagonal=k) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def vander( + x: torch.tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + # torch.vander hasn't been used as it produces 0 gradients + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + ret = torch.pow( + torch.transpose(torch.unsqueeze(x, 0), 0, 1), + torch.arange(start, stop, step), + out=out, + ) + if ret.dtype != x.dtype: + return ret.to(x.dtype) + return ret + + @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -543,3 +576,6 @@ def vector_to_skew_symmetric_matrix( row3 = torch.cat((-a2s, a1s, zs), -1) # BS x 3 x 3 return torch.cat((row1, row2, row3), -2, out=out) + + +vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/torch/manipulation.py b/ivy/functional/backends/torch/manipulation.py index 1b5efcec1cf06..fe8b37b4dcafd 100644 --- a/ivy/functional/backends/torch/manipulation.py +++ b/ivy/functional/backends/torch/manipulation.py @@ -14,42 +14,12 @@ from . import backend_version -# --- Helpers --- # -# --------------- # - - def _reshape_fortran_torch(x, shape): if len(x.shape) > 0: x = x.permute(*reversed(range(len(x.shape)))) return x.reshape(shape[::-1]).permute(list(range(len(shape)))[::-1]) -# --- Main --- # -# ------------ # - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version -) -def clip( - x: torch.Tensor, - x_min: Union[Number, torch.Tensor], - x_max: Union[Number, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if hasattr(x_min, "dtype"): - x_min = torch.asarray(x_min, device=x.device) - x_max = torch.asarray(x_max, device=x.device) - promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) - promoted_type = torch.promote_types(promoted_type, x.dtype) - x_min = x_min.to(promoted_type) - x_max = x_max.to(promoted_type) - x = x.to(promoted_type) - return torch.clamp(x, x_min, x_max, out=out) - - # Array API Standard # # -------------------# @@ -73,26 +43,7 @@ def concat( return torch.cat(xs, dim=axis, out=out) -def constant_pad( - x: torch.Tensor, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if 0 in x.shape: - new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] - return torch.ones(new_shape, dtype=x.dtype) * value - if x.shape == (): - x = x.unsqueeze(0) - if isinstance(pad_width, torch.Tensor): - pad_width = pad_width.detach().cpu().numpy().tolist() - pad_width_flat: List[int] = list() - for pad_width_sec in reversed(pad_width): - for item in pad_width_sec: - pad_width_flat.append(item) - return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) +concat.support_native_out = True def expand_dims( @@ -142,23 +93,6 @@ def permute_dims( return torch.permute(x, axes) -@with_unsupported_dtypes( - {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version -) -def repeat( - x: torch.Tensor, - /, - repeats: Union[int, Iterable[int]], - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if len(x.shape) == 0 and axis in [0, -1]: - axis = None - repeats = torch.tensor(repeats) - return torch.repeat_interleave(x, repeats, axis) - - def reshape( x: torch.Tensor, /, @@ -198,6 +132,69 @@ def roll( return torch.roll(x, shift, axis) +def squeeze( + x: torch.Tensor, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if isinstance(axis, int): + if x.size(dim=axis) > 1: + raise ValueError( + "Expected dimension of size [{}, {}], but found " + "dimension size {}".format(-x.dim(), x.dim(), axis) + ) + if x.shape[axis] != 1: + raise ivy.utils.exceptions.IvyException( + f"Expected size of axis to be 1 but was {x.shape[axis]}" + ) + return torch.squeeze(x, axis) + if axis is None: + if copy: + newarr = torch.clone(x) + return torch.squeeze(newarr) + return torch.squeeze(x) + newarr = torch.clone(x) + if isinstance(axis, tuple): + axis = list(axis) + normalise_axis = [ + (len(x.shape) - abs(element)) if element < 0 else element for element in axis + ] + normalise_axis.sort() + axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] + dim = x.dim() + for i in axis_updated_after_squeeze: + shape = x.shape[i] + if shape > 1 and (shape < -dim or dim <= shape): + raise ValueError( + "Expected dimension of size [{}, {}], " + "but found dimension size {}".format(-dim, dim, shape) + ) + else: + if copy: + newarr = torch.squeeze(newarr, i) + else: + x = torch.squeeze(x, i) + if copy: + return newarr + return x + + +def stack( + arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], + /, + *, + axis: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.stack(arrays, axis, out=out) + + +stack.support_native_out = True + + # Extra # # ------# @@ -247,64 +244,61 @@ def split( return list(torch.split(x, num_or_size_splits, axis)) -def squeeze( +@with_unsupported_dtypes( + {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version +) +def repeat( x: torch.Tensor, /, + repeats: Union[int, Iterable[int]], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, + axis: Optional[int] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if isinstance(axis, int): - if x.size(dim=axis) > 1: - raise ValueError( - "Expected dimension of size [{}, {}], but found " - "dimension size {}".format(-x.dim(), x.dim(), axis) - ) - if x.shape[axis] != 1: - raise ivy.utils.exceptions.IvyException( - f"Expected size of axis to be 1 but was {x.shape[axis]}" - ) - return torch.squeeze(x, axis) - if axis is None: - if copy: - newarr = torch.clone(x) - return torch.squeeze(newarr) - return torch.squeeze(x) - newarr = torch.clone(x) - if isinstance(axis, tuple): - axis = list(axis) - normalise_axis = [ - (len(x.shape) - abs(element)) if element < 0 else element for element in axis - ] - normalise_axis.sort() - axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] - dim = x.dim() - for i in axis_updated_after_squeeze: - shape = x.shape[i] - if shape > 1 and (shape < -dim or dim <= shape): - raise ValueError( - "Expected dimension of size [{}, {}], " - "but found dimension size {}".format(-dim, dim, shape) - ) - else: - if copy: - newarr = torch.squeeze(newarr, i) - else: - x = torch.squeeze(x, i) - if copy: - return newarr - return x + if len(x.shape) == 0 and axis in [0, -1]: + axis = None + repeats = torch.tensor(repeats) + return torch.repeat_interleave(x, repeats, axis) -def stack( - arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], +def tile( + x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + if isinstance(repeats, torch.Tensor): + repeats = repeats.detach().cpu().numpy().tolist() + return x.repeat(repeats) + + +def constant_pad( + x: torch.Tensor, /, + pad_width: List[List[int]], *, - axis: int = 0, + value: Number = 0.0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.stack(arrays, axis, out=out) + if 0 in x.shape: + new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] + return torch.ones(new_shape, dtype=x.dtype) * value + if x.shape == (): + x = x.unsqueeze(0) + if isinstance(pad_width, torch.Tensor): + pad_width = pad_width.detach().cpu().numpy().tolist() + pad_width_flat: List[int] = list() + for pad_width_sec in reversed(pad_width): + for item in pad_width_sec: + pad_width_flat.append(item) + return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) + + +def zero_pad( + x: torch.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[torch.Tensor] = None, +): + return constant_pad(x, pad_width, value=0.0) def swapaxes( @@ -319,12 +313,29 @@ def swapaxes( return torch.transpose(x, axis0, axis1) -def tile( - x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes( + {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version +) +def clip( + x: torch.Tensor, + x_min: Union[Number, torch.Tensor], + x_max: Union[Number, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if isinstance(repeats, torch.Tensor): - repeats = repeats.detach().cpu().numpy().tolist() - return x.repeat(repeats) + if hasattr(x_min, "dtype"): + x_min = torch.asarray(x_min, device=x.device) + x_max = torch.asarray(x_max, device=x.device) + promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) + promoted_type = torch.promote_types(promoted_type, x.dtype) + x_min = x_min.to(promoted_type) + x_max = x_max.to(promoted_type) + x = x.to(promoted_type) + return torch.clamp(x, x_min, x_max, out=out) + + +clip.support_native_out = True def unstack( @@ -344,18 +355,3 @@ def unstack( if keepdims: return [r.unsqueeze(axis) for r in ret] return ret - - -def zero_pad( - x: torch.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[torch.Tensor] = None, -): - return constant_pad(x, pad_width, value=0.0) - - -concat.support_native_out = True -stack.support_native_out = True -clip.support_native_out = True diff --git a/ivy/functional/backends/torch/random.py b/ivy/functional/backends/torch/random.py index 12f7a19686be8..55495e074b067 100644 --- a/ivy/functional/backends/torch/random.py +++ b/ivy/functional/backends/torch/random.py @@ -14,10 +14,52 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, torch.Tensor] = 0.0, + high: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + device: torch.device, + seed: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + shape = _check_bounds_and_get_shape(low, high, shape).shape + rand_range = high - low + if seed: + torch.manual_seed(seed) + if torch.is_tensor(shape): + shape = shape.tolist() + return ( + torch.rand(shape, device=device, dtype=torch.float) * rand_range + low + ).type(dtype) + + +def random_normal( + *, + mean: Union[float, torch.Tensor] = 0.0, + std: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + seed: Optional[int] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + dtype = ivy.as_native_dtype(dtype) + if seed: + torch.manual_seed(seed) + if isinstance(mean, (int, float)) and isinstance(std, (int, float)): + return torch.normal(mean, std, shape, out=out).type(dtype).to(device) + return torch.normal(mean, std, out=out).type(dtype).to(device) + random_normal.support_native_out = True -multinomial.support_native_out = True -shuffle.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) @@ -48,6 +90,9 @@ def multinomial( return torch.multinomial(probs.float(), num_samples, replace, out=out).to(device) +multinomial.support_native_out = True + + def randint( low: Union[int, torch.Tensor], high: Union[int, torch.Tensor], @@ -70,51 +115,6 @@ def randint( return (torch.rand(shape, device=device) * rand_range + low).to(dtype) -def random_normal( - *, - mean: Union[float, torch.Tensor] = 0.0, - std: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - seed: Optional[int] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - dtype = ivy.as_native_dtype(dtype) - if seed: - torch.manual_seed(seed) - if isinstance(mean, (int, float)) and isinstance(std, (int, float)): - return torch.normal(mean, std, shape, out=out).type(dtype).to(device) - return torch.normal(mean, std, out=out).type(dtype).to(device) - - -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, torch.Tensor] = 0.0, - high: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - device: torch.device, - seed: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - shape = _check_bounds_and_get_shape(low, high, shape).shape - rand_range = high - low - if seed: - torch.manual_seed(seed) - if torch.is_tensor(shape): - shape = shape.tolist() - return ( - torch.rand(shape, device=device, dtype=torch.float) * rand_range + low - ).type(dtype) - - def seed(*, seed_value: int = 0) -> None: torch.manual_seed(seed_value) torch.cuda.manual_seed(seed_value) @@ -140,3 +140,6 @@ def shuffle( if seed: torch.manual_seed(seed) return torch.index_select(x, 0, torch.randperm(batch_size), out=out) + + +shuffle.support_native_out = True diff --git a/ivy/functional/backends/torch/searching.py b/ivy/functional/backends/torch/searching.py index 778a0a3b5d141..702100b3b5ab1 100644 --- a/ivy/functional/backends/torch/searching.py +++ b/ivy/functional/backends/torch/searching.py @@ -69,19 +69,6 @@ def argmin( return ret -# Extra # -# ----- # - - -def argwhere( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.argwhere(x) - - def nonzero( x: torch.Tensor, /, @@ -120,3 +107,16 @@ def where( if condition.dtype is not torch.bool: condition = condition == 1.0 return ivy.astype(torch.where(condition, x1, x2), x1.dtype, copy=False) + + +# Extra # +# ----- # + + +def argwhere( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.argwhere(x) diff --git a/ivy/functional/backends/torch/sorting.py b/ivy/functional/backends/torch/sorting.py index f5e085b0088a3..6fea3ca9794c3 100644 --- a/ivy/functional/backends/torch/sorting.py +++ b/ivy/functional/backends/torch/sorting.py @@ -8,12 +8,6 @@ from . import backend_version -argsort.support_native_out = True -sort.support_native_out = True -msort.support_native_out = True -searchsorted.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def argsort( x: torch.Tensor, @@ -32,6 +26,30 @@ def argsort( return sorted_indices +argsort.support_native_out = True + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def sort( + x: torch.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if out is not None: + out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) + sorted_tensor, _ = torch.sort( + x, dim=axis, descending=descending, stable=stable, out=out + ) + return sorted_tensor + + +sort.support_native_out = True + + # msort @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def msort( @@ -40,6 +58,9 @@ def msort( return torch.msort(a, out=out) +msort.support_native_out = True + + @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def searchsorted( x: torch.Tensor, @@ -92,19 +113,4 @@ def searchsorted( return torch.searchsorted(x, v, sorter=sorter, side=side).to(ret_dtype) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def sort( - x: torch.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if out is not None: - out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) - sorted_tensor, _ = torch.sort( - x, dim=axis, descending=descending, stable=stable, out=out - ) - return sorted_tensor +searchsorted.support_native_out = True diff --git a/ivy/functional/backends/torch/statistical.py b/ivy/functional/backends/torch/statistical.py index 930858211703a..04694a0bdca3f 100644 --- a/ivy/functional/backends/torch/statistical.py +++ b/ivy/functional/backends/torch/statistical.py @@ -1,3 +1,5 @@ +# global +torch_scatter = None from typing import Union, Optional, Sequence import torch @@ -8,130 +10,30 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# global -torch_scatter = None - - -# --- Helpers --- # -# --------------- # - - -def _infer_dtype(dtype: torch.dtype) -> torch.dtype: - default_dtype = ivy.infer_default_dtype(dtype) - if default_dtype in ivy.valid_dtypes: - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return ivy.as_native_dtype(default_dtype) - return ivy.as_native_dtype(dtype) - - -# --- Main --- # -# ------------ # - - -# Extra # -# ----- # +# Array API Standard # +# -------------------# -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - }, - backend_version, -) -def cumprod( +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def min( x: torch.Tensor, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - - if not (exclusive or reverse): - return torch.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - ret = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumprod(x, -1, dtype=dtype) - ret = torch.transpose(x, axis, -1) - else: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - ret = torch.flip(x, dims=(axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "1.12.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cumsum( - x: torch.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[torch.dtype] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - res = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumsum(x, -1, dtype=dtype) - res = torch.transpose(x, axis, -1) - else: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - res = torch.flip(x, dims=(axis,)) + if axis == (): if ivy.exists(out): - return ivy.inplace_update(out, res) - return res - return torch.cumsum(x, axis, dtype=dtype, out=out) + return ivy.inplace_update(out, x) + else: + return x + if not keepdims and not axis and axis != 0: + return torch.amin(input=x, out=out) + return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def einsum( - equation: str, - *operands: torch.Tensor, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = _get_promoted_type_of_operands(operands) - return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) +min.support_native_out = True @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -153,6 +55,9 @@ def max( return torch.amax(input=x, dim=axis, keepdim=keepdims, out=out) +max.support_native_out = True + + def mean( x: torch.Tensor, /, @@ -176,27 +81,15 @@ def mean( return torch.mean(x, dim=axis, keepdim=keepdims, out=out) -# Array API Standard # -# -------------------# +mean.support_native_out = True -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def min( - x: torch.Tensor, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if axis == (): - if ivy.exists(out): - return ivy.inplace_update(out, x) - else: - return x - if not keepdims and not axis and axis != 0: - return torch.amin(input=x, out=out) - return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) +def _infer_dtype(dtype: torch.dtype) -> torch.dtype: + default_dtype = ivy.infer_default_dtype(dtype) + if default_dtype in ivy.valid_dtypes: + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return ivy.as_native_dtype(default_dtype) + return ivy.as_native_dtype(dtype) # Function does support uint8, but allowing support for unsigned will cause @@ -321,8 +214,113 @@ def var( ).to(x.dtype) -min.support_native_out = True -max.support_native_out = True -mean.support_native_out = True +# Extra # +# ----- # + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + }, + backend_version, +) +def cumprod( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + + if not (exclusive or reverse): + return torch.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + ret = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumprod(x, -1, dtype=dtype) + ret = torch.transpose(x, axis, -1) + else: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + ret = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + cumprod.support_native_out = True + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "1.12.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cumsum( + x: torch.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + res = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumsum(x, -1, dtype=dtype) + res = torch.transpose(x, axis, -1) + else: + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + res = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res + return torch.cumsum(x, axis, dtype=dtype, out=out) + + cumsum.support_native_out = True + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def einsum( + equation: str, + *operands: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = _get_promoted_type_of_operands(operands) + return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) diff --git a/ivy/functional/backends/torch/utility.py b/ivy/functional/backends/torch/utility.py index 5d62f84cb7479..ca643ce0d4095 100644 --- a/ivy/functional/backends/torch/utility.py +++ b/ivy/functional/backends/torch/utility.py @@ -3,10 +3,6 @@ from typing import Union, Optional, Sequence -all.support_native_out = True -any.support_native_out = True - - def all( x: torch.Tensor, /, @@ -29,6 +25,9 @@ def all( return x +all.support_native_out = True + + def any( x: torch.Tensor, /, @@ -49,3 +48,6 @@ def any( for i, a in enumerate(axis): x = torch.any(x, dim=a if keepdims else a - i, keepdim=keepdims, out=out) return x + + +any.support_native_out = True diff --git a/ivy/functional/frontends/torch/comparison_ops.py b/ivy/functional/frontends/torch/comparison_ops.py index eeaadadfbf864..b743b38d135f0 100644 --- a/ivy/functional/frontends/torch/comparison_ops.py +++ b/ivy/functional/frontends/torch/comparison_ops.py @@ -290,7 +290,7 @@ def topk(input, k, dim=None, largest=True, sorted=True, *, out=None): gt = greater -ne = not_equal ge = greater_equal le = less_equal lt = less +ne = not_equal diff --git a/ivy/functional/ivy/activations.py b/ivy/functional/ivy/activations.py index d55d9dfcde068..f19180ab7a529 100644 --- a/ivy/functional/ivy/activations.py +++ b/ivy/functional/ivy/activations.py @@ -18,10 +18,6 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # - - def _gelu_jax_like( x: Union[ivy.Array, ivy.NativeArray], /, @@ -35,87 +31,6 @@ def _gelu_jax_like( return fn_original(x, approximate=True, out=out) -def _leaky_relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original: Optional[Callable] = None, - alpha: float = 0.2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.astype(x * alpha, x.dtype), - x, - ) - - -def _relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.array(0.0, dtype=x.dtype), - x, - ) - - -def _softplus_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, - out: Optional[ivy.Array] = None, -): - if beta is not None: - x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) - else: - x_beta = x - amax = ivy.relu(x_beta) - res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) - res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) - res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( - x.dtype - ) * ivy.astype(1j, x.dtype) - if beta is not None: - res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) - if threshold is not None: - res = ivy.where( - ivy.real(x_beta) < threshold, - res, - x, - ).astype(x.dtype) - return res - - -def _wrap_between(y, a): - """Wrap y between [-a, a]""" - a = ivy.array(a, dtype=y.dtype) - a2 = ivy.array(2 * a, dtype=y.dtype) - zero = ivy.array(0, dtype=y.dtype) - rem = ivy.remainder(ivy.add(y, a), a2) - rem = ivy.where(rem < zero, rem + a2, rem) - a - return rem - - -# --- Main --- # -# ------------ # - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -184,52 +99,26 @@ def gelu( return current_backend(x).gelu(x, approximate=approximate, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -def hardswish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply the hardswish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the hardswish activation of each element in ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 0., 4.]) - >>> y = ivy.hardswish(x) - >>> y - ivy.array([0., 0., 4.]) +gelu.jax_like = _gelu_jax_like - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) - >>> x = ivy.hardswish(x, out=x) - >>> x - { - a: ivy.array([-0., 4., 5.]), - b: ivy.array([0., 5.]) - } - """ - return current_backend(x).hardswish(x, out=out) +def _leaky_relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original: Optional[Callable] = None, + alpha: float = 0.2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.astype(x * alpha, x.dtype), + x, + ) @handle_exceptions @@ -310,6 +199,9 @@ def leaky_relu( return current_backend(x).leaky_relu(x, alpha=alpha, out=out) +leaky_relu.jax_like = _leaky_relu_jax_like + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -384,60 +276,22 @@ def log_softmax( return current_backend(x).log_softmax(x, axis=axis, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def mish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +def _relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Apply the mish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the mish activation of each element in - ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.mish(x) - >>> print(y) - ivy.array([-0.30340147, 0. , 0.86509842]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.mish(x, out = y) - >>> print(y) - ivy.array([ 1.40337825, 0.56114835, -0.20788449]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.mish(x) - >>> print(x) - { - a: ivy.array([0.86509842, -0.30883577]), - b: ivy.array([0.28903052, -0.10714479]) - } - """ - return current_backend(x).mish(x, out=out) + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.array(0.0, dtype=x.dtype), + x, + ) @handle_exceptions @@ -509,6 +363,9 @@ def relu( return current_backend(x).relu(x, out=out) +relu.jax_like = _relu_jax_like + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -639,6 +496,46 @@ def softmax( return current_backend(x).softmax(x, axis=axis, out=out) +def _wrap_between(y, a): + """Wrap y between [-a, a]""" + a = ivy.array(a, dtype=y.dtype) + a2 = ivy.array(2 * a, dtype=y.dtype) + zero = ivy.array(0, dtype=y.dtype) + rem = ivy.remainder(ivy.add(y, a), a2) + rem = ivy.where(rem < zero, rem + a2, rem) - a + return rem + + +def _softplus_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, + out: Optional[ivy.Array] = None, +): + if beta is not None: + x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) + else: + x_beta = x + amax = ivy.relu(x_beta) + res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) + res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) + res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( + x.dtype + ) * ivy.astype(1j, x.dtype) + if beta is not None: + res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) + if threshold is not None: + res = ivy.where( + ivy.real(x_beta) < threshold, + res, + x, + ).astype(x.dtype) + return res + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -710,7 +607,108 @@ def softplus( return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out) -gelu.jax_like = _gelu_jax_like -leaky_relu.jax_like = _leaky_relu_jax_like -relu.jax_like = _relu_jax_like softplus.jax_like = _softplus_jax_like + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def mish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply the mish activation function element-wise. + + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the mish activation of each element in + ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.mish(x) + >>> print(y) + ivy.array([-0.30340147, 0. , 0.86509842]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.mish(x, out = y) + >>> print(y) + ivy.array([ 1.40337825, 0.56114835, -0.20788449]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.mish(x) + >>> print(x) + { + a: ivy.array([0.86509842, -0.30883577]), + b: ivy.array([0.28903052, -0.10714479]) + } + """ + return current_backend(x).mish(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def hardswish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply the hardswish activation function element-wise. + + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the hardswish activation of each element in ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 0., 4.]) + >>> y = ivy.hardswish(x) + >>> y + ivy.array([0., 0., 4.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) + >>> x = ivy.hardswish(x, out=x) + >>> x + { + a: ivy.array([-0., 4., 5.]), + b: ivy.array([0., 5.]) + } + """ + return current_backend(x).hardswish(x, out=out) diff --git a/ivy/functional/ivy/constants.py b/ivy/functional/ivy/constants.py index 461877e75b3b2..a830e3df30d66 100644 --- a/ivy/functional/ivy/constants.py +++ b/ivy/functional/ivy/constants.py @@ -1,55 +1,64 @@ # global import math -atto = 1e-18 -centi = 1e-2 -deci = 1e-1 -deka = 1e1 + # Array API Standard # # -------------------# e = math.e -exa = 1e18 -exbi = 2**60 -femto = 1e-15 -gibi = 2**30 -giga = 1e9 +"""IEEE 754 floating-point representation of Euler's constant.""" + +pi = math.pi +"""IEEE 754 floating-point representation of the mathematical constant π.""" + +nan = math.nan +"""IEEE 754 floating-point representation of Not a Number (NaN).""" + +inf = math.inf +"""IEEE 754 floating-point representation of (positive) infinity.""" + +newaxis = None +"""An alias for None which is useful for indexing arrays.""" + + # Mathematical constants # # ------# golden = golden_ratio = (1 + math.sqrt(5)) / 2 +quetta = 1e30 +ronna = 1e27 +yotta = 1e24 +zetta = 1e21 +exa = 1e18 +peta = 1e15 +tera = 1e12 +giga = 1e9 +mega = 1e6 +kilo = 1e3 hecto = 1e2 -inf = math.inf +deka = 1e1 +deci = 1e-1 +centi = 1e-2 +milli = 1e-3 +micro = 1e-6 +nano = 1e-9 +pico = 1e-12 +femto = 1e-15 +atto = 1e-18 +zepto = 1e-21 +yocto = 1e-24 +ronto = 1e-27 +quecto = 1e-30 + + # Binary prefixes # # ------# kibi = 2**10 -kilo = 1e3 mebi = 2**20 -mega = 1e6 -micro = 1e-6 -milli = 1e-3 -nan = math.nan -nano = 1e-9 -newaxis = None -pebi = 2**50 -peta = 1e15 -pi = math.pi -pico = 1e-12 -quecto = 1e-30 -quetta = 1e30 -ronna = 1e27 -ronto = 1e-27 +gibi = 2**30 tebi = 2**40 -tera = 1e12 -yobi = 2**80 -yocto = 1e-24 -yotta = 1e24 +pebi = 2**50 +exbi = 2**60 zebi = 2**70 -zepto = 1e-21 -zetta = 1e21 -"""IEEE 754 floating-point representation of Euler's constant.""" -"""IEEE 754 floating-point representation of the mathematical constant π.""" -"""IEEE 754 floating-point representation of Not a Number (NaN).""" -"""IEEE 754 floating-point representation of (positive) infinity.""" -"""An alias for None which is useful for indexing arrays.""" +yobi = 2**80 diff --git a/ivy/functional/ivy/control_flow_ops.py b/ivy/functional/ivy/control_flow_ops.py index e41773bbcf366..4c2d78dc13c39 100644 --- a/ivy/functional/ivy/control_flow_ops.py +++ b/ivy/functional/ivy/control_flow_ops.py @@ -9,100 +9,6 @@ ) -# --- Helpers --- # -# --------------- # - - -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) - - -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} - - -# --- Main --- # -# ------------ # - - -def cast_bool(x): - return bool(x) - - -# todo (nightcrab) find a better place for these cmp functions - - -def cmp_is(left, right): - return left is right - - -def cmp_isnot(left, right): - return left is not right - - -def for_loop( - iterable: Iterable[Any], - body_fn: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - """ - Loops over an iterable, passing the current iteration along with a tuple of - variables into the provided body function. - - Parameters - ---------- - iterable - The iterable to loop over. - body_fn - A function to call each iteration, first taking the iterator value - and then a tuple of extra parameters. - vars - Extra parameters to be passed to body_fn. - - Returns - ------- - ret - The loop's return value (if any). - - Example - ---- - ``` - def body_fn(k, args): - print(k+1) - return args - - lst = [5,6] - - ivy.for_loop(lst, body_fn, ()) - >>> 5 - >>> 6 - ``` - """ - iterator = iterable.__iter__() - - vars_dict = _tuple_to_dict(vars) - - def test_fn(iterator, original_body, vars_dict): - try: - val = iterator.__next__() - except StopIteration: - return False - - vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) - - for k in range(len(vars_tuple)): - vars_dict[k] = vars_tuple[k] - - return True - - def empty_function(iterator, original_body, vars_dict): - return (iterator, original_body, vars_dict) - - packed_vars = (iterator, body_fn, vars_dict) - - return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) - - def if_else( cond: Callable, body_fn: Callable, @@ -164,17 +70,6 @@ def _if_else(cond, body_fn, orelse_fn, vars): return _if_else(cond, body_fn, orelse_fn, vars) -def try_except( - body1: Callable, - body2: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - try: - return body1(*vars) - except Exception as e: - return body2(*vars, e) - - def while_loop( test_fn: Callable, body_fn: Callable, @@ -229,3 +124,100 @@ def _while_loop(test_fn, body_fn, vars): body_fn = to_ivy_arrays_and_back(body_fn) return _while_loop(test_fn, body_fn, vars) + + +def for_loop( + iterable: Iterable[Any], + body_fn: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + """ + Loops over an iterable, passing the current iteration along with a tuple of + variables into the provided body function. + + Parameters + ---------- + iterable + The iterable to loop over. + body_fn + A function to call each iteration, first taking the iterator value + and then a tuple of extra parameters. + vars + Extra parameters to be passed to body_fn. + + Returns + ------- + ret + The loop's return value (if any). + + Example + ---- + ``` + def body_fn(k, args): + print(k+1) + return args + + lst = [5,6] + + ivy.for_loop(lst, body_fn, ()) + >>> 5 + >>> 6 + ``` + """ + iterator = iterable.__iter__() + + vars_dict = _tuple_to_dict(vars) + + def test_fn(iterator, original_body, vars_dict): + try: + val = iterator.__next__() + except StopIteration: + return False + + vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) + + for k in range(len(vars_tuple)): + vars_dict[k] = vars_tuple[k] + + return True + + def empty_function(iterator, original_body, vars_dict): + return (iterator, original_body, vars_dict) + + packed_vars = (iterator, body_fn, vars_dict) + + return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) + + +def try_except( + body1: Callable, + body2: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + try: + return body1(*vars) + except Exception as e: + return body2(*vars, e) + + +# todo (nightcrab) find a better place for these cmp functions + + +def cmp_is(left, right): + return left is right + + +def cmp_isnot(left, right): + return left is not right + + +def cast_bool(x): + return bool(x) + + +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} + + +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) diff --git a/ivy/functional/ivy/creation.py b/ivy/functional/ivy/creation.py index f12f82585431e..0273b8f4b0e55 100644 --- a/ivy/functional/ivy/creation.py +++ b/ivy/functional/ivy/creation.py @@ -35,32 +35,44 @@ handle_backend_invalid, ) +# Helpers # +# --------# -# Type hints # -# -----------# - -SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") -_T_co = TypeVar("_T_co", covariant=True) +def asarray_handle_nestable(fn: Callable) -> Callable: + fn_name = fn.__name__ -class NestedSequence(Protocol[_T_co]): - def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: - ... + @functools.wraps(fn) + def _asarray_handle_nestable(*args, **kwargs): + """ + Call `fn` with the *nestable* property of the function correctly handled. This + means mapping the function to the container leaves if any containers are passed + in the input. - def __len__(self, /) -> int: - ... + Parameters + ---------- + args + The arguments to be passed to the function. + kwargs + The keyword arguments to be passed to the function. -# --- Helpers --- # -# --------------- # + Returns + ------- + The return of the function, with the nestable property handled correctly. + """ + # This decorator should only be applied to ivy.asarray, so we know where + # the container must be if there is one. + cont_fn = getattr(ivy.Container, "static_" + fn_name) + if isinstance(args[0], ivy.Container): + return cont_fn(*args, **kwargs) + # if the passed arguments does not contain a container, the function using + # the passed arguments, returning an ivy or a native array. + return fn(*args, **kwargs) -def _flatten_nest(xs): - for x in xs: - if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): - yield from _flatten_nest(x) - else: - yield x + _asarray_handle_nestable.handle_nestable = True + return _asarray_handle_nestable def _ivy_to_native(x): @@ -80,15 +92,6 @@ def _ivy_to_native(x): return x -def _remove_np_bfloat16(obj): - # unlike other frameworks, torch and paddle do not support creating tensors - # from numpy arrays that have bfloat16 dtype using any extension because - # bfloat16 in not supported natively by numpy (as of version <=1.25) - if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": - return obj.tolist() - return obj - - def _shape_to_native(x): # checks the first element of the leaf list and # converts it to a native array if it is an ivy array @@ -106,8 +109,168 @@ def _shape_to_native(x): return x -# --- Main --- # -# ------------ # +def _flatten_nest(xs): + for x in xs: + if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): + yield from _flatten_nest(x) + else: + yield x + + +def _remove_np_bfloat16(obj): + # unlike other frameworks, torch and paddle do not support creating tensors + # from numpy arrays that have bfloat16 dtype using any extension because + # bfloat16 in not supported natively by numpy (as of version <=1.25) + if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": + return obj.tolist() + return obj + + +def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): + """ + Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances + and return arrays are all converted to `ivy.Array` instances. + + This wrapper is specifically for the backend implementations of + asarray. + + It assumes either all the elements in a leaf list are ivy arrays + or none of them are. It checks the first element of all the leaf + list. If it is an ivy array, it converts all the elements in the + leaf list to native otherwise it skips that leaf list. + """ + new_arg = _ivy_to_native(args[0]) + new_args = (new_arg,) + args[1:] + if dtype is not None: + dtype = ivy.default_dtype(dtype=dtype, as_native=True) + return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) + + return _asarray_to_native_arrays_and_back + + +def asarray_infer_dtype(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_dtype(*args, dtype=None, **kwargs): + """ + Determine the correct `dtype`, and then calls the function with the `dtype` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. + + Parameters + ---------- + args + The arguments to be passed to the function. + + dtype + The dtype for the function. + + kwargs + The keyword arguments to be passed to the function. + + Returns + ------- + The return of the function, with `dtype` passed explicitly. + """ + + def _infer_dtype(obj): + if isinstance(obj, ivy.NativeShape): + obj = list(obj) + if hasattr(obj, "dtype"): + return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype + else: + return ivy.default_dtype(item=obj) + + if not ivy.exists(dtype): + arr = args[0] + # get default dtypes for all elements + dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] + # flatten the nested structure + dtype_list = _flatten_nest(dtype_list) + # keep unique dtypes + dtype_list = list(set(dtype_list)) + if len(dtype_list) != 0: # handle the case of empty input + # promote all dtypes to a single dtype + dtype = dtype_list[0] + # we disable precise mode to avoid wider than necessary casting + # that might result from the mixing of int32 and float32 + with ivy.PreciseMode(False): + for dt in dtype_list[1:]: + dtype = ivy.promote_types(dtype, dt) + else: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + # call the function with dtype provided explicitly + return fn(*args, dtype=dtype, **kwargs) + + _asarray_infer_dtype.infer_dtype = True + return _asarray_infer_dtype + + +def asarray_infer_device(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_device(*args, device=None, **kwargs): + """ + Determine the correct `device`, and then calls the function with the `device` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. + + Parameters + ---------- + args + The arguments to be passed to the function. + + device + The device for the function. + + kwargs + The keyword arguments to be passed to the function. + + Returns + ------- + The return of the function, with `device` passed explicitly. + """ + if isinstance(args[0], list): + return fn( + *args, device=ivy.default_device(device, as_native=True), **kwargs + ) + + # find the first array argument, if required + arr = None if ivy.exists(device) else args[0] + # infer the correct device + device = ivy.default_device(device, item=arr, as_native=True) + # call the function with device provided explicitly + return fn(*args, device=device, **kwargs) + + _asarray_infer_device.infer_device = True + return _asarray_infer_device + + +def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: + @functools.wraps(fn) + def _inputs_to_native_shapes(*args, **kwargs): + new_arg = _shape_to_native(args[0]) + new_args = (new_arg,) + args[1:] + return fn(*new_args, **kwargs) + + _inputs_to_native_shapes.inputs_to_native_shapes = True + return _inputs_to_native_shapes + + +# Type hints # +# -----------# + +SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") +_T_co = TypeVar("_T_co", covariant=True) + + +class NestedSequence(Protocol[_T_co]): + def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: + ... + + def __len__(self, /) -> int: + ... # Array API Standard # @@ -309,199 +472,434 @@ def asarray( ) -def asarray_handle_nestable(fn: Callable) -> Callable: - fn_name = fn.__name__ +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def zeros( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with zeros. - @functools.wraps(fn) - def _asarray_handle_nestable(*args, **kwargs): - """ - Call `fn` with the *nestable* property of the function correctly handled. This - means mapping the function to the container leaves if any containers are passed - in the input. + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type must + be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - Parameters - ---------- - args - The arguments to be passed to the function. + Returns + ------- + ret + an array containing zeros. - kwargs - The keyword arguments to be passed to the function. - Returns - ------- - The return of the function, with the nestable property handled correctly. - """ - # This decorator should only be applied to ivy.asarray, so we know where - # the container must be if there is one. - cont_fn = getattr(ivy.Container, "static_" + fn_name) - if isinstance(args[0], ivy.Container): - return cont_fn(*args, **kwargs) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - # if the passed arguments does not contain a container, the function using - # the passed arguments, returning an ivy or a native array. - return fn(*args, **kwargs) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - _asarray_handle_nestable.handle_nestable = True - return _asarray_handle_nestable + Examples + -------- + With :class:`ivy.NativeShape` input: + >>> shape = (3, 5) + >>> x = ivy.zeros(shape) + >>> print(x) + ivy.array([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + >>> x = ivy.zeros(5) + >>> print(x) + ivy.array([0., 0., 0., 0., 0.]) + """ + return current_backend().zeros(shape, dtype=dtype, device=device, out=out) -def asarray_infer_device(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_device(*args, device=None, **kwargs): - """ - Determine the correct `device`, and then calls the function with the `device` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. - Parameters - ---------- - args - The arguments to be passed to the function. +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with ones. - device - The device for the function. + .. note:: - kwargs - The keyword arguments to be passed to the function. + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal to + zero (i.e., ``1 + 0j``). - Returns - ------- - The return of the function, with `device` passed explicitly. - """ - if isinstance(args[0], list): - return fn( - *args, device=ivy.default_device(device, as_native=True), **kwargs - ) + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - # find the first array argument, if required - arr = None if ivy.exists(device) else args[0] - # infer the correct device - device = ivy.default_device(device, item=arr, as_native=True) - # call the function with device provided explicitly - return fn(*args, device=device, **kwargs) + Returns + ------- + ret + an array containing ones. - _asarray_infer_device.infer_device = True - return _asarray_infer_device + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. -def asarray_infer_dtype(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_dtype(*args, dtype=None, **kwargs): - """ - Determine the correct `dtype`, and then calls the function with the `dtype` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - Parameters - ---------- - args - The arguments to be passed to the function. + Examples + -------- + With :class:`ivy.Shape` input: - dtype - The dtype for the function. + >>> shape = (2,2) + >>> x = ivy.ones(shape) + >>> print(x) + ivy.array([[1., 1.], + [1., 1.]]) - kwargs - The keyword arguments to be passed to the function. + With :class:`ivy.Dtype` input: - Returns - ------- - The return of the function, with `dtype` passed explicitly. - """ + >>> shape = (3,2) + >>> d_type = ivy.int64 + >>> y = ivy.ones(shape, dtype=d_type) + >>> print(y) + ivy.array([[1, 1], + [1, 1], + [1, 1]]) - def _infer_dtype(obj): - if isinstance(obj, ivy.NativeShape): - obj = list(obj) - if hasattr(obj, "dtype"): - return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype - else: - return ivy.default_dtype(item=obj) + With :class:`ivy.Device` input: - if not ivy.exists(dtype): - arr = args[0] - # get default dtypes for all elements - dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] - # flatten the nested structure - dtype_list = _flatten_nest(dtype_list) - # keep unique dtypes - dtype_list = list(set(dtype_list)) - if len(dtype_list) != 0: # handle the case of empty input - # promote all dtypes to a single dtype - dtype = dtype_list[0] - # we disable precise mode to avoid wider than necessary casting - # that might result from the mixing of int32 and float32 - with ivy.PreciseMode(False): - for dt in dtype_list[1:]: - dtype = ivy.promote_types(dtype, dt) - else: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - # call the function with dtype provided explicitly - return fn(*args, dtype=dtype, **kwargs) + >>> shape = (3,2) + >>> y = ivy.ones(shape, device="cpu") + >>> print(y) + ivy.array([[1., 1.], + [1., 1.], + [1., 1.]]) - _asarray_infer_dtype.infer_dtype = True - return _asarray_infer_dtype + With :class:`ivy.Array` input: + + >>> shape = (1, 5, 2) + >>> x = ivy.zeros(shape) + >>> ivy.ones(shape, out=x) + >>> print(x) + ivy.array([[[1., 1.], + [1., 1.], + [1., 1.], + [1., 1.], + [1., 1.]]]) + """ + return current_backend().ones(shape, dtype=dtype, device=device, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def full_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + fill_value: Number, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ``fill_value`` and having the same ``shape`` as an + input array ``x`` . + + Parameters + ---------- + x + input array from which to derive the output array shape. + fill_value + Scalar fill value + dtype + output array data type. If ``dtype`` is `None`, the output array data type must + be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and where every element is equal to + ``fill_value``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :code:`int` datatype: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> fill_value = 1 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> fill_value = 0.000123 + >>> x = ivy.ones(5) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With float datatype: + + >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([3.0, 8.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x,fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123]) + + >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([[0.000123, 0.000123, 0.000123], + [0.000123, 0.000123, 0.000123]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), + ... b=ivy.array([4.123, 5.23, 6.23])) + >>> fill_value = 15.0 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + { + a: ivy.array([15., 15., 15.]), + b: ivy.array([15., 15., 15.]) + } + """ + return current_backend(x).full_like( + x, fill_value, dtype=dtype, device=device, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ones and having the same shape as an input array + ``x``. + + .. note:: + + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal + to zero (i.e., ``1 + 0j``). + + Parameters + ---------- + x + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default ``None``. + device + device on which to place the created array. If device is ``None``, the output + array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and filled with ``ones``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1., 1., 1.], + [1., 1., 1.]]) + + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.zeros(3) + >>> ivy.ones_like(x, out=y) + >>> print(y) + ivy.array([1., 1., 1.]) + + With :class:`ivy.NativeArray` input: + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1, 1, 1], + [1, 1, 1]]) -def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: - @functools.wraps(fn) - def _inputs_to_native_shapes(*args, **kwargs): - new_arg = _shape_to_native(args[0]) - new_args = (new_arg,) + args[1:] - return fn(*new_args, **kwargs) + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) - _inputs_to_native_shapes.inputs_to_native_shapes = True - return _inputs_to_native_shapes + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.ones_like(x) + >>> print(y) + { + a: ivy.array([1, 1, 1]), + b: ivy.array([1, 1, 1]) + } -def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): - """ - Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances - and return arrays are all converted to `ivy.Array` instances. + With :class:`ivy.Array` input: - This wrapper is specifically for the backend implementations of - asarray. + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.ones_like() + >>> print(y) + ivy.array([1, 1, 1, 1, 1]) - It assumes either all the elements in a leaf list are ivy arrays - or none of them are. It checks the first element of all the leaf - list. If it is an ivy array, it converts all the elements in the - leaf list to native otherwise it skips that leaf list. - """ - new_arg = _ivy_to_native(args[0]) - new_args = (new_arg,) + args[1:] - if dtype is not None: - dtype = ivy.default_dtype(dtype=dtype, as_native=True) - return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) + With :class:'ivy.Container' input: - return _asarray_to_native_arrays_and_back + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.ones_like() + >>> print(y) + { + a: ivy.array([1., 1.]), + b: ivy.array([1., 1.]) + } + """ + return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function +@infer_dtype @handle_device_shifting -def copy_array( +@infer_device +def zeros_like( x: Union[ivy.Array, ivy.NativeArray], /, *, - to_ivy_array: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Copy an array. + Return a new array filled with zeros and having the same ``shape`` as an input array + ``x``. Parameters ---------- x - array, input array containing elements to copy. - to_ivy_array - boolean, if True the returned array will be an ivy.Array object otherwise - returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., - depending on the backend), defaults to True. + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -509,83 +907,195 @@ def copy_array( Returns ------- ret - a copy of the input array ``x``. + an array having the same shape as ``x`` and filled with ``zeros``. - Examples - -------- - With one :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.copy_array(x) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.zeros_like(x) >>> print(y) - ivy.array([-1, 0, 1]) + ivy.array([0, 0, 0, 0, 0, 0]) - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = ivy.copy_array(x) + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.zeros_like(x) >>> print(y) - ivy.array([1, 0, 1, 1]) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) - >>> x = ivy.array([1, 0, 1, -1]) - >>> y = ivy.zeros((1, 4)) - >>> ivy.copy_array(x, out=y) + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.ones(3) + >>> ivy.zeros_like(x, out=y) >>> print(y) - ivy.array([1, 0, 1, -1]) + ivy.array([0., 0., 0.]) - >>> x = ivy.array([1, 0, 1, 1]) - >>> ivy.copy_array(x, out=x) - >>> print(x) - ivy.array([1, 0, 1, 1]) + With :class:`ivy.NativeArray` input: - With one :class:`ivy.Container` input: + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([[0, 0, 0],[0, 0, 0]]) - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.copy_array(x) + + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) >>> print(y) - { - a: ivy.array([-1, 0, 1]) - } + ivy.array([0, 0, 0, 0, 0, 0]) - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.copy_array(x) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.zeros_like(x) >>> print(y) { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) + a: ivy.array([0, 0, 0]), + b: ivy.array([0, 0, 0]) } - With one :class:`ivy.Container` static method: - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.Container.static_copy_array(x) + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.zeros_like() + >>> print(y) + ivy.array([0, 0, 0, 0, 0]) + + With :class:'ivy.Container' input: + + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.zeros_like() >>> print(y) { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) + a: ivy.array([0., 0.]), + b: ivy.array([0., 0.]) } + """ + return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) - With one :class:`ivy.Array` instance method: - >>> x = ivy.array([-1, 0, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([-1, 0, 1]) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def tril( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the lower triangular part of a matrix (or a stack of matrices) ``x``. + + .. note:: + + The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` + on the interval ``[0, min(M, N) - 1]``. + + Parameters + ---------- + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. + k + diagonal above which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the lower triangular part(s). The returned array must have + the same shape and data type as x. All elements above the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + return current_backend(x).tril(x, k=k, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def triu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the upper triangular part of a matrix (or a stack of matrices) ``x``. + + .. note:: + + The upper triangular part of the matrix is defined as the elements + on and above the specified diagonal ``k``. + + Parameters + ---------- + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. *, + k + diagonal below which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the upper triangular part(s). The returned array must have + the same shape and data type as x. All elements below the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([1, 0, 1, 1]) - With :class:`ivy.Container` instance method: + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) - >>> y = x.copy_array() - >>> print(y) - { - a: ivy.array([1, 0, 1]), - b: ivy.array([-1, 0, 1, 1]) - } + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. """ - return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) + return current_backend(x).triu(x, k=k, out=out) @handle_backend_invalid @@ -842,18 +1352,43 @@ def eye( @handle_out_argument @to_native_arrays_and_back @handle_array_function +@infer_dtype @handle_device_shifting -def from_dlpack( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +@infer_device +def linspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], + /, + num: int, + *, + axis: Optional[int] = None, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array containing the data from another (array) object with a - ``__dlpack__`` method. + Generate a certain number of evenly-spaced values in an interval along a given axis. + + See :math:`arange` that allows to specify the step size of evenly spaced values in + an interval. Parameters ---------- - x object - input (array) object. + start + First entry in the range. + stop + Final entry in the range. + num + Number of values to generate. + axis + Axis along which the operation is performed. + endpoint + If True, stop is the last sample. Otherwise, it is not included. + dtype + output array data type. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -861,87 +1396,186 @@ def from_dlpack( Returns ------- ret - an array containing the data in `x`. - - .. admonition:: Note - :class: note - - The returned array may be either a copy or a view. See - :ref:`data-interchange` for details. + Tensor of evenly-spaced values. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.linspace.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. + + Functional Examples + ------------------- + + With float input: + + >>> x = ivy.linspace(1, 2, 3) + >>> print(x) + ivy.array([1. , 1.5, 2. ]) + + >>> x = ivy.linspace(1, 2, 4, endpoint=False) + >>> print(x) + ivy.array([1., 1.25, 1.5 , 1.75]) + + >>> x = ivy.linspace(1, 10, 4, dtype="int32") + >>> print(x) + ivy.array([ 1, 4, 7, 10]) + + >>> x = ivy.linspace(1, 2, 4, device= "cpu") + >>> print(x) + ivy.array([1., 1.33333337, 1.66666663, 2.]) + + >>> y = ivy.array([0,0,0,0]) + >>> ivy.linspace(1, 2, 4, out= y) + >>> print(y) + ivy.array([1, 1, 1, 2]) + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2]) + >>> y = ivy.array([4,5]) + >>> z = ivy.linspace(x, y, 4, axis = 0) + >>> print(z) + ivy.array([[1, 2], + [2, 3], + [3, 4], + [4, 5]]) """ - return current_backend(x).from_dlpack(x, out=out) + return current_backend(start).linspace( + start, + stop, + num, + axis=axis, + endpoint=endpoint, + dtype=dtype, + device=device, + out=out, + ) +@handle_backend_invalid @handle_nestable -@outputs_to_ivy_arrays -def frombuffer( - buffer: bytes, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> ivy.Array: - r""" - Interpret a buffer as a 1-dimensional array. - - .. note:: - Note that either of the following must be true: - 1. count is a positive non-zero number, and the total number of bytes - in the buffer is equal or greater than offset plus count times the size - (in bytes) of dtype. - 2. count is negative, and the length (number of bytes) of the buffer - subtracted by the offset is a multiple of the size (in bytes) of dtype. +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def meshgrid( + *arrays: Union[ivy.Array, ivy.NativeArray], + sparse: bool = False, + indexing: str = "xy", + out: Optional[ivy.Array] = None, +) -> List[ivy.Array]: + """ + Return coordinate matrices from coordinate vectors. Parameters ---------- - buffer - An object that exposes the buffer interface. - dtype - Data-type of the returned array; default: float. - count - Number of items to read. -1 means all data in the buffer. - offset - Start reading the buffer from this offset (in bytes); default: 0. + arrays + an arbitrary number of one-dimensional arrays representing grid coordinates. + Each array should have the same numeric data type. + sparse + if True, a sparse grid is returned in order to conserve memory. + Default: ``False``. + indexing + Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or + one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, + respectively), the ``indexing`` keyword has no effect and should be ignored. + Default: ``'xy'``. Returns ------- - out - 1-dimensional array. + ret + list of N arrays, where ``N`` is the number of provided one-dimensional input + arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional + arrays having lengths ``Ni = len(xi)``, - Examples - -------- - With :class:`bytes` inputs: + - if matrix indexing ``ij``, then each returned array must have the shape + ``(N1, N2, N3, ..., Nn)``. + - if Cartesian indexing ``xy``, then each returned array must have shape + ``(N2, N1, N3, ..., Nn)``. + + Accordingly, for the two-dimensional case with input one-dimensional arrays of + length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must + have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned + array must have shape ``(N, M)``. + + Similarly, for the three-dimensional case with input one-dimensional arrays of + length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned + array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then + each returned array must have shape ``(N, M, P)``. + + Each returned array should have the same data type as the input arrays. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of + the `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) + + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) + + >>> x = ivy.array([1, 2, 5]) + >>> y = ivy.array([4, 1]) + >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') + >>> print(xv) + ivy.array([[1, 1], + [2, 2], + [5, 5]]) + + >>> print(yv) + ivy.array([[4, 1], + [4, 1], + [4, 1]]) + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> xv, yv = ivy.meshgrid(x, y, sparse=True) + >>> print(xv) + ivy.array([[1, 2, 3]]) - >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' - >>> y = ivy.frombuffer(x, dtype=ivy.float64) - >>> print(y) - ivy.array([1., 2.]) + >>> print(yv) + ivy.array([[4], [5], [6]]) - >>> x = b'\x01\x02\x03\x04' - >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) - >>> print(y) - ivy.array([2, 3, 4]) + With :class:`ivy.NativeArray` input: - >>> x = b'\x00<\x00@\x00B\x00D\x00E' - >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) - >>> print(y) - ivy.array([2., 3., 4., 5.]) + >>> x = ivy.native_array([1, 2]) + >>> y = ivy.native_array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) + + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) """ - return current_backend().frombuffer( - buffer, - dtype=dtype, - count=count, - offset=offset, + return current_backend().meshgrid( + *arrays, sparse=sparse, indexing=indexing, out=out ) @@ -1043,239 +1677,17 @@ def full( >>> fill_value = ivy.Container(a=0.99, b=False) >>> dtype = ivy.Container(a=ivy.float64, b=ivy.bool) >>> device = ivy.Container(a=ivy.NativeDevice('cpu'), b=ivy.Device('cpu')) - >>> x = ivy.full(shape, fill_value, dtype=dtype, device=device) - >>> print(x) - { - a: ivy.array([[0.99], - [0.99]]), - b: ivy.array([[[False, False]], - [[False, False]]]) - } - """ - return current_backend().full( - shape, fill_value, dtype=dtype, device=device, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def full_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - fill_value: Number, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ``fill_value`` and having the same ``shape`` as an - input array ``x`` . - - Parameters - ---------- - x - input array from which to derive the output array shape. - fill_value - Scalar fill value - dtype - output array data type. If ``dtype`` is `None`, the output array data type must - be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array having the same shape as ``x`` and where every element is equal to - ``fill_value``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :code:`int` datatype: - - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> fill_value = 1 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) - - >>> fill_value = 0.000123 - >>> x = ivy.ones(5) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) - - With float datatype: - - >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([3.0, 8.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x,fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123]) - - >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([[0.000123, 0.000123, 0.000123], - [0.000123, 0.000123, 0.000123]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), - ... b=ivy.array([4.123, 5.23, 6.23])) - >>> fill_value = 15.0 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - { - a: ivy.array([15., 15., 15.]), - b: ivy.array([15., 15., 15.]) - } - """ - return current_backend(x).full_like( - x, fill_value, dtype=dtype, device=device, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def linspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], - /, - num: int, - *, - axis: Optional[int] = None, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Generate a certain number of evenly-spaced values in an interval along a given axis. - - See :math:`arange` that allows to specify the step size of evenly spaced values in - an interval. - - Parameters - ---------- - start - First entry in the range. - stop - Final entry in the range. - num - Number of values to generate. - axis - Axis along which the operation is performed. - endpoint - If True, stop is the last sample. Otherwise, it is not included. - dtype - output array data type. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Tensor of evenly-spaced values. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With float input: - - >>> x = ivy.linspace(1, 2, 3) - >>> print(x) - ivy.array([1. , 1.5, 2. ]) - - >>> x = ivy.linspace(1, 2, 4, endpoint=False) - >>> print(x) - ivy.array([1., 1.25, 1.5 , 1.75]) - - >>> x = ivy.linspace(1, 10, 4, dtype="int32") - >>> print(x) - ivy.array([ 1, 4, 7, 10]) - - >>> x = ivy.linspace(1, 2, 4, device= "cpu") - >>> print(x) - ivy.array([1., 1.33333337, 1.66666663, 2.]) - - >>> y = ivy.array([0,0,0,0]) - >>> ivy.linspace(1, 2, 4, out= y) - >>> print(y) - ivy.array([1, 1, 1, 2]) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2]) - >>> y = ivy.array([4,5]) - >>> z = ivy.linspace(x, y, 4, axis = 0) - >>> print(z) - ivy.array([[1, 2], - [2, 3], - [3, 4], - [4, 5]]) + >>> x = ivy.full(shape, fill_value, dtype=dtype, device=device) + >>> print(x) + { + a: ivy.array([[0.99], + [0.99]]), + b: ivy.array([[[False, False]], + [[False, False]]]) + } """ - return current_backend(start).linspace( - start, - stop, - num, - axis=axis, - endpoint=endpoint, - dtype=dtype, - device=device, - out=out, + return current_backend().full( + shape, fill_value, dtype=dtype, device=device, out=out ) @@ -1285,237 +1697,163 @@ def linspace( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@infer_dtype -@infer_device -def logspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], - /, - num: int, - *, - base: float = 10.0, - axis: int = 0, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, +@handle_device_shifting +def from_dlpack( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None ) -> ivy.Array: """ - Generate a certain number of evenly-spaced values in log space, in an interval along - a given axis. + Return a new array containing the data from another (array) object with a + ``__dlpack__`` method. Parameters ---------- - start - First value in the range in log space. base ** start is the starting value in - the sequence. Can be an array or a float. - stop - Last value in the range in log space. base ** stop is the final value in the - sequence. Can be an array or a float. - num - Number of values to generate. - base - The base of the log space. Default is 10.0 - axis - Axis along which the operation is performed. Relevant only if start or stop are - array-like. Default is 0. - endpoint - If True, stop is the last sample. Otherwise, it is not included. Default is - True. - dtype - The data type of the output tensor. If None, the dtype of on_value is used or if - that is None, the dtype of off_value is used, or if that is None, defaults to - float32. Default is None. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is - None. + x object + input (array) object. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. Default is None. + inputs broadcast to. Returns ------- ret - Tensor of evenly-spaced values in log space. + an array containing the data in `x`. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + .. admonition:: Note + :class: note - Functional Examples - ------------------- - With float input: + The returned array may be either a copy or a view. See + :ref:`data-interchange` for details. - >>> print(ivy.logspace(1, 2, 4)) - ivy.array([ 10., 21.5443469, 46.41588834, 100.]) - >>> print(ivy.logspace(1, 2, 4, endpoint=False)) - ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> print(ivy.logspace(1, 2, 4, dtype= int)) - ivy.array([ 10., 10., 10., 100.]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + return current_backend(x).from_dlpack(x, out=out) - >>> out = ivy.array([0,0,0,0]) - >>> ivy.logspace(1, 2, 4, out = out) - >>> print(out) - ivy.array([ 10, 21, 46, 100]) - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[1.e+01, 1.e+02], - [1.e+02, 1.e+03], - [1.e+03, 1.e+04], - [1.e+04, 1.e+05]) +# Extra # +# ------# - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4, axis = 1)) - ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], - [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[ 10., 100.], - [ 100., 100.], - [ 1000., 1000.], - [10000., 10000.]]) - """ - result = base ** linspace( - start, - stop, - num, - endpoint=endpoint, - axis=axis, - dtype=dtype, - device=device, - ) - if ivy.exists(out): - return ivy.inplace_update(out, result) - return result +array = asarray @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function @handle_device_shifting -def meshgrid( - *arrays: Union[ivy.Array, ivy.NativeArray], - sparse: bool = False, - indexing: str = "xy", +def copy_array( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + to_ivy_array: bool = True, out: Optional[ivy.Array] = None, -) -> List[ivy.Array]: +) -> ivy.Array: """ - Return coordinate matrices from coordinate vectors. + Copy an array. Parameters ---------- - arrays - an arbitrary number of one-dimensional arrays representing grid coordinates. - Each array should have the same numeric data type. - sparse - if True, a sparse grid is returned in order to conserve memory. - Default: ``False``. - indexing - Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or - one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, - respectively), the ``indexing`` keyword has no effect and should be ignored. - Default: ``'xy'``. + x + array, input array containing elements to copy. + to_ivy_array + boolean, if True the returned array will be an ivy.Array object otherwise + returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., + depending on the backend), defaults to True. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - list of N arrays, where ``N`` is the number of provided one-dimensional input - arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional - arrays having lengths ``Ni = len(xi)``, - - - if matrix indexing ``ij``, then each returned array must have the shape - ``(N1, N2, N3, ..., Nn)``. - - if Cartesian indexing ``xy``, then each returned array must have shape - ``(N2, N1, N3, ..., Nn)``. - - Accordingly, for the two-dimensional case with input one-dimensional arrays of - length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must - have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned - array must have shape ``(N, M)``. - - Similarly, for the three-dimensional case with input one-dimensional arrays of - length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned - array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then - each returned array must have shape ``(N, M, P)``. - - Each returned array should have the same data type as the input arrays. + a copy of the input array ``x``. + Examples + -------- + With one :class:`ivy.Array` input: - This function conforms to the `Array API Standard - `_. This docstring is an extension of - the `docstring `_ - in the standard. + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.copy_array(x) + >>> print(y) + ivy.array([-1, 0, 1]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = ivy.copy_array(x) + >>> print(y) + ivy.array([1, 0, 1, 1]) - Functional Examples - ------------------- + >>> x = ivy.array([1, 0, 1, -1]) + >>> y = ivy.zeros((1, 4)) + >>> ivy.copy_array(x, out=y) + >>> print(y) + ivy.array([1, 0, 1, -1]) - With :class:`ivy.Array` input: + >>> x = ivy.array([1, 0, 1, 1]) + >>> ivy.copy_array(x, out=x) + >>> print(x) + ivy.array([1, 0, 1, 1]) - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) + With one :class:`ivy.Container` input: - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.copy_array(x) + >>> print(y) + { + a: ivy.array([-1, 0, 1]) + } - >>> x = ivy.array([1, 2, 5]) - >>> y = ivy.array([4, 1]) - >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') - >>> print(xv) - ivy.array([[1, 1], - [2, 2], - [5, 5]]) + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.copy_array(x) + >>> print(y) + { + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) + } - >>> print(yv) - ivy.array([[4, 1], - [4, 1], - [4, 1]]) + With one :class:`ivy.Container` static method: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> xv, yv = ivy.meshgrid(x, y, sparse=True) - >>> print(xv) - ivy.array([[1, 2, 3]]) + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.Container.static_copy_array(x) + >>> print(y) + { + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) + } - >>> print(yv) - ivy.array([[4], [5], [6]]) + With one :class:`ivy.Array` instance method: - With :class:`ivy.NativeArray` input: + >>> x = ivy.array([-1, 0, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([-1, 0, 1]) - >>> x = ivy.native_array([1, 2]) - >>> y = ivy.native_array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([1, 0, 1, 1]) - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) + With :class:`ivy.Container` instance method: + + >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) + >>> y = x.copy_array() + >>> print(y) + { + a: ivy.array([1, 0, 1]), + b: ivy.array([-1, 0, 1, 1]) + } """ - return current_backend().meshgrid( - *arrays, sparse=sparse, indexing=indexing, out=out - ) + return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) @handle_backend_invalid @@ -1615,215 +1953,7 @@ def one_hot( The data type of the output tensor. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if - None. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Tensor of zeros with the same shape and type as a, unless dtype provided which - overrides. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([3, 1]) - >>> y = 5 - >>> z = x.one_hot(5) - >>> print(z) - ivy.array([[0., 0., 0., 1., 0.], - ... [0., 1., 0., 0., 0.]]) - - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, y) - ivy.array([[1., 0., 0., 0., 0.]]) - - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, 5, out=z) - ivy.array([[1., 0., 0., 0., 0.]]) - >>> print(z) - ivy.array([[1., 0., 0., 0., 0.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1, 2]), \ - b=ivy.array([3, 1]), c=ivy.array([2, 3])) - >>> y = 5 - >>> z = x.one_hot(y) - >>> print(z) - { - a: ivy.array([[0., 1., 0., 0., 0.], - [0., 0., 1., 0., 0.]]), - b: ivy.array([[0., 0., 0., 1., 0.], - [0., 1., 0., 0., 0.]]), - c: ivy.array([[0., 0., 1., 0., 0.], - [0., 0., 0., 1., 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([2]), \ - b=ivy.array([]), c=ivy.native_array([4])) - >>> y = 7 - >>> z = x.one_hot(y) - >>> print(z) - { - a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), - b: ivy.array([], shape=(0, 7)), - c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) - } - """ - return current_backend(indices).one_hot( - indices, - depth, - on_value=on_value, - off_value=off_value, - axis=axis, - dtype=dtype, - device=device, - out=out, - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with ones. - - .. note:: - - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal to - zero (i.e., ``1 + 0j``). - - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing ones. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Shape` input: - - >>> shape = (2,2) - >>> x = ivy.ones(shape) - >>> print(x) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Dtype` input: - - >>> shape = (3,2) - >>> d_type = ivy.int64 - >>> y = ivy.ones(shape, dtype=d_type) - >>> print(y) - ivy.array([[1, 1], - [1, 1], - [1, 1]]) - - With :class:`ivy.Device` input: - - >>> shape = (3,2) - >>> y = ivy.ones(shape, device="cpu") - >>> print(y) - ivy.array([[1., 1.], - [1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input: - - >>> shape = (1, 5, 2) - >>> x = ivy.zeros(shape) - >>> ivy.ones(shape, out=x) - >>> print(x) - ivy.array([[[1., 1.], - [1., 1.], - [1., 1.], - [1., 1.], - [1., 1.]]]) - """ - return current_backend().ones(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ones and having the same shape as an input array - ``x``. - - .. note:: - - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal - to zero (i.e., ``1 + 0j``). - - Parameters - ---------- - x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default ``None``. - device - device on which to place the created array. If device is ``None``, the output - array device must be inferred from ``x``. Default: ``None``. + None. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1831,82 +1961,69 @@ def ones_like( Returns ------- ret - an array having the same shape as ``x`` and filled with ``ones``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) - - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1., 1., 1.], - [1., 1., 1.]]) - - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.zeros(3) - >>> ivy.ones_like(x, out=y) - >>> print(y) - ivy.array([1., 1., 1.]) + Tensor of zeros with the same shape and type as a, unless dtype provided which + overrides. + + Examples + -------- + With :class:`ivy.Array` inputs: - With :class:`ivy.NativeArray` input: + >>> x = ivy.array([3, 1]) + >>> y = 5 + >>> z = x.one_hot(5) + >>> print(z) + ivy.array([[0., 0., 0., 1., 0.], + ... [0., 1., 0., 0., 0.]]) - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1, 1, 1], - [1, 1, 1]]) + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, y) + ivy.array([[1., 0., 0., 0., 0.]]) - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, 5, out=z) + ivy.array([[1., 0., 0., 0., 0.]]) + >>> print(z) + ivy.array([[1., 0., 0., 0., 0.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.ones_like(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([1, 2]), \ + b=ivy.array([3, 1]), c=ivy.array([2, 3])) + >>> y = 5 + >>> z = x.one_hot(y) + >>> print(z) { - a: ivy.array([1, 1, 1]), - b: ivy.array([1, 1, 1]) + a: ivy.array([[0., 1., 0., 0., 0.], + [0., 0., 1., 0., 0.]]), + b: ivy.array([[0., 0., 0., 1., 0.], + [0., 1., 0., 0., 0.]]), + c: ivy.array([[0., 0., 1., 0., 0.], + [0., 0., 0., 1., 0.]]) } - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.ones_like() - >>> print(y) - ivy.array([1, 1, 1, 1, 1]) - - With :class:'ivy.Container' input: - - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.ones_like() - >>> print(y) + >>> x = ivy.Container(a=ivy.array([2]), \ + b=ivy.array([]), c=ivy.native_array([4])) + >>> y = 7 + >>> z = x.one_hot(y) + >>> print(z) { - a: ivy.array([1., 1.]), - b: ivy.array([1., 1.]) + a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), + b: ivy.array([], shape=(0, 7)), + c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) } """ - return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) + return current_backend(indices).one_hot( + indices, + depth, + on_value=on_value, + off_value=off_value, + axis=axis, + dtype=dtype, + device=device, + out=out, + ) @handle_backend_invalid @@ -1915,110 +2032,178 @@ def ones_like( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def tril( - x: Union[ivy.Array, ivy.NativeArray], +@infer_dtype +@infer_device +def logspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], /, + num: int, *, - k: int = 0, + base: float = 10.0, + axis: int = 0, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the lower triangular part of a matrix (or a stack of matrices) ``x``. - - .. note:: - - The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` - on the interval ``[0, min(M, N) - 1]``. + Generate a certain number of evenly-spaced values in log space, in an interval along + a given axis. Parameters ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. - k - diagonal above which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. + start + First value in the range in log space. base ** start is the starting value in + the sequence. Can be an array or a float. + stop + Last value in the range in log space. base ** stop is the final value in the + sequence. Can be an array or a float. + num + Number of values to generate. + base + The base of the log space. Default is 10.0 + axis + Axis along which the operation is performed. Relevant only if start or stop are + array-like. Default is 0. + endpoint + If True, stop is the last sample. Otherwise, it is not included. Default is + True. + dtype + The data type of the output tensor. If None, the dtype of on_value is used or if + that is None, the dtype of off_value is used, or if that is None, defaults to + float32. Default is None. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is + None. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + inputs broadcast to. Default is None. Returns ------- ret - an array containing the lower triangular part(s). The returned array must have - the same shape and data type as x. All elements above the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + Tensor of evenly-spaced values in log space. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. + + Functional Examples + ------------------- + With float input: + + >>> print(ivy.logspace(1, 2, 4)) + ivy.array([ 10., 21.5443469, 46.41588834, 100.]) + + >>> print(ivy.logspace(1, 2, 4, endpoint=False)) + ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) + + >>> print(ivy.logspace(1, 2, 4, dtype= int)) + ivy.array([ 10., 10., 10., 100.]) + + >>> out = ivy.array([0,0,0,0]) + >>> ivy.logspace(1, 2, 4, out = out) + >>> print(out) + ivy.array([ 10, 21, 46, 100]) + + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[1.e+01, 1.e+02], + [1.e+02, 1.e+03], + [1.e+03, 1.e+04], + [1.e+04, 1.e+05]) + + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4, axis = 1)) + ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], + [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) + + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[ 10., 100.], + [ 100., 100.], + [ 1000., 1000.], + [10000., 10000.]]) """ - return current_backend(x).tril(x, k=k, out=out) + result = base ** linspace( + start, + stop, + num, + endpoint=endpoint, + axis=axis, + dtype=dtype, + device=device, + ) + if ivy.exists(out): + return ivy.inplace_update(out, result) + return result -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def triu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, +@outputs_to_ivy_arrays +def frombuffer( + buffer: bytes, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + count: Optional[int] = -1, + offset: Optional[int] = 0, ) -> ivy.Array: - """ - Return the upper triangular part of a matrix (or a stack of matrices) ``x``. + r""" + Interpret a buffer as a 1-dimensional array. .. note:: - - The upper triangular part of the matrix is defined as the elements - on and above the specified diagonal ``k``. + Note that either of the following must be true: + 1. count is a positive non-zero number, and the total number of bytes + in the buffer is equal or greater than offset plus count times the size + (in bytes) of dtype. + 2. count is negative, and the length (number of bytes) of the buffer + subtracted by the offset is a multiple of the size (in bytes) of dtype. Parameters ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. *, - k - diagonal below which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + buffer + An object that exposes the buffer interface. + dtype + Data-type of the returned array; default: float. + count + Number of items to read. -1 means all data in the buffer. + offset + Start reading the buffer from this offset (in bytes); default: 0. Returns ------- - ret - an array containing the upper triangular part(s). The returned array must have - the same shape and data type as x. All elements below the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. + out + 1-dimensional array. + Examples + -------- + With :class:`bytes` inputs: - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' + >>> y = ivy.frombuffer(x, dtype=ivy.float64) + >>> print(y) + ivy.array([1., 2.]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = b'\x01\x02\x03\x04' + >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) + >>> print(y) + ivy.array([2, 3, 4]) + + >>> x = b'\x00<\x00@\x00B\x00D\x00E' + >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) + >>> print(y) + ivy.array([2., 3., 4., 5.]) """ - return current_backend(x).triu(x, k=k, out=out) + return current_backend().frombuffer( + buffer, + dtype=dtype, + count=count, + offset=offset, + ) @handle_exceptions @@ -2115,193 +2300,3 @@ def triu_indices( (ivy.array([0, 0, 0, 0, 1, 1, 1, 1]), ivy.array([0, 1, 2, 3, 0, 1, 2, 3])) """ return current_backend().triu_indices(n_rows, n_cols, k, device=device) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def zeros( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with zeros. - - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type must - be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing zeros. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.NativeShape` input: - >>> shape = (3, 5) - >>> x = ivy.zeros(shape) - >>> print(x) - ivy.array([[0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.]]) - - >>> x = ivy.zeros(5) - >>> print(x) - ivy.array([0., 0., 0., 0., 0.]) - """ - return current_backend().zeros(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def zeros_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with zeros and having the same ``shape`` as an input array - ``x``. - - Parameters - ---------- - x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array having the same shape as ``x`` and filled with ``zeros``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.zeros_like(x) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) - - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.zeros_like(x) - >>> print(y) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) - - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.ones(3) - >>> ivy.zeros_like(x, out=y) - >>> print(y) - ivy.array([0., 0., 0.]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.zeros_like(x) - >>> print(y) - ivy.array([[0, 0, 0],[0, 0, 0]]) - - - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.zeros_like(x) - >>> print(y) - { - a: ivy.array([0, 0, 0]), - b: ivy.array([0, 0, 0]) - } - - - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.zeros_like() - >>> print(y) - ivy.array([0, 0, 0, 0, 0]) - - With :class:'ivy.Container' input: - - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.zeros_like() - >>> print(y) - { - a: ivy.array([0., 0.]), - b: ivy.array([0., 0.]) - } - """ - return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) - - -# Extra # -# ------# - - -array = asarray diff --git a/ivy/functional/ivy/data_type.py b/ivy/functional/ivy/data_type.py index b5114cfda5ca2..4a084d541e98c 100644 --- a/ivy/functional/ivy/data_type.py +++ b/ivy/functional/ivy/data_type.py @@ -26,188 +26,52 @@ from ivy.utils.exceptions import handle_exceptions -# Array API Standard # -# -------------------# - -Finfo = None -Iinfo = None -default_complex_dtype_stack = list() -# Extra # -# ------# - -default_dtype_stack = list() -default_float_dtype_stack = list() -default_int_dtype_stack = list() -default_uint_dtype_stack = list() - - -class DefaultDtype: - """Ivy's DefaultDtype class.""" - - def __init__(self, dtype: ivy.Dtype): - self._dtype = dtype - - def __enter__(self): - set_default_dtype(self._dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultFloatDtype: - """Ivy's DefaultFloatDtype class.""" - - def __init__(self, float_dtype: ivy.Dtype): - self._float_dtype = float_dtype - - def __enter__(self): - set_default_float_dtype(self._float_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_float_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultIntDtype: - """Ivy's DefaultIntDtype class.""" - - def __init__(self, int_dtype: ivy.Dtype): - self._int_dtype = int_dtype - - def __enter__(self): - set_default_int_dtype(self._int_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_int_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultUintDtype: - """Ivy's DefaultUintDtype class.""" - - def __init__(self, uint_dtype: ivy.UintDtype): - self._uint_dtype = uint_dtype - - def __enter__(self): - set_default_uint_dtype(self._uint_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_uint_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultComplexDtype: - """Ivy's DefaultComplexDtype class.""" - - def __init__(self, complex_dtype: ivy.Dtype): - self._complex_dtype = complex_dtype - - def __enter__(self): - set_default_complex_dtype(self._complex_dtype) - return self +# Helpers # +# --------# - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_complex_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -# --- Helpers --- # -# --------------- # - - -def _check_complex128(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "complex128" - elif isinstance(input, np.ndarray): - return str(input.dtype) == "complex128" - if hasattr(input, "real") and hasattr(input, "imag"): - return _check_float64(input.real) and _check_float64(input.imag) - return False - - -def _check_float64(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "float64" - if math.isfinite(input): - m, e = math.frexp(input) - return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) - return False - - -# Get the list of dtypes supported by the function -# by default returns the supported dtypes -def _get_dtypes(fn, complement=True): - supported = set(ivy.valid_dtypes) - # We only care about getting dtype info from the base function - # if we do need to at some point use dtype information from the parent function - # we can comment out the following condition - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") - if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: - if complement: - supported = set(ivy.all_dtypes).difference(supported) - return supported +def _is_valid_dtypes_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): + fn_supported_dtypes = fn.supported_dtypes + fn_unsupported_dtypes = fn.unsupported_dtypes + if isinstance(fn_supported_dtypes, dict): + if isinstance(fn_unsupported_dtypes, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_dtypes + and backend_str in fn_unsupported_dtypes + ): + return False + else: + if isinstance(fn_unsupported_dtypes, tuple): + return False + return True - # Their values are formatted like either - # 1. fn.supported_dtypes = ("float16",) - # Could also have the "all" value for the framework - basic = [ - ("supported_dtypes", set.intersection, ivy.valid_dtypes), - ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), - ] - # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. - typesets = { - "valid": ivy.valid_dtypes, - "numeric": ivy.valid_numeric_dtypes, - "float": ivy.valid_float_dtypes, - "integer": ivy.valid_int_dtypes, - "unsigned": ivy.valid_uint_dtypes, - "complex": ivy.valid_complex_dtypes, - } +def _handle_nestable_dtype_info(fn): + def _handle_nestable_dtype_info_wrapper(type): + if isinstance(type, ivy.Container): + type = type.cont_map(lambda x, kc: fn(x)) + type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) + type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) + return type + return fn(type) - for key, merge_fn, base in basic: - if hasattr(fn, key): - dtypes = getattr(fn, key) - # only einops allowed to be a dictionary - if isinstance(dtypes, dict): - dtypes = dtypes.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(dtypes, tuple) - dtypes = list(dtypes) - typeset_list = [] - for i, dtype in reversed(list(enumerate(dtypes))): - if dtype in typesets: - typeset_list.extend(typesets[dtype]) - dtypes.pop(i) - dtypes = dtypes + typeset_list - supported = merge_fn(supported, set(dtypes)) + return _handle_nestable_dtype_info_wrapper - if complement: - supported = set(ivy.all_dtypes).difference(supported) - return tuple(supported) +# Unindent every line in the source such that +# class methods can be compiled as normal methods +def _lstrip_lines(source: str) -> str: + # Separate all lines + source = source.split("\n") + # Check amount of indent before first character + indent = len(source[0]) - len(source[0].lstrip()) + # Remove same spaces from all lines + for i in range(len(source)): + source[i] = source[i][indent:] + source = "\n".join(source) + return source # Get the list of function used the function @@ -262,50 +126,6 @@ def _get_functions_from_string(func_names, module): return ret -def _handle_nestable_dtype_info(fn): - def _handle_nestable_dtype_info_wrapper(type): - if isinstance(type, ivy.Container): - type = type.cont_map(lambda x, kc: fn(x)) - type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) - type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) - return type - return fn(type) - - return _handle_nestable_dtype_info_wrapper - - -def _is_valid_dtypes_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): - fn_supported_dtypes = fn.supported_dtypes - fn_unsupported_dtypes = fn.unsupported_dtypes - if isinstance(fn_supported_dtypes, dict): - if isinstance(fn_unsupported_dtypes, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_dtypes - and backend_str in fn_unsupported_dtypes - ): - return False - else: - if isinstance(fn_unsupported_dtypes, tuple): - return False - return True - - -# Unindent every line in the source such that -# class methods can be compiled as normal methods -def _lstrip_lines(source: str) -> str: - # Separate all lines - source = source.split("\n") - # Check amount of indent before first character - indent = len(source[0]) - len(source[0].lstrip()) - # Remove same spaces from all lines - for i in range(len(source)): - source[i] = source[i][indent:] - source = "\n".join(source) - return source - - # Get dtypes/device of nested functions, used for unsupported and supported dtypes # IMPORTANT: a few caveats: # 1. The base functions must be defined in ivy or the same module @@ -361,44 +181,67 @@ def _nested_get(f, base_set, merge_fn, get_fn, wrapper=set): return out -# --- Main --- # -# ------------ # - - -@handle_exceptions -def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: - """ - Convert native data type to string representation. +# Get the list of dtypes supported by the function +# by default returns the supported dtypes +def _get_dtypes(fn, complement=True): + supported = set(ivy.valid_dtypes) - Parameters - ---------- - dtype_in - The data type to convert to string. + # We only care about getting dtype info from the base function + # if we do need to at some point use dtype information from the parent function + # we can comment out the following condition + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") + if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: + if complement: + supported = set(ivy.all_dtypes).difference(supported) + return supported - Returns - ------- - ret - data type string 'float32' - """ - return current_backend(None).as_ivy_dtype(dtype_in) + # Their values are formatted like either + # 1. fn.supported_dtypes = ("float16",) + # Could also have the "all" value for the framework + basic = [ + ("supported_dtypes", set.intersection, ivy.valid_dtypes), + ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), + ] + # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. + typesets = { + "valid": ivy.valid_dtypes, + "numeric": ivy.valid_numeric_dtypes, + "float": ivy.valid_float_dtypes, + "integer": ivy.valid_int_dtypes, + "unsigned": ivy.valid_uint_dtypes, + "complex": ivy.valid_complex_dtypes, + } -@handle_exceptions -def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: - """ - Convert data type string representation to native data type. + for key, merge_fn, base in basic: + if hasattr(fn, key): + dtypes = getattr(fn, key) + # only einops allowed to be a dictionary + if isinstance(dtypes, dict): + dtypes = dtypes.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(dtypes, tuple) + dtypes = list(dtypes) + typeset_list = [] + for i, dtype in reversed(list(enumerate(dtypes))): + if dtype in typesets: + typeset_list.extend(typesets[dtype]) + dtypes.pop(i) + dtypes = dtypes + typeset_list + supported = merge_fn(supported, set(dtypes)) - Parameters - ---------- - dtype_in - The data type string to convert to native data type. + if complement: + supported = set(ivy.all_dtypes).difference(supported) - Returns - ------- - ret - data type e.g. ivy.float32. - """ - return current_backend(None).as_native_dtype(dtype_in) + return tuple(supported) + + +# Array API Standard # +# -------------------# + +Finfo = None +Iinfo = None @handle_exceptions @@ -745,251 +588,488 @@ def can_cast( @handle_exceptions -def check_float(x: Any) -> bool: +@handle_backend_invalid +@inputs_to_native_arrays +@handle_device_shifting +def finfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Finfo: """ - Check if the input is a float or a float-like object. + Machine limits for floating-point data types. Parameters ---------- - x - Input to check. + type + the kind of floating-point data-type about which to get information. Returns ------- ret - "True" if the input is a float or a float-like object, otherwise "False". - """ - return isinstance(x, (int, float)) and x is not bool + an object having the followng attributes: + - **bits**: *int* -@handle_exceptions -def closest_valid_dtype(type: Union[ivy.Dtype, str, None], /) -> Union[ivy.Dtype, str]: - """ - Determine the closest valid datatype to the datatype passed as input. + number of bits occupied by the floating-point data type. - Parameters - ---------- - type - The data type for which to check the closest valid type for. + - **eps**: *float* - Returns - ------- - ret - The closest valid data type as a native ivy.Dtype + difference between 1.0 and the next smallest representable floating-point + number larger than 1.0 according to the IEEE-754 standard. + + - **max**: *float* + + largest representable number. + + - **min**: *float* + + smallest representable number. + + - **smallest_normal**: *float* + + smallest positive floating-point number with full precision. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Examples -------- With :class:`ivy.Dtype` input: - >>> xType = ivy.float16 - >>> yType = ivy.closest_valid_dtype(xType) - >>> print(yType) - float16 + >>> y = ivy.finfo(ivy.float32) + >>> print(y) + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - With :class:`ivy.NativeDtype` inputs: + With :code:`str` input: - >>> xType = ivy.native_uint16 - >>> yType = ivy.closest_valid_dtype(xType) - >>> print(yType) - uint16 + >>> y = ivy.finfo('float32') + >>> print(y) + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - With :code:`str` input: + With :class:`ivy.Array` input: - >>> xType = 'int32' - >>> yType = ivy.closest_valid_dtype(xType) - >>> print(yType) - int32 + >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) + >>> print(ivy.finfo(x)) + finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) + + >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) + >>> print(ivy.finfo(x)) + finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) + + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), + ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) + >>> print(ivy.finfo(c)) + { + x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), + y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) + } """ - return current_backend(type).closest_valid_dtype(type) + return current_backend(None).finfo(type) @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays +@handle_backend_invalid +@inputs_to_native_arrays @handle_device_shifting -def default_complex_dtype( - *, - input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, - as_native: bool = False, -) -> Union[ivy.Dtype, str, ivy.NativeDtype]: +def iinfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Iinfo: """ + Machine limits for integer data types. + Parameters ---------- - input - Number or array for inferring the complex dtype. - complex_dtype - The complex dtype to be returned. - as_native - Whether to return the complex dtype as native dtype. + type + the kind of integer data-type about which to get information. Returns ------- - Return ``complex_dtype`` as native or ivy dtype if provided, else - if ``input`` is given, return its complex dtype, otherwise return the - global default complex dtype. + ret + a class with that encapsules the following attributes: + + - **bits**: *int* + + number of bits occupied by the type. + + - **max**: *int* + + largest representable number. + + - **min**: *int* + + smallest representable number. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Examples -------- - >>> ivy.default_complex_dtype() - 'complex64' + With :class:`ivy.Dtype` input: - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' + >>> ivy.iinfo(ivy.int32) + iinfo(min=-2147483648, max=2147483647, dtype=int32) - >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) - 'complex128' + With :code:`str` input: - >>> ivy.default_complex_dtype(input=4294.967346) - 'complex64' + >>> ivy.iinfo('int32') + iinfo(min=-2147483648, max=2147483647, dtype=int32) - >>> x = ivy.array([9.8,8.9], dtype="complex128") - >>> ivy.default_complex_dtype(input=x) - 'complex128' - """ - global default_complex_dtype_stack - if ivy.exists(complex_dtype): - if as_native is True: - return ivy.as_native_dtype(complex_dtype) - return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - as_native = ivy.default(as_native, False) - if ivy.exists(input): - if ivy.is_array(input): - ret = ivy.dtype(input) - elif isinstance(input, np.ndarray): - ret = str(input.dtype) - elif isinstance(input, (list, tuple, dict)): - if ivy.nested_argwhere( - input, lambda x: _check_complex128(x), stop_after_n_found=1 - ): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - elif isinstance(input, Number): - if _check_complex128(input): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - if as_native: - return ivy.as_native_dtype(ret) - return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) + With :class:`ivy.Array` input: + + >>> x = ivy.array([13,21,34], dtype=ivy.int8) + >>> ivy.iinfo(x) + iinfo(min=-128, max=127, dtype=int8) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) + >>> ivy.iinfo(x) + iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) + + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), + ... y=ivy.array([76,81,16], dtype=ivy.uint32)) + >>> ivy.iinfo(c) + { + x: iinfo(min=0, max=65535, dtype=uint16), + y: iinfo(min=0, max=4294967295, dtype=uint32) + } + """ + return current_backend(None).iinfo(type) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def result_type( + *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] +) -> ivy.Dtype: + """ + Return the dtype that results from applying the type promotion rules (see + :ref:`type-promotion`) to the arguments. + + .. note:: + If provided mixed dtypes (e.g., integer and floating-point), the returned dtype + will be implementation-specific. + + Parameters + ---------- + arrays_and_dtypes + an arbitrary number of input arrays and/or dtypes. + + Returns + ------- + ret + the dtype resulting from an operation involving the input arrays and dtypes. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.array([3., 4., 5.]) + >>> d = ivy.result_type(x, y) + >>> print(d) + float32 + + With :class:`ivy.Dtype` input: + + >>> d = ivy.result_type(ivy.uint8, ivy.uint64) + >>> print(d) + uint64 + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = x.a.dtype + >>> print(d) + int32 + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = ivy.result_type(x, ivy.float64) + >>> print(d) + { + a: float64 + } + """ + return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) + + +# Extra # +# ------# + +default_dtype_stack = list() +default_float_dtype_stack = list() +default_int_dtype_stack = list() +default_uint_dtype_stack = list() +default_complex_dtype_stack = list() + + +class DefaultDtype: + """Ivy's DefaultDtype class.""" + + def __init__(self, dtype: ivy.Dtype): + self._dtype = dtype + + def __enter__(self): + set_default_dtype(self._dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultFloatDtype: + """Ivy's DefaultFloatDtype class.""" + + def __init__(self, float_dtype: ivy.Dtype): + self._float_dtype = float_dtype + + def __enter__(self): + set_default_float_dtype(self._float_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_float_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultIntDtype: + """Ivy's DefaultIntDtype class.""" + + def __init__(self, int_dtype: ivy.Dtype): + self._int_dtype = int_dtype + + def __enter__(self): + set_default_int_dtype(self._int_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_int_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultUintDtype: + """Ivy's DefaultUintDtype class.""" + + def __init__(self, uint_dtype: ivy.UintDtype): + self._uint_dtype = uint_dtype + + def __enter__(self): + set_default_uint_dtype(self._uint_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_uint_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultComplexDtype: + """Ivy's DefaultComplexDtype class.""" + + def __init__(self, complex_dtype: ivy.Dtype): + self._complex_dtype = complex_dtype + + def __enter__(self): + set_default_complex_dtype(self._complex_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_complex_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +@handle_exceptions +def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: + """ + Get the number of bits used for representing the input data type. + + Parameters + ---------- + dtype_in + The data type to determine the number of bits for. + + Returns + ------- + ret + The number of bits used to represent the data type. + + Examples + -------- + With :class:`ivy.Dtype` inputs: + + >>> x = ivy.dtype_bits(ivy.float32) + >>> print(x) + 32 + + >>> x = ivy.dtype_bits('int64') + >>> print(x) + 64 + + With :class:`ivy.NativeDtype` inputs: + + >>> x = ivy.dtype_bits(ivy.native_bool) + >>> print(x) + 1 + """ + return current_backend(dtype_in).dtype_bits(dtype_in) + + +@handle_exceptions +def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: + """ + Check if the given data type is hashable or not. + + Parameters + ---------- + dtype_in + The data type to check. + + Returns + ------- + ret + True if data type is hashable else False + """ + try: + hash(dtype_in) + return True + except TypeError: + return False @handle_exceptions -@inputs_to_ivy_arrays -def default_dtype( - *, - dtype: Optional[Union[ivy.Dtype, str]] = None, - item: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - as_native: bool = False, -) -> Union[ivy.Dtype, ivy.NativeDtype, str]: +def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: """ + Convert native data type to string representation. + Parameters ---------- - item - Number or array for inferring the dtype. - dtype - The dtype to be returned. - as_native - Whether to return the dtype as native dtype. + dtype_in + The data type to convert to string. Returns ------- - Return ``dtype`` as native or ivy dtype if provided, else - if ``item`` is given, return its dtype, otherwise return the - global default dtype. + ret + data type string 'float32' + """ + return current_backend(None).as_ivy_dtype(dtype_in) - Examples - -------- - >>> ivy.default_dtype() - 'float32' - >>> ivy.set_default_dtype(ivy.bool) - >>> ivy.default_dtype() - 'bool' +@handle_exceptions +def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: + """ + Convert data type string representation to native data type. - >>> ivy.set_default_dtype(ivy.int16) - >>> ivy.default_dtype() - 'int16' + Parameters + ---------- + dtype_in + The data type string to convert to native data type. - >>> ivy.set_default_dtype(ivy.float64) - >>> ivy.default_dtype() - 'float64' + Returns + ------- + ret + data type e.g. ivy.float32. + """ + return current_backend(None).as_native_dtype(dtype_in) - >>> ivy.default_dtype(dtype="int32") - 'int32' - >>> ivy.default_dtype(dtype=ivy.float16) - 'float16' +def _check_float64(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "float64" + if math.isfinite(input): + m, e = math.frexp(input) + return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) + return False - >>> ivy.default_dtype(item=53.234) - 'float64' - >>> ivy.default_dtype(item=[1, 2, 3]) - 'int32' +def _check_complex128(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "complex128" + elif isinstance(input, np.ndarray): + return str(input.dtype) == "complex128" + if hasattr(input, "real") and hasattr(input, "imag"): + return _check_float64(input.real) and _check_float64(input.imag) + return False - >>> x = ivy.array([5.2, 9.7], dtype="complex128") - >>> ivy.default_dtype(item=x) - 'complex128' + +@handle_exceptions +def closest_valid_dtype(type: Union[ivy.Dtype, str, None], /) -> Union[ivy.Dtype, str]: """ - if ivy.exists(dtype): - if as_native is True: - return ivy.as_native_dtype(dtype) - return ivy.as_ivy_dtype(dtype) - as_native = ivy.default(as_native, False) - if ivy.exists(item): - if hasattr(item, "override_dtype_check"): - return item.override_dtype_check() - elif isinstance(item, (list, tuple, dict)) and len(item) == 0: - pass - elif ivy.is_complex_dtype(item): - return ivy.default_complex_dtype(input=item, as_native=as_native) - elif ivy.is_float_dtype(item): - return ivy.default_float_dtype(input=item, as_native=as_native) - elif ivy.is_uint_dtype(item): - return ivy.default_int_dtype(input=item, as_native=as_native) - elif ivy.is_int_dtype(item): - return ivy.default_int_dtype(input=item, as_native=as_native) - elif as_native: - return ivy.as_native_dtype("bool") - else: - return "bool" - global default_dtype_stack - if not default_dtype_stack: - global default_float_dtype_stack - if default_float_dtype_stack: - ret = default_float_dtype_stack[-1] - else: - ret = "float32" - else: - ret = default_dtype_stack[-1] - if as_native: - return ivy.as_native_dtype(ret) - return ivy.as_ivy_dtype(ret) + Determine the closest valid datatype to the datatype passed as input. + + Parameters + ---------- + type + The data type for which to check the closest valid type for. + + Returns + ------- + ret + The closest valid data type as a native ivy.Dtype + + Examples + -------- + With :class:`ivy.Dtype` input: + + >>> xType = ivy.float16 + >>> yType = ivy.closest_valid_dtype(xType) + >>> print(yType) + float16 + + With :class:`ivy.NativeDtype` inputs: + + >>> xType = ivy.native_uint16 + >>> yType = ivy.closest_valid_dtype(xType) + >>> print(yType) + uint16 + + With :code:`str` input: + + >>> xType = 'int32' + >>> yType = ivy.closest_valid_dtype(xType) + >>> print(yType) + int32 + """ + return current_backend(type).closest_valid_dtype(type) @handle_exceptions @@ -1084,7 +1164,148 @@ def default_float_dtype( ret = default_float_dtype_stack[-1] if as_native: return ivy.as_native_dtype(ret) - return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) + return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) + + +@handle_exceptions +def infer_default_dtype( + dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: + """ + Summary. + + Parameters + ---------- + dtype + + as_native + (Default value = False) + + Returns + ------- + Return the default data type for the “kind” (integer or floating-point) of dtype + + Examples + -------- + >>> ivy.set_default_int_dtype("int32") + >>> ivy.infer_default_dtype("int8") + 'int8' + + >>> ivy.set_default_float_dtype("float64") + >>> ivy.infer_default_dtype("float32") + 'float64' + + >>> ivy.set_default_uint_dtype("uint32") + >>> x = ivy.array([0], dtype="uint64") + >>> ivy.infer_default_dtype(x.dtype) + 'uint32' + """ + if ivy.is_complex_dtype(dtype): + default_dtype = ivy.default_complex_dtype(as_native=as_native) + elif ivy.is_float_dtype(dtype): + default_dtype = ivy.default_float_dtype(as_native=as_native) + elif ivy.is_uint_dtype(dtype): + default_dtype = ivy.default_uint_dtype(as_native=as_native) + elif ivy.is_int_dtype(dtype): + default_dtype = ivy.default_int_dtype(as_native=as_native) + elif as_native: + default_dtype = ivy.as_native_dtype("bool") + else: + default_dtype = ivy.as_ivy_dtype("bool") + return default_dtype + + +@handle_exceptions +@inputs_to_ivy_arrays +def default_dtype( + *, + dtype: Optional[Union[ivy.Dtype, str]] = None, + item: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + as_native: bool = False, +) -> Union[ivy.Dtype, ivy.NativeDtype, str]: + """ + Parameters + ---------- + item + Number or array for inferring the dtype. + dtype + The dtype to be returned. + as_native + Whether to return the dtype as native dtype. + + Returns + ------- + Return ``dtype`` as native or ivy dtype if provided, else + if ``item`` is given, return its dtype, otherwise return the + global default dtype. + + Examples + -------- + >>> ivy.default_dtype() + 'float32' + + >>> ivy.set_default_dtype(ivy.bool) + >>> ivy.default_dtype() + 'bool' + + >>> ivy.set_default_dtype(ivy.int16) + >>> ivy.default_dtype() + 'int16' + + >>> ivy.set_default_dtype(ivy.float64) + >>> ivy.default_dtype() + 'float64' + + >>> ivy.default_dtype(dtype="int32") + 'int32' + + >>> ivy.default_dtype(dtype=ivy.float16) + 'float16' + + >>> ivy.default_dtype(item=53.234) + 'float64' + + >>> ivy.default_dtype(item=[1, 2, 3]) + 'int32' + + >>> x = ivy.array([5.2, 9.7], dtype="complex128") + >>> ivy.default_dtype(item=x) + 'complex128' + """ + if ivy.exists(dtype): + if as_native is True: + return ivy.as_native_dtype(dtype) + return ivy.as_ivy_dtype(dtype) + as_native = ivy.default(as_native, False) + if ivy.exists(item): + if hasattr(item, "override_dtype_check"): + return item.override_dtype_check() + elif isinstance(item, (list, tuple, dict)) and len(item) == 0: + pass + elif ivy.is_complex_dtype(item): + return ivy.default_complex_dtype(input=item, as_native=as_native) + elif ivy.is_float_dtype(item): + return ivy.default_float_dtype(input=item, as_native=as_native) + elif ivy.is_uint_dtype(item): + return ivy.default_int_dtype(input=item, as_native=as_native) + elif ivy.is_int_dtype(item): + return ivy.default_int_dtype(input=item, as_native=as_native) + elif as_native: + return ivy.as_native_dtype("bool") + else: + return "bool" + global default_dtype_stack + if not default_dtype_stack: + global default_float_dtype_stack + if default_float_dtype_stack: + ret = default_float_dtype_stack[-1] + else: + ret = "float32" + else: + ret = default_dtype_stack[-1] + if as_native: + return ivy.as_native_dtype(ret) + return ivy.as_ivy_dtype(ret) @handle_exceptions @@ -1301,177 +1522,149 @@ def default_uint_dtype( @handle_exceptions -@handle_backend_invalid @handle_nestable -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_device_shifting -def dtype( - x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: +def default_complex_dtype( + *, + input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, + as_native: bool = False, +) -> Union[ivy.Dtype, str, ivy.NativeDtype]: """ - Get the data type for input array x. - Parameters ---------- - x - Tensor for which to get the data type. + input + Number or array for inferring the complex dtype. + complex_dtype + The complex dtype to be returned. as_native - Whether or not to return the dtype in string format. Default is ``False``. - - Returns - ------- - ret - Data type of the array. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) - >>> y = ivy.dtype(x1) - >>> print(y) - float32 - - With :class:`ivy.NativeArray` inputs: - - >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) - >>> y = ivy.dtype(x1) - >>> print(y) - int32 - - With :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), - ... b=ivy.native_array([1, 0, 0, 0, 1])) - >>> y = ivy.dtype(x.a) - >>> print(y) - float32 - """ - return current_backend(x).dtype(x, as_native=as_native) - - -@handle_exceptions -def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: - """ - Get the number of bits used for representing the input data type. - - Parameters - ---------- - dtype_in - The data type to determine the number of bits for. + Whether to return the complex dtype as native dtype. Returns ------- - ret - The number of bits used to represent the data type. + Return ``complex_dtype`` as native or ivy dtype if provided, else + if ``input`` is given, return its complex dtype, otherwise return the + global default complex dtype. Examples -------- - With :class:`ivy.Dtype` inputs: + >>> ivy.default_complex_dtype() + 'complex64' - >>> x = ivy.dtype_bits(ivy.float32) - >>> print(x) - 32 + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' - >>> x = ivy.dtype_bits('int64') - >>> print(x) - 64 + >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) + 'complex128' - With :class:`ivy.NativeDtype` inputs: + >>> ivy.default_complex_dtype(input=4294.967346) + 'complex64' - >>> x = ivy.dtype_bits(ivy.native_bool) - >>> print(x) - 1 + >>> x = ivy.array([9.8,8.9], dtype="complex128") + >>> ivy.default_complex_dtype(input=x) + 'complex128' """ - return current_backend(dtype_in).dtype_bits(dtype_in) + global default_complex_dtype_stack + if ivy.exists(complex_dtype): + if as_native is True: + return ivy.as_native_dtype(complex_dtype) + return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) + as_native = ivy.default(as_native, False) + if ivy.exists(input): + if ivy.is_array(input): + ret = ivy.dtype(input) + elif isinstance(input, np.ndarray): + ret = str(input.dtype) + elif isinstance(input, (list, tuple, dict)): + if ivy.nested_argwhere( + input, lambda x: _check_complex128(x), stop_after_n_found=1 + ): + ret = ivy.complex128 + else: + if not default_complex_dtype_stack: + def_dtype = default_dtype() + if ivy.is_complex_dtype(def_dtype): + ret = def_dtype + else: + ret = "complex64" + else: + ret = default_complex_dtype_stack[-1] + elif isinstance(input, Number): + if _check_complex128(input): + ret = ivy.complex128 + else: + if not default_complex_dtype_stack: + def_dtype = default_dtype() + if ivy.is_complex_dtype(def_dtype): + ret = def_dtype + else: + ret = "complex64" + else: + ret = default_complex_dtype_stack[-1] + else: + if not default_complex_dtype_stack: + def_dtype = default_dtype() + if ivy.is_complex_dtype(def_dtype): + ret = def_dtype + else: + ret = "complex64" + else: + ret = default_complex_dtype_stack[-1] + if as_native: + return ivy.as_native_dtype(ret) + return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) @handle_exceptions @handle_backend_invalid +@handle_nestable @inputs_to_native_arrays @handle_device_shifting -def finfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Finfo: +def dtype( + x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: """ - Machine limits for floating-point data types. + Get the data type for input array x. Parameters ---------- - type - the kind of floating-point data-type about which to get information. + x + Tensor for which to get the data type. + as_native + Whether or not to return the dtype in string format. Default is ``False``. Returns ------- ret - an object having the followng attributes: - - - **bits**: *int* - - number of bits occupied by the floating-point data type. - - - **eps**: *float* - - difference between 1.0 and the next smallest representable floating-point - number larger than 1.0 according to the IEEE-754 standard. - - - **max**: *float* - - largest representable number. - - - **min**: *float* - - smallest representable number. - - - **smallest_normal**: *float* - - smallest positive floating-point number with full precision. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + Data type of the array. Examples -------- - With :class:`ivy.Dtype` input: - - >>> y = ivy.finfo(ivy.float32) - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - - With :code:`str` input: - - >>> y = ivy.finfo('float32') - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) + With :class:`ivy.Array` inputs: - With :class:`ivy.Array` input: + >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) + >>> y = ivy.dtype(x1) + >>> print(y) + float32 - >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) - >>> print(ivy.finfo(x)) - finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) + With :class:`ivy.NativeArray` inputs: - >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) - >>> print(ivy.finfo(x)) - finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) + >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) + >>> y = ivy.dtype(x1) + >>> print(y) + int32 - With :class:`ivy.Container` input: + With :class:`ivy.Container` inputs: - >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), - ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) - >>> print(ivy.finfo(c)) - { - x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), - y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) - } + >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), + ... b=ivy.native_array([1, 0, 0, 0, 1])) + >>> y = ivy.dtype(x.a) + >>> print(y) + float32 """ - return current_backend(None).finfo(type) + return current_backend(x).dtype(x, as_native=as_native) @handle_exceptions @@ -1575,131 +1768,6 @@ def function_unsupported_dtypes( ) -@handle_exceptions -@handle_backend_invalid -@inputs_to_native_arrays -@handle_device_shifting -def iinfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Iinfo: - """ - Machine limits for integer data types. - - Parameters - ---------- - type - the kind of integer data-type about which to get information. - - Returns - ------- - ret - a class with that encapsules the following attributes: - - - **bits**: *int* - - number of bits occupied by the type. - - - **max**: *int* - - largest representable number. - - - **min**: *int* - - smallest representable number. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Dtype` input: - - >>> ivy.iinfo(ivy.int32) - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :code:`str` input: - - >>> ivy.iinfo('int32') - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([13,21,34], dtype=ivy.int8) - >>> ivy.iinfo(x) - iinfo(min=-128, max=127, dtype=int8) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) - >>> ivy.iinfo(x) - iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) - - With :class:`ivy.Container` input: - - >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), - ... y=ivy.array([76,81,16], dtype=ivy.uint32)) - >>> ivy.iinfo(c) - { - x: iinfo(min=0, max=65535, dtype=uint16), - y: iinfo(min=0, max=4294967295, dtype=uint32) - } - """ - return current_backend(None).iinfo(type) - - -@handle_exceptions -def infer_default_dtype( - dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: - """ - Summary. - - Parameters - ---------- - dtype - - as_native - (Default value = False) - - Returns - ------- - Return the default data type for the “kind” (integer or floating-point) of dtype - - Examples - -------- - >>> ivy.set_default_int_dtype("int32") - >>> ivy.infer_default_dtype("int8") - 'int8' - - >>> ivy.set_default_float_dtype("float64") - >>> ivy.infer_default_dtype("float32") - 'float64' - - >>> ivy.set_default_uint_dtype("uint32") - >>> x = ivy.array([0], dtype="uint64") - >>> ivy.infer_default_dtype(x.dtype) - 'uint32' - """ - if ivy.is_complex_dtype(dtype): - default_dtype = ivy.default_complex_dtype(as_native=as_native) - elif ivy.is_float_dtype(dtype): - default_dtype = ivy.default_float_dtype(as_native=as_native) - elif ivy.is_uint_dtype(dtype): - default_dtype = ivy.default_uint_dtype(as_native=as_native) - elif ivy.is_int_dtype(dtype): - default_dtype = ivy.default_int_dtype(as_native=as_native) - elif as_native: - default_dtype = ivy.as_native_dtype("bool") - else: - default_dtype = ivy.as_ivy_dtype("bool") - return default_dtype - - @handle_exceptions def invalid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -1761,144 +1829,23 @@ def is_bool_dtype( if ivy.is_array(dtype_in): dtype_in = ivy.dtype(dtype_in) elif isinstance(dtype_in, np.ndarray): - return "bool" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return ( - True - if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) - else False - ) - elif isinstance(dtype_in, (list, tuple, dict)): - return ( - True - if ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (bool, np.bool)) and x is not int, - ) - else False - ) - return "bool" in ivy.as_ivy_dtype(dtype_in) - - -@handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -def is_complex_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], - /, -) -> bool: - """ - Determine whether the input data type is a complex dtype. - - Parameters - ---------- - dtype_in - The array or data type to check - - Returns - ------- - ret - Whether or not the array or data type is of a complex dtype - - Examples - -------- - >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) - True - - >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) - True - - >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) - False - """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "complex" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return isinstance(dtype_in, (complex, np.complexfloating)) - elif isinstance(dtype_in, (list, tuple, dict)): - return ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (complex, np.complexfloating)) - or (ivy.is_array(x) and "complex" in ivy.dtype(x)), - ) - return "complex" in as_ivy_dtype(dtype_in) - - -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -def is_float_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], - /, -) -> bool: - """ - Determine whether the input data type is a float dtype. - - Parameters - ---------- - dtype_in - The array or data type to check - - Returns - ------- - ret - Whether or not the array or data type is of a floating point dtype - - Examples - -------- - >>> x = ivy.is_float_dtype(ivy.float32) - >>> print(x) - True - - >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) - >>> print(ivy.is_float_dtype(arr)) - True - """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "float" in dtype_in.dtype.name + return "bool" in dtype_in.dtype.name elif isinstance(dtype_in, Number): - return True if isinstance(dtype_in, (float, np.floating)) else False + return ( + True + if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) + else False + ) elif isinstance(dtype_in, (list, tuple, dict)): return ( True if ivy.nested_argwhere( dtype_in, - lambda x: isinstance(x, (float, np.floating)) - or (ivy.is_array(x) and "float" in ivy.dtype(x)), + lambda x: isinstance(x, (bool, np.bool)) and x is not int, ) else False ) - return "float" in as_ivy_dtype(dtype_in) - - -@handle_exceptions -def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: - """ - Check if the given data type is hashable or not. - - Parameters - ---------- - dtype_in - The data type to check. - - Returns - ------- - ret - True if data type is hashable else False - """ - try: - hash(dtype_in) - return True - except TypeError: - return False + return "bool" in ivy.as_ivy_dtype(dtype_in) @handle_exceptions @@ -1990,34 +1937,72 @@ def is_int_dtype( @handle_exceptions -def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: +def check_float(x: Any) -> bool: """ - Determine whether the input dtype is a Native dtype. + Check if the input is a float or a float-like object. + + Parameters + ---------- + x + Input to check. + + Returns + ------- + ret + "True" if the input is a float or a float-like object, otherwise "False". + """ + return isinstance(x, (int, float)) and x is not bool + + +@handle_exceptions +@handle_nestable +@inputs_to_native_arrays +def is_float_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: + """ + Determine whether the input data type is a float dtype. Parameters ---------- dtype_in - Determine whether the input data type is a native data type object. + The array or data type to check Returns ------- ret - Boolean, whether or not dtype_in is a native data type. + Whether or not the array or data type is of a floating point dtype Examples -------- - >>> ivy.set_backend('numpy') - >>> ivy.is_native_dtype(np.int32) + >>> x = ivy.is_float_dtype(ivy.float32) + >>> print(x) True - >>> ivy.set_backend('numpy') - >>> ivy.is_native_array(ivy.float64) - False + >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) + >>> print(ivy.is_float_dtype(arr)) + True """ - try: - return current_backend(None).is_native_dtype(dtype_in) - except ValueError: - return False + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() + elif isinstance(dtype_in, np.ndarray): + return "float" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return True if isinstance(dtype_in, (float, np.floating)) else False + elif isinstance(dtype_in, (list, tuple, dict)): + return ( + True + if ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (float, np.floating)) + or (ivy.is_array(x) and "float" in ivy.dtype(x)), + ) + else False + ) + return "float" in as_ivy_dtype(dtype_in) @handle_exceptions @@ -2069,203 +2054,101 @@ def is_uint_dtype( @handle_exceptions -def promote_types( - type1: Union[ivy.Dtype, ivy.NativeDtype], - type2: Union[ivy.Dtype, ivy.NativeDtype], - /, - *, - array_api_promotion: bool = False, -) -> ivy.Dtype: - """ - Promote the datatypes type1 and type2, returning the data type they promote to. - - Parameters - ---------- - type1 - the first of the two types to promote - type2 - the second of the two types to promote - array_api_promotion - whether to only use the array api promotion rules - - Returns - ------- - ret - The type that both input types promote to - """ - query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] - query.sort(key=lambda x: str(x)) - query = tuple(query) - - def _promote(query): - if array_api_promotion: - return ivy.array_api_promotion_table[query] - return ivy.promotion_table[query] - - try: - ret = _promote(query) - except KeyError: - # try again with the dtypes swapped - query = (query[1], query[0]) - try: - ret = _promote(query) - except KeyError: - raise ivy.utils.exceptions.IvyDtypePromotionError( - "these dtypes ({} and {}) are not type promotable, ".format( - type1, type2 - ) - ) - return ret - - -@handle_exceptions -def promote_types_of_inputs( - x1: Union[ivy.NativeArray, Number, Iterable[Number]], - x2: Union[ivy.NativeArray, Number, Iterable[Number]], - /, - *, - array_api_promotion: bool = False, -) -> Tuple[ivy.NativeArray, ivy.NativeArray]: - """ - Promote the dtype of the given native array inputs to a common dtype based on type - promotion rules. - - While passing float or integer values or any other non-array input - to this function, it should be noted that the return will be an - array-like object. Therefore, outputs from this function should be - used as inputs only for those functions that expect an array-like or - tensor-like objects, otherwise it might give unexpected results. - """ - - def _special_case(a1, a2): - # check for float number and integer array case - return isinstance(a1, float) and "int" in str(a2.dtype) - - def _get_target_dtype(scalar, arr): - # identify a good dtype to give the scalar value, - # based on it's own type and that of the arr value - if _special_case(scalar, arr): - return "float64" - elif arr.dtype == bool and not isinstance(scalar, bool): - return None # let ivy infer a dtype - elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): - return "complex128" - else: - return arr.dtype - - if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): - device = ivy.default_device(item=x1, as_native=True) - x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) - elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): - device = ivy.default_device(item=x2, as_native=True) - x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) - elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): - x1 = ivy.asarray(x1) - x2 = ivy.asarray(x2) - - if x1.dtype != x2.dtype: - promoted = promote_types( - x1.dtype, x2.dtype, array_api_promotion=array_api_promotion - ) - x1 = ivy.astype(x1, promoted, copy=False) - x2 = ivy.astype(x2, promoted, copy=False) - - ivy.utils.assertions._check_jax_x64_flag(x1.dtype) - return ivy.to_native(x1), ivy.to_native(x2) - - -@handle_exceptions -@handle_backend_invalid @handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def result_type( - *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] -) -> ivy.Dtype: +@inputs_to_ivy_arrays +def is_complex_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: """ - Return the dtype that results from applying the type promotion rules (see - :ref:`type-promotion`) to the arguments. - - .. note:: - If provided mixed dtypes (e.g., integer and floating-point), the returned dtype - will be implementation-specific. + Determine whether the input data type is a complex dtype. Parameters ---------- - arrays_and_dtypes - an arbitrary number of input arrays and/or dtypes. - - Returns - ------- - ret - the dtype resulting from an operation involving the input arrays and dtypes. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.array([3., 4., 5.]) - >>> d = ivy.result_type(x, y) - >>> print(d) - float32 - - With :class:`ivy.Dtype` input: - - >>> d = ivy.result_type(ivy.uint8, ivy.uint64) - >>> print(d) - uint64 + dtype_in + The array or data type to check - With :class:`ivy.Container` input: + Returns + ------- + ret + Whether or not the array or data type is of a complex dtype - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = x.a.dtype - >>> print(d) - int32 + Examples + -------- + >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) + True - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = ivy.result_type(x, ivy.float64) - >>> print(d) - { - a: float64 - } + >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) + True + + >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) + False """ - return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() + elif isinstance(dtype_in, np.ndarray): + return "complex" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return isinstance(dtype_in, (complex, np.complexfloating)) + elif isinstance(dtype_in, (list, tuple, dict)): + return ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (complex, np.complexfloating)) + or (ivy.is_array(x) and "complex" in ivy.dtype(x)), + ) + return "complex" in as_ivy_dtype(dtype_in) @handle_exceptions -def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): +def promote_types( + type1: Union[ivy.Dtype, ivy.NativeDtype], + type2: Union[ivy.Dtype, ivy.NativeDtype], + /, + *, + array_api_promotion: bool = False, +) -> ivy.Dtype: """ - Set the 'complex_dtype' as the default data type. + Promote the datatypes type1 and type2, returning the data type they promote to. Parameters ---------- - complex_dtype - The complex data type to be set as the default. + type1 + the first of the two types to promote + type2 + the second of the two types to promote + array_api_promotion + whether to only use the array api promotion rules - Examples - -------- - With :class: `ivy.Dtype` input: + Returns + ------- + ret + The type that both input types promote to + """ + query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] + query.sort(key=lambda x: str(x)) + query = tuple(query) - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' + def _promote(query): + if array_api_promotion: + return ivy.array_api_promotion_table[query] + return ivy.promotion_table[query] - >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) - >>> ivy.default_complex_dtype() - 'complex128' - """ - complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - ivy.utils.assertions._check_jax_x64_flag(complex_dtype) - global default_complex_dtype_stack - default_complex_dtype_stack.append(complex_dtype) + try: + ret = _promote(query) + except KeyError: + # try again with the dtypes swapped + query = (query[1], query[0]) + try: + ret = _promote(query) + except KeyError: + raise ivy.utils.exceptions.IvyDtypePromotionError( + "these dtypes ({} and {}) are not type promotable, ".format( + type1, type2 + ) + ) + return ret @handle_exceptions @@ -2386,6 +2269,34 @@ def set_default_uint_dtype(uint_dtype: Union[ivy.Dtype, str], /): default_uint_dtype_stack.append(uint_dtype) +@handle_exceptions +def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): + """ + Set the 'complex_dtype' as the default data type. + + Parameters + ---------- + complex_dtype + The complex data type to be set as the default. + + Examples + -------- + With :class: `ivy.Dtype` input: + + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' + + >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) + >>> ivy.default_complex_dtype() + 'complex128' + """ + complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) + ivy.utils.assertions._check_jax_x64_flag(complex_dtype) + global default_complex_dtype_stack + default_complex_dtype_stack.append(complex_dtype) + + @handle_exceptions def type_promote_arrays( x1: Union[ivy.Array, ivy.NativeArray], @@ -2412,27 +2323,6 @@ def type_promote_arrays( return ivy.astype(x1, new_type), ivy.astype(x2, new_type) -@handle_exceptions -def unset_default_complex_dtype(): - """ - Reset the current default complex dtype to the previous state. - - Examples - -------- - >>> ivy.set_default_complex_dtype(ivy.complex64) - >>> ivy.set_default_complex_dtype(ivy.complex128) - >>> ivy.default_complex_dtype_stack - ['complex64','complex128'] - - >>> ivy.unset_default_complex_dtype() - >>> ivy.default_complex_dtype_stack - ['complex64'] - """ - global default_complex_dtype_stack - if default_complex_dtype_stack: - default_complex_dtype_stack.pop(-1) - - @handle_exceptions def unset_default_dtype(): """ @@ -2523,6 +2413,27 @@ def unset_default_uint_dtype(): default_uint_dtype_stack.pop(-1) +@handle_exceptions +def unset_default_complex_dtype(): + """ + Reset the current default complex dtype to the previous state. + + Examples + -------- + >>> ivy.set_default_complex_dtype(ivy.complex64) + >>> ivy.set_default_complex_dtype(ivy.complex128) + >>> ivy.default_complex_dtype_stack + ['complex64','complex128'] + + >>> ivy.unset_default_complex_dtype() + >>> ivy.default_complex_dtype_stack + ['complex64'] + """ + global default_complex_dtype_stack + if default_complex_dtype_stack: + default_complex_dtype_stack.pop(-1) + + @handle_exceptions def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -2555,3 +2466,90 @@ def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bo if dtype_in is None: return True return ivy.as_ivy_dtype(dtype_in) in ivy.valid_dtypes + + +@handle_exceptions +def promote_types_of_inputs( + x1: Union[ivy.NativeArray, Number, Iterable[Number]], + x2: Union[ivy.NativeArray, Number, Iterable[Number]], + /, + *, + array_api_promotion: bool = False, +) -> Tuple[ivy.NativeArray, ivy.NativeArray]: + """ + Promote the dtype of the given native array inputs to a common dtype based on type + promotion rules. + + While passing float or integer values or any other non-array input + to this function, it should be noted that the return will be an + array-like object. Therefore, outputs from this function should be + used as inputs only for those functions that expect an array-like or + tensor-like objects, otherwise it might give unexpected results. + """ + + def _special_case(a1, a2): + # check for float number and integer array case + return isinstance(a1, float) and "int" in str(a2.dtype) + + def _get_target_dtype(scalar, arr): + # identify a good dtype to give the scalar value, + # based on it's own type and that of the arr value + if _special_case(scalar, arr): + return "float64" + elif arr.dtype == bool and not isinstance(scalar, bool): + return None # let ivy infer a dtype + elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): + return "complex128" + else: + return arr.dtype + + if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): + device = ivy.default_device(item=x1, as_native=True) + x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) + elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): + device = ivy.default_device(item=x2, as_native=True) + x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) + elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): + x1 = ivy.asarray(x1) + x2 = ivy.asarray(x2) + + if x1.dtype != x2.dtype: + promoted = promote_types( + x1.dtype, x2.dtype, array_api_promotion=array_api_promotion + ) + x1 = ivy.astype(x1, promoted, copy=False) + x2 = ivy.astype(x2, promoted, copy=False) + + ivy.utils.assertions._check_jax_x64_flag(x1.dtype) + return ivy.to_native(x1), ivy.to_native(x2) + + +@handle_exceptions +def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: + """ + Determine whether the input dtype is a Native dtype. + + Parameters + ---------- + dtype_in + Determine whether the input data type is a native data type object. + + Returns + ------- + ret + Boolean, whether or not dtype_in is a native data type. + + Examples + -------- + >>> ivy.set_backend('numpy') + >>> ivy.is_native_dtype(np.int32) + True + + >>> ivy.set_backend('numpy') + >>> ivy.is_native_array(ivy.float64) + False + """ + try: + return current_backend(None).is_native_dtype(dtype_in) + except ValueError: + return False diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py index 058b9dd513163..1fee8c85ac25b 100644 --- a/ivy/functional/ivy/device.py +++ b/ivy/functional/ivy/device.py @@ -10,21 +10,6 @@ import types from typing import Type, Optional, Tuple -# nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. - -from typing import Union, Callable, Iterable, Any - -# local -import ivy -from ivy.func_wrapper import ( - handle_out_argument, - to_native_arrays_and_back, - handle_nestable, - handle_array_like_without_promotion, - handle_backend_invalid, -) -from ivy.utils.exceptions import handle_exceptions - # noinspection PyUnresolvedReferences try: import pynvml @@ -39,11 +24,26 @@ " of the Ivy's device module will be limited. Please install pynvml if" " you wish to use GPUs with Ivy." ) + # nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. + +from typing import Union, Callable, Iterable, Any + +# local +import ivy +from ivy.func_wrapper import ( + handle_out_argument, + to_native_arrays_and_back, + handle_nestable, + handle_array_like_without_promotion, + handle_backend_invalid, +) +from ivy.utils.exceptions import handle_exceptions default_device_stack = list() -max_chunk_sizes = dict() soft_device_mode_stack = list() +dev_handles = dict() split_factors = dict() +max_chunk_sizes = dict() # Extra # @@ -136,87 +136,10 @@ def __exit__( return self -# Profiler # - - -class Profiler(abc.ABC): - """ - The profiler class is used to profile the execution of some code. - - Parameters - ---------- - save_dir - The directory to save the profile data to. - """ - - def __init__(self, save_dir: str): - self._save_dir = save_dir - - @abc.abstractmethod - def start(self): - """ - Start the profiler. - - This should be called before the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException - - @abc.abstractmethod - def stop(self): - """ - Stop the profiler. - - This should be called after the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException - - @abc.abstractmethod - def __enter__(self): - raise ivy.utils.exceptions.IvyNotImplementedException - - @abc.abstractmethod - def __exit__(self, exc_type, exc_val, exc_tb): - raise ivy.utils.exceptions.IvyNotImplementedException - - -# --- Helpers --- # -# --------------- # - - -def _get_devices(fn: Callable, complement: bool = True) -> Tuple: - valid_devices = ivy.valid_devices - invalid_devices = ivy.invalid_devices - all_devices = ivy.all_devices - - supported = set(ivy.valid_devices) - - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - supported = set(all_devices).difference(supported) - return supported - - # Their values are formated like either - # 1. fn.supported_devices = ("cpu",) - # Could also have the "all" value for the framework - basic = [ - ("supported_devices", set.intersection, valid_devices), - ("unsupported_devices", set.difference, invalid_devices), - ] - for key, merge_fn, base in basic: - if hasattr(fn, key): - v = getattr(fn, key) - if "einops" in fn.__name__ and isinstance(v, dict): - v = v.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(v, tuple) - supported = merge_fn(supported, set(v)) - - if complement: - supported = set(all_devices).difference(supported) - - return tuple(supported) +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return ivy.current_backend().handle_soft_device_variable( + *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs + ) # Helpers # @@ -232,24 +155,6 @@ def _get_nvml_gpu_handle(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: return handle -def _is_valid_devices_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): - fn_supported_devices = fn.supported_devices - fn_unsupported_devices = fn.unsupported_devices - if isinstance(fn_supported_devices, dict): - if isinstance(fn_unsupported_devices, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_devices - and backend_str in fn_unsupported_devices - ): - return False - else: - if isinstance(fn_unsupported_devices, tuple): - return False - return True - - def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kwargs): with ivy.ArrayMode(False): default_device = ivy.default_device(device_shifting_dev, as_native=True) @@ -264,177 +169,176 @@ def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kw return args, kwargs, default_device -# --- Main --- # -# ------------ # - +# Device Queries # -# Conversions +# Array Printing @handle_exceptions -def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: +def get_all_ivy_arrays_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, +) -> ivy.Container: """ - Convert device to string representation. + Get all ivy arrays which are currently alive on the specified device. Parameters ---------- device - The device handle to convert to string. + The device handle from which to get the arrays Returns ------- ret - Device string e.g. 'cuda:0'. + Container with the arrays found for the specified device [identity, array] Examples -------- - >>> y = ivy.as_ivy_dev('cpu') - >>> print(y) - cpu + >>> x = ivy.array([1,0,2]) + >>> y = ivy.dev(x) + >>> z = ivy.get_all_ivy_arrays_on_dev(y) + >>> print(z) + {139740789224448:ivy.array([1,0,2])}, """ - return ivy.current_backend().as_ivy_dev(device) + device = ivy.as_ivy_dev(device) + all_arrays = list() + for obj in gc.get_objects(): + if ( + obj is ivy.data_classes.array.array.Array + and ivy.is_ivy_array(obj) + and ivy.dev(obj) == device + ): + all_arrays.append(obj) + + return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) @handle_exceptions -def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: +def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: """ - Convert device string representation to native device type. + Return the number of arrays which are currently alive on the specified device. Parameters ---------- device - The device string to convert to native device handle. - A native device handle can be passed in instead - in this case - the unmodified parameter is returned. + The device handle from which to count the arrays Returns ------- ret - Native device handle. + Number of arrays on the specified device Examples -------- - With :class:`ivy.Device` input: - - >>> ivy.set_backend("numpy") - >>> ivy.as_native_dev("cpu") - 'cpu' - - >>> ivy.set_backend("tensorflow") - >>> ivy.as_native_dev("tpu:3") - '/TPU:3' - - With :class:`ivy.NativeDevice` input: + >>> x1 = ivy.array([-1, 0, 5.2]) + >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 2 - >>> import torch - >>> device = torch.device("cuda") - >>> device - device(type='cuda') + >>> x1 = ivy.native_array([-1, 0, 5.2]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 0 - >>> ivy.as_native_dev(device) - device(type='cuda') + >>> x = ivy.Container(x1=ivy.array([-1]), + ... x2=ivy.native_array([-1])) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 1 """ - return ivy.current_backend().as_native_dev(device) - - -# Memory + return len(ivy.get_all_ivy_arrays_on_dev(device)) @handle_exceptions -def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +@handle_nestable +def print_all_ivy_arrays_on_dev( + *, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + attr_only: bool = True, +) -> None: """ - Clear memory cache on target device. + Print the shape and dtype for all ivy arrays which are currently alive on the + specified device. Parameters ---------- device - The device string to convert to native device handle or native device handle. + The device on which to print the arrays + + attr_only + Whether or not to only print the `shape` and `dtype` attributes of the array Examples -------- - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.clear_cached_mem_on_dev(device) + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y) + ((3,), 'int32') + ((3,), 'int32') + + + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) + [1,0,2] + [3,2,1] """ - ivy.current_backend().clear_cached_mem_on_dev(device) + arrs = ivy.get_all_ivy_arrays_on_dev(device).values() + if attr_only: + [print((arr.shape, arr.dtype)) for arr in arrs] + else: + [print(arr) for arr in arrs] -# Default Device # +ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False -# noinspection PyShadowingNames @handle_exceptions -def default_device( - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - /, - *, - item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, - as_native: bool = None, -) -> Union[ivy.Device, ivy.NativeDevice]: +def set_soft_device_mode(mode: bool) -> None: """ - Return the input device or the default device. If the as_native flag is set, the - device will be converted to a native device. If the item is provided, the item's - device is returned. If the device is not provided, the last default device is - returned. If a default device has not been set, the first gpu is returned if - available, otherwise the cpu is returned. - - Parameters - ---------- - device - The device to be returned or converted. - item - The item to get the device from. - as_native - Whether to convert the device to a native device. - - Returns - ------- - ret - Device handle or string. + Set the mode of whether to move input arrays to `ivy.default_device()` before + performing an operation. + Parameter + --------- + mode + boolean whether to move input arrays Examples -------- - >>> ivy.default_device() - device(type='cpu') - - >>> ivy.default_device("gpu:0") - 'gpu:0' - - >>> ivy.default_device(item=[], as_native=False) - 'cpu' + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode + False + >>> ivy.set_soft_device_mode(True) + >>> ivy.soft_device_mode + True + """ + global soft_device_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + soft_device_mode_stack.append(mode) + ivy.__setattr__("soft_device_mode", mode, True) - >>> ivy.default_device(item=(), as_native=True) - device(type='cpu') - >>> ivy.default_device(item={"a": 1}, as_native=True) - device(type='cpu') +@handle_exceptions +def unset_soft_device_mode() -> None: + """ + Reset the mode of moving input arrays to `ivy.default_device()` before performing an + operation. - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'gpu:0') - >>> ivy.default_device(item=x, as_native=True) - device(type='gpu', id=0) + Examples + -------- + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode + False + >>> ivy.unset_soft_device_mode() + >>> ivy.soft_device_mode + True """ - if ivy.exists(device): - if as_native is True: - return ivy.as_native_dev(device) - elif as_native is False: - return ivy.as_ivy_dev(device) - return device - as_native = ivy.default(as_native, False) - if ivy.exists(item): - if isinstance(item, (list, tuple, dict)) and len(item) == 0: - pass - elif ivy.is_array(item): - return ivy.dev(item, as_native=as_native) - global default_device_stack - if not default_device_stack: - ret = "gpu:0" if ivy.gpu_is_available() else "cpu" - else: - ret = default_device_stack[-1] - if as_native: - return ivy.as_native_dev(ret) - return ivy.as_ivy_dev(ret) + global soft_device_mode_stack + if soft_device_mode_stack: + soft_device_mode_stack.pop(-1) + mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False + ivy.__setattr__("soft_device_mode", mode, True) # Retrieval @@ -481,301 +385,195 @@ def dev( return ivy.current_backend(x).dev(x, as_native=as_native) -# Utilization +# Conversions @handle_exceptions -def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: """ - Get the current utilization (%) for a given device. + Convert device to string representation. Parameters ---------- device - The device string of the device to query utilization for. + The device handle to convert to string. Returns ------- ret - The device utilization (%) + Device string e.g. 'cuda:0'. - Example - ------- - >>> ivy.dev_util('cpu') - 13.4 - >>> ivy.dev_util('gpu:0') - 7.8 - >>> ivy.dev_util('cpu') - 93.4 - >>> ivy.dev_util('gpu:2') - 57.4 - >>> ivy.dev_util('cpu') - 84.2 + Examples + -------- + >>> y = ivy.as_ivy_dev('cpu') + >>> print(y) + cpu """ - if device == "cpu": - return psutil.cpu_percent() - elif "gpu" in device: - handle = _get_nvml_gpu_handle(device) - return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + return ivy.current_backend().as_ivy_dev(device) @handle_exceptions -@handle_nestable -def function_supported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: +def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: """ - Return the supported devices of the current backend's function. The function returns - a dict containing the supported devices for the compositional and primary - implementations in case of partial mixed functions. + Convert device string representation to native device type. Parameters ---------- - fn - The function to check for the supported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + device + The device string to convert to native device handle. + A native device handle can be passed in instead - in this case + the unmodified parameter is returned. Returns ------- ret - Tuple or dict containing the supported devices of the function + Native device handle. Examples -------- - >>> import ivy - >>> print(ivy.function_supported_devices(ivy.ones)) - ('cpu', 'gpu') - """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=False), - } - else: - supported_devices = set(_get_devices(fn, complement=False)) - if recurse: - supported_devices = ivy.functional.data_type._nested_get( - fn, supported_devices, set.intersection, function_supported_devices - ) - - return ( - supported_devices - if isinstance(supported_devices, dict) - else tuple(supported_devices) - ) + With :class:`ivy.Device` input: + >>> ivy.set_backend("numpy") + >>> ivy.as_native_dev("cpu") + 'cpu' -@handle_exceptions -@handle_nestable -def function_unsupported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: - """ - Return the unsupported devices of the current backend's function. The function - returns a dict containing the unsupported devices for the compositional and primary - implementations in case of partial mixed functions. + >>> ivy.set_backend("tensorflow") + >>> ivy.as_native_dev("tpu:3") + '/TPU:3' - Parameters - ---------- - fn - The function to check for the unsupported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + With :class:`ivy.NativeDevice` input: - Returns - ------- - ret - Tuple or dict containing the unsupported devices of the function + >>> import torch + >>> device = torch.device("cuda") + >>> device + device(type='cuda') - Examples - -------- - >>> print(ivy.function_unsupported_devices(ivy.ones)) - ('tpu',) + >>> ivy.as_native_dev(device) + device(type='cuda') """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=True), - } - else: - unsupported_devices = set(_get_devices(fn, complement=True)) - if recurse: - unsupported_devices = ivy.functional.data_type._nested_get( - fn, unsupported_devices, set.union, function_unsupported_devices - ) - return ( - unsupported_devices - if isinstance(unsupported_devices, dict) - else tuple(unsupported_devices) - ) - + return ivy.current_backend().as_native_dev(device) -# Device Queries # -# Array Printing +# Memory @handle_exceptions -def get_all_ivy_arrays_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, -) -> ivy.Container: +def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: """ - Get all ivy arrays which are currently alive on the specified device. + Clear memory cache on target device. Parameters ---------- device - The device handle from which to get the arrays - - Returns - ------- - ret - Container with the arrays found for the specified device [identity, array] + The device string to convert to native device handle or native device handle. Examples -------- - >>> x = ivy.array([1,0,2]) - >>> y = ivy.dev(x) - >>> z = ivy.get_all_ivy_arrays_on_dev(y) - >>> print(z) - {139740789224448:ivy.array([1,0,2])}, - """ - device = ivy.as_ivy_dev(device) - all_arrays = list() - for obj in gc.get_objects(): - if ( - obj is ivy.data_classes.array.array.Array - and ivy.is_ivy_array(obj) - and ivy.dev(obj) == device - ): - all_arrays.append(obj) - - return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) - - -# Availability - - -@handle_exceptions -def gpu_is_available() -> bool: - """ - Determine whether a GPU is available to use, with the backend framework. - - Returns - ------- - ret - Boolean, as to whether a gpu is available. - - Examples - -------- - >>> print(ivy.gpu_is_available()) - False + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.clear_cached_mem_on_dev(device) """ - return ivy.current_backend().gpu_is_available() - - -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return ivy.current_backend().handle_soft_device_variable( - *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs - ) + ivy.current_backend().clear_cached_mem_on_dev(device) @handle_exceptions -def num_cpu_cores(*, logical: bool = True) -> int: +def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Determine the number of cores available in the cpu. + Get the total amount of memory (in GB) for a given device string. In case of CPU, + the total RAM is returned. Parameters ---------- - logical - Whether request is for number of physical or logical cores available in CPU + device + The device string to convert to native device handle. Returns ------- ret - Number of cores available in CPU + The total memory on the device in GB. Examples -------- - >>> print(ivy.num_cpu_cores(logical=False)) - 2 - """ - if logical: - return psutil.cpu_count(logical=logical) - else: - return psutil.cpu_count(logical=False) - - -@handle_exceptions -def num_gpus() -> int: - """ - Determine the number of available GPUs, with the backend framework. - - Returns - ------- - ret - Number of available GPUs. + >>> x = ivy.total_mem_on_dev("cpu") + >>> print(x) + 53.66700032 - Examples - -------- - >>> print(ivy.num_gpus()) - 1 + >>> x = ivy.total_mem_on_dev("gpu:0") + >>> print(x) + 8.589934592 """ - return ivy.current_backend().num_gpus() + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.total / 1e9 + elif device == "cpu": + return psutil.virtual_memory().total / 1e9 + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions -def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: +def used_mem_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + process_specific: bool = False, +) -> float: """ - Return the number of arrays which are currently alive on the specified device. + Get the used memory (in GB) for a given device string. In case of CPU, the used RAM + is returned. Parameters ---------- device - The device handle from which to count the arrays + The device string to convert to native device handle. + process_specific + Whether to check the memory used by this python process alone. Default is + False. Returns ------- ret - Number of arrays on the specified device + The used memory on the device in GB. Examples -------- - >>> x1 = ivy.array([-1, 0, 5.2]) - >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 2 + >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) + >>> print(x) + 6.219563008 - >>> x1 = ivy.native_array([-1, 0, 5.2]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 0 + >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) + >>> print(x) + 0.902400346 - >>> x = ivy.Container(x1=ivy.array([-1]), - ... x2=ivy.native_array([-1])) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) >>> print(y) - 1 + 0.525205504 """ - return len(ivy.get_all_ivy_arrays_on_dev(device)) + ivy.clear_cached_mem_on_dev(device) + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + if process_specific: + pid = os.getpid() + for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): + if process.pid == pid: + return process.usedGpuMemory / 1e9 + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.used / 1e9 + elif device == "cpu": + if process_specific: + return psutil.Process(os.getpid()).memory_info().rss / 1e9 + vm = psutil.virtual_memory() + return (vm.total - vm.available) / 1e9 + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions @@ -839,151 +637,308 @@ def percent_used_mem_on_dev( ) +# Utilization + + @handle_exceptions -@handle_nestable -def print_all_ivy_arrays_on_dev( - *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - attr_only: bool = True, -) -> None: +def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Print the shape and dtype for all ivy arrays which are currently alive on the - specified device. + Get the current utilization (%) for a given device. Parameters ---------- device - The device on which to print the arrays - - attr_only - Whether or not to only print the `shape` and `dtype` attributes of the array - - Examples - -------- - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y) - ((3,), 'int32') - ((3,), 'int32') + The device string of the device to query utilization for. + Returns + ------- + ret + The device utilization (%) - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) - [1,0,2] - [3,2,1] + Example + ------- + >>> ivy.dev_util('cpu') + 13.4 + >>> ivy.dev_util('gpu:0') + 7.8 + >>> ivy.dev_util('cpu') + 93.4 + >>> ivy.dev_util('gpu:2') + 57.4 + >>> ivy.dev_util('cpu') + 84.2 """ - arrs = ivy.get_all_ivy_arrays_on_dev(device).values() - if attr_only: - [print((arr.shape, arr.dtype)) for arr in arrs] + if device == "cpu": + return psutil.cpu_percent() + elif "gpu" in device: + handle = _get_nvml_gpu_handle(device) + return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu else: - [print(arr) for arr in arrs] + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) + + +# Availability @handle_exceptions -def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +def gpu_is_available() -> bool: """ - Set the default device to given device instance. + Determine whether a GPU is available to use, with the backend framework. - Parameters - ---------- - device - The device to set as the default device + Returns + ------- + ret + Boolean, as to whether a gpu is available. Examples -------- - >>> ivy.set_default_device("cpu") - >>> ivy.default_device() - 'cpu' - - >>> ivy.set_backend("torch") - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device(as_native=True) - device(type='cuda', index=0) - - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_default_device(device) - >>> ivy.default_device(as_native=True) - device(type='cuda') + >>> print(ivy.gpu_is_available()) + False """ - global default_device_stack - default_device_stack.append(device) + return ivy.current_backend().gpu_is_available() @handle_exceptions -def set_soft_device_mode(mode: bool) -> None: +def num_cpu_cores(*, logical: bool = True) -> int: """ - Set the mode of whether to move input arrays to `ivy.default_device()` before - performing an operation. + Determine the number of cores available in the cpu. + + Parameters + ---------- + logical + Whether request is for number of physical or logical cores available in CPU + + Returns + ------- + ret + Number of cores available in CPU - Parameter - --------- - mode - boolean whether to move input arrays Examples -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode + >>> print(ivy.num_cpu_cores(logical=False)) + 2 + """ + if logical: + return psutil.cpu_count(logical=logical) + else: + return psutil.cpu_count(logical=False) + + +@handle_exceptions +def num_gpus() -> int: + """ + Determine the number of available GPUs, with the backend framework. + + Returns + ------- + ret + Number of available GPUs. + + Examples + -------- + >>> print(ivy.num_gpus()) + 1 + """ + return ivy.current_backend().num_gpus() + + +@handle_exceptions +def tpu_is_available() -> bool: + """ + Determine whether a TPU is available to use, with the backend framework. + + Returns + ------- + ret + Boolean, as to whether a tpu is available. + + Examples + -------- + >>> print(ivy.tpu_is_available()) False - >>> ivy.set_soft_device_mode(True) - >>> ivy.soft_device_mode - True """ - global soft_device_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - soft_device_mode_stack.append(mode) - ivy.__setattr__("soft_device_mode", mode, True) + return ivy.current_backend().tpu_is_available() +# Default Device # + + +# noinspection PyShadowingNames @handle_exceptions -def set_split_factor( - factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None -) -> None: +def default_device( + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + /, + *, + item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, + as_native: bool = None, +) -> Union[ivy.Device, ivy.NativeDevice]: """ - Set the global split factor for a given device, which can be used to scale batch - splitting chunk sizes for the device across the codebase. + Return the input device or the default device. If the as_native flag is set, the + device will be converted to a native device. If the item is provided, the item's + device is returned. If the device is not provided, the last default device is + returned. If a default device has not been set, the first gpu is returned if + available, otherwise the cpu is returned. Parameters ---------- - factor - The factor to set the device-specific split factor to. device - The device to set the split factor for. Sets the default device by default. + The device to be returned or converted. + item + The item to get the device from. + as_native + Whether to convert the device to a native device. + + Returns + ------- + ret + Device handle or string. Examples -------- - >>> print(ivy.default_device()) - cpu + >>> ivy.default_device() + device(type='cpu') - >>> ivy.set_split_factor(0.5) - >>> print(ivy.split_factors) - {'cpu': 0.5} + >>> ivy.default_device("gpu:0") + 'gpu:0' - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.3, device=device) - >>> print(ivy.split_factors) - {device(type='cuda'): 0.3} + >>> ivy.default_device(item=[], as_native=False) + 'cpu' - >>> ivy.set_split_factor(0.4, device="tpu") - >>> print(ivy.split_factors) - {'tpu': 0.4} + >>> ivy.default_device(item=(), as_native=True) + device(type='cpu') + + >>> ivy.default_device(item={"a": 1}, as_native=True) + device(type='cpu') + + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'gpu:0') + >>> ivy.default_device(item=x, as_native=True) + device(type='gpu', id=0) + """ + if ivy.exists(device): + if as_native is True: + return ivy.as_native_dev(device) + elif as_native is False: + return ivy.as_ivy_dev(device) + return device + as_native = ivy.default(as_native, False) + if ivy.exists(item): + if isinstance(item, (list, tuple, dict)) and len(item) == 0: + pass + elif ivy.is_array(item): + return ivy.dev(item, as_native=as_native) + global default_device_stack + if not default_device_stack: + ret = "gpu:0" if ivy.gpu_is_available() else "cpu" + else: + ret = default_device_stack[-1] + if as_native: + return ivy.as_native_dev(ret) + return ivy.as_ivy_dev(ret) + + +@handle_exceptions +def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: + """ + Set the default device to given device instance. + + Parameters + ---------- + device + The device to set as the default device + + Examples + -------- + >>> ivy.set_default_device("cpu") + >>> ivy.default_device() + 'cpu' + + >>> ivy.set_backend("torch") + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device(as_native=True) + device(type='cuda', index=0) >>> import torch >>> ivy.set_backend("torch") >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.2) - >>> ivy.set_split_factor(0.3, device='gpu') - >>> print(ivy.split_factors) - {'cpu': 0.2, 'gpu': 0.3} + >>> ivy.set_default_device(device) + >>> ivy.default_device(as_native=True) + device(type='cuda') """ - ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) - global split_factors - device = ivy.default(device, default_device()) - split_factors[device] = factor + global default_device_stack + default_device_stack.append(device) + + +@handle_exceptions +def unset_default_device() -> None: + """ + Reset the default device to "cpu". + + Examples + -------- + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device() + "gpu:0" + >>> ivy.unset_default_device() + >>> ivy.default_device() + "cpu" + """ + global default_device_stack + if default_device_stack: + default_device_stack.pop(-1) + + +# Device Allocation # + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def to_device( + x: Union[ivy.Array, ivy.NativeArray], + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + stream: Optional[Union[int, Any]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Move the input array x to the desired device, specified by device string. + + Parameters + ---------- + x + input array to be moved to the desired device + device + device to move the input array `x` to + stream + stream object to use during copy. In addition to the types supported in + array.__dlpack__(), implementations may choose to support any library-specific + stream object with the caveat that any code using such an object would not be + portable. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + input array x placed on the desired device + + Examples + -------- + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'cpu') + >>> print(x.device) + cpu + """ + return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) # Function Splitting # @@ -1029,9 +984,58 @@ def split_factor( @handle_exceptions -def split_func_call( - func: Callable, - inputs: Union[ivy.Array, ivy.NativeArray], +def set_split_factor( + factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None +) -> None: + """ + Set the global split factor for a given device, which can be used to scale batch + splitting chunk sizes for the device across the codebase. + + Parameters + ---------- + factor + The factor to set the device-specific split factor to. + device + The device to set the split factor for. Sets the default device by default. + + Examples + -------- + >>> print(ivy.default_device()) + cpu + + >>> ivy.set_split_factor(0.5) + >>> print(ivy.split_factors) + {'cpu': 0.5} + + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_split_factor(0.3, device=device) + >>> print(ivy.split_factors) + {device(type='cuda'): 0.3} + + >>> ivy.set_split_factor(0.4, device="tpu") + >>> print(ivy.split_factors) + {'tpu': 0.4} + + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_split_factor(0.2) + >>> ivy.set_split_factor(0.3, device='gpu') + >>> print(ivy.split_factors) + {'cpu': 0.2, 'gpu': 0.3} + """ + ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) + global split_factors + device = ivy.default(device, default_device()) + split_factors[device] = factor + + +@handle_exceptions +def split_func_call( + func: Callable, + inputs: Union[ivy.Array, ivy.NativeArray], mode: str, /, *, @@ -1162,213 +1166,200 @@ def split_func_call( return ret[0] if len(ret) == 1 else ret -# Device Allocation # +def _is_valid_devices_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): + fn_supported_devices = fn.supported_devices + fn_unsupported_devices = fn.unsupported_devices + if isinstance(fn_supported_devices, dict): + if isinstance(fn_unsupported_devices, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_devices + and backend_str in fn_unsupported_devices + ): + return False + else: + if isinstance(fn_unsupported_devices, tuple): + return False + return True -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def to_device( - x: Union[ivy.Array, ivy.NativeArray], - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - stream: Optional[Union[int, Any]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Move the input array x to the desired device, specified by device string. +def _get_devices(fn: Callable, complement: bool = True) -> Tuple: + valid_devices = ivy.valid_devices + invalid_devices = ivy.invalid_devices + all_devices = ivy.all_devices - Parameters - ---------- - x - input array to be moved to the desired device - device - device to move the input array `x` to - stream - stream object to use during copy. In addition to the types supported in - array.__dlpack__(), implementations may choose to support any library-specific - stream object with the caveat that any code using such an object would not be - portable. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + supported = set(ivy.valid_devices) - Returns - ------- - ret - input array x placed on the desired device + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + supported = set(all_devices).difference(supported) + return supported - Examples - -------- - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'cpu') - >>> print(x.device) - cpu - """ - return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) + # Their values are formated like either + # 1. fn.supported_devices = ("cpu",) + # Could also have the "all" value for the framework + basic = [ + ("supported_devices", set.intersection, valid_devices), + ("unsupported_devices", set.difference, invalid_devices), + ] + for key, merge_fn, base in basic: + if hasattr(fn, key): + v = getattr(fn, key) + if "einops" in fn.__name__ and isinstance(v, dict): + v = v.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(v, tuple) + supported = merge_fn(supported, set(v)) + + if complement: + supported = set(all_devices).difference(supported) + + return tuple(supported) @handle_exceptions -def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +@handle_nestable +def function_supported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: """ - Get the total amount of memory (in GB) for a given device string. In case of CPU, - the total RAM is returned. + Return the supported devices of the current backend's function. The function returns + a dict containing the supported devices for the compositional and primary + implementations in case of partial mixed functions. Parameters ---------- - device - The device string to convert to native device handle. + fn + The function to check for the supported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. Returns ------- ret - The total memory on the device in GB. + Tuple or dict containing the supported devices of the function Examples -------- - >>> x = ivy.total_mem_on_dev("cpu") - >>> print(x) - 53.66700032 - - >>> x = ivy.total_mem_on_dev("gpu:0") - >>> print(x) - 8.589934592 + >>> import ivy + >>> print(ivy.function_supported_devices(ivy.ones)) + ('cpu', 'gpu') """ - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.total / 1e9 - elif device == "cpu": - return psutil.virtual_memory().total / 1e9 + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=False), + } else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + supported_devices = set(_get_devices(fn, complement=False)) + if recurse: + supported_devices = ivy.functional.data_type._nested_get( + fn, supported_devices, set.intersection, function_supported_devices + ) + + return ( + supported_devices + if isinstance(supported_devices, dict) + else tuple(supported_devices) + ) @handle_exceptions -def tpu_is_available() -> bool: +@handle_nestable +def function_unsupported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: """ - Determine whether a TPU is available to use, with the backend framework. + Return the unsupported devices of the current backend's function. The function + returns a dict containing the unsupported devices for the compositional and primary + implementations in case of partial mixed functions. + + Parameters + ---------- + fn + The function to check for the unsupported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. Returns ------- ret - Boolean, as to whether a tpu is available. - - Examples - -------- - >>> print(ivy.tpu_is_available()) - False - """ - return ivy.current_backend().tpu_is_available() - - -@handle_exceptions -def unset_default_device() -> None: - """ - Reset the default device to "cpu". + Tuple or dict containing the unsupported devices of the function Examples -------- - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device() - "gpu:0" - >>> ivy.unset_default_device() - >>> ivy.default_device() - "cpu" + >>> print(ivy.function_unsupported_devices(ivy.ones)) + ('tpu',) """ - global default_device_stack - if default_device_stack: - default_device_stack.pop(-1) - + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=True), + } + else: + unsupported_devices = set(_get_devices(fn, complement=True)) + if recurse: + unsupported_devices = ivy.functional.data_type._nested_get( + fn, unsupported_devices, set.union, function_unsupported_devices + ) + return ( + unsupported_devices + if isinstance(unsupported_devices, dict) + else tuple(unsupported_devices) + ) -@handle_exceptions -def unset_soft_device_mode() -> None: - """ - Reset the mode of moving input arrays to `ivy.default_device()` before performing an - operation. - Examples - -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode - False - >>> ivy.unset_soft_device_mode() - >>> ivy.soft_device_mode - True - """ - global soft_device_mode_stack - if soft_device_mode_stack: - soft_device_mode_stack.pop(-1) - mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False - ivy.__setattr__("soft_device_mode", mode, True) +# Profiler # -@handle_exceptions -def used_mem_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - process_specific: bool = False, -) -> float: +class Profiler(abc.ABC): """ - Get the used memory (in GB) for a given device string. In case of CPU, the used RAM - is returned. + The profiler class is used to profile the execution of some code. Parameters ---------- - device - The device string to convert to native device handle. - process_specific - Whether to check the memory used by this python process alone. Default is - False. + save_dir + The directory to save the profile data to. + """ - Returns - ------- - ret - The used memory on the device in GB. + def __init__(self, save_dir: str): + self._save_dir = save_dir - Examples - -------- - >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) - >>> print(x) - 6.219563008 + @abc.abstractmethod + def start(self): + """ + Start the profiler. - >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) - >>> print(x) - 0.902400346 + This should be called before the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException - >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) - >>> print(y) - 0.525205504 - """ - ivy.clear_cached_mem_on_dev(device) - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - if process_specific: - pid = os.getpid() - for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): - if process.pid == pid: - return process.usedGpuMemory / 1e9 - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.used / 1e9 - elif device == "cpu": - if process_specific: - return psutil.Process(os.getpid()).memory_info().rss / 1e9 - vm = psutil.virtual_memory() - return (vm.total - vm.available) / 1e9 - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + @abc.abstractmethod + def stop(self): + """ + Stop the profiler. + + This should be called after the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException + @abc.abstractmethod + def __enter__(self): + raise ivy.utils.exceptions.IvyNotImplementedException -ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False -dev_handles = dict() + @abc.abstractmethod + def __exit__(self, exc_type, exc_val, exc_tb): + raise ivy.utils.exceptions.IvyNotImplementedException diff --git a/ivy/functional/ivy/elementwise.py b/ivy/functional/ivy/elementwise.py index e5280d85a6f3e..5fd532ccfb9c9 100644 --- a/ivy/functional/ivy/elementwise.py +++ b/ivy/functional/ivy/elementwise.py @@ -18,20 +18,6 @@ from ivy.utils.exceptions import handle_exceptions -pow.unsupported_gradients = {"torch": ["float16"]} -trunc_divide.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - "handle_backend_invalid", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - # Array API Standard # # -------------------# @@ -512,51 +498,6 @@ def add( return ivy.current_backend(x1, x2).add(x1, x2, alpha=alpha, out=out) -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def angle( - z: Union[ivy.Array, ivy.NativeArray], - /, - *, - deg: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate Element-wise the angle for an array of complex numbers(x+yj). - - Parameters - ---------- - z - Array-like input. - deg - optional bool. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Returns an array of angles for each complex number in the input. - If deg is False(default), angle is calculated in radian and if - deg is True, then angle is calculated in degrees. - - Examples - -------- - >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) - >>> z - ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) - >>> ivy.angle(z) - ivy.array([ 2.35619449, 2.35619449, -0.78539816]) - >>> ivy.angle(z,deg=True) - ivy.array([135., 135., -45.]) - """ - return ivy.current_backend(z).angle(z, deg=deg, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -1954,85 +1895,6 @@ def cosh( return ivy.current_backend(x).cosh(x, out=out) -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def deg2rad( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Convert the input from degrees to radians. - - Parameters - ---------- - x - input array whose elements are each expressed in degrees. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array with each element in ``x`` converted from degrees to radians. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x=ivy.array([0,90,180,270,360]) - >>> y=ivy.deg2rad(x) - >>> print(y) - ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) - - >>> x=ivy.array([0,-1.5,-50,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.deg2rad(x,out=y) - >>> print(y) - ivy.array([ 0., -0.02617994, -0.87266463, nan]) - - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.deg2rad(x, out=x) - >>> print(x) - ivy.array([[ 0.01919862, 0.03839725, 0.05759586], - [-0.07679449, -0.09599311, -0.11519173]]) - - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.deg2rad(x,out=y) - >>> print(y) - ivy.array([0., 0.35081118, nan]) - - With :class:`ivy.Container` input: - - >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), - ... b=ivy.array([0,90,180,270,360])) - >>> y=ivy.deg2rad(x) - >>> print(y) - { - a: ivy.array([0., 0.35081118, -0.88139129, nan]), - b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) - } - - >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), - ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) - >>> y=ivy.deg2rad(x) - >>> print(y) - { - a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), - b: ivy.array([0., -0.02617994, -0.87266463, nan]) - } - """ - return ivy.current_backend(x).deg2rad(x, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2342,49 +2204,6 @@ def equal( return ivy.current_backend(x1, x2).equal(x1, x2, out=out) -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def erf( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Gauss error function of ``x`` element-wise. - - Parameters - ---------- - x - Value to compute exponential for. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The Gauss error function of x. - - Examples - -------- - >>> x = ivy.array([0, 0.3, 0.7, 1.0]) - >>> ivy.erf(x) - ivy.array([0., 0.328, 0.677, 0.842]) - """ - return ivy.current_backend(x).erf(x, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2532,70 +2351,210 @@ def exp( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def exp2( - x: Union[ivy.Array, float, list, tuple], +def imag( + val: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate 2**p for all p in the input array. + Return the imaginary part of a complex number for each element ``x_i`` of the input + array ``val``. Parameters ---------- - x - Array-like input. + val + input array. Should have a complex floating-point data type. out optional output array, for writing the result to. Returns ------- ret - Element-wise 2 to the power x. This is a scalar if x is a scalar. + Returns an array with the imaginary part of complex numbers. + The returned arrau must have a floating-point data type determined by + the precision of ``val`` (e.g., if ``val`` is ``complex64``, + the returned array must be ``float32``). + + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.exp2(x) - ivy.array([2., 4., 8.]) - >>> x = [5, 6, 7] - >>> ivy.exp2(x) - ivy.array([32., 64., 128.]) + >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) + >>> b + ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) + >>> ivy.imag(b) + ivy.array([2., 4., 6.]) """ - return ivy.current_backend(x).exp2(x, out=out) + return ivy.current_backend(val).imag(val, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def expm1( - x: Union[ivy.Array, ivy.NativeArray], +def angle( + z: Union[ivy.Array, ivy.NativeArray], /, *, + deg: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain - ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element - ``x_i`` of the input array ``x``. + Calculate Element-wise the angle for an array of complex numbers(x+yj). - .. note:: - The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when - ``x`` is close to zero. Accordingly, conforming implementations should avoid - implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other - IEEE 754-2019 compliant mathematical library, for a potential reference - implementation. + Parameters + ---------- + z + Array-like input. + deg + optional bool. + out + optional output array, for writing the result to. - .. note:: - For complex floating-point operands, ``expm1(conj(x))`` - must equal ``conj(expm1(x))``. + Returns + ------- + ret + Returns an array of angles for each complex number in the input. + If deg is False(default), angle is calculated in radian and if + deg is True, then angle is calculated in degrees. - .. note:: + Examples + -------- + >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> z + ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) + >>> ivy.angle(z) + ivy.array([ 2.35619449, 2.35619449, -0.78539816]) + >>> ivy.angle(z,deg=True) + ivy.array([135., 135., -45.]) + """ + return ivy.current_backend(z).angle(z, deg=deg, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def gcd( + x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the greatest common divisor of |x1| and |x2|. + + Parameters + ---------- + x1 + First array-like input. + x2 + Second array-input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Element-wise gcd of |x1| and |x2|. + + Examples + -------- + >>> x1 = ivy.array([1, 2, 3]) + >>> x2 = ivy.array([4, 5, 6]) + >>> ivy.gcd(x1, x2) + ivy.array([1., 1., 3.]) + >>> x1 = ivy.array([1, 2, 3]) + >>> ivy.gcd(x1, 10) + ivy.array([1., 2., 1.]) + """ + return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def exp2( + x: Union[ivy.Array, float, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Calculate 2**p for all p in the input array. + + Parameters + ---------- + x + Array-like input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Element-wise 2 to the power x. This is a scalar if x is a scalar. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> ivy.exp2(x) + ivy.array([2., 4., 8.]) + >>> x = [5, 6, 7] + >>> ivy.exp2(x) + ivy.array([32., 64., 128.]) + """ + return ivy.current_backend(x).exp2(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def expm1( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain + ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element + ``x_i`` of the input array ``x``. + + .. note:: + The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when + ``x`` is close to zero. Accordingly, conforming implementations should avoid + implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other + IEEE 754-2019 compliant mathematical library, for a potential reference + implementation. + + .. note:: + For complex floating-point operands, ``expm1(conj(x))`` + must equal ``conj(expm1(x))``. + + .. note:: The exponential function is an entire function in the complex plane and has no branch cuts. @@ -3016,92 +2975,6 @@ def fmin( return ivy.current_backend(x1, x2).fmin(x1, x2, out=out) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fmod( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Compute the element-wise remainder of divisions of two arrays. - - Parameters - ---------- - x1 - First input array. - x2 - Second input array - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array with element-wise remainder of divisions. - - Examples - -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmod(x1, x2) - ivy.array([ 0, 3, 0]) - - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmod(x1, x2) - ivy.array([ nan, nan, nan]) - """ - return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def gcd( - x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the greatest common divisor of |x1| and |x2|. - - Parameters - ---------- - x1 - First array-like input. - x2 - Second array-input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Element-wise gcd of |x1| and |x2|. - - Examples - -------- - >>> x1 = ivy.array([1, 2, 3]) - >>> x2 = ivy.array([4, 5, 6]) - >>> ivy.gcd(x1, x2) - ivy.array([1., 1., 3.]) - >>> x1 = ivy.array([1, 2, 3]) - >>> ivy.gcd(x1, 10) - ivy.array([1., 2., 1.]) - """ - return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -3291,57 +3164,277 @@ def greater_equal( return ivy.current_backend(x1, x2).greater_equal(x1, x2, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def imag( - val: Union[ivy.Array, ivy.NativeArray], +def less_equal( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the imaginary part of a complex number for each element ``x_i`` of the input - array ``val``. + Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 + with the respective element x2_i of the input array x2. Parameters ---------- - val - input array. Should have a complex floating-point data type. + x1 + first input array. May have any data type. + x2 + second input array. Must be compatible with x1 (with Broadcasting). May have any + data type. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - ret - Returns an array with the imaginary part of complex numbers. - The returned arrau must have a floating-point data type determined by - the precision of ``val`` (e.g., if ``val`` is ``complex64``, - the returned array must be ``float32``). + ret + an array containing the element-wise results. The returned array must have a + data type of bool. - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.less_equal.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- - >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) - >>> b - ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) - >>> ivy.imag(b) - ivy.array([2., 4., 6.]) + With :class:`ivy.Array` input: + + >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) + >>> print(x) + ivy.array([True, True, False]) + + >>> x = ivy.array([[10.1, 2.3, -3.6]]) + >>> y = ivy.array([[4.8], [5.2], [6.1]]) + >>> shape = (3,3) + >>> fill_value = False + >>> z = ivy.full(shape, fill_value) + >>> ivy.less_equal(x, y, out=z) + >>> print(z) + ivy.array([[False, True, True], + [False, True, True], + [False, True, True]]) + + >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) + >>> y = ivy.array([[8.4], [2.5], [1.6]]) + >>> ivy.less_equal(x, y, out=x) + >>> print(x) + ivy.array([[[1.], + [0.], + [1.]]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) + >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) + >>> z = ivy.less_equal(x, y) + >>> print(z) + { + a: ivy.array([False, False, False]), + b: ivy.array([True, True, True]) + } """ - return ivy.current_backend(val).imag(val, out=out) + return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def multiply( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Calculate the product for each element x1_i of the input array x1 with the + respective element x2_i of the input array x2. + + .. note:: + Floating-point multiplication is not always associative due to finite precision. + + **Special Cases** + + For real-valued floating-point operands, + + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` and ``x2_i`` have the same mathematical sign, + the result has a positive mathematical sign, unless the result is ``NaN``. + If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` and ``x2_i`` have different mathematical signs, + the result has a negative mathematical sign, + unless the result is ``NaN``. If the result is ``NaN``, + the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, + the result is a signed infinity with the mathematical sign determined by + the rule already stated above. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` + is a nonzero finite number, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - If ``x1_i`` is a nonzero finite number and ``x2_i`` + is either ``+infinity`` or ``-infinity``, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - In the remaining cases, where neither ``infinity`` nor ``NaN`` + is involved, the product must be computed and rounded to the nearest + representable value according to IEEE 754-2019 and a supported + rounding mode. If the magnitude is too large to represent, + the result is an `infinity` of appropriate mathematical sign. + If the magnitude is too small to represent, the result is a zero of + appropriate mathematical sign. + + For complex floating-point operands, multiplication is defined according to the + following table. For real components ``a`` and ``c`` and + imaginary components ``b`` and ``d``, + + +------------+----------------+-----------------+--------------------------+ + | | c | dj | c + dj | + +============+================+=================+==========================+ + | **a** | a * c | (a*d)j | (a*c) + (a*d)j | + +------------+----------------+-----------------+--------------------------+ + | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | + +------------+----------------+-----------------+--------------------------+ + | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | + +------------+----------------+-----------------+--------------------------+ + + In general, for complex floating-point operands, real-valued floating-point + special cases must independently apply to the real and imaginary component + operations involving real numbers as described in the above table. + + When ``a``, ``b``, ``c``, or ``d`` are all finite numbers + (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), + multiplication of complex floating-point operands should be computed + as if calculated according to the textbook formula for complex number multiplication + + .. math:: + (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j + + When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, + ``+infinity``, or ``-infinity``, + + - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, the result is implementation dependent. + + .. note:: + For complex floating-point operands, the results of special cases may be + implementation dependent depending on how an implementation chooses + to model complex numbers and complex infinity + (e.g., complex plane versus Riemann sphere). + For those implementations following C99 and its one-infinity model, + when at least one component is infinite, + even if the other component is ``NaN``, + the complex value is infinite, and the usual arithmetic + rules do not apply to complex-complex multiplication. + In the interest of performance, other implementations + may want to avoid the complex branching logic necessary + to implement the one-infinity model and choose to implement + all complex-complex multiplication according to the textbook formula. + Accordingly, special case behavior is unlikely + to be consistent across implementations. + + Parameters + ---------- + x1 + first input array. Should have a numeric data type. + + x2 + second input array. Must be compatible with ``x1`` + (see :ref'`broadcasting`). Should have a numeric data type + + out + optional output array, for writing the array result to. + It must have a shape that the inputs broadcast to. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Returns + ------- + ret + an array containing the element-wise products. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. + + Examples + -------- + With :code:`ivy.Array` inputs: + + >>> x1 = ivy.array([3., 5., 7.]) + >>> x2 = ivy.array([4., 6., 8.]) + >>> y = ivy.multiply(x1, x2) + >>> print(y) + ivy.array([12., 30., 56.]) + + With :code:`ivy.NativeArray` inputs: + + >>> x1 = ivy.native_array([1., 3., 9.]) + >>> x2 = ivy.native_array([4., 7.2, 1.]) + >>> y = ivy.multiply(x1, x2) + >>> print(y) + ivy.array([ 4. , 21.6, 9. ]) + + With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: + + >>> x1 = ivy.array([8., 6., 7.]) + >>> x2 = ivy.native_array([1., 2., 3.]) + >>> y = ivy.multiply(x1, x2) + >>> print(y) + ivy.array([ 8., 12., 21.]) + + With :code:`ivy.Container` inputs: + + >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) + >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) + >>> y = ivy.multiply(x1, x2) + >>> print(y) + { + a: ivy.array([12.,12.,24.]), + b: ivy.array([9.,3.,10.]) + } + + With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: + + >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) + >>> x2 = ivy.array([1.,2.,3.]) + >>> y = ivy.multiply(x1, x2) + >>> print(y) + { + a: ivy.array([3.,8.,15.]), + b: ivy.array([2.,4.,3.]) + } + """ + return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) @handle_exceptions @@ -3681,27 +3774,28 @@ def isnan( @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def isreal( - x: Union[ivy.Array, ivy.NativeArray], +def less( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Test each element ``x_i`` of the input array ``x`` to determine whether the element - is real number. Returns a bool array, where True if input element is real. If - element has complex type with zero complex part, the return value for that element - is True. + Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. Parameters ---------- - x - input array. + x1 + first input array. Should have a numeric data type. + x2 + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -3709,122 +3803,12 @@ def isreal( Returns ------- ret - an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is - real number and ``False`` otherwise. The returned array should have a data type - of ``bool``. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances in place of - :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints - and also the examples below. + an array containing the element-wise results. The returned array must have a + data type of ``bool``. Examples -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([[[True], [True], [True]]]) - - >>> x = ivy.array([1-0j, 3j, 7+5j]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([ True, False, False]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ - b=ivy.array([5j, 5-6j, 3])) - >>> z = ivy.isreal(x) - >>> print(z) - { - a: ivy.array([False, True, True]), - b: ivy.array([False, False, True]) - } - """ - return ivy.current_backend(x).isreal(x, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lcm( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the element-wise least common multiple (LCM) of x1 and x2. - - Parameters - ---------- - x1 - first input array, must be integers - x2 - second input array, must be integers - out - optional output array, for writing the result to. - - Returns - ------- - ret - an array that includes the element-wise least common multiples of x1 and x2 - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x1=ivy.array([2, 3, 4]) - >>> x2=ivy.array([5, 7, 15]) - >>> x1.lcm(x1, x2) - ivy.array([10, 21, 60]) - """ - return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def less( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. - - Parameters - ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. - - Examples - -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` input: >>> x = ivy.less(ivy.array([1,2,3]),ivy.array([2,2,2])) >>> print(x) @@ -3877,93 +3861,6 @@ def less( return ivy.current_backend(x1).less(x1, x2, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def less_equal( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 - with the respective element x2_i of the input array x2. - - Parameters - ---------- - x1 - first input array. May have any data type. - x2 - second input array. Must be compatible with x1 (with Broadcasting). May have any - data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the element-wise results. The returned array must have a - data type of bool. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) - >>> print(x) - ivy.array([True, True, False]) - - >>> x = ivy.array([[10.1, 2.3, -3.6]]) - >>> y = ivy.array([[4.8], [5.2], [6.1]]) - >>> shape = (3,3) - >>> fill_value = False - >>> z = ivy.full(shape, fill_value) - >>> ivy.less_equal(x, y, out=z) - >>> print(z) - ivy.array([[False, True, True], - [False, True, True], - [False, True, True]]) - - >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) - >>> y = ivy.array([[8.4], [2.5], [1.6]]) - >>> ivy.less_equal(x, y, out=x) - >>> print(x) - ivy.array([[[1.], - [0.], - [1.]]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) - >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) - >>> z = ivy.less_equal(x, y) - >>> print(z) - { - a: ivy.array([False, False, False]), - b: ivy.array([True, True, True]) - } - """ - return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -4882,125 +4779,98 @@ def logical_xor( return ivy.current_backend(x1, x2).logical_xor(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def maximum( - x1: Union[ivy.Array, ivy.NativeArray, Number], - x2: Union[ivy.Array, ivy.NativeArray, Number], +def nan_to_num( + x: Union[ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. + Replace NaN with zero and infinity with large finite numbers (default behaviour) or + with the numbers defined by the user using the nan, posinf and/or neginf keywords. Parameters ---------- - x1 - Input array containing elements to maximum threshold. - x2 - Tensor containing maximum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum - is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. + x + Array input. + copy + Whether to create a copy of x (True) or to replace values in-place (False). + The in-place operation only occurs if casting to an array does not require + a copy. Default is True. + nan + Value to be used to fill NaN values. If no value is passed then NaN values + will be replaced with 0.0. + posinf + Value to be used to fill positive infinity values. If no value is passed + then positive infinity values will be replaced with a very large number. + neginf + Value to be used to fill negative infinity values. + If no value is passed then negative infinity values + will be replaced with a very small (or negative) number. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - An array with the elements of x1, but clipped to not be lower than the x2 - values. + Array with the non-finite values replaced. + If copy is False, this may be x itself. Examples -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.maximum(x, y) - >>> print(z) - ivy.array([9, 9, 5]) - - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.maximum(x, y, out=z) - >>> print(z) - ivy.array([[9., 9., 9., 9., 9., 9.], - [3., 5., 9., 8., 3., 7.], - [2., 5., 9., 8., 3., 7.]]) - - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.maximum(x, y, out=x) - >>> print(x) - ivy.array([[7, 7]]) - - With one :class:`ivy.Container` input: - - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]), - ... b=ivy.array([-5, 9])) - >>> z = ivy.maximum(x, y) - >>> print(z) - { - a: ivy.array([[1, 3], - [2, 4], - [3, 7]]), - b: ivy.array([[1, 9], - [2, 9], - [3, 9]]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) - >>> z = ivy.maximum(x, y) - >>> print(z) - { - a: ivy.array([1, 5, 6]), - b: ivy.array([5, 9, 7]) - } - """ - return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) + >>> x = ivy.array([1, 2, 3, nan]) + >>> ivy.nan_to_num(x) + ivy.array([1., 1., 3., 0.0]) + >>> x = ivy.array([1, 2, 3, inf]) + >>> ivy.nan_to_num(x, posinf=5e+100) + ivy.array([1., 2., 3., 5e+100]) + """ + return ivy.current_backend(x).nan_to_num( + x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out + ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def minimum( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def negative( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. + Return a new array with the negative value of each element in ``x``. + + .. note:: + For signed integer data types, the numerical negative of + the minimum representable integer is implementation-dependent. + + .. note:: + If ``x`` has a complex floating-point data type, + both the real and imaginary components for each ``x_i`` + must be negated (a result which follows from the rules of + complex number multiplication). Parameters ---------- - x1 - Input array containing elements to minimum threshold. - x2 - Tensor containing minimum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum - is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. + x + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5008,62 +4878,53 @@ def minimum( Returns ------- ret - An array with the elements of x1, but clipped to not exceed the x2 values. + A new array with the negative value of each element in ``x``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input: - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.minimum(x, y) - >>> print(z) - ivy.array([7, 3, 2]) + >>> x = ivy.array([0,1,1,2]) + >>> y = ivy.negative(x) + >>> print(y) + ivy.array([ 0, -1, -1, -2]) - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.minimum(x, y, out=z) - >>> print(z) - ivy.array([[1.,5.,9.,8.,3.,7.], - [1.,3.,3.,3.,3.,3.], - [1.,2.,2.,2.,2.,2.]]) + >>> x = ivy.array([0,-1,-0.5,2,3]) + >>> y = ivy.zeros(5) + >>> ivy.negative(x, out=y) + >>> print(y) + ivy.array([-0. , 1. , 0.5, -2. , -3. ]) - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.minimum(x, y, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.negative(x,out=x) >>> print(x) - ivy.array([[0, 3]]) - - With one :class:`ivy.Container` input: - - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) - >>> z = ivy.minimum(x, y) - >>> print(z) - { - a: ivy.array([[1, 0], - [1, 0], - [1, 0]]), - b: ivy.array([[-5, 3], - [-5, 4], - [-5, 7]]) - } + ivy.array([[-1.1, -2.2, -3.3], + [4.4, 5.5, 6.6]]) - With multiple :class:`ivy.Container` inputs: + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 3, 1]), - ... b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]), - ... b=ivy.array([5, 9, 7])) - >>> z = ivy.minimum(x, y) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.negative(x) + >>> print(y) { - a: ivy.array([1, 3, 1]), - b: ivy.array([2, 8, 5]) + a: ivy.array([-0., -1., -2.]), + b: ivy.array([-3., -4., 5.]) } """ - return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) + return ivy.current_backend(x).negative(x, out=out) @handle_exceptions @@ -5073,277 +4934,350 @@ def minimum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def multiply( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def not_equal( + x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], + x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate the product for each element x1_i of the input array x1 with the - respective element x2_i of the input array x2. - - .. note:: - Floating-point multiplication is not always associative due to finite precision. + """ + Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. **Special Cases** For real-valued floating-point operands, - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` and ``x2_i`` have the same mathematical sign, - the result has a positive mathematical sign, unless the result is ``NaN``. - If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` and ``x2_i`` have different mathematical signs, - the result has a negative mathematical sign, - unless the result is ``NaN``. If the result is ``NaN``, - the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, - the result is a signed infinity with the mathematical sign determined by - the rule already stated above. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` - is a nonzero finite number, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - If ``x1_i`` is a nonzero finite number and ``x2_i`` - is either ``+infinity`` or ``-infinity``, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - In the remaining cases, where neither ``infinity`` nor ``NaN`` - is involved, the product must be computed and rounded to the nearest - representable value according to IEEE 754-2019 and a supported - rounding mode. If the magnitude is too large to represent, - the result is an `infinity` of appropriate mathematical sign. - If the magnitude is too small to represent, the result is a zero of - appropriate mathematical sign. - - For complex floating-point operands, multiplication is defined according to the - following table. For real components ``a`` and ``c`` and - imaginary components ``b`` and ``d``, - - +------------+----------------+-----------------+--------------------------+ - | | c | dj | c + dj | - +============+================+=================+==========================+ - | **a** | a * c | (a*d)j | (a*c) + (a*d)j | - +------------+----------------+-----------------+--------------------------+ - | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | - +------------+----------------+-----------------+--------------------------+ - | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | - +------------+----------------+-----------------+--------------------------+ - - In general, for complex floating-point operands, real-valued floating-point - special cases must independently apply to the real and imaginary component - operations involving real numbers as described in the above table. - - When ``a``, ``b``, ``c``, or ``d`` are all finite numbers - (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), - multiplication of complex floating-point operands should be computed - as if calculated according to the textbook formula for complex number multiplication - - .. math:: - (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j - - When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, - ``+infinity``, or ``-infinity``, + - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, + the result is ``True``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, + the result is ``True``. + - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, + and ``x1_i`` does not equal ``x2_i``, the result is ``True``. + - In the remaining cases, the result is ``False``. - - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, the result is implementation dependent. + For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, + ``c = real(x2_i)``, ``d = imag(x2_i)``, and - .. note:: - For complex floating-point operands, the results of special cases may be - implementation dependent depending on how an implementation chooses - to model complex numbers and complex infinity - (e.g., complex plane versus Riemann sphere). - For those implementations following C99 and its one-infinity model, - when at least one component is infinite, - even if the other component is ``NaN``, - the complex value is infinite, and the usual arithmetic - rules do not apply to complex-complex multiplication. - In the interest of performance, other implementations - may want to avoid the complex branching logic necessary - to implement the one-infinity model and choose to implement - all complex-complex multiplication according to the textbook formula. - Accordingly, special case behavior is unlikely - to be consistent across implementations. + - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. + - In the remaining cases, the result is the logical OR of + the equality comparison between the real values ``a`` and ``c`` + (real components) and between the real values ``b`` and ``d`` + (imaginary components), as described above for real-valued floating-point operands + (i.e., ``a != c OR b != d``). Parameters ---------- x1 first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` - (see :ref'`broadcasting`). Should have a numeric data type - + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. out - optional output array, for writing the array result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of ``bool``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.not_equal.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - Returns - ------- - ret - an array containing the element-wise products. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. + Functional Examples + ------------------ - Examples - -------- - With :code:`ivy.Array` inputs: + With :class:`ivy.Array` inputs: + + >>> x1 = ivy.array([1, 0, 1, 1]) + >>> x2 = ivy.array([1, 0, 0, -1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False, True, True]) + + >>> x1 = ivy.array([1, 0, 1, 0]) + >>> x2 = ivy.array([0, 1, 0, 1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([True, True, True, True]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.not_equal(x1, x2, out=x1) + >>> print(y) + ivy.array([1, 0, 0, 1]) + + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + + >>> x1 = ivy.native_array([1, 2]) + >>> x2 = ivy.array([1, 2]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False]) - >>> x1 = ivy.array([3., 5., 7.]) - >>> x2 = ivy.array([4., 6., 8.]) - >>> y = ivy.multiply(x1, x2) + >>> x1 = ivy.native_array([1, -1]) + >>> x2 = ivy.array([0, 1]) + >>> y = ivy.not_equal(x1, x2) >>> print(y) - ivy.array([12., 30., 56.]) + ivy.array([True, True]) - With :code:`ivy.NativeArray` inputs: + >>> x1 = ivy.native_array([1, -1, 1, -1]) + >>> x2 = ivy.native_array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) - >>> x1 = ivy.native_array([1., 3., 9.]) - >>> x2 = ivy.native_array([4., 7.2, 1.]) - >>> y = ivy.multiply(x1, x2) + >>> x1 = ivy.native_array([1, 2, 3, 4]) + >>> x2 = ivy.native_array([0, 2, 3, 4]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) >>> print(y) - ivy.array([ 4. , 21.6, 9. ]) + ivy.array([1., 0., 0., 0.]) - With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: + With :class:`ivy.Container` input: - >>> x1 = ivy.array([8., 6., 7.]) - >>> x2 = ivy.native_array([1., 2., 3.]) - >>> y = ivy.multiply(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) - ivy.array([ 8., 12., 21.]) - - With :code:`ivy.Container` inputs: + { + a: ivy.array([False, True, False]), + b: ivy.array([False, False, False]), + c: ivy.array([False, False, False]) + } - >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) - >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) - >>> y = ivy.multiply(x1, x2) + >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1.0, 2.0, 4.0])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.native_array([1.1, 2.1, 3.1]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) { - a: ivy.array([12.,12.,24.]), - b: ivy.array([9.,3.,10.]) + a: ivy.array([True, True, True]), + b: ivy.array([True, True, True]), + c: ivy.array([False, False, False]) } - With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) - >>> x2 = ivy.array([1.,2.,3.]) - >>> y = ivy.multiply(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 3, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 4, 5])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) { - a: ivy.array([3.,8.,15.]), - b: ivy.array([2.,4.,3.]) + a: ivy.array([False, False, False]), + b: ivy.array([False, True, False]) + } + + >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), + ... b=ivy.array([1, 4, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), + ... b=ivy.array([1.0, 4.0, 5.0])) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + { + a: ivy.array([False, False, False]), + b: ivy.array([False, False, False]) } """ - return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) + return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def nan_to_num( - x: Union[ivy.Array, ivy.NativeArray], +def positive( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Replace NaN with zero and infinity with large finite numbers (default behaviour) or - with the numbers defined by the user using the nan, posinf and/or neginf keywords. + Return a new array with the positive value of each element in ``x``. Parameters ---------- x - Array input. - copy - Whether to create a copy of x (True) or to replace values in-place (False). - The in-place operation only occurs if casting to an array does not require - a copy. Default is True. - nan - Value to be used to fill NaN values. If no value is passed then NaN values - will be replaced with 0.0. - posinf - Value to be used to fill positive infinity values. If no value is passed - then positive infinity values will be replaced with a very large number. - neginf - Value to be used to fill negative infinity values. - If no value is passed then negative infinity values - will be replaced with a very small (or negative) number. + Input array. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Array with the non-finite values replaced. - If copy is False, this may be x itself. + A new array with the positive value of each element in ``x``. - Examples - -------- - >>> x = ivy.array([1, 2, 3, nan]) - >>> ivy.nan_to_num(x) - ivy.array([1., 1., 3., 0.0]) - >>> x = ivy.array([1, 2, 3, inf]) - >>> ivy.nan_to_num(x, posinf=5e+100) - ivy.array([1., 2., 3., 5e+100]) + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3 ,5, 7]) + >>> y = ivy.positive(x) + >>> print(y) + ivy.array([2, 3, 5, 7]) + + >>> x = ivy.array([0, -1, -0.5, 2, 3]) + >>> y = ivy.zeros(5) + >>> ivy.positive(x, out=y) + >>> print(y) + ivy.array([0., -1., -0.5, 2., 3.]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.positive(x,out=x) + >>> print(x) + ivy.array([[ 1.1, 2.2, 3.3], + [-4.4, -5.5, -6.6]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.positive(x) + >>> print(y) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., -5.]) + } """ - return ivy.current_backend(x).nan_to_num( - x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out - ) + return ivy.current_backend(x).positive(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def negative( - x: Union[float, ivy.Array, ivy.NativeArray], +def pow( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the negative value of each element in ``x``. + Calculate an implementation-dependent approximation of exponentiation by raising + each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` + (the exponent), where ``x2_i`` is the corresponding element of the input array + ``x2``. .. note:: - For signed integer data types, the numerical negative of - the minimum representable integer is implementation-dependent. + If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when + ``x2_i`` is negative (i.e., less than zero) is unspecified and thus + implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a + floating-point data type, behavior is implementation-dependent (type promotion + between data type "kinds" (integer versus floating-point) is unspecified). + + **Special cases** + + For floating-point operands, + + - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result + is ``+infinity``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result + is ``+0``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. + - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is + ``+0``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is + ``+0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an + odd integer value, the result is ``-infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not + an odd integer value, the result is ``+infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an + odd integer value, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer + value, the result is ``-infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+infinity``. + - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite + number, and ``x2_i`` is not an integer value, the result is ``NaN``. + + For complex floating-point operands, special cases should be handled + as if the operation is implemented as ``exp(x2*log(x1))``. .. note:: - If ``x`` has a complex floating-point data type, - both the real and imaginary components for each ``x_i`` - must be negated (a result which follows from the rules of - complex number multiplication). + Conforming implementations are allowed to treat special cases involving + complex floating-point operands more carefully than as described + in this specification. Parameters - ---------- - x - Input array. + ---------- + x1 + first input array whose elements correspond to the exponentiation base. Should + have a numeric data type. + x2 + second input array whose elements correspond to the exponentiation exponent. + Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric + data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5351,103 +5285,80 @@ def negative( Returns ------- ret - A new array with the negative value of each element in ``x``. + an array containing the element-wise results. The returned array must have a + data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.pow.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0,1,1,2]) - >>> y = ivy.negative(x) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.pow(x, 3) >>> print(y) - ivy.array([ 0, -1, -1, -2]) + ivy.array([1, 8, 27]) - >>> x = ivy.array([0,-1,-0.5,2,3]) - >>> y = ivy.zeros(5) - >>> ivy.negative(x, out=y) + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.pow(x, 2, out=y) >>> print(y) - ivy.array([-0. , 1. , 0.5, -2. , -3. ]) + ivy.array([2.25, 0.64, 0.09]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.negative(x,out=x) + >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) + >>> ivy.pow(x, 2.3, out=x) >>> print(x) - ivy.array([[-1.1, -2.2, -3.3], - [4.4, 5.5, 6.6]]) + ivy.array([[ 1.52095687, 4.92457771, 13.49372482], + [ 1. , 8.22738838, 156.5877228 ]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.negative(x) + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.pow(x, 3) >>> print(y) { - a: ivy.array([-0., -1., -2.]), - b: ivy.array([-3., -4., 5.]) + a:ivy.array([0,1]), + b:ivy.array([8,27]) } """ - return ivy.current_backend(x).negative(x, out=out) + return ivy.current_backend(x1, x2).pow(x1, x2, out=out) + + +pow.unsupported_gradients = {"torch": ["float16"]} @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def not_equal( - x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], - x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], +def real( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. - - **Special Cases** - - For real-valued floating-point operands, - - - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, - the result is ``True``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, - the result is ``True``. - - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, - and ``x1_i`` does not equal ``x2_i``, the result is ``True``. - - In the remaining cases, the result is ``False``. - - For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, - ``c = real(x2_i)``, ``d = imag(x2_i)``, and - - - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. - - In the remaining cases, the result is the logical OR of - the equality comparison between the real values ``a`` and ``c`` - (real components) and between the real values ``b`` and ``d`` - (imaginary components), as described above for real-valued floating-point operands - (i.e., ``a != c OR b != d``). + Test each element ``x_i`` of the input array ``x`` to take only real part from it. + Returns a float array, where it only contains . If element has complex type with + zero complex part, the return value will be that element, else it only returns real + part. Parameters ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. + x + input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5455,133 +5366,175 @@ def not_equal( Returns ------- ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. + an array containing test results. An element ``out_i`` is + ``real number`` if ``x_i`` contain real number part only + and if it is ``real number with complex part also`` then it + returns the real number part. + The returned array must have a floating-point data type with the + same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, + the returned array must have the floating-point precision of ``float32``). + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + Examples + -------- + With :class:`ivy.Array` inputs: - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = ivy.array([[[1.1], [2], [-6.3]]]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([[[1.1], [2.], [-6.3]]]) - Functional Examples - ------------------ + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([4.2, 0., 7.]) - With :class:`ivy.Array` inputs: + With :class:`ivy.Container` input: - >>> x1 = ivy.array([1, 0, 1, 1]) - >>> x2 = ivy.array([1, 0, 0, -1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False, True, True]) + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ + b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.real(x) + >>> print(z) + { + a: ivy.array([-6.7, 0.314, 1.23]), + b: ivy.array([0., 5.32, 3.001]) + } + """ + return ivy.current_backend(x).real(x, out=out) - >>> x1 = ivy.array([1, 0, 1, 0]) - >>> x2 = ivy.array([0, 1, 0, 1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([True, True, True, True]) - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def remainder( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + modulus: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the remainder of division for each element ``x1_i`` of the input array ``x1`` + and the respective element ``x2_i`` of the input array ``x2``. - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.not_equal(x1, x2, out=x1) - >>> print(y) - ivy.array([1, 0, 0, 1]) + .. note:: + This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For + input arrays which promote to an integer data type, the result of division by + zero is unspecified and thus implementation-defined. In general, similar to + Python’s ``%`` operator, this function is not recommended for floating-point + operands as semantics do not follow IEEE 754. That this function is specified + to accept floating-point operands is primarily for reasons of backward + compatibility. - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + **Special Cases** - >>> x1 = ivy.native_array([1, 2]) - >>> x2 = ivy.array([1, 2]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False]) + For floating-point operands, - >>> x1 = ivy.native_array([1, -1]) - >>> x2 = ivy.array([0, 1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([True, True]) + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either + ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, + the result is ``NaN``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x2_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x2_i``. (note: this results matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - In the remaining cases, the result must match that of the Python ``%`` operator. - >>> x1 = ivy.native_array([1, -1, 1, -1]) - >>> x2 = ivy.native_array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) + Parameters + ---------- + x1 + dividend input array. Should have a numeric data type. + x2 + divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). + Should have a numeric data type. + modulus + whether to compute the modulus instead of the remainder. Default is ``True``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. Each element-wise result must have + the same sign as the respective element ``x2_i``. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. - >>> x1 = ivy.native_array([1, 2, 3, 4]) - >>> x2 = ivy.native_array([0, 2, 3, 4]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 0.]) - With :class:`ivy.Container` input: + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - { - a: ivy.array([False, True, False]), - b: ivy.array([False, False, False]), - c: ivy.array([False, False, False]) - } + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments - >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1.0, 2.0, 4.0])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.native_array([1.1, 2.1, 3.1]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x1 = ivy.array([2., 5., 15.]) + >>> x2 = ivy.array([3., 2., 4.]) + >>> y = ivy.remainder(x1, x2) >>> print(y) - { - a: ivy.array([True, True, True]), - b: ivy.array([True, True, True]), - c: ivy.array([False, False, False]) - } + ivy.array([2., 1., 3.]) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 3, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 4, 5])) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.array([23., 1., 6.]) + >>> x2 = ivy.native_array([11., 2., 4.]) + >>> y = ivy.remainder(x1, x2) >>> print(y) - { - a: ivy.array([False, False, False]), - b: ivy.array([False, True, False]) - } + ivy.array([1., 1., 2.]) - >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), - ... b=ivy.array([1, 4, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), - ... b=ivy.array([1.0, 4.0, 5.0])) - >>> y = ivy.not_equal(x1, x2) + With :class:`ivy.Container` inputs: + + >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) + >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) + >>> y = ivy.remainder(x1, x2) >>> print(y) { - a: ivy.array([False, False, False]), - b: ivy.array([False, False, False]) + a: ivy.array([0., 0., 1.]), + b: ivy.array([0., 2., 1.]) } """ - return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) + return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) @handle_exceptions @@ -5592,19 +5545,55 @@ def not_equal( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def positive( - x: Union[float, ivy.Array, ivy.NativeArray], +def round( + x: Union[ivy.Array, ivy.NativeArray], /, *, + decimals: Optional[int] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the positive value of each element in ``x``. + Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued + number. + + .. note:: + For complex floating-point operands, real and imaginary components + must be independently rounded to the nearest integer-valued number. + + Rounded real and imaginary components must be equal + to their equivalent rounded real-valued floating-point + counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` + must equal ``round(real(x)))`` and ``imag(round(x))`` must equal + ``round(imag(x))``). + + **Special cases** + + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + + For floating-point operands, + + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If two integers are equally close to ``x_i``, the result is + the even integer closest to ``x_i``. + + .. note:: + For complex floating-point operands, the following special + cases apply to real and imaginary components independently + (e.g., if ``real(x_i)`` is ``NaN``, the rounded + real component is ``NaN``). + + - If ``x_i`` is already integer-valued, the result is ``x_i``. Parameters ---------- x - Input array. + input array containing elements to round. + decimals + number of decimal places to round to. Default is ``0``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5612,145 +5601,121 @@ def positive( Returns ------- ret - A new array with the positive value of each element in ``x``. + An array of the same shape and type as x, with the elements rounded to integers. + + Note: PyTorch supports an additional argument :code:`decimals` for the + `round function `_. + It has been deliberately omitted here due to the imprecise + nature of the argument in :code:`torch.round`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.round.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - Functional Examples - ------------------- + instances in place of any of the arguments. + Examples + -------- With :class:`ivy.Array` input: - >>> x = ivy.array([2, 3 ,5, 7]) - >>> y = ivy.positive(x) + >>> x = ivy.array([1.2, 2.4, 3.6]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([2, 3, 5, 7]) + ivy.array([1.,2.,4.]) - >>> x = ivy.array([0, -1, -0.5, 2, 3]) - >>> y = ivy.zeros(5) - >>> ivy.positive(x, out=y) + >>> x = ivy.array([-0, 5, 4.5]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([0., -1., -0.5, 2., 3.]) + ivy.array([0.,5.,4.]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.positive(x,out=x) + >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) + >>> y = ivy.zeros(4) + >>> ivy.round(x, out=y) + >>> print(y) + ivy.array([2.,2.,15.,-5.]) + + >>> x = ivy.array([[0, 5.433, -343.3, 1.5], + ... [-5.5, 44.2, 11.5, 12.01]]) + >>> ivy.round(x, out=x) >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-4.4, -5.5, -6.6]]) + ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.positive(x) + >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), + ... b=ivy.array([-300.9, -527.3, 4.5])) + >>> y = ivy.round(x) >>> print(y) { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., -5.]) + a:ivy.array([4.,9.,7.,0.]), + b:ivy.array([-301.,-527.,4.]) } """ - return ivy.current_backend(x).positive(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def pow( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate an implementation-dependent approximation of exponentiation by raising - each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` - (the exponent), where ``x2_i`` is the corresponding element of the input array - ``x2``. - - .. note:: - If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when - ``x2_i`` is negative (i.e., less than zero) is unspecified and thus - implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a - floating-point data type, behavior is implementation-dependent (type promotion - between data type "kinds" (integer versus floating-point) is unspecified). - - **Special cases** - - For floating-point operands, - - - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result - is ``+infinity``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result - is ``+0``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. - - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is - ``+0``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is - ``+0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an - odd integer value, the result is ``-infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not - an odd integer value, the result is ``+infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an - odd integer value, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer - value, the result is ``-infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+infinity``. - - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite - number, and ``x2_i`` is not an integer value, the result is ``NaN``. + return ivy.current_backend(x).round(x, decimals=decimals, out=out) - For complex floating-point operands, special cases should be handled - as if the operation is implemented as ``exp(x2*log(x1))``. - .. note:: - Conforming implementations are allowed to treat special cases involving - complex floating-point operands more carefully than as described - in this specification. +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def sign( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + np_variant: Optional[bool] = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Return an indication of the sign of a number for each element ``x_i`` of the input + array ``x``. + + The sign function (also known as the **signum function**) + of a number :math:`x_{i}` is defined as + + .. math:: + \operatorname{sign}(x_i) = \begin{cases} + 0 & \textrm{if } x_i = 0 \\ + \frac{x}{|x|} & \textrm{otherwise} + \end{cases} + + where :math:`|x_i|` is the absolute value of :math:`x_i`. + + **Special cases** + + - If ``x_i`` is less than ``0``, the result is ``-1``. + - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. + - If ``x_i`` is greater than ``0``, the result is ``+1``. + - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``-0`` or ``+0`` and ``b`` is + either ``-0`` or ``+0``, the result is ``0 + 0j``. + - If ``a`` is ``NaN`` or ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, special cases must be handled + according to the rules of complex number division. Parameters ---------- - x1 - first input array whose elements correspond to the exponentiation base. Should - have a numeric data type. - x2 - second input array whose elements correspond to the exponentiation exponent. - Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric - data type. + x + input array. Should have a numeric data type. + np_variant + Handles complex numbers like numpy does If ``True``, + ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. + otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, + otherwise y = 0.`` + out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5758,54 +5723,53 @@ def pow( Returns ------- ret - an array containing the element-wise results. The returned array must have a - data type determined by :ref:`type-promotion`. + an array containing the evaluated result for each element in ``x``. The returned + array must have the same data type as ``x``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sign.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.pow(x, 3) + >>> x = ivy.array([8.3, -0, 6.8, 0.07]) + >>> y = ivy.sign(x) >>> print(y) - ivy.array([1, 8, 27]) + ivy.array([1., 0., 1., 1.]) - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.pow(x, 2, out=y) + >>> x = ivy.array([[5.78, -4., -6.9, 0], + ... [-.4, 0.5, 8, -0.01]]) + >>> y = ivy.sign(x) >>> print(y) - ivy.array([2.25, 0.64, 0.09]) - - >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) - >>> ivy.pow(x, 2.3, out=x) - >>> print(x) - ivy.array([[ 1.52095687, 4.92457771, 13.49372482], - [ 1. , 8.22738838, 156.5877228 ]]) + ivy.array([[ 1., -1., -1., 0.], + [-1., 1., 1., -1.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.pow(x, 3) + >>> x = ivy.Container(a=ivy.array([0., -0.]), + ... b=ivy.array([1.46, 5.9, -0.0]), + ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) + >>> y = ivy.sign(x) >>> print(y) { - a:ivy.array([0,1]), - b:ivy.array([8,27]) + a: ivy.array([0., 0.]), + b: ivy.array([1., 1., 0.]), + c: ivy.array([-1., -1., -1., 1.]) } """ - return ivy.current_backend(x1, x2).pow(x1, x2, out=out) + return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -5813,19 +5777,37 @@ def pow( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def rad2deg( +def sin( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Convert the input from radians to degrees. + r""" + Calculate an implementation-dependent approximation to the sine, having domain + ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of + the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. + + .. note:: + The sine is an entire function on the complex plane and has no branch cuts. + + **Special cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + + For complex floating-point operands, special cases + must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. Parameters ---------- x - input array whose elements are each expressed in radians. + input array whose elements are each expressed in radians. Should have a + floating-point data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5833,56 +5815,53 @@ def rad2deg( Returns ------- ret - an array with each element in ``x`` converted from radians to degrees. + an array containing the sine of each element in ``x``. The returned array must + have a floating-point data type determined by :ref:`type-promotion`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) - >>> y=ivy.rad2deg(x) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.sin(x) >>> print(y) - ivy.array([ 0., 90., 180., 270., 360.]) + ivy.array([0., 0.841, 0.909]) - >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.rad2deg(x,out=y) + >>> x = ivy.array([0., 1.2, -2.3, 3.6]) + >>> y = ivy.zeros(4) + >>> ivy.sin(x, out=y) >>> print(y) - ivy.array([ 0. , -1.5, -50. , nan]) + ivy.array([0., 0.932, -0.746, -0.443]) - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.rad2deg(x, out=x) + >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) + >>> ivy.sin(x, out=x) >>> print(x) - ivy.array([[ 63., 126., 189.], - [-252., -315., -378.]]) - - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.rad2deg(x,out=y) - >>> print(y) - ivy.array([ 0., 1150., nan]) + ivy.array([[0.841, 0.909, 0.141], + [0.757, 0.959, 0.279]]) With :class:`ivy.Container` input: - >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), - ... b=ivy.array([0., 1., 2., 3., 4.])) - >>> y=ivy.rad2deg(x) - >>> print(y) - { - a: ivy.array([0., 1150., -2890., nan]), - b: ivy.array([0., 57.3, 115., 172., 229.]) - } - - >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), - ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) - >>> y=ivy.rad2deg(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), + ... b=ivy.array([-4., -5., -6., -7.])) + >>> y = ivy.sin(x) >>> print(y) { - a: ivy.array([0., 573., 10300., 487., 344.]), - b: ivy.array([0., -85.9, 28.6, nan]) + a: ivy.array([0., 0.841, 0.909, 0.141]), + b: ivy.array([0.757, 0.959, 0.279, -0.657]) } """ - return ivy.current_backend(x).rad2deg(x, out=out) + return ivy.current_backend(x).sin(x, out=out) @handle_exceptions @@ -5891,23 +5870,76 @@ def rad2deg( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def real( +def sinh( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Test each element ``x_i`` of the input array ``x`` to take only real part from it. - Returns a float array, where it only contains . If element has complex type with - zero complex part, the return value will be that element, else it only returns real - part. + r""" + Calculate an implementation-dependent approximation to the hyperbolic sine, having + domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each + element ``x_i`` of the input array ``x``. + + .. math:: + \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} + + .. note:: + The hyperbolic sine is an entire function in the + complex plane and has no branch cuts. + The function is periodic, with period + :math:`2\pi j`, with respect to the imaginary component. + + **Special cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + .. note:: + For complex floating-point operands, ``sinh(conj(x))`` + must equal ``conj(sinh(x))``. + + - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. + - If ``a`` is ``+0`` and ``b`` is ``+infinity``, + the result is ``0 + NaN j`` (sign of the real component is unspecified). + - If ``a`` is ``+0`` and ``b`` is ``NaN``, + the result is ``0 + NaN j`` (sign of the real component is unspecified). + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is ``+0``, + the result is ``+infinity + 0j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, + the result is ``+infinity * cis(b)``. + - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, + the result is ``infinity + NaN j`` (sign of the real component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``infinity + NaN j`` + (sign of the real component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + + where ``cis(v)`` is ``cos(v) + sin(v)*1j``. Parameters ---------- x - input array. + input array whose elements each represent a hyperbolic angle. Should have a + floating-point data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5915,45 +5947,45 @@ def real( Returns ------- ret - an array containing test results. An element ``out_i`` is - ``real number`` if ``x_i`` contain real number part only - and if it is ``real number with complex part also`` then it - returns the real number part. - The returned array must have a floating-point data type with the - same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, - the returned array must have the floating-point precision of ``float32``). + an array containing the hyperbolic sine of each element in ``x``. The returned + array must have a floating-point data type determined by :ref:`type-promotion`. - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input: - >>> x = ivy.array([[[1.1], [2], [-6.3]]]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([[[1.1], [2.], [-6.3]]]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.sinh(x) + >>> print(y) + ivy.array([1.18, 3.63, 10.]) - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([4.2, 0., 7.]) + >>> x = ivy.array([0.23, 3., -1.2]) + >>> ivy.sinh(x, out=x) + >>> print(x) + ivy.array([0.232, 10., -1.51]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ - b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.real(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) + >>> y = ivy.sinh(x) + >>> print(y) { - a: ivy.array([-6.7, 0.314, 1.23]), - b: ivy.array([0., 5.32, 3.001]) + a: ivy.array([0.232, -0.253, 1.18]), + b: ivy.array([10., -27.3, 1.62]) } """ - return ivy.current_backend(x).real(x, out=out) + return ivy.current_backend(x).sinh(x, out=out) @handle_exceptions @@ -5964,114 +5996,77 @@ def real( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def reciprocal( - x: Union[float, ivy.Array, ivy.NativeArray], +def sqrt( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Return a new array with the reciprocal of each element in ``x``. - - Parameters - ---------- - x - Input array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + r""" + Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, + +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, + each result must be indistinguishable from the infinitely precise result (as + required by IEEE 754). - Returns - ------- - ret - A new array with the positive value of each element in ``x``. + .. note:: + After rounding, each result must be indistinguishable + from the infinitely precise result (as required by IEEE 754). - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.reciprocal(x) - >>> print(y) - ivy.array([1. , 0.5 , 0.33333333]) - """ - return ivy.current_backend(x).reciprocal(x, out=out) + .. note:: + For complex floating-point operands, ``sqrt(conj(x))`` + must equal ``conj(sqrt(x))``. + .. note:: + By convention, the branch cut of the square root is + the negative real axis :math:`(-\infty, 0)`. -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def remainder( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - modulus: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the remainder of division for each element ``x1_i`` of the input array ``x1`` - and the respective element ``x2_i`` of the input array ``x2``. + The square root is a continuous function from above + the branch cut, taking into account the sign of the imaginary component. - .. note:: - This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For - input arrays which promote to an integer data type, the result of division by - zero is unspecified and thus implementation-defined. In general, similar to - Python’s ``%`` operator, this function is not recommended for floating-point - operands as semantics do not follow IEEE 754. That this function is specified - to accept floating-point operands is primarily for reasons of backward - compatibility. + Accordingly, for complex arguments, the function returns + the square root in the range of the right half-plane, + including the imaginary axis (i.e., the plane defined by + :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` + along the imaginary axis). - **Special Cases** + **Special cases** For floating-point operands, - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either - ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, - the result is ``NaN``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x2_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x2_i``. (note: this results matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - In the remaining cases, the result must match that of the Python ``%`` operator. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is less than ``0``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, + the result is ``+0 + 0j``. + - If ``a`` is any value (including ``NaN``) and ``b`` is + ``+infinity``, the result is ``+infinity + infinity j``. + - If ``a`` is a finite number and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - If ``a`` ``-infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. + - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, + the result is ``NaN + infinity j`` + (sign of the imaginary component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``+infinity + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is any value, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + Parameters ---------- - x1 - dividend input array. Should have a numeric data type. - x2 - divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). - Should have a numeric data type. - modulus - whether to compute the modulus instead of the remainder. Default is ``True``. + x + input array. Should have a floating-point data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6079,15 +6074,14 @@ def remainder( Returns ------- ret - an array containing the element-wise results. Each element-wise result must have - the same sign as the respective element ``x2_i``. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. + an array containing the square root of each element in ``x``. The returned array + must have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sqrt.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6096,34 +6090,34 @@ def remainder( Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input: - >>> x1 = ivy.array([2., 5., 15.]) - >>> x2 = ivy.array([3., 2., 4.]) - >>> y = ivy.remainder(x1, x2) + >>> x = ivy.array([0, 4., 8.]) + >>> y = ivy.sqrt(x) >>> print(y) - ivy.array([2., 1., 3.]) + ivy.array([0., 2., 2.83]) - With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + >>> x = ivy.array([1, 2., 4.]) + >>> y = ivy.zeros(3) + >>> ivy.sqrt(x, out=y) + ivy.array([1., 1.41, 2.]) - >>> x1 = ivy.array([23., 1., 6.]) - >>> x2 = ivy.native_array([11., 2., 4.]) - >>> y = ivy.remainder(x1, x2) - >>> print(y) - ivy.array([1., 1., 2.]) + >>> X = ivy.array([40., 24., 100.]) + >>> ivy.sqrt(x, out=x) + >>> ivy.array([6.32455532, 4.89897949, 10.]) - With :class:`ivy.Container` inputs: + With :class:`ivy.Container` input: - >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) - >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) - >>> y = ivy.remainder(x1, x2) + >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa + >>> y = ivy.sqrt(x) >>> print(y) { - a: ivy.array([0., 0., 1.]), - b: ivy.array([0., 2., 1.]) + a: ivy.array([6.63, 7.48, 13.]), + b: ivy.array([[7., 1.], + [0., 4.47]]) } """ - return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) + return ivy.current_backend(x).sqrt(x, out=out) @handle_exceptions @@ -6134,55 +6128,19 @@ def remainder( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def round( +def square( x: Union[ivy.Array, ivy.NativeArray], /, *, - decimals: Optional[int] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued - number. - - .. note:: - For complex floating-point operands, real and imaginary components - must be independently rounded to the nearest integer-valued number. - - Rounded real and imaginary components must be equal - to their equivalent rounded real-valued floating-point - counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` - must equal ``round(real(x)))`` and ``imag(round(x))`` must equal - ``round(imag(x))``). - - **Special cases** - - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - - For floating-point operands, - - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If two integers are equally close to ``x_i``, the result is - the even integer closest to ``x_i``. - - .. note:: - For complex floating-point operands, the following special - cases apply to real and imaginary components independently - (e.g., if ``real(x_i)`` is ``NaN``, the rounded - real component is ``NaN``). - - - If ``x_i`` is already integer-valued, the result is ``x_i``. + Each element ``x_i`` of the input array ``x``. Parameters ---------- x - input array containing elements to round. - decimals - number of decimal places to round to. Default is ``0``. + Input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6190,18 +6148,13 @@ def round( Returns ------- ret - An array of the same shape and type as x, with the elements rounded to integers. - + an array containing the evaluated result for each element in ``x``. - Note: PyTorch supports an additional argument :code:`decimals` for the - `round function `_. - It has been deliberately omitted here due to the imprecise - nature of the argument in :code:`torch.round`. - This function conforms to the `Array API Standard + This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.square.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6212,99 +6165,63 @@ def round( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1.2, 2.4, 3.6]) - >>> y = ivy.round(x) - >>> print(y) - ivy.array([1.,2.,4.]) - - >>> x = ivy.array([-0, 5, 4.5]) - >>> y = ivy.round(x) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.square(x) >>> print(y) - ivy.array([0.,5.,4.]) + ivy.array([1, 4, 9]) - >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) - >>> y = ivy.zeros(4) - >>> ivy.round(x, out=y) + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.square(x, out=y) >>> print(y) - ivy.array([2.,2.,15.,-5.]) + ivy.array([2.25, 0.64, 0.09]) - >>> x = ivy.array([[0, 5.433, -343.3, 1.5], - ... [-5.5, 44.2, 11.5, 12.01]]) - >>> ivy.round(x, out=x) + >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) + >>> ivy.square(x, out=x) >>> print(x) - ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) + ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), - ... b=ivy.array([-300.9, -527.3, 4.5])) - >>> y = ivy.round(x) + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.square(x) >>> print(y) { - a:ivy.array([4.,9.,7.,0.]), - b:ivy.array([-301.,-527.,4.]) + a:ivy.array([0,1]), + b:ivy.array([4,9]) } """ - return ivy.current_backend(x).round(x, decimals=decimals, out=out) + return ivy.current_backend(x).square(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sign( - x: Union[ivy.Array, ivy.NativeArray], +def subtract( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, - np_variant: Optional[bool] = True, + alpha: Optional[Union[int, float]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Return an indication of the sign of a number for each element ``x_i`` of the input - array ``x``. - - The sign function (also known as the **signum function**) - of a number :math:`x_{i}` is defined as - - .. math:: - \operatorname{sign}(x_i) = \begin{cases} - 0 & \textrm{if } x_i = 0 \\ - \frac{x}{|x|} & \textrm{otherwise} - \end{cases} - - where :math:`|x_i|` is the absolute value of :math:`x_i`. - - **Special cases** - - - If ``x_i`` is less than ``0``, the result is ``-1``. - - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. - - If ``x_i`` is greater than ``0``, the result is ``+1``. - - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` - - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and - - - If ``a`` is either ``-0`` or ``+0`` and ``b`` is - either ``-0`` or ``+0``, the result is ``0 + 0j``. - - If ``a`` is ``NaN`` or ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, special cases must be handled - according to the rules of complex number division. + """ + Calculate the difference for each element ``x1_i`` of the input array ``x1`` with + the respective element ``x2_i`` of the input array ``x2``. Parameters ---------- - x - input array. Should have a numeric data type. - np_variant - Handles complex numbers like numpy does If ``True``, - ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. - otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, - otherwise y = 0.`` - + x1 + first input array. Should have a numeric data type. + x2 + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. + alpha + optional scalar multiplier for ``x2``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6312,14 +6229,13 @@ def sign( Returns ------- ret - an array containing the evaluated result for each element in ``x``. The returned - array must have the same data type as ``x``. + an array containing the element-wise differences. - This function conforms to the `Array API Standard + This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.subtract.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6328,34 +6244,19 @@ def sign( Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([8.3, -0, 6.8, 0.07]) - >>> y = ivy.sign(x) - >>> print(y) - ivy.array([1., 0., 1., 1.]) - - >>> x = ivy.array([[5.78, -4., -6.9, 0], - ... [-.4, 0.5, 8, -0.01]]) - >>> y = ivy.sign(x) - >>> print(y) - ivy.array([[ 1., -1., -1., 0.], - [-1., 1., 1., -1.]]) - - With :class:`ivy.Container` input: + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y) + >>> print(z) + ivy.array([ 1, 5, -3]) - >>> x = ivy.Container(a=ivy.array([0., -0.]), - ... b=ivy.array([1.46, 5.9, -0.0]), - ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) - >>> y = ivy.sign(x) - >>> print(y) - { - a: ivy.array([0., 0.]), - b: ivy.array([1., 1., 0.]), - c: ivy.array([-1., -1., -1., 1.]) - } + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y, alpha=2) + >>> print(z) + ivy.array([-1, 4, -9]) """ - return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) + return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) @handle_exceptions @@ -6366,19 +6267,28 @@ def sign( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sin( +def tan( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the sine, having domain - ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of - the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. - + r"""Calculate an implementation-dependent approximation to the tangent, having + domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each + element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be + expressed in radians. .. note:: - The sine is an entire function on the complex plane and has no branch cuts. + Tangent is an analytical function on the complex plane + and has no branch cuts. The function is periodic, + with period :math:`\pi j`, with respect to the real + component and has first order poles along the real + line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. + However, IEEE 754 binary floating-point representation + cannot represent the value :math:`\pi / 2` exactly, and, + thus, no argument value is possible for + which a pole error occurs. + + where :math:`{tanh}` is the hyperbolic tangent. **Special cases** @@ -6389,68 +6299,68 @@ def sin( - If ``x_i`` is ``-0``, the result is ``-0``. - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - For complex floating-point operands, special cases - must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. + For complex floating-point operands, special cases must + be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. Parameters ---------- x - input array whose elements are each expressed in radians. Should have a + input array whose elements are expressed in radians. Should have a floating-point data type. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - an array containing the sine of each element in ``x``. The returned array must - have a floating-point data type determined by :ref:`type-promotion`. + an array containing the tangent of each element in ``x``. The return must have a + floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.tan.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.sin(x) + >>> y = ivy.tan(x) >>> print(y) - ivy.array([0., 0.841, 0.909]) + ivy.array([0., 1.56, -2.19]) - >>> x = ivy.array([0., 1.2, -2.3, 3.6]) - >>> y = ivy.zeros(4) - >>> ivy.sin(x, out=y) + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tan(x, out=y) >>> print(y) - ivy.array([0., 0.932, -0.746, -0.443]) + ivy.array([0.546, -0.842, -0.916]) - >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) - >>> ivy.sin(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tan(x, out=x) >>> print(x) - ivy.array([[0.841, 0.909, 0.141], - [0.757, 0.959, 0.279]]) + ivy.array([[1.96, -1.37, 0.16], + [-3.1, 0.996, -0.328]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), - ... b=ivy.array([-4., -5., -6., -7.])) - >>> y = ivy.sin(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.tan(x) >>> print(y) { - a: ivy.array([0., 0.841, 0.909, 0.141]), - b: ivy.array([0.757, 0.959, 0.279, -0.657]) + a: ivy.array([0., 1.56, -2.19]), + b: ivy.array([-0.143, 1.16, -3.38]) } """ - return ivy.current_backend(x).sin(x, out=out) + return ivy.current_backend(x).tan(x, out=out) @handle_exceptions @@ -6461,25 +6371,18 @@ def sin( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sinh( +@handle_complex_input +def tanh( x: Union[ivy.Array, ivy.NativeArray], /, *, + complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the hyperbolic sine, having - domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each - element ``x_i`` of the input array ``x``. - - .. math:: - \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} - - .. note:: - The hyperbolic sine is an entire function in the - complex plane and has no branch cuts. - The function is periodic, with period - :math:`2\pi j`, with respect to the imaginary component. + """ + Calculate an implementation-dependent approximation to the hyperbolic tangent, + having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element + ``x_i`` of the input array ``x``. **Special cases** @@ -6488,62 +6391,80 @@ def sinh( - If ``x_i`` is ``NaN``, the result is ``NaN``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+infinity``, the result is ``+1``. + - If ``x_i`` is ``-infinity``, the result is ``-1``. For complex floating-point operands, let ``a = real(x_i)``, ``b = imag(x_i)``, and .. note:: - For complex floating-point operands, ``sinh(conj(x))`` - must equal ``conj(sinh(x))``. + For complex floating-point operands, ``tanh(conj(x))`` + must equal ``conj(tanh(x))``. - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. + - If ``a`` is a nonzero finite number and ``b`` is + ``+infinity``, the result is ``NaN + NaN j``. - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). + the result is ``+0 + NaN j``. + - If ``a`` is a nonzero finite number and ``b`` + is ``NaN``, the result is ``NaN + NaN j``. - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is ``+0``, - the result is ``+infinity + 0j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, - the result is ``+infinity * cis(b)``. + the result is ``+0 + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``infinity + NaN j`` (sign of the real component is unspecified). + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``infinity + NaN j`` - (sign of the real component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, + the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero number, the result is ``NaN + NaN j``. - If ``a`` is ``NaN`` and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - where ``cis(v)`` is ``cos(v) + sin(v)*1j``. + .. warning:: + For historical reasons stemming from the C standard, + array libraries may not return the expected + result when ``a`` is ``+0`` and ``b`` is either + ``+infinity`` or ``NaN``. The result should be + ``+0 + NaN j`` in both cases; however, for libraries + compiled against older C versions, the result may be + ``NaN + NaN j``. + + Array libraries are not required to patch these older + C versions, and, thus, users are advised that results + may vary across array library implementations for + these special cases. + Parameters ---------- x input array whose elements each represent a hyperbolic angle. Should have a - floating-point data type. + real-valued floating-point data + type. + complex_mode + optional specifier for how to handle complex data types. See + ``ivy.func_wrapper.handle_complex_input`` for more detail. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - an array containing the hyperbolic sine of each element in ``x``. The returned - array must have a floating-point data type determined by :ref:`type-promotion`. + an array containing the hyperbolic tangent of each element in ``x``. + The returned array must have a real-valued floating-point data type + determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.tanh.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6554,108 +6475,131 @@ def sinh( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.sinh(x) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.tanh(x) >>> print(y) - ivy.array([1.18, 3.63, 10.]) + ivy.array([0., 0.762, 0.964]) - >>> x = ivy.array([0.23, 3., -1.2]) - >>> ivy.sinh(x, out=x) + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tanh(x, out=y) + >>> print(y) + ivy.array([0.462, -0.604, 0.984]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tanh(x, out=x) >>> print(x) - ivy.array([0.232, 10., -1.51]) + ivy.array([[0.8, 0.976, 0.997], + [-1., -1., -1.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) - >>> y = ivy.sinh(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.tanh(x) >>> print(y) { - a: ivy.array([0.232, -0.253, 1.18]), - b: ivy.array([10., -27.3, 1.62]) + a: ivy.array([0., 0.762, 0.964]), + b: ivy.array([0.995, 0.999, 1.]) } """ - return ivy.current_backend(x).sinh(x, out=out) + return ivy.current_backend(x).tanh(x, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def sqrt( - x: Union[ivy.Array, ivy.NativeArray], +def trapz( + y: ivy.Array, /, *, + x: Optional[ivy.Array] = None, + dx: float = 1.0, + axis: int = -1, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, - +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, - each result must be indistinguishable from the infinitely precise result (as - required by IEEE 754). + """ + Integrate along the given axis using the composite trapezoidal rule. - .. note:: - After rounding, each result must be indistinguishable - from the infinitely precise result (as required by IEEE 754). + If x is provided, the integration happens in sequence along its elements + - they are not sorted.. - .. note:: - For complex floating-point operands, ``sqrt(conj(x))`` - must equal ``conj(sqrt(x))``. + Parameters + ---------- + y + The array that should be integrated. + x + The sample points corresponding to the input array values. + If x is None, the sample points are assumed to be evenly spaced + dx apart. The default is None. + dx + The spacing between sample points when x is None. The default is 1. + axis + The axis along which to integrate. + out + optional output array, for writing the result to. - .. note:: - By convention, the branch cut of the square root is - the negative real axis :math:`(-\infty, 0)`. + Returns + ------- + ret + Definite integral of n-dimensional array as approximated along + a single axis by the trapezoidal rule. If the input array is a + 1-dimensional array, then the result is a float. If n is greater + than 1, then the result is an n-1 dimensional array. - The square root is a continuous function from above - the branch cut, taking into account the sign of the imaginary component. + Examples + -------- + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3]) + 4.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], x=[4, 6, 8]) + 8.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], dx=2) + 8.0 + """ + return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) - Accordingly, for complex arguments, the function returns - the square root in the range of the right half-plane, - including the imaginary axis (i.e., the plane defined by - :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` - along the imaginary axis). + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def trunc( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Round each element x_i of the input array x to the integer-valued number that is + closest to but no greater than x_i. **Special cases** + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + For floating-point operands, - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is less than ``0``, the result is ``NaN``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and - - - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, - the result is ``+0 + 0j``. - - If ``a`` is any value (including ``NaN``) and ``b`` is - ``+infinity``, the result is ``+infinity + infinity j``. - - If ``a`` is a finite number and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - If ``a`` ``-infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. - - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, - the result is ``NaN + infinity j`` - (sign of the imaginary component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``+infinity + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is any value, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - + - If ``x_i`` is ``NaN``, the result is ``NaN``. Parameters ---------- x - input array. Should have a floating-point data type. + input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6663,14 +6607,14 @@ def sqrt( Returns ------- ret - an array containing the square root of each element in ``x``. The returned array - must have a floating-point data type determined by :ref:`type-promotion`. + an array containing the rounded result for each element in ``x``. + The returned array must have the same data type as ``x``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.trunc.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6681,32 +6625,38 @@ def sqrt( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0, 4., 8.]) - >>> y = ivy.sqrt(x) + >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) + >>> y = ivy.trunc(x) >>> print(y) - ivy.array([0., 2., 2.83]) + ivy.array([-1., 0., 3., -0.]) - >>> x = ivy.array([1, 2., 4.]) - >>> y = ivy.zeros(3) - >>> ivy.sqrt(x, out=y) - ivy.array([1., 1.41, 2.]) + >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) + >>> ivy.trunc(x, out=x) + >>> print(x) + ivy.array([ 0., 7., -23., -0.]) - >>> X = ivy.array([40., 24., 100.]) - >>> ivy.sqrt(x, out=x) - >>> ivy.array([6.32455532, 4.89897949, 10.]) + >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) + >>> y = ivy.zeros([2,3]) + >>> ivy.trunc(x, out=y) + >>> print(y) + ivy.array([[ 0., -8., 0.], + [ 0., 0., 2.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa - >>> y = ivy.sqrt(x) + >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) + >>> y = ivy.trunc(x) >>> print(y) { - a: ivy.array([6.63, 7.48, 13.]), - b: ivy.array([[7., 1.], - [0., 4.47]]) + a: ivy.array([-0., 4., 1.]), + b: ivy.array([12., -3., 1.]) } """ - return ivy.current_backend(x).sqrt(x, out=out) + return ivy.current_backend(x).trunc(x, out=out) + + +# Extra # +# ------# @handle_exceptions @@ -6717,19 +6667,19 @@ def sqrt( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def square( +def erf( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Each element ``x_i`` of the input array ``x``. + Compute the Gauss error function of ``x`` element-wise. Parameters ---------- x - Input array. Should have a numeric data type. + Value to compute exponential for. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6737,50 +6687,15 @@ def square( Returns ------- ret - an array containing the evaluated result for each element in ``x``. - - - This method conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The Gauss error function of x. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.square(x) - >>> print(y) - ivy.array([1, 4, 9]) - - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.square(x, out=y) - >>> print(y) - ivy.array([2.25, 0.64, 0.09]) - - >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) - >>> ivy.square(x, out=x) - >>> print(x) - ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.square(x) - >>> print(y) - { - a:ivy.array([0,1]), - b:ivy.array([4,9]) - } + >>> x = ivy.array([0, 0.3, 0.7, 1.0]) + >>> ivy.erf(x) + ivy.array([0., 0.328, 0.677, 0.842]) """ - return ivy.current_backend(x).square(x, out=out) + return ivy.current_backend(x).erf(x, out=out) @handle_exceptions @@ -6790,27 +6705,26 @@ def square( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def subtract( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def maximum( + x1: Union[ivy.Array, ivy.NativeArray, Number], + x2: Union[ivy.Array, ivy.NativeArray, Number], /, *, - alpha: Optional[Union[int, float]] = None, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the difference for each element ``x1_i`` of the input array ``x1`` with - the respective element ``x2_i`` of the input array ``x2``. + Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. Parameters ---------- x1 - first input array. Should have a numeric data type. + Input array containing elements to maximum threshold. x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - alpha - optional scalar multiplier for ``x2``. + Tensor containing maximum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum + is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6818,138 +6732,154 @@ def subtract( Returns ------- ret - an array containing the element-wise differences. + An array with the elements of x1, but clipped to not be lower than the x2 + values. + Examples + -------- + With :class:`ivy.Array` inputs: - This method conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.maximum(x, y) + >>> print(z) + ivy.array([9, 9, 5]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.maximum(x, y, out=z) + >>> print(z) + ivy.array([[9., 9., 9., 9., 9., 9.], + [3., 5., 9., 8., 3., 7.], + [2., 5., 9., 8., 3., 7.]]) - Examples - -------- - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y) + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.maximum(x, y, out=x) + >>> print(x) + ivy.array([[7, 7]]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]), + ... b=ivy.array([-5, 9])) + >>> z = ivy.maximum(x, y) >>> print(z) - ivy.array([ 1, 5, -3]) + { + a: ivy.array([[1, 3], + [2, 4], + [3, 7]]), + b: ivy.array([[1, 9], + [2, 9], + [3, 9]]) + } - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y, alpha=2) + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) + >>> z = ivy.maximum(x, y) >>> print(z) - ivy.array([-1, 4, -9]) + { + a: ivy.array([1, 5, 6]), + b: ivy.array([5, 9, 7]) + } """ - return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) + return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def tan( - x: Union[ivy.Array, ivy.NativeArray], +def minimum( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r"""Calculate an implementation-dependent approximation to the tangent, having - domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each - element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be - expressed in radians. - .. note:: - Tangent is an analytical function on the complex plane - and has no branch cuts. The function is periodic, - with period :math:`\pi j`, with respect to the real - component and has first order poles along the real - line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. - However, IEEE 754 binary floating-point representation - cannot represent the value :math:`\pi / 2` exactly, and, - thus, no argument value is possible for - which a pole error occurs. - - where :math:`{tanh}` is the hyperbolic tangent. - - **Special cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - For complex floating-point operands, special cases must - be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. + """ + Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. Parameters ---------- - x - input array whose elements are expressed in radians. Should have a - floating-point data type. + x1 + Input array containing elements to minimum threshold. + x2 + Tensor containing minimum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum + is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the tangent of each element in ``x``. The return must have a - floating-point data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + An array with the elements of x1, but clipped to not exceed the x2 values. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.minimum(x, y) + >>> print(z) + ivy.array([7, 3, 2]) + + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.minimum(x, y, out=z) + >>> print(z) + ivy.array([[1.,5.,9.,8.,3.,7.], + [1.,3.,3.,3.,3.,3.], + [1.,2.,2.,2.,2.,2.]]) - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tan(x) - >>> print(y) - ivy.array([0., 1.56, -2.19]) + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.minimum(x, y, out=x) + >>> print(x) + ivy.array([[0, 3]]) - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tan(x, out=y) - >>> print(y) - ivy.array([0.546, -0.842, -0.916]) + With one :class:`ivy.Container` input: - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tan(x, out=x) - >>> print(x) - ivy.array([[1.96, -1.37, 0.16], - [-3.1, 0.996, -0.328]]) + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) + >>> z = ivy.minimum(x, y) + >>> print(z) + { + a: ivy.array([[1, 0], + [1, 0], + [1, 0]]), + b: ivy.array([[-5, 3], + [-5, 4], + [-5, 7]]) + } - With :class:`ivy.Container` input: + With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.tan(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([1, 3, 1]), + ... b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]), + ... b=ivy.array([5, 9, 7])) + >>> z = ivy.minimum(x, y) + >>> print(z) { - a: ivy.array([0., 1.56, -2.19]), - b: ivy.array([-0.143, 1.16, -3.38]) + a: ivy.array([1, 3, 1]), + b: ivy.array([2, 8, 5]) } """ - return ivy.current_backend(x).tan(x, out=out) + return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) @handle_exceptions @@ -6960,140 +6890,36 @@ def tan( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -@handle_complex_input -def tanh( - x: Union[ivy.Array, ivy.NativeArray], +def reciprocal( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, - complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate an implementation-dependent approximation to the hyperbolic tangent, - having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element - ``x_i`` of the input array ``x``. - - **Special cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+1``. - - If ``x_i`` is ``-infinity``, the result is ``-1``. - - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and - - .. note:: - For complex floating-point operands, ``tanh(conj(x))`` - must equal ``conj(tanh(x))``. - - - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. - - If ``a`` is a nonzero finite number and ``b`` is - ``+infinity``, the result is ``NaN + NaN j``. - - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``+0 + NaN j``. - - If ``a`` is a nonzero finite number and ``b`` - is ``NaN``, the result is ``NaN + NaN j``. - - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``+0 + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. - - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, - the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero number, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - .. warning:: - For historical reasons stemming from the C standard, - array libraries may not return the expected - result when ``a`` is ``+0`` and ``b`` is either - ``+infinity`` or ``NaN``. The result should be - ``+0 + NaN j`` in both cases; however, for libraries - compiled against older C versions, the result may be - ``NaN + NaN j``. - - Array libraries are not required to patch these older - C versions, and, thus, users are advised that results - may vary across array library implementations for - these special cases. - + Return a new array with the reciprocal of each element in ``x``. Parameters ---------- x - input array whose elements each represent a hyperbolic angle. Should have a - real-valued floating-point data - type. - complex_mode - optional specifier for how to handle complex data types. See - ``ivy.func_wrapper.handle_complex_input`` for more detail. + Input array. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the hyperbolic tangent of each element in ``x``. - The returned array must have a real-valued floating-point data type - determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + A new array with the positive value of each element in ``x``. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tanh(x) - >>> print(y) - ivy.array([0., 0.762, 0.964]) - - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tanh(x, out=y) - >>> print(y) - ivy.array([0.462, -0.604, 0.984]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tanh(x, out=x) - >>> print(x) - ivy.array([[0.8, 0.976, 0.997], - [-1., -1., -1.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.tanh(x) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.reciprocal(x) >>> print(y) - { - a: ivy.array([0., 0.762, 0.964]), - b: ivy.array([0.995, 0.999, 1.]) - } + ivy.array([1. , 0.5 , 0.33333333]) """ - return ivy.current_backend(x).tanh(x, out=out) + return ivy.current_backend(x).reciprocal(x, out=out) @handle_backend_invalid @@ -7101,61 +6927,80 @@ def tanh( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def trapz( - y: ivy.Array, +def deg2rad( + x: Union[ivy.Array, ivy.NativeArray], /, *, - x: Optional[ivy.Array] = None, - dx: float = 1.0, - axis: int = -1, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Integrate along the given axis using the composite trapezoidal rule. - - If x is provided, the integration happens in sequence along its elements - - they are not sorted.. + Convert the input from degrees to radians. Parameters ---------- - y - The array that should be integrated. x - The sample points corresponding to the input array values. - If x is None, the sample points are assumed to be evenly spaced - dx apart. The default is None. - dx - The spacing between sample points when x is None. The default is 1. - axis - The axis along which to integrate. + input array whose elements are each expressed in degrees. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Definite integral of n-dimensional array as approximated along - a single axis by the trapezoidal rule. If the input array is a - 1-dimensional array, then the result is a float. If n is greater - than 1, then the result is an n-1 dimensional array. + an array with each element in ``x`` converted from degrees to radians. Examples -------- - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3]) - 4.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], x=[4, 6, 8]) - 8.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], dx=2) - 8.0 + With :class:`ivy.Array` input: + + >>> x=ivy.array([0,90,180,270,360]) + >>> y=ivy.deg2rad(x) + >>> print(y) + ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + + >>> x=ivy.array([0,-1.5,-50,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([ 0., -0.02617994, -0.87266463, nan]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.deg2rad(x, out=x) + >>> print(x) + ivy.array([[ 0.01919862, 0.03839725, 0.05759586], + [-0.07679449, -0.09599311, -0.11519173]]) + + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([0., 0.35081118, nan]) + + With :class:`ivy.Container` input: + + >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), + ... b=ivy.array([0,90,180,270,360])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 0.35081118, -0.88139129, nan]), + b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + } + + >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), + ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), + b: ivy.array([0., -0.02617994, -0.87266463, nan]) + } """ - return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) + return ivy.current_backend(x).deg2rad(x, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -7163,32 +7008,19 @@ def trapz( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def trunc( +def rad2deg( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Round each element x_i of the input array x to the integer-valued number that is - closest to but no greater than x_i. - - **Special cases** - - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - - For floating-point operands, - - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. + Convert the input from radians to degrees. Parameters ---------- x - input array. Should have a numeric data type. + input array whose elements are each expressed in radians. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -7196,52 +7028,56 @@ def trunc( Returns ------- ret - an array containing the rounded result for each element in ``x``. - The returned array must have the same data type as ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + an array with each element in ``x`` converted from radians to degrees. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) - >>> y = ivy.trunc(x) + >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) + >>> y=ivy.rad2deg(x) >>> print(y) - ivy.array([-1., 0., 3., -0.]) + ivy.array([ 0., 90., 180., 270., 360.]) - >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) - >>> ivy.trunc(x, out=x) + >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.rad2deg(x,out=y) + >>> print(y) + ivy.array([ 0. , -1.5, -50. , nan]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.rad2deg(x, out=x) >>> print(x) - ivy.array([ 0., 7., -23., -0.]) + ivy.array([[ 63., 126., 189.], + [-252., -315., -378.]]) - >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) - >>> y = ivy.zeros([2,3]) - >>> ivy.trunc(x, out=y) + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.rad2deg(x,out=y) >>> print(y) - ivy.array([[ 0., -8., 0.], - [ 0., 0., 2.]]) + ivy.array([ 0., 1150., nan]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) - >>> y = ivy.trunc(x) + >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), + ... b=ivy.array([0., 1., 2., 3., 4.])) + >>> y=ivy.rad2deg(x) >>> print(y) { - a: ivy.array([-0., 4., 1.]), - b: ivy.array([12., -3., 1.]) + a: ivy.array([0., 1150., -2890., nan]), + b: ivy.array([0., 57.3, 115., 172., 229.]) + } + + >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), + ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) + >>> y=ivy.rad2deg(x) + >>> print(y) + { + a: ivy.array([0., 573., 10300., 487., 344.]), + b: ivy.array([0., -85.9, 28.6, nan]) } """ - return ivy.current_backend(x).trunc(x, out=out) + return ivy.current_backend(x).rad2deg(x, out=out) @handle_exceptions @@ -7288,3 +7124,169 @@ def trunc_divide( ivy.array([ 0., -1., 14.]) """ return ivy.trunc(ivy.divide(x1, x2), out=out) + + +trunc_divide.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + "handle_backend_invalid", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def isreal( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Test each element ``x_i`` of the input array ``x`` to determine whether the element + is real number. Returns a bool array, where True if input element is real. If + element has complex type with zero complex part, the return value for that element + is True. + + Parameters + ---------- + x + input array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is + real number and ``False`` otherwise. The returned array should have a data type + of ``bool``. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances in place of + :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints + and also the examples below. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([[[True], [True], [True]]]) + + >>> x = ivy.array([1-0j, 3j, 7+5j]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([ True, False, False]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ + b=ivy.array([5j, 5-6j, 3])) + >>> z = ivy.isreal(x) + >>> print(z) + { + a: ivy.array([False, True, True]), + b: ivy.array([False, False, True]) + } + """ + return ivy.current_backend(x).isreal(x, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fmod( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Compute the element-wise remainder of divisions of two arrays. + + Parameters + ---------- + x1 + First input array. + x2 + Second input array + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array with element-wise remainder of divisions. + + Examples + -------- + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmod(x1, x2) + ivy.array([ 0, 3, 0]) + + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmod(x1, x2) + ivy.array([ nan, nan, nan]) + """ + return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lcm( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the element-wise least common multiple (LCM) of x1 and x2. + + Parameters + ---------- + x1 + first input array, must be integers + x2 + second input array, must be integers + out + optional output array, for writing the result to. + + Returns + ------- + ret + an array that includes the element-wise least common multiples of x1 and x2 + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x1=ivy.array([2, 3, 4]) + >>> x2=ivy.array([5, 7, 15]) + >>> x1.lcm(x1, x2) + ivy.array([10, 21, 60]) + """ + return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) diff --git a/ivy/functional/ivy/experimental/activations.py b/ivy/functional/ivy/experimental/activations.py index 30de525bedc40..34203cffbaf6d 100644 --- a/ivy/functional/ivy/experimental/activations.py +++ b/ivy/functional/ivy/experimental/activations.py @@ -17,68 +17,6 @@ ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -def elu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - alpha: float = 1.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Apply the elu unit function element-wise. - - Parameters - ---------- - x - Input array. - alpha - scaler for controlling the slope of the function for x <= 0 Default: 1.0 - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The input array with elu applied element-wise. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([0.39, -0.85]) - >>> y = ivy.elu(x) - >>> print(y) - ivy.array([ 0.38999999, -0.57258511]) - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.elu(x, out=y) - >>> print(y) - ivy.array([ 1.5, 0.69999999, -0.90928203]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.elu(x, out=x) - >>> print(x) - ivy.array([[ 1.10000002, 2.20000005, 3.29999995], - [-0.98772264, -0.99591321, -0.99863964]]) - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.elu(x, out=x) - >>> print(x) - { - a: ivy.array([0., -0.69880581]), - b: ivy.array([0.40000001, -0.18126924]) - } - """ - return current_backend(x).elu(x, alpha=alpha, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -129,57 +67,6 @@ def logit( return current_backend(x).logit(x, eps=eps, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def logsigmoid( - input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply element-wise Log-sigmoid of x. - - logsigmoid(x) = log(1 / (1 + exp(-x)). - - Parameters - ---------- - input - Input array. - - Returns - ------- - Array with same shape as input with Log-sigmoid applied to every element. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-1.31326175, -0.69314718, -0.31326169]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-0.20141329, -0.40318608, -2.48683619]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.logsigmoid(x) - >>> print(x) - { - a: ivy.array([-0.31326169, -1.46328247]), - b: ivy.array([-0.59813893, -0.43748799]) - } - """ - return ivy.current_backend(input).logsigmoid(input, out=out) - - @handle_exceptions @handle_nestable @handle_array_like_without_promotion @@ -238,6 +125,67 @@ def prelu( raise IvyException +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def thresholded_relu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + threshold: Union[int, float] = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Apply the rectified linear unit function with custom threshold. + + Parameters + ---------- + x + input array + threshold + threshold value above which the activation is linear. Default: ``0``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the rectified linear unit activation of each element in + ``x``. with custom ``threshold``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.thresholded_relu(x, threshold=0.5) + >>> print(y) + ivy.array([0., 0. , 1.]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.thresholded_relu(x, threshold=1, out = y) + >>> print(y) + ivy.array([ 1.5, 0., 0.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.thresholded_relu(x, threshold=0.5) + >>> print(x) + { + a: ivy.array([1., 0.]), + b: ivy.array([0., 0.6]) + } + """ + return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -284,6 +232,57 @@ def relu6( return current_backend(x).relu6(x, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def logsigmoid( + input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply element-wise Log-sigmoid of x. + + logsigmoid(x) = log(1 / (1 + exp(-x)). + + Parameters + ---------- + input + Input array. + + Returns + ------- + Array with same shape as input with Log-sigmoid applied to every element. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-1.31326175, -0.69314718, -0.31326169]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-0.20141329, -0.40318608, -2.48683619]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.logsigmoid(x) + >>> print(x) + { + a: ivy.array([-0.31326169, -1.46328247]), + b: ivy.array([-0.59813893, -0.43748799]) + } + """ + return ivy.current_backend(input).logsigmoid(input, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -343,40 +342,6 @@ def selu( return current_backend(x).selu(x, out=out) -def sequence_length( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.int64: - """ - Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy - array input. - - Parameters - ---------- - x - Can be a sequence of any tensor type: bool, complex128, - complex64, double, float, float16, int16, int32, int64, - int8, string, uint16, uint32, uint64, uint8 - - Returns - ------- - length - Length of the input sequence, as a scalar (empty shape tensor). - - Examples - -------- - >>> x = ivy.array([True, False, True]) - >>> y = ivy.sequence_length(x) - >>> print(y) - 3 - - >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] - >>> y = ivy.sequence_length(x) - >>> print(y) - 5 - """ - return current_backend(x).sequence_length(x, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -433,23 +398,23 @@ def silu( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def thresholded_relu( +@handle_array_function +def elu( x: Union[ivy.Array, ivy.NativeArray], /, *, - threshold: Union[int, float] = 0, + alpha: float = 1.0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply the rectified linear unit function with custom threshold. + Apply the elu unit function element-wise. Parameters ---------- x - input array - threshold - threshold value above which the activation is linear. Default: ``0``. + Input array. + alpha + scaler for controlling the slope of the function for x <= 0 Default: 1.0 out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -457,32 +422,67 @@ def thresholded_relu( Returns ------- ret - an array containing the rectified linear unit activation of each element in - ``x``. with custom ``threshold``. + The input array with elu applied element-wise. Examples -------- With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.thresholded_relu(x, threshold=0.5) + >>> x = ivy.array([0.39, -0.85]) + >>> y = ivy.elu(x) >>> print(y) - ivy.array([0., 0. , 1.]) - + ivy.array([ 0.38999999, -0.57258511]) >>> x = ivy.array([1.5, 0.7, -2.4]) >>> y = ivy.zeros(3) - >>> ivy.thresholded_relu(x, threshold=1, out = y) + >>> ivy.elu(x, out=y) >>> print(y) - ivy.array([ 1.5, 0., 0.]) - + ivy.array([ 1.5, 0.69999999, -0.90928203]) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.elu(x, out=x) + >>> print(x) + ivy.array([[ 1.10000002, 2.20000005, 3.29999995], + [-0.98772264, -0.99591321, -0.99863964]]) With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.thresholded_relu(x, threshold=0.5) + >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.elu(x, out=x) >>> print(x) { - a: ivy.array([1., 0.]), - b: ivy.array([0., 0.6]) + a: ivy.array([0., -0.69880581]), + b: ivy.array([0.40000001, -0.18126924]) } """ - return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) + return current_backend(x).elu(x, alpha=alpha, out=out) + + +def sequence_length( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.int64: + """ + Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy + array input. + + Parameters + ---------- + x + Can be a sequence of any tensor type: bool, complex128, + complex64, double, float, float16, int16, int32, int64, + int8, string, uint16, uint32, uint64, uint8 + + Returns + ------- + length + Length of the input sequence, as a scalar (empty shape tensor). + + Examples + -------- + >>> x = ivy.array([True, False, True]) + >>> y = ivy.sequence_length(x) + >>> print(y) + 3 + + >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] + >>> y = ivy.sequence_length(x) + >>> print(y) + 5 + """ + return current_backend(x).sequence_length(x, out=out) diff --git a/ivy/functional/ivy/experimental/creation.py b/ivy/functional/ivy/experimental/creation.py index cf96237fa41af..bb3fc06d68e66 100644 --- a/ivy/functional/ivy/experimental/creation.py +++ b/ivy/functional/ivy/experimental/creation.py @@ -21,31 +21,56 @@ ) -# --- Helpers --- # -# --------------- # +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def vorbis_window( + window_length: Union[ivy.Array, ivy.NativeArray], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array that contains a vorbis power complementary window of size + window_length. + Parameters + ---------- + window_length + the length of the vorbis window. + dtype + data type of the returned array. By default float32. + out + optional output array, for writing the result to. -def _iter_product(*args, repeat=1): - # itertools.product - pools = [tuple(pool) for pool in args] * repeat - result = [[]] - for pool in pools: - result = [x + [y] for x in result for y in pool] - for prod in result: - yield tuple(prod) + Returns + ------- + ret + Input array with the vorbis window. + Examples + -------- + >>> ivy.vorbis_window(3) + ivy.array([0.38268346, 1. , 0.38268352]) -# --- Main --- # -# ------------ # + >>> ivy.vorbis_window(5) + ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) + """ + return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @infer_dtype @handle_device_shifting -def blackman_window( +def hann_window( size: int, *, periodic: bool = True, @@ -53,14 +78,13 @@ def blackman_window( out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Generate a Blackman window. The Blackman window is a taper formed by using the first - three terms of a summation of cosines. It was designed to have close to the minimal - leakage possible. It is close to optimal, only slightly worse than a Kaiser window. + Generate a Hann window. The Hanning window is a taper formed by using a weighted + cosine. Parameters ---------- - window_length - the window_length of the returned window. + size + the size of the returned window. periodic If True, returns a window to be used as periodic function. If False, return a symmetric window. @@ -73,111 +97,122 @@ def blackman_window( ------- ret The array containing the window. + Functional Examples ------------------- - >>> ivy.blackman_window(4, periodic = True) - ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) - >>> ivy.blackman_window(7, periodic = False) - ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, - 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) + >>> ivy.hann_window(4, periodic = True) + ivy.array([0. , 0.5, 1. , 0.5]) + + >>> ivy.hann_window(7, periodic = False) + ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) """ - return ivy.current_backend().blackman_window( + return ivy.current_backend().hann_window( size, periodic=periodic, dtype=dtype, out=out ) @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays +@to_native_arrays_and_back @infer_dtype -@infer_device -def eye_like( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, *, - k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the - same ``shape`` as the first and last dim of input array ``x``. input array ``x`` - should to be 2D. + Compute the Kaiser window with window length window_length and shape beta. Parameters ---------- - x - input array from which to derive the output array shape. - k - index of the diagonal. A positive value refers to an upper diagonal, a negative - value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. + window_length + an int defining the length of the window. + periodic + If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta + a float used as shape parameter for the window. dtype - output array data type. If dtype is None, the output array data type must be the - default floating-point data type. Default: ``None``. - device - the device on which to place the created array. + data type of the returned array. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array having the same shape as ``x`` and filled with ``ones`` in - diagonal ``k`` and ``zeros`` elsewhere. - + The array containing the window. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances as a replacement to any of the arguments. + Examples + -------- + >>> ivy.kaiser_window(5) + ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) + >>> ivy.kaiser_window(5, True, 5) + ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) + >>> ivy.kaiser_window(5, False, 5) + ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) + """ + return ivy.current_backend().kaiser_window( + window_length, periodic, beta, dtype=dtype, out=out + ) - Functional Examples - ------------------- - With :class:`ivy.Array` input: +@handle_exceptions +@handle_nestable +@handle_out_argument +@infer_dtype +def kaiser_bessel_derived_window( + window_length: int, + beta: float = 12.0, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Kaiser bessel derived window with window length window_length and shape + beta. - >>> x1 = ivy.array([[0, 1],[2, 3]]) - >>> y1 = ivy.eye_like(x1) - >>> print(y1) - ivy.array([[1., 0.], - [0., 1.]]) + Parameters + ---------- + window_length + an int defining the length of the window. + beta + a float used as shape parameter for the window. + dtype + data type of the returned array + out + optional output array, for writing the result to. - >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) - >>> y1 = ivy.eye_like(x1, k=1) - >>> print(y1) - ivy.array([[0., 1., 0.], - [0., 0., 1.], - [0., 0., 0.]]) + Returns + ------- + ret + The array containing the window. - With :class:`ivy.Container` input: + Functional Examples + ------------------- + >>> ivy.kaiser_bessel_derived_window(5) + ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) - >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) - >>> y = x.eye_like() - >>> print(y) - { - a: ivy.array([[1., 0.], - [0., 1.]]), - b: ivy.array([[1., 0.], - [0., 1.]]) - } + >>> ivy.kaiser_bessel_derived_window(5, 5) + ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) """ - shape = ivy.shape(x, as_array=True) - dim = len(shape) - if dim <= 1: - cols = dim - else: - cols = int(shape[-1]) - rows = 0 if dim < 1 else int(shape[0]) - return ivy.eye( - rows, - cols, - k=k, - dtype=dtype, - device=device, - out=out, - ) + if window_length < 2: + result = ivy.array([], dtype=dtype) + if ivy.exists(out): + ivy.inplace_update(out, result) + return result + half_len = window_length // 2 + kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) + kaiser_w_csum = ivy.cumsum(kaiser_w) + half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) + window = ivy.concat((half_w, half_w[::-1]), axis=0) + result = window.astype(dtype) + return result @handle_exceptions @@ -237,211 +272,214 @@ def hamming_window( return result +hamming_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "handle_device_shifting", + ), + "to_skip": (), +} + + @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def hann_window( - size: int, +@outputs_to_ivy_arrays +@infer_device +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, *, - periodic: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, +) -> Tuple[ivy.Array, ...]: """ - Generate a Hann window. The Hanning window is a taper formed by using a weighted - cosine. - - Parameters - ---------- - size - the size of the returned window. - periodic - If True, returns a window to be used as periodic function. - If False, return a symmetric window. - dtype - The data type to produce. Must be a floating point type. - out - optional output array, for writing the result to. - - Returns - ------- - ret - The array containing the window. - - Functional Examples - ------------------- - >>> ivy.hann_window(4, periodic = True) - ivy.array([0. , 0.5, 1. , 0.5]) - - >>> ivy.hann_window(7, periodic = False) - ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) - """ - return ivy.current_backend().hann_window( - size, periodic=periodic, dtype=dtype, out=out - ) + Return the indices of the lower triangular part of a row by col matrix in a 2-by-N + shape (tuple of two N dimensional arrays), where the first row contains row + coordinates of all indices and the second row contains column coordinates. Indices + are ordered based on rows and then columns. The lower triangular part of the matrix + is defined as the elements on and below the diagonal. The argument k controls which + diagonal to consider. If k = 0, all elements on and below the main diagonal are + retained. A positive value excludes just as many diagonals below the main diagonal, + and similarly a negative value includes just as many diagonals above the main + diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, + n_cols}−1]. + Notes + ----- + Primary purpose of this function is to slice an array of shape (n,m). See + https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html + for examples -@handle_exceptions -def indices( - dimensions: Sequence[int], - *, - dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, - sparse: bool = False, -) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: - """ - Return an array representing the indices of a grid. + Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices Parameters ---------- - dimensions - The shape of the grid. - dtype - The data type of the result. - sparse - Return a sparse representation of the grid instead of a dense representation. + n_rows + number of rows in the 2-d matrix. + n_cols + number of columns in the 2-d matrix. If None n_cols will be the same as n_rows + k + number of shifts from the main diagonal. k = 0 includes main diagonal, + k > 0 moves downward and k < 0 moves upward + device + device on which to place the created array. Default: ``None``. Returns ------- ret - If sparse is False, returns one grid indices array of shape - (len(dimensions),) + tuple(dimensions). - If sparse is True, returns a tuple of arrays each of shape - (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. + an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) + contains row coordinates of all indices and the second subarray (i.e ret[1]) + contains columns indices. + + Function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> ivy.indices((3, 2)) - ivy.array([[[0 0] - [1 1] - [2 2]] - [[0 1] - [0 1] - [0 1]]]) - >>> ivy.indices((3, 2), sparse=True) - (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) + >>> x = ivy.tril_indices(4,4,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,1) + >>> print(x) + (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,-2) + >>> print(x) + (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,2,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 0, 1])) + + >>> x = ivy.tril_indices(2,4,0) + >>> print(x) + (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,-4,0) + >>> print(x) + (ivy.array([]), ivy.array([])) + + >>> x = ivy.tril_indices(4,4,100) + >>> print(x) + (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(2,4,-100) + >>> print(x) + (ivy.array([]), ivy.array([])) """ - if sparse: - return tuple( - ivy.arange(dim) - .expand_dims( - axis=[j for j in range(len(dimensions)) if i != j], - ) - .astype(dtype) - for i, dim in enumerate(dimensions) - ) - else: - grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") - return ivy.stack(grid, axis=0).astype(dtype) + return current_backend().tril_indices(n_rows, n_cols, k, device=device) @handle_exceptions @handle_nestable +@handle_array_like_without_promotion @handle_out_argument +@inputs_to_ivy_arrays @infer_dtype -def kaiser_bessel_derived_window( - window_length: int, - beta: float = 12.0, +@infer_device +def eye_like( + x: Union[ivy.Array, ivy.NativeArray], *, + k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kaiser bessel derived window with window length window_length and shape - beta. + Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the + same ``shape`` as the first and last dim of input array ``x``. input array ``x`` + should to be 2D. Parameters ---------- - window_length - an int defining the length of the window. - beta - a float used as shape parameter for the window. + x + input array from which to derive the output array shape. + k + index of the diagonal. A positive value refers to an upper diagonal, a negative + value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. dtype - data type of the returned array + output array data type. If dtype is None, the output array data type must be the + default floating-point data type. Default: ``None``. + device + the device on which to place the created array. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The array containing the window. + an array having the same shape as ``x`` and filled with ``ones`` in + diagonal ``k`` and ``zeros`` elsewhere. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances as a replacement to any of the arguments. Functional Examples ------------------- - >>> ivy.kaiser_bessel_derived_window(5) - ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) - - >>> ivy.kaiser_bessel_derived_window(5, 5) - ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) - """ - if window_length < 2: - result = ivy.array([], dtype=dtype) - if ivy.exists(out): - ivy.inplace_update(out, result) - return result - half_len = window_length // 2 - kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) - kaiser_w_csum = ivy.cumsum(kaiser_w) - half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) - window = ivy.concat((half_w, half_w[::-1]), axis=0) - result = window.astype(dtype) - return result + With :class:`ivy.Array` input: -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Kaiser window with window length window_length and shape beta. + >>> x1 = ivy.array([[0, 1],[2, 3]]) + >>> y1 = ivy.eye_like(x1) + >>> print(y1) + ivy.array([[1., 0.], + [0., 1.]]) - Parameters - ---------- - window_length - an int defining the length of the window. - periodic - If True, returns a periodic window suitable for use in spectral analysis. - If False, returns a symmetric window suitable for use in filter design. - beta - a float used as shape parameter for the window. - dtype - data type of the returned array. - out - optional output array, for writing the result to. + >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) + >>> y1 = ivy.eye_like(x1, k=1) + >>> print(y1) + ivy.array([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) - Returns - ------- - ret - The array containing the window. + With :class:`ivy.Container` input: - Examples - -------- - >>> ivy.kaiser_window(5) - ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) - >>> ivy.kaiser_window(5, True, 5) - ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) - >>> ivy.kaiser_window(5, False, 5) - ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) + >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) + >>> y = x.eye_like() + >>> print(y) + { + a: ivy.array([[1., 0.], + [0., 1.]]), + b: ivy.array([[1., 0.], + [0., 1.]]) + } """ - return ivy.current_backend().kaiser_window( - window_length, periodic, beta, dtype=dtype, out=out + shape = ivy.shape(x, as_array=True) + dim = len(shape) + if dim <= 1: + cols = dim + else: + cols = int(shape[-1]) + rows = 0 if dim < 1 else int(shape[0]) + return ivy.eye( + rows, + cols, + k=k, + dtype=dtype, + device=device, + out=out, ) +def _iter_product(*args, repeat=1): + # itertools.product + pools = [tuple(pool) for pool in args] * repeat + result = [[]] + for pool in pools: + result = [x + [y] for x in result for y in pool] + for prod in result: + yield tuple(prod) + + @handle_exceptions @inputs_to_ivy_arrays def ndenumerate( @@ -515,60 +553,185 @@ def ndindex( @handle_exceptions -@handle_nestable -@infer_dtype -def random_cp( - shape: Sequence[int], - rank: int, - /, +def indices( + dimensions: Sequence[int], *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - full: Optional[bool] = False, - orthogonal: Optional[bool] = False, - seed: Optional[int] = None, - normalise_factors: Optional[bool] = True, -) -> Union[ivy.CPTensor, ivy.Array]: + dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, + sparse: bool = False, +) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: """ - Generate a random CP tensor. + Return an array representing the indices of a grid. Parameters ---------- - shape - shape of the tensor to generate - rank - rank of the CP decomposition - full - if True, a full tensor is returned - otherwise, the decomposed tensor is returned - orthogonal - if True, creates a tensor with orthogonal components - seed - seed for generating random numbers + dimensions + The shape of the grid. + dtype + The data type of the result. + sparse + Return a sparse representation of the grid instead of a dense representation. Returns ------- - ivy.CPTensor + ret + If sparse is False, returns one grid indices array of shape + (len(dimensions),) + tuple(dimensions). + If sparse is True, returns a tuple of arrays each of shape + (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. + + Examples + -------- + >>> ivy.indices((3, 2)) + ivy.array([[[0 0] + [1 1] + [2 2]] + [[0 1] + [0 1] + [0 1]]]) + >>> ivy.indices((3, 2), sparse=True) + (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) """ - rank = ivy.CPTensor.validate_cp_rank(shape, rank) - if (rank > min(shape)) and orthogonal: - warnings.warn( - "Can only construct orthogonal tensors when rank <= min(shape) but got " - f"a tensor with min(shape)={min(shape)} < rank={rank}" + if sparse: + return tuple( + ivy.arange(dim) + .expand_dims( + axis=[j for j in range(len(dimensions)) if i != j], + ) + .astype(dtype) + for i, dim in enumerate(dimensions) ) + else: + grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") + return ivy.stack(grid, axis=0).astype(dtype) - factors = [ - (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape - ] - weights = ivy.ones((rank,), dtype=dtype) - if orthogonal: - factors = [ivy.qr(factor)[0] for factor in factors] - if full: - return ivy.CPTensor.cp_to_tensor((weights, factors)) - elif normalise_factors: - return ivy.CPTensor.cp_normalize((weights, factors)) - else: - return ivy.CPTensor((weights, factors)) +indices.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_min( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the minimum along segments of an array. Segments are defined by an integer + array of segment IDs. + + Note + ---- + If the given segment ID `i` is negative, then the corresponding + value is dropped, and will not be included in the result. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented min operation. + For each segment, it computes the min value in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_sum( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the sum of elements along segments of an array. Segments are defined by an + integer array of segment IDs. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented sum operation. + For each segment, it computes the sum of values in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def blackman_window( + size: int, + *, + periodic: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a Blackman window. The Blackman window is a taper formed by using the first + three terms of a summation of cosines. It was designed to have close to the minimal + leakage possible. It is close to optimal, only slightly worse than a Kaiser window. + + Parameters + ---------- + window_length + the window_length of the returned window. + periodic + If True, returns a window to be used as periodic function. + If False, return a symmetric window. + dtype + The data type to produce. Must be a floating point type. + out + optional output array, for writing the result to. + + Returns + ------- + ret + The array containing the window. + Functional Examples + ------------------- + >>> ivy.blackman_window(4, periodic = True) + ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) + >>> ivy.blackman_window(7, periodic = False) + ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, + 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) + """ + return ivy.current_backend().blackman_window( + size, periodic=periodic, dtype=dtype, out=out + ) @handle_exceptions @@ -643,96 +806,59 @@ def random_tucker( @handle_exceptions @handle_nestable -@outputs_to_ivy_arrays -@infer_device -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, +@infer_dtype +def random_cp( + shape: Sequence[int], + rank: int, + /, *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, -) -> Tuple[ivy.Array, ...]: + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + full: Optional[bool] = False, + orthogonal: Optional[bool] = False, + seed: Optional[int] = None, + normalise_factors: Optional[bool] = True, +) -> Union[ivy.CPTensor, ivy.Array]: """ - Return the indices of the lower triangular part of a row by col matrix in a 2-by-N - shape (tuple of two N dimensional arrays), where the first row contains row - coordinates of all indices and the second row contains column coordinates. Indices - are ordered based on rows and then columns. The lower triangular part of the matrix - is defined as the elements on and below the diagonal. The argument k controls which - diagonal to consider. If k = 0, all elements on and below the main diagonal are - retained. A positive value excludes just as many diagonals below the main diagonal, - and similarly a negative value includes just as many diagonals above the main - diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, - n_cols}−1]. - - Notes - ----- - Primary purpose of this function is to slice an array of shape (n,m). See - https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html - for examples - - Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices + Generate a random CP tensor. Parameters ---------- - n_rows - number of rows in the 2-d matrix. - n_cols - number of columns in the 2-d matrix. If None n_cols will be the same as n_rows - k - number of shifts from the main diagonal. k = 0 includes main diagonal, - k > 0 moves downward and k < 0 moves upward - device - device on which to place the created array. Default: ``None``. + shape + shape of the tensor to generate + rank + rank of the CP decomposition + full + if True, a full tensor is returned + otherwise, the decomposed tensor is returned + orthogonal + if True, creates a tensor with orthogonal components + seed + seed for generating random numbers Returns ------- - ret - an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) - contains row coordinates of all indices and the second subarray (i.e ret[1]) - contains columns indices. - - Function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - >>> x = ivy.tril_indices(4,4,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) - - >>> x = ivy.tril_indices(4,4,1) - >>> print(x) - (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) - - >>> x = ivy.tril_indices(4,4,-2) - >>> print(x) - (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) - - >>> x = ivy.tril_indices(4,2,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 0, 1])) - - >>> x = ivy.tril_indices(2,4,0) - >>> print(x) - (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) - - >>> x = ivy.tril_indices(4,-4,0) - >>> print(x) - (ivy.array([]), ivy.array([])) + ivy.CPTensor + """ + rank = ivy.CPTensor.validate_cp_rank(shape, rank) + if (rank > min(shape)) and orthogonal: + warnings.warn( + "Can only construct orthogonal tensors when rank <= min(shape) but got " + f"a tensor with min(shape)={min(shape)} < rank={rank}" + ) - >>> x = ivy.tril_indices(4,4,100) - >>> print(x) - (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + factors = [ + (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape + ] + weights = ivy.ones((rank,), dtype=dtype) + if orthogonal: + factors = [ivy.qr(factor)[0] for factor in factors] - >>> x = ivy.tril_indices(2,4,-100) - >>> print(x) - (ivy.array([]), ivy.array([])) - """ - return current_backend().tril_indices(n_rows, n_cols, k, device=device) + if full: + return ivy.CPTensor.cp_to_tensor((weights, factors)) + elif normalise_factors: + return ivy.CPTensor.cp_normalize((weights, factors)) + else: + return ivy.CPTensor((weights, factors)) @handle_nestable @@ -784,135 +910,3 @@ def trilu( instances in place of any of the arguments. """ return current_backend(x).trilu(x, k=k, upper=upper, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@to_native_arrays_and_back -def unsorted_segment_min( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the minimum along segments of an array. Segments are defined by an integer - array of segment IDs. - - Note - ---- - If the given segment ID `i` is negative, then the corresponding - value is dropped, and will not be included in the result. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented min operation. - For each segment, it computes the min value in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_nestable -@to_native_arrays_and_back -def unsorted_segment_sum( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the sum of elements along segments of an array. Segments are defined by an - integer array of segment IDs. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented sum operation. - For each segment, it computes the sum of values in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def vorbis_window( - window_length: Union[ivy.Array, ivy.NativeArray], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array that contains a vorbis power complementary window of size - window_length. - - Parameters - ---------- - window_length - the length of the vorbis window. - dtype - data type of the returned array. By default float32. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Input array with the vorbis window. - - Examples - -------- - >>> ivy.vorbis_window(3) - ivy.array([0.38268346, 1. , 0.38268352]) - - >>> ivy.vorbis_window(5) - ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) - """ - return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) - - -hamming_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "handle_device_shifting", - ), - "to_skip": (), -} -indices.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} diff --git a/ivy/functional/ivy/experimental/elementwise.py b/ivy/functional/ivy/experimental/elementwise.py index e8c672460e7f1..613fe70acc6b2 100644 --- a/ivy/functional/ivy/experimental/elementwise.py +++ b/ivy/functional/ivy/experimental/elementwise.py @@ -17,214 +17,227 @@ from ivy.utils.exceptions import handle_exceptions -lerp.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def allclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +@handle_array_function +def lgamma( + x: Union[ivy.Array, ivy.NativeArray], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> bool: +) -> ivy.Array: """ - Return a True if the two arrays are element-wise equal within given tolerance; - otherwise False. - - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(x2)) and the absolute difference - atol are added together to compare against the absolute difference - between x1 and x2. - The default atol is not appropriate for comparing numbers that are - much smaller than one + Compute the natural logarithm of the absolute value of the gamma function on x. Parameters ---------- - x1 - First input array. - x2 - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in x1 will be - considered equal to NaN's in x2 in the output array. + x + input array. Should have a floating-point data type. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Returns True if the two arrays are equal within the given tolerance; - False otherwise. + an array containing the natural log of Gamma(x) of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. Examples -------- - >>> x1 = ivy.array([1e10, 1e-7]) - >>> x2 = ivy.array([1.00001e10, 1e-8]) - >>> y = ivy.allclose(x1, x2) + >>> x = ivy.array([1.6, 2.6, 3.5]) + >>> y = x.lgamma() >>> print(y) - ivy.array(False) + ivy.array([-0.11259177, 0.3574118 , 1.20097363]) - >>> x1 = ivy.array([1.0, ivy.nan]) - >>> x2 = ivy.array([1.0, ivy.nan]) - >>> y = ivy.allclose(x1, x2, equal_nan=True) + >>> x = ivy.array([1., 2., 3. ]) + >>> y = x.lgamma() >>> print(y) - ivy.array(True) + ivy.array([0. ,0. ,0.69314718]) - >>> x1 = ivy.array([1e-10, 1e-10]) - >>> x2 = ivy.array([1.00001e-10, 1e-10]) - >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) - >>> print(y) - ivy.array(True) + >>> x = ivy.array([4.5, -4, -5.6]) + >>> x.lgamma(out = x) + >>> print(x) + ivy.array([2.45373654, inf, -4.6477685 ]) """ - return ivy.current_backend().allclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) + return ivy.current_backend(x).lgamma(x, out=out) @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays -def binarizer( +@to_native_arrays_and_back +@handle_device_shifting +def sinc( x: Union[ivy.Array, ivy.NativeArray], /, *, - threshold: float = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Map the values of the input tensor to either 0 or 1, element-wise, based on the - outcome of a comparison against a threshold value. + Calculate an implementation-dependent approximation of the principal value of the + normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain + ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element + ``x_i`` is assumed to be expressed in radians. + + **Special cases** + + For floating-point operands, + + - If x_i is NaN, the result is NaN. + - If ``x_i`` is ``0``, the result is ``1``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. Parameters ---------- x - Data to be binarized - threshold - Values greater than this are - mapped to 1, others to 0. + input array. Should have a floating-point data type. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Binarized output data + an array containing the normalized sinc function of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = x.sinc() + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + >>> x = ivy.array([1.5, 0.5, -1.5]) + >>> y = ivy.zeros(3) + >>> ivy.sinc(x, out=y) + >>> print(y) + ivy.array([-0.212,0.637,-0.212]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = ivy.sinc(x) + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), + ... b=ivy.array([3.5, 4.5, 5.5])) + >>> y = x.sinc() + >>> print(y) + { + a: ivy.array([0.637,-0.212,0.127]), + b: ivy.array([-0.0909,0.0707,-0.0579]) + } """ - xc = ivy.copy_array(x, out=out) - if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": - xc = ivy.astype(xc, ivy.default_float_dtype()) - if ivy.is_complex_dtype(xc): - xc = ivy.abs(xc) - ret = ivy.where(xc > threshold, 1, 0) - return ret + return ivy.current_backend(x).sinc(x, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def conj( - x: Union[ivy.Array, ivy.NativeArray], +def fmax( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Return the complex conjugate for each element ``x_i`` of the input array ``x``. - - For complex number of the form + Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the + case where one of the elements is NaN. ivy.maximum returns the NaN element while + ivy.fmax returns the non-NaN element. - .. math:: - a + bj + Parameters + ---------- + x1 + First input array. + x2 + Second input array. + out + optional output array, for writing the result to. - the complex conjugate is defined as + Returns + ------- + ret + Array with element-wise maximums. - .. math:: - a - bj + Examples + -------- + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmax(x1, x2) + ivy.array([ 2., 5., 4.]) - Hence, the returned conjugates must be computed by negating - the imaginary component of each element ``x_i`` + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmax(x1, x2) + ivy.array([ 0., 0., nan]) + """ + return ivy.current_backend().fmax(x1, x2, out=out) - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the - `docstring `_ - in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def float_power( + x1: Union[ivy.Array, float, list, tuple], + x2: Union[ivy.Array, float, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must + be broadcastable to the same shape. This differs from the power function in that + integers, float16, and float32 are promoted to floats with a minimum precision of + float64 so that the result is always inexact. Parameters ---------- - x - input array. + x1 + Array-like with elements to raise in power. + x2 + Array-like of exponents. If x1.shape != x2.shape, + they must be broadcastable to a common shape + (which becomes the shape of the output). out optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. Returns ------- ret - an arrray of the same dtype as the input array with - the complex conjugates of the complex values present - in the input array. If x is a scalar then a scalar - will be returned. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. - + The bases in x1 raised to the exponents in x2. + This is a scalar if both x1 and x2 are scalars Examples -------- - With :class:`ivy.Array` inputs: - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.conj(x) - >>> print(z) - ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) - - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), - ... b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.conj(x) - >>> print(z) - { - a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), - b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) - } + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> ivy.float_power(x1, 3) + ivy.array([1., 8., 27., 64., 125.]) + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> x2 = ivy.array([2, 3, 3, 2, 1]) + >>> ivy.float_power(x1, x2) + ivy.array([1., 8., 27., 16., 5.]) """ - return ivy.current_backend(x).conj(x, out=out) + return ivy.current_backend().float_power(x1, x2, out=out) @handle_exceptions @@ -328,9 +341,208 @@ def count_nonzero( >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) ivy.array([[[3, 4]]]) """ - return ivy.current_backend().count_nonzero( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out - ) + return ivy.current_backend().count_nonzero( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def nansum( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as + zero. + + Parameters + ---------- + x + Input array. + axis + Axis or axes along which the sum is computed. + The default is to compute the sum of the flattened array. + dtype + The type of the returned array and of the accumulator in + which the elements are summed. By default, the dtype of input is used. + keepdims + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + A new array holding the result is returned unless out is specified, + in which it is returned. + + Examples + -------- + >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) + >>> ivy.nansum(a) + 10.0 + >>> ivy.nansum(a, axis=0) + ivy.array([2.1, 5.8, 2.1]) + >>> ivy.nansum(a, axis=1) + ivy.array([5.5, 4.5]) + """ + return ivy.current_backend().nansum( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out + ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def isclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a boolean array where two arrays are element-wise equal within a tolerance. + + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(b)) and the absolute difference + atol are added together to compare against the absolute difference + between a and b. + The default atol is not appropriate for comparing numbers that are + much smaller than one + + Parameters + ---------- + a + First input array. + b + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in a will be + considered equal to NaN's in b in the output array. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + Returns a boolean array of where a and b are equal within the given + tolerance. If both a and b are scalars, returns a single boolean value. + + Examples + -------- + >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) + ivy.array([True, False]) + >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) + ivy.array([True, True]) + >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) + ivy.array([False, False]) + >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) + ivy.array([False, True]) + """ + return ivy.current_backend().isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def signbit( + x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return element-wise True where signbit is set (less than zero). + + Parameters + ---------- + x + Array-like input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Output array, or reference to out if that was supplied. + This is a scalar if x is a scalar. + + Examples + -------- + >>> x = ivy.array([1, -2, 3]) + >>> ivy.signbit(x) + ivy.array([False, True, False]) + """ + return ivy.current_backend(x).signbit(x, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def hypot( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Return the hypotenuse given the two sides of a right angle triangle. + + Parameters + ---------- + x1 + The first input array + x2 + The second input array + + Returns + ------- + ret + An array with the hypotenuse + + Examples + -------- + >>> a = ivy.array([3.0, 4.0, 5.0]) + >>> b = ivy.array([4.0, 5.0, 6.0]) + >>> ivy.hypot(a, b) + ivy.array([5.0, 6.4031, 7.8102]) + """ + return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) @handle_backend_invalid @@ -390,26 +602,45 @@ def diff( @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_out_argument +@handle_array_like_without_promotion @to_native_arrays_and_back -def digamma( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def allclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> bool: """ - Compute the logarithmic derivative of the gamma function at x. + Return a True if the two arrays are element-wise equal within given tolerance; + otherwise False. - Note - ---- - The Ivy version only accepts real-valued inputs. + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(x2)) and the absolute difference + atol are added together to compare against the absolute difference + between x1 and x2. + The default atol is not appropriate for comparing numbers that are + much smaller than one Parameters ---------- - x - Input array. + x1 + First input array. + x2 + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in x1 will be + considered equal to NaN's in x2 in the output array. out Alternate output array in which to place the result. The default is None. @@ -417,16 +648,32 @@ def digamma( Returns ------- ret - Array with values computed from digamma function from - input arrays' values, element-wise. + Returns True if the two arrays are equal within the given tolerance; + False otherwise. Examples -------- - >>> x = ivy.array([.9, 3, 3.2]) - >>> y = ivy.digamma(x) - ivy.array([-0.7549271 0.92278427 0.9988394]) + >>> x1 = ivy.array([1e10, 1e-7]) + >>> x2 = ivy.array([1.00001e10, 1e-8]) + >>> y = ivy.allclose(x1, x2) + >>> print(y) + ivy.array(False) + + >>> x1 = ivy.array([1.0, ivy.nan]) + >>> x2 = ivy.array([1.0, ivy.nan]) + >>> y = ivy.allclose(x1, x2, equal_nan=True) + >>> print(y) + ivy.array(True) + + >>> x1 = ivy.array([1e-10, 1e-10]) + >>> x2 = ivy.array([1.00001e-10, 1e-10]) + >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) + >>> print(y) + ivy.array(True) """ - return ivy.current_backend(x).digamma(x, out=out) + return ivy.current_backend().allclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) @handle_backend_invalid @@ -467,70 +714,22 @@ def fix( return ivy.current_backend(x).fix(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def float_power( - x1: Union[ivy.Array, float, list, tuple], - x2: Union[ivy.Array, float, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must - be broadcastable to the same shape. This differs from the power function in that - integers, float16, and float32 are promoted to floats with a minimum precision of - float64 so that the result is always inexact. - - Parameters - ---------- - x1 - Array-like with elements to raise in power. - x2 - Array-like of exponents. If x1.shape != x2.shape, - they must be broadcastable to a common shape - (which becomes the shape of the output). - out - optional output array, for writing the result to. - - Returns - ------- - ret - The bases in x1 raised to the exponents in x2. - This is a scalar if both x1 and x2 are scalars - - Examples - -------- - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> ivy.float_power(x1, 3) - ivy.array([1., 8., 27., 64., 125.]) - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> x2 = ivy.array([2, 3, 3, 2, 1]) - >>> ivy.float_power(x1, x2) - ivy.array([1., 8., 27., 16., 5.]) - """ - return ivy.current_backend().float_power(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fmax( +def nextafter( x1: Union[ivy.Array, ivy.NativeArray], x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> bool: """ - Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the - case where one of the elements is NaN. ivy.maximum returns the NaN element while - ivy.fmax returns the non-NaN element. + Return the next floating-point value after x1 towards x2, element-wise. Parameters ---------- @@ -539,26 +738,22 @@ def fmax( x2 Second input array. out - optional output array, for writing the result to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - Array with element-wise maximums. + The next representable values of x1 in the direction of x2. Examples -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmax(x1, x2) - ivy.array([ 2., 5., 4.]) - - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmax(x1, x2) - ivy.array([ 0., 0., nan]) + >>> x1 = ivy.array([1.0e-50, 2.0e+50]) + >>> x2 = ivy.array([2.0, 1.0]) + >>> ivy.nextafter(x1, x2) + ivy.array([1.4013e-45., 3.4028e+38]) """ - return ivy.current_backend().fmax(x1, x2, out=out) + return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) @handle_exceptions @@ -568,35 +763,41 @@ def fmax( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def frexp( +def zeta( x: Union[ivy.Array, ivy.NativeArray], + q: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> bool: """ - Decompose the elements of x into mantissa and twos exponent. + Compute the Hurwitz zeta function elementwisely with each pair of floats in two + arrays. Parameters ---------- x - Input array. + First input array. + q + Second input array, must have the same shape as the first input array out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - A tuple of two arrays, the mantissa and the twos exponent. + Array with values computed from zeta function from + input arrays' values. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.frexp(x) - (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) + >>> x = ivy.array([5.0, 3.0]) + >>> q = ivy.array([2.0, 2.0]) + >>> ivy.zeta(x, q) + ivy.array([0.0369, 0.2021]) """ - return ivy.current_backend(x).frexp(x, out=out) + return ivy.current_backend(x, q).zeta(x, q, out=out) @handle_exceptions @@ -679,41 +880,91 @@ def gradient( ) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def hypot( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def xlogy( + x: Union[ivy.Array, ivy.NativeArray], + y: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> bool: """ - Return the hypotenuse given the two sides of a right angle triangle. + Compute x*log(y) element-wise so that the result is 0 if x = 0. Parameters ---------- - x1 - The first input array - x2 - The second input array + x + First input array. + y + Second input array. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + The next representable values of x1 in the direction of x2. + + Examples + -------- + >>> x = ivy.zeros(3) + >>> y = ivy.array([-1.0, 0.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([0.0, 0.0, 0.0]) + + >>> x = ivy.array([1.0, 2.0, 3.0]) + >>> y = ivy.array([3.0, 2.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([1.0986, 1.3863, 0.0000]) + """ + return ivy.current_backend(x, y).xlogy(x, y, out=out) + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_ivy_arrays +def binarizer( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + threshold: float = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Map the values of the input tensor to either 0 or 1, element-wise, based on the + outcome of a comparison against a threshold value. + + Parameters + ---------- + x + Data to be binarized + threshold + Values greater than this are + mapped to 1, others to 0. + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - An array with the hypotenuse - - Examples - -------- - >>> a = ivy.array([3.0, 4.0, 5.0]) - >>> b = ivy.array([4.0, 5.0, 6.0]) - >>> ivy.hypot(a, b) - ivy.array([5.0, 6.4031, 7.8102]) + Binarized output data """ - return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) + xc = ivy.copy_array(x, out=out) + if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": + xc = ivy.astype(xc, ivy.default_float_dtype()) + if ivy.is_complex_dtype(xc): + xc = ivy.abs(xc) + ret = ivy.where(xc > threshold, 1, 0) + return ret @handle_exceptions @@ -723,63 +974,80 @@ def hypot( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def isclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def conj( + x: Union[ivy.Array, ivy.NativeArray], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a boolean array where two arrays are element-wise equal within a tolerance. + Return the complex conjugate for each element ``x_i`` of the input array ``x``. - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(b)) and the absolute difference - atol are added together to compare against the absolute difference - between a and b. - The default atol is not appropriate for comparing numbers that are - much smaller than one + For complex number of the form + + .. math:: + a + bj + + the complex conjugate is defined as + + .. math:: + a - bj + + Hence, the returned conjugates must be computed by negating + the imaginary component of each element ``x_i`` + + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Parameters ---------- - a - First input array. - b - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in a will be - considered equal to NaN's in b in the output array. + x + input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - Returns a boolean array of where a and b are equal within the given - tolerance. If both a and b are scalars, returns a single boolean value. + an arrray of the same dtype as the input array with + the complex conjugates of the complex values present + in the input array. If x is a scalar then a scalar + will be returned. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. + Examples -------- - >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) - ivy.array([True, False]) - >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) - ivy.array([True, True]) - >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) - ivy.array([False, False]) - >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) - ivy.array([False, True]) + With :class:`ivy.Array` inputs: + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.conj(x) + >>> print(z) + ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) + + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), + ... b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.conj(x) + >>> print(z) + { + a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), + b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) + } """ - return ivy.current_backend().isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) + return ivy.current_backend(x).conj(x, out=out) @handle_exceptions @@ -928,166 +1196,30 @@ def lerp( end = ivy.array([end]) if ( ivy.dtype(input) not in input_end_allowed_types - or ivy.dtype(end) not in input_end_allowed_types - ): - input = ivy.astype(input, "float64") - end = ivy.astype(end, "float64") - - if ivy.is_array(weight): - if ivy.dtype(weight) not in weight_allowed_types: - weight = ivy.astype(weight, "float64") - else: - if not isinstance(weight, float): - weight = ivy.astype(ivy.array([weight]), "float64") - - return ivy.add(input, ivy.multiply(weight, ivy.subtract(end, input)), out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -def lgamma( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the natural logarithm of the absolute value of the gamma function on x. - - Parameters - ---------- - x - input array. Should have a floating-point data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the natural log of Gamma(x) of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. - - Examples - -------- - >>> x = ivy.array([1.6, 2.6, 3.5]) - >>> y = x.lgamma() - >>> print(y) - ivy.array([-0.11259177, 0.3574118 , 1.20097363]) - - >>> x = ivy.array([1., 2., 3. ]) - >>> y = x.lgamma() - >>> print(y) - ivy.array([0. ,0. ,0.69314718]) - - >>> x = ivy.array([4.5, -4, -5.6]) - >>> x.lgamma(out = x) - >>> print(x) - ivy.array([2.45373654, inf, -4.6477685 ]) - """ - return ivy.current_backend(x).lgamma(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def modf( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: - """ - Decompose the elements of x into fractional and integral parts. - - Parameters - ---------- - x - Input array. - out - Optional output array for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - A tuple of two arrays, the fractional and integral parts. - - Examples - -------- - >>> x = ivy.array([1.5, 2.7, 3.9]) - >>> ivy.modf(x) - (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) - """ - return ivy.current_backend(x).modf(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nansum( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as - zero. - - Parameters - ---------- - x - Input array. - axis - Axis or axes along which the sum is computed. - The default is to compute the sum of the flattened array. - dtype - The type of the returned array and of the accumulator in - which the elements are summed. By default, the dtype of input is used. - keepdims - If this is set to True, the axes which are reduced are left - in the result as dimensions with size one. - out - Alternate output array in which to place the result. - The default is None. + or ivy.dtype(end) not in input_end_allowed_types + ): + input = ivy.astype(input, "float64") + end = ivy.astype(end, "float64") - Returns - ------- - ret - A new array holding the result is returned unless out is specified, - in which it is returned. + if ivy.is_array(weight): + if ivy.dtype(weight) not in weight_allowed_types: + weight = ivy.astype(weight, "float64") + else: + if not isinstance(weight, float): + weight = ivy.astype(ivy.array([weight]), "float64") - Examples - -------- - >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) - >>> ivy.nansum(a) - 10.0 - >>> ivy.nansum(a, axis=0) - ivy.array([2.1, 5.8, 2.1]) - >>> ivy.nansum(a, axis=1) - ivy.array([5.5, 4.5]) - """ - return ivy.current_backend().nansum( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out - ) + return ivy.add(input, ivy.multiply(weight, ivy.subtract(end, input)), out=out) + + +lerp.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} @handle_exceptions @@ -1097,154 +1229,112 @@ def nansum( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def nextafter( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def frexp( + x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> bool: + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Return the next floating-point value after x1 towards x2, element-wise. + Decompose the elements of x into mantissa and twos exponent. Parameters ---------- - x1 - First input array. - x2 - Second input array. + x + Input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - The next representable values of x1 in the direction of x2. + A tuple of two arrays, the mantissa and the twos exponent. Examples -------- - >>> x1 = ivy.array([1.0e-50, 2.0e+50]) - >>> x2 = ivy.array([2.0, 1.0]) - >>> ivy.nextafter(x1, x2) - ivy.array([1.4013e-45., 3.4028e+38]) + >>> x = ivy.array([1, 2, 3]) + >>> ivy.frexp(x) + (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) """ - return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) + return ivy.current_backend(x).frexp(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def signbit( - x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], +def modf( + x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Return element-wise True where signbit is set (less than zero). + Decompose the elements of x into fractional and integral parts. Parameters ---------- x - Array-like input. + Input array. out - optional output array, for writing the result to. + Optional output array for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - Output array, or reference to out if that was supplied. - This is a scalar if x is a scalar. + A tuple of two arrays, the fractional and integral parts. Examples -------- - >>> x = ivy.array([1, -2, 3]) - >>> ivy.signbit(x) - ivy.array([False, True, False]) + >>> x = ivy.array([1.5, 2.7, 3.9]) + >>> ivy.modf(x) + (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) """ - return ivy.current_backend(x).signbit(x, out=out) + return ivy.current_backend(x).modf(x, out=out) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def sinc( +def digamma( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate an implementation-dependent approximation of the principal value of the - normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain - ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element - ``x_i`` is assumed to be expressed in radians. - - **Special cases** - - For floating-point operands, + Compute the logarithmic derivative of the gamma function at x. - - If x_i is NaN, the result is NaN. - - If ``x_i`` is ``0``, the result is ``1``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + Note + ---- + The Ivy version only accepts real-valued inputs. Parameters ---------- x - input array. Should have a floating-point data type. + Input array. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - an array containing the normalized sinc function of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. + Array with values computed from digamma function from + input arrays' values, element-wise. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = x.sinc() - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - >>> x = ivy.array([1.5, 0.5, -1.5]) - >>> y = ivy.zeros(3) - >>> ivy.sinc(x, out=y) - >>> print(y) - ivy.array([-0.212,0.637,-0.212]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = ivy.sinc(x) - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), - ... b=ivy.array([3.5, 4.5, 5.5])) - >>> y = x.sinc() - >>> print(y) - { - a: ivy.array([0.637,-0.212,0.127]), - b: ivy.array([-0.0909,0.0707,-0.0579]) - } + >>> x = ivy.array([.9, 3, 3.2]) + >>> y = ivy.digamma(x) + ivy.array([-0.7549271 0.92278427 0.9988394]) """ - return ivy.current_backend(x).sinc(x, out=out) + return ivy.current_backend(x).digamma(x, out=out) @handle_exceptions @@ -1298,93 +1388,3 @@ def sparsify_tensor( tensor = ivy.concat([ivy.zeros(len(x) - card, dtype=x.dtype), x[-card:]], axis=0) return ivy.reshape(tensor, _shape, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def xlogy( - x: Union[ivy.Array, ivy.NativeArray], - y: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> bool: - """ - Compute x*log(y) element-wise so that the result is 0 if x = 0. - - Parameters - ---------- - x - First input array. - y - Second input array. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - The next representable values of x1 in the direction of x2. - - Examples - -------- - >>> x = ivy.zeros(3) - >>> y = ivy.array([-1.0, 0.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([0.0, 0.0, 0.0]) - - >>> x = ivy.array([1.0, 2.0, 3.0]) - >>> y = ivy.array([3.0, 2.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([1.0986, 1.3863, 0.0000]) - """ - return ivy.current_backend(x, y).xlogy(x, y, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def zeta( - x: Union[ivy.Array, ivy.NativeArray], - q: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> bool: - """ - Compute the Hurwitz zeta function elementwisely with each pair of floats in two - arrays. - - Parameters - ---------- - x - First input array. - q - Second input array, must have the same shape as the first input array - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - Array with values computed from zeta function from - input arrays' values. - - Examples - -------- - >>> x = ivy.array([5.0, 3.0]) - >>> q = ivy.array([2.0, 2.0]) - >>> ivy.zeta(x, q) - ivy.array([0.0369, 0.2021]) - """ - return ivy.current_backend(x, q).zeta(x, q, out=out) diff --git a/ivy/functional/ivy/experimental/general.py b/ivy/functional/ivy/experimental/general.py index faa8df1e69b2a..0cb749b476ffa 100644 --- a/ivy/functional/ivy/experimental/general.py +++ b/ivy/functional/ivy/experimental/general.py @@ -13,10 +13,6 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # - - def _correct_ivy_callable(func): # get the current backend of the given ivy callable if ivy.nested_any( @@ -29,10 +25,6 @@ def _correct_ivy_callable(func): return func -# --- Main --- # -# ------------ # - - @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py index 60cb61ec63724..120ebe0279784 100644 --- a/ivy/functional/ivy/experimental/layers.py +++ b/ivy/functional/ivy/experimental/layers.py @@ -21,1084 +21,1104 @@ from ivy.functional.ivy.experimental.general import _correct_ivy_callable from ivy.utils.exceptions import handle_exceptions - -# --- Helpers --- # -# --------------- # +_min = builtins.min +_slice = builtins.slice +_max = builtins.max -def _cast_init(init, dtype): - if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): - if ivy.is_float_dtype(dtype): - info = ivy.finfo(dtype) - else: - info = ivy.iinfo(dtype) - if "float64" not in str(dtype): - init = info.max if init > 0 else info.min - return ivy.array(init, dtype=dtype) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool1d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 1-D max pool given 3-D input x. + Parameters + ---------- + x + Input image *[batch_size, w, d_in]* if data_format is "NWC". + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) + data_format + "NWC" or "NCW". Defaults to "NWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. -def _compute_idx(in_size, out_size, device): - out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) - i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) - maxlength = in_size // out_size + 1 - in_size_mod = in_size % out_size - # adaptive = True iff there are kernels with different lengths - adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) - if adaptive: - maxlength += 1 - elif in_size_mod == 0: - maxlength -= 1 - range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) - idx = ivy.expand_dims(i0, axis=-1) + range_max - if adaptive: - maxval = ivy.full_like(idx, fill_value=in_size - 1) - idx = ivy.minimum(idx, maxval) - i1 = ivy.trunc_divide( - (out_range + 1) * in_size + out_size - 1, out_size - ).astype(ivy.int64) - length = i1 - i0 - else: - length = maxlength - return idx, length, range_max, adaptive + Returns + ------- + ret + The result of the pooling operation. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. -def _compute_weight_mat( - input_size, - output_size, - scale, - align_corners, - kernel_fn, - antialias: bool, - dim_scale_factor, -): - inv_scale = 1.0 / scale - kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 - if not align_corners: - sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 - x = ( - ivy.abs( - ivy.expand_dims(sample_f) - - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) - / kernel_scale - ) - else: - sample_f = ivy.arange(output_size) * dim_scale_factor - x = ivy.abs( - ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) / (kernel_scale) - weights = kernel_fn(x) - total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) - weights = ivy.where( - ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), - ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), - 0, - ) - input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 - return ivy.where( - ivy.expand_dims( - ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) - ), - weights, - 0, - ) + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 4., 5., 6., 7.]], -def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): - def _pad(arr, pads, pad_value): - out = ivy.astype( - ivy.pad( - arr, - ivy.maximum(0, pads).to_list(), - mode="constant", - constant_values=ivy.to_scalar(pad_value), - ), - arr.dtype, - ) - slices = tuple( - _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) - for (lo, hi), dim in zip(pads, arr.shape) - ) - return out[slices] + [[16., 17., 18., 19.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa + ivy.array([[[ 1., 3.], + [ 5., 7.], + [ 9., 11.]], - if ( - _min(lhs.ndim, len(rhs_shape)) < 2 - or lhs.ndim != len(rhs_shape) - or lhs.shape[1] != rhs_shape[1] - ): - raise ValueError("Dimension mismatch") - if len(window_strides) != len(rhs_shape) - 2: - raise ValueError("Wrong number of strides for spatial dimensions") - if len(pads) != len(rhs_shape) - 2: - raise ValueError("Wrong number of pads for spatial dimensions") + [[13., 15.], + [17., 19.], + [21., 23.]]]) + """ + return ivy.current_backend(x).max_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) - lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) - in_shape = lhs.shape[2:] - filter_shape = rhs_shape[2:] - dim = len(filter_shape) - out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() - view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_unpool1d( + x: ivy.Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 1-D max unpooling given the 1-D pooled input x and its indices. - out_shape = [ - (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) - ] - view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] + Parameters + ---------- + x + Pooled input image *[batch_size, w, d_in]*. + indices + Indices obtained from the corresponding max pooling operation. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + out + optional output array, for writing the result to. - view = ivy.as_strided(lhs, view_shape, view_strides) + Returns + ------- + ret + The result of the unpooling operation. - view_axes = list(range(view.ndim)) - sum_axes = view_axes[-dim - 1 :] - rhs_axes = [view.ndim] + sum_axes - out_axes = [0, view.ndim] + list(range(1, dim + 1)) + Both the description and the type hints above assume an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - return view, view_axes, rhs_axes, out_axes + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') + >>> print(pool_result) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') + >>> print(unpool_result) + ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], + [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], -def _cubic_kernel(x): - out = ((1.5 * x - 2.5) * x) * x + 1.0 - out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) - return ivy.where(x >= 2.0, 0.0, out) + [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], + [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) + """ + return ivy.current_backend(x).max_unpool1d( + x, indices, kernel, strides, padding, data_format=data_format, out=out + ) -def _dilate(operand, factors, fill_value): - outspace = list(operand.shape[:2]) + [ - shape + (factors[i] - 1) * (shape - 1) - for i, shape in enumerate(operand.shape[2:]) - ] - out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) - lhs_slices = tuple(_slice(None, None, step) for step in factors) - out[(_slice(None),) * 2 + lhs_slices] = operand - return out +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool2d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D max pool given 4-D input x. + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NHWC" or "NCHW". Defaults to "NHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. -def _dim_scale_factor(input_size, output_size, align_corners, scales): - if align_corners: - if output_size > 1: - dim_scale_factor = (input_size - 1) / (output_size - 1) - else: - dim_scale_factor = 0.0 - else: - dim_scale_factor = ( - input_size / (input_size * scales) - if scales is not None - else input_size / output_size - ) - return dim_scale_factor + Returns + ------- + ret + The result of the pooling operation. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. -def _expand_to_dim(x, dim): - for _ in range(dim - len(x.shape)): - x = ivy.expand_dims(x, axis=-1) - return x + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 2., 3.], + [ 4., 5.], + [ 4., 5.]]], -def _get_identity(func, dtype, init): - func_name = func.__name__ - if func_name in identities: - identity = identities[func_name] - return _cast_init(identity, dtype) - return init + [[[ 8., 9.], + [10., 11.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[16., 17.]], -def _get_size(scale_factor, size, dims, x_shape): - if scale_factor is not None: - if isinstance(scale_factor, (float, int)): - scale_factor = [scale_factor] * dims - elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: - scale_factor = [scale_factor[0]] * dims + [[22., 23.]]], - size = tuple( - [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] - ) - else: - size = (size,) * dims if isinstance(size, int) else tuple(size) - return size + [[[40., 41.]], -def _interpolate_with_kernel( - x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode -): - spatial_dims = [2 + i for i in range(dims)] - equation = generate_einsum_equation(dims) - kernel_func = get_interpolate_kernel(mode) - output_shape = tuple(input_shape[:2]) + size - operands = [] - for i, d in enumerate(spatial_dims): - m = input_shape[d] - n = output_shape[d] - dim_scale_factor = _dim_scale_factor( - m, - n, - align_corners, - scale_factor[i] if scale_factor is not None else None, - ) - w = _compute_weight_mat( - m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor - ).astype(x.dtype) - operands.append(w) - return ivy.einsum(equation, x, *operands) + [[46., 47.]]]]) + """ + return ivy.current_backend(x).max_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) -def _lanczos_kernel(radius, x): - y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) - out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) - return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool3d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 3-D max pool given 5-D input x. + Parameters + ---------- + x + Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) + data_format + "NDHWC" or "NCDHW". Defaults to "NDHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. -def _mask(vals, length, range_max, dim, mask_value=0.0): - if isinstance(length, int): - return vals, length - else: - assert dim < 0 - mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) - if dim == -2: - mask = _expand_to_dim(mask, 4) - vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) - length = _expand_to_dim(length, -dim) - return vals, length + Returns + ------- + ret + The result of the pooling operation. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. -def _mitchellcubic_kernel(x): - absx = abs(x) - if absx < 1: - return (7 * absx**3 - 12 * absx**2 + 6) / 6 - elif absx < 2: - return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 - else: - return 0 + Examples + -------- + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) + ivy.array([[[[[14., 15.]]]], -def _output_ceil_shape(w, f, p, s): - return math.ceil((w - f + p) / s) + 1 + [[[[38., 39.]]]]]) + >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) + ivy.array([[[[[14., 15.]]], -def _padding_ceil_mode(w, f, p, s, return_added_padding=False): - remaining_pixels = (w - f + sum(p)) % s - added_padding = 0 - if s > 1 and remaining_pixels != 0 and f > 1: - input_size = w + sum(p) - # making sure that the remaining pixels are supposed - # to be covered by the window - # they won't be covered if stride is big enough to skip them - if input_size - remaining_pixels - (f - 1) + s > input_size: - return p - output_shape = _output_ceil_shape( - w, - f, - sum(p), - s, - ) - # calculating new padding with ceil_output_shape - new_pad = (output_shape - 1) * s + f - w - # updating pad_list with new padding by adding it to the end - added_padding = new_pad - sum(p) - p = ( - p[0], - p[1] + added_padding, - ) - if return_added_padding: - return p, added_padding - return p - - -def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): - if padding.upper() == "SAME": - out_shape = [ - math.ceil(in_size / stride) - for in_size, stride in zip(in_shape, window_strides) - ] - pad_sizes = [ - _max((out_size - 1) * stride + filter_size - in_size, 0) - for out_size, stride, filter_size, in_size in zip( - out_shape, window_strides, filter_shape, in_shape - ) - ] - return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] - else: - return [(0, 0)] * len(in_shape) + [[[22., 23.]]]], -def _sum_tensors(ts): - return _reduce(ivy.add, ts) -def _tf_area_dim_scale(index, starting_index, scale, ending_index): - if index < starting_index: - dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index - else: - dim_scale = ending_index - index if index + 1 > ending_index else 1.0 - return dim_scale + [[[[38., 39.]]], -def _tf_area_indices(dim_index, scale): - starting_index = dim_index * scale - ending_index = (dim_index + 1) * scale - rounded_indices = ( - int(starting_index), - math.ceil(ending_index), + [[[46., 47.]]]]]) + """ + return ivy.current_backend(x).max_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, ) - return starting_index, ending_index, rounded_indices - - -def _tf_area_interpolate(x, size, dims): - ret = ivy.zeros((x.shape[:2] + size)) - scale = ivy.divide(ivy.shape(x)[2:], size) - area = 1.0 / ivy.prod(scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) - sum_data = ivy.zeros( - ( - d_index[1] - d_index[0], - h_index[1] - h_index[0], - w_index[1] - w_index[0], - ) - ) - for d_ind in range(d_index[0], d_index[1]): - scale_z = _tf_area_dim_scale( - d_ind, d_in, scale[0], d_in1 - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale( - h_ind, h_in, scale[1], h_in1 - ) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[2], w_in1 - ) - sum_data[ - d_ind - d_index[0], - h_ind - h_index[0], - w_ind - w_index[0], - ] = ( - ivy.array(ch[d_ind, h_ind, w_ind]) - * scale_x - * scale_y - * scale_z - * area - ) - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) - sum_data = ivy.zeros( - (h_index[1] - h_index[0], w_index[1] - w_index[0]) - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[1], w_in1 - ) - sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( - ivy.array(ch[h_ind, w_ind]) - * scale_x - * scale_y - * area - ) - ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) - else: - for w_dim in range(size[0]): - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) - sum_data = ivy.zeros((w_index[1] - w_index[0],)) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) - sum_data[w_ind - w_index[0]] = ( - ivy.array(ch[w_ind]) * scale_x * area - ) - ret[i, j, w_dim] = ivy.sum(sum_data) - return ret -def _triangle_kernel(x): - return ivy.maximum(0, 1 - ivy.abs(x)) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def avg_pool1d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + division_override: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 1-D avg pool given 3-D input x. + Parameters + ---------- + x + Input image *[batch_size, w, d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + division_override + If specified, it will be used as the divisor, + otherwise kernel_size will be used. + out + optional output array, for writing the result to. -def _upsample_bicubic2d_default( - a, - output_size, - align_corners, - scale_h=None, - scale_w=None, -): - N, C, iH, iW = a.shape - oH, oW = output_size + Returns + ------- + ret + The result of the pooling operation. - def compute_scale(in_size, out_size, align_corners, scale=None): - if align_corners: - return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 - else: - return 1 / scale if scale is not None and scale > 0 else in_size / out_size + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - def compute_source_index(scale, dst_index, align_corners): - if align_corners: - return scale * dst_index - else: - return scale * (dst_index + 0.5) - 0.5 + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 2., 3., 4., 5.], + [ 8., 9., 10., 11.]], - height_scale = compute_scale(iH, oH, align_corners, scale_h) - width_scale = compute_scale(iW, oW, align_corners, scale_w) + [[14., 15., 16., 17.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 2., 3., 4., 5.]], - N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) - C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) - out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) - out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) + [[14., 15., 16., 17.]]]) + """ + return ivy.current_backend(x).avg_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + division_override=division_override, + out=out, + ) - real_x = compute_source_index(width_scale, out_x, align_corners) - in_x = ivy.floor(real_x) - t_x = real_x - in_x - ix = ivy.astype(in_x, ivy.int64) - real_y = compute_source_index(height_scale, out_y, align_corners) - in_y = ivy.floor(real_y) - t_y = real_y - in_y - iy = ivy.astype(in_y, ivy.int64) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def avg_pool2d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, + /, + *, + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D average pool given 4-D input x. - iys_ofs = (iy - 1, iy, iy + 1, iy + 2) - ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) - - def load_bounded(ys, xs): - y_idx = ivy.clip(ys, 0, iH - 1) - x_idx = ivy.clip(xs, 0, iW - 1) - return a[N_idx, C_idx, y_idx, x_idx] - - def get_x_interp(y): - coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) - return _upsample_cubic_interp1d(coeffs_x, t_x) - - coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) - result = _upsample_cubic_interp1d(coeffs_y, t_y) + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimensio paddings. + data_format + NHWC" or "NCHW". Defaults to "NHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + out + optional output array, for writing the result to. - return result + Returns + ------- + ret + The result of the pooling operation. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. -def _upsample_cubic_convolution1(x, A): - return ((A + 2) * x - (A + 3)) * x * x + 1 + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], -def _upsample_cubic_convolution2(x, A): - return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[ 8., 9.]], + [[14., 15.]]], -def _upsample_cubic_interp1d(coeffs, ts): - coeffs2 = _upsample_get_cubic_coefficients(ts) - return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) + [[[32., 33.]], -def _upsample_get_cubic_coefficients(t): - A = -0.75 - return ( - _upsample_cubic_convolution2(t + 1.0, A), - _upsample_cubic_convolution1(t, A), - _upsample_cubic_convolution1(1.0 - t, A), - _upsample_cubic_convolution2(2.0 - t, A), + [[38., 39.]]]]) + """ + return ivy.current_backend(x).avg_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, ) -# --- Main --- # -# ------------ # - - +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_avg_pool1d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: int, +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def avg_pool3d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, + /, + *, + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 1D adaptive average pooling over an input signal composed of several input - planes. + Compute a 3-D avg pool given 5-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, L_in) or (C, L_in) where N is - the batch dimension, C is the feature dimension, and L_in is the spatial - dimension. - output_size - Spatial output size. + x + Input volume *[batch_size,d,h,w,d_in]*. + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension + paddings. + data_format + NDHWC" or "NCDHW". Defaults to "NDHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + divisor_override + If specified, it will be used as divisor, otherwise kernel_size will be used. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - The result of the pooling operation. Will have shape (N, C, L_out) or - (C, L_out), where L_out = `output_size` - """ - squeeze = False - if input.ndim == 2: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 3: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", - ) + ret + The result of the pooling operation. - if input.shape[-1] % output_size == 0: - stride = input.shape[-1] // output_size - kernel_size = input.shape[-1] - (output_size - 1) * stride - pooled_output = ivy.avg_pool1d( - input, kernel_size, stride, "VALID", data_format="NCW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size, input.device - ) + Examples + -------- + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.avg_pool3d(x,2,2,'VALID')) + ivy.array([[[[[ 7., 8.]]]], - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., idxw]) - if not adaptive_w: - ret = ivy.mean(vals, axis=-1) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + [[[[31., 32.]]]]]) + >>> print(ivy.avg_pool3d(x,2,2,'SAME')) + ivy.array([[[[[ 7., 8.]]], - ret = None - for i in range(vals.shape[-1]): - if ret is None: - ret = vals[..., i] - else: - ret = ret + vals[..., i] - pooled_output = ret / length_w.astype(ret.dtype) - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[[19., 20.]]]], -@handle_exceptions + + [[[[31., 32.]]], + + + [[[43., 44.]]]]]) + """ + return ivy.current_backend(x).avg_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, + ) + + +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def adaptive_avg_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], +@handle_out_argument +@to_native_arrays_and_back +def pool( + x: Union[ivy.Array, ivy.NativeArray], + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, + *, + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 2D adaptive average pooling over an input signal composed of several input - planes. + Perform an N-D pooling operation. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + Input array to pool over. + window_shape + Shape of the pooling window. + pool_type + Type of pooling operation, either 'MAX' or 'AVG'. + strides + Strides of the pooling operation. + padding + Padding type, either 'VALID' or 'SAME'. + data_format + Data format of the input and output data, either 'NCHW' or 'NHWC'. + dilations + Dilation rate of the pooling operation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) - - if isinstance(output_size, int): - output_size = (output_size, output_size) - - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.avg_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output + ret + The result of the pooling operation. - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) + ivy.array([[[[ 8., 9.]], + [[14., 15.]]], + [[[32., 33.]], + [[38., 39.]]]]) + """ + return ivy.current_backend(x).pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ceil_mode=ceil_mode, + out=out, ) - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) - - if not adaptive_h and not adaptive_w: - ret = ivy.mean(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret - - vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) - - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ret + vals[..., i, :, j] - pooled_output = ret / (length_h * length_w).astype(vals.dtype) - - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output - +@handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_max_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], -): +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dct( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Apply a 2D adaptive maximum pooling over an input signal composed of several input - planes. + Compute the 1D Discrete Cosine Tranformation of a given signal. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + The input signal. + type + The type of the dct. Must be 1, 2, 3 or 4. + n + The lenght of the transform. If n is less than the input signal lenght, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the DCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". + out + optional output array, for writing the result to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) + ret + Array containing the transformed input. - if isinstance(output_size, int): - output_size = (output_size, output_size) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.max_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output + Examples + -------- + With :class:`ivy.Array` input: - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device - ) + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.dct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, + 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array( - input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device - ) + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], - if not adaptive_h and not adaptive_w: - ret = ivy.max(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) - vals, length_h = _mask( - vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") - ) - vals, length_w = _mask( - vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") - ) + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.dct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ivy.maximum(ret, vals[..., i, :, j]) - pooled_output = ret.astype(vals.dtype) + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.dct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + With one :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.dct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475, 5.19664812, -1.95411837]), + b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, + -1.25980937, 0.64958102, -0.2442648]) + } -def area_interpolate(x, dims, size, scale): - ret = ivy.zeros((x.shape[:2] + size)) - inv_scale = ivy.divide(1.0, scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_index = ( - int(d_dim * inv_scale[0]), - math.ceil((d_dim + 1) * inv_scale[0]), - ) - h_index = ( - int(h_dim * inv_scale[1]), - math.ceil((h_dim + 1) * inv_scale[1]), - ) - w_index = ( - int(w_dim * scale[2]), - math.ceil((w_dim + 1) * inv_scale[2]), - ) - scale_z = d_index[1] - d_index[0] - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_z * scale_y * scale_x - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( - ch[ - d_index[0] : d_index[1], - h_index[0] : h_index[1], - w_index[0] : w_index[1], - ] - ) * (1 / area) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_index = ( - int(h_dim * inv_scale[0]), - math.ceil((h_dim + 1) * inv_scale[0]), - ) - w_index = ( - int(w_dim * inv_scale[1]), - math.ceil((w_dim + 1) * inv_scale[1]), - ) - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_y * scale_x - ret[i, j, h_dim, w_dim] = ivy.sum( - ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] - ) * (1 / area) - else: - for w_dim in range(size[0]): - w_index = ( - int(w_dim * inv_scale[0]), - math.ceil((w_dim + 1) * inv_scale[0]), - ) - scale_x = w_index[1] - w_index[0] - ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( - 1 / scale_x - ) - return ret + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, + -26.00041008, 19.75149155, -16.97056389, 10.87819386, + -5.89381361]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } + """ + return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) -@handle_backend_invalid +@handle_exceptions @handle_nestable @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def avg_pool1d( +def idct( x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, /, *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - division_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Compute a 1-D avg pool given 3-D input x. + Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. Parameters ---------- x - Input image *[batch_size, w, d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - division_override - If specified, it will be used as the divisor, - otherwise kernel_size will be used. + The input signal. + type + The type of the idct. Must be 1, 2, 3 or 4. + n + The length of the transform. If n is less than the input signal length, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the IDCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". out optional output array, for writing the result to. Returns ------- ret - The result of the pooling operation. + Array containing the transformed input. - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 2., 3., 4., 5.], - [ 8., 9., 10., 11.]], + With :class:`ivy.Array` input: - [[14., 15., 16., 17.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 2., 3., 4., 5.]], + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.idct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475 , 5.19664812, -1.95411837]) - [[14., 15., 16., 17.]]]) + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], + + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) + + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.idct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) + + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.idct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.idct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, + -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, + -8.80319249e-08, -4.05617893e-01]), + b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, + -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, + -1.10039906e-08, -5.07022366e-02]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, + -16.18951607, 18.06697273, -17.57439613, 11.68861485, + -4.41308832]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } """ - return ivy.current_backend(x).avg_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - division_override=division_override, - out=out, - ) + return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) + + +idct.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def avg_pool2d( +def fft( x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, + dim: int, /, *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute a 2-D average pool given 4-D input x. + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. Parameters ---------- x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimensio paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT. + dim + The dimension along which to take the one dimensional FFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing FFT. + Should be a integer greater than 1. out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The result of the pooling operation. - - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + The result of the FFT operation. Examples -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], - - - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[ 8., 9.]], - - [[14., 15.]]], - - - [[[32., 33.]], - - [[38., 39.]]]]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, + 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, + 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, + 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, + 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, + 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, + 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, + 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, + 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, + 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) """ - return ivy.current_backend(x).avg_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, - ) + return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def avg_pool3d( +def dropout1d( x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, + prob: float, /, *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, + training: bool = True, + data_format: str = "NWC", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 3-D avg pool given 5-D input x. + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D + feature map. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]*. - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension - paddings. + a 2D or 3D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout1d is performed during training or ignored + during testing. data_format - NDHWC" or "NCDHW". Defaults to "NDHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - divisor_override - If specified, it will be used as divisor, otherwise kernel_size will be used. + "NWC" or "NCW". Defaults to "NWC". out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - The result of the pooling operation. + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.avg_pool3d(x,2,2,'VALID')) - ivy.array([[[[[ 7., 8.]]]], - - - - [[[[31., 32.]]]]]) - >>> print(ivy.avg_pool3d(x,2,2,'SAME')) - ivy.array([[[[[ 7., 8.]]], - - - [[[19., 20.]]]], - - + With :class:`ivy.Array` input: - [[[[31., 32.]]], + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + ivy.array([[[2., 0, 2.]]]) + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") + >>> print(y) + ivy.array([[[1, 1, 1]]]) - [[[43., 44.]]]]]) + With one :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), + ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + { + a: ivy.array([[[200., 400., 0.]]]), + b: ivy.array([[[0., 0., 0.]]]) + } """ - return ivy.current_backend(x).avg_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, + return ivy.current_backend(x).dropout1d( + x, prob, training=training, data_format=data_format, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dct( +def dropout2d( x: Union[ivy.Array, ivy.NativeArray], + prob: float, /, *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + training: bool = True, + data_format: str = "NHWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Compute the 1D Discrete Cosine Tranformation of a given signal. + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D + feature map. Parameters ---------- x - The input signal. - type - The type of the dct. Must be 1, 2, 3 or 4. - n - The lenght of the transform. If n is less than the input signal lenght, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the DCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". + a 3D or 4D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout2d is performed during training or ignored + during testing. + data_format + "NHWC" or "NCHW". Defaults to "NHWC". out optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - Array containing the transformed input. + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1108,147 +1128,19 @@ def dct( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.dct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, - 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) - - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], - - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) - - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.dct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) - - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.dct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) - - With one :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.dct(x, type=3, n=None, norm='ortho') + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 0.5) >>> print(y) - { - a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475, 5.19664812, -1.95411837]), - b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, - -1.25980937, 0.64958102, -0.2442648]) - } - - With multiple :class:`ivy.Container` inputs: + ivy.array([[0., 2., 2.]]) - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") >>> print(y) - { - a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, - -26.00041008, 19.75149155, -16.97056389, 10.87819386, - -5.89381361]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } + ivy.array([[1, 1, 1]]) """ - return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) - - -@handle_exceptions -@handle_nestable -@handle_out_argument -@inputs_to_ivy_arrays -def dft( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = 1, - inverse: bool = False, - onesided: bool = False, - dft_length: Optional[Union[int, Tuple[int]]] = None, - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the discrete Fourier transform of input. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT. - axis - The axis on which to perform the DFT. By default this - value is set to 1, which corresponds to the first dimension - after the batch index. - inverse - Whether to perform the inverse discrete fourier transform. - By default this value is set to False. - onesided - If onesided is True, only values for w in [0, 1, 2, …, floor(n_fft/2) + 1] - are returned because the real-to-complex Fourier transform satisfies the - conjugate symmetry, i.e., X[m, w] = X[m,w]=X[m,n_fft-w]*. Note if the - input or window tensors are complex, then onesided output is not possible. - Enabling onesided with real inputs performs a Real-valued fast Fourier - transform (RFFT). When invoked with real or complex valued input, the - default value is False. Values can be True or False. - dft_length - The length of the signal.If greater than the axis dimension, - the signal will be zero-padded up to dft_length. If less than - the axis dimension, only the first dft_length values will be - used as the signal. It’s an optional value. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be - "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by 1/sqrt(n). - "forward" indicates normalization by 1/n. - out - Optional output array, for writing the result to. It must - have a shape that the inputs broadcast to. - - Returns - ------- - ret - The Fourier Transform of the input vector.If onesided is False, - the following shape is expected: [batch_idx][signal_dim1][signal_dim2] - …[signal_dimN][2]. If axis=0 and onesided is True, the following shape - is expected: [batch_idx][floor(signal_dim1/2)+1][signal_dim2]…[signal_dimN][2]. - If axis=1 and onesided is True, the following shape is expected: - [batch_idx][signal_dim1][floor(signal_dim2/2)+1]…[signal_dimN][2]. - If axis=N-1 and onesided is True, the following shape is expected: - [batch_idx][signal_dim1][signal_dim2]…[floor(signal_dimN/2)+1][2]. - The signal_dim at the specified axis is equal to the dft_length. - """ - if inverse: - res = ivy.ifft(x, axis, norm=norm, n=dft_length, out=out) - else: - res = ivy.fft(x, axis, norm=norm, n=dft_length, out=out) - - if onesided: - slices = [slice(0, a) for a in res.shape] - slices[axis] = slice(0, ivy.shape(res, as_array=True)[axis] // 2 + 1) - res = res[tuple(slices)] - return res + return ivy.current_backend(x).dropout2d( + x, prob, training=training, data_format=data_format, out=out + ) @handle_exceptions @@ -1258,32 +1150,32 @@ def dft( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout1d( +def dropout3d( x: Union[ivy.Array, ivy.NativeArray], prob: float, /, *, training: bool = True, - data_format: str = "NWC", + data_format: str = "NDHWC", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ Randomly zero out entire channels with probability prob using samples from a Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D + case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D feature map. Parameters ---------- x - a 2D or 3D input array. Should have a floating-point data type. + a 4D or 5D input array. Should have a floating-point data type. prob probability of a channel to be zero-ed. training - controls whether dropout1d is performed during training or ignored + controls whether dropout3d is performed during training or ignored during testing. data_format - "NWC" or "NCW". Defaults to "NWC". + "NDHWC" or "NCDHW". Defaults to "NDHWC". out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1297,32 +1189,8 @@ def dropout1d( Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - ivy.array([[[2., 0, 2.]]]) - - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[[1, 1, 1]]]) - - With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), - ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - { - a: ivy.array([[[200., 400., 0.]]]), - b: ivy.array([[[0., 0., 0.]]]) - } """ - return ivy.current_backend(x).dropout1d( + return ivy.current_backend(x).dropout3d( x, prob, training=training, data_format=data_format, out=out ) @@ -1334,115 +1202,66 @@ def dropout1d( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout2d( +def ifft( x: Union[ivy.Array, ivy.NativeArray], - prob: float, - /, + dim: int, *, - training: bool = True, - data_format: str = "NHWC", + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D - feature map. + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. Parameters ---------- x - a 3D or 4D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout2d is performed during training or ignored - during testing. - data_format - "NHWC" or "NCHW". Defaults to "NHWC". + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs IFFT. + dim + The dimension along which to take the one dimensional IFFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing IFFT. + Should be a integer greater than 1. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The result of the IFFT operation. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 0.5) - >>> print(y) - ivy.array([[0., 2., 2.]]) - - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[1, 1, 1]]) - """ - return ivy.current_backend(x).dropout2d( - x, prob, training=training, data_format=data_format, out=out - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def dropout3d( - x: Union[ivy.Array, ivy.NativeArray], - prob: float, - /, - *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D - feature map. - - Parameters - ---------- - x - a 4D or 5D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout3d is performed during training or ignored - during testing. - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, + 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, + 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, + 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, + 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, + 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, + 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, + 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, + 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, + 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) """ - return ivy.current_backend(x).dropout3d( - x, prob, training=training, data_format=data_format, out=out - ) + return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) @handle_exceptions @@ -1500,497 +1319,529 @@ def embedding( @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fft( +@inputs_to_ivy_arrays +def dft( x: Union[ivy.Array, ivy.NativeArray], - dim: int, /, *, + axis: int = 1, + inverse: bool = False, + onesided: bool = False, + dft_length: Optional[Union[int, Tuple[int]]] = None, norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. + """ + Compute the discrete Fourier transform of input. Parameters ---------- x Input volume *[...,d_in,...]*, where d_in indicates the dimension that needs FFT. - dim - The dimension along which to take the one dimensional FFT. + axis + The axis on which to perform the DFT. By default this + value is set to 1, which corresponds to the first dimension + after the batch index. + inverse + Whether to perform the inverse discrete fourier transform. + By default this value is set to False. + onesided + If onesided is True, only values for w in [0, 1, 2, …, floor(n_fft/2) + 1] + are returned because the real-to-complex Fourier transform satisfies the + conjugate symmetry, i.e., X[m, w] = X[m,w]=X[m,n_fft-w]*. Note if the + input or window tensors are complex, then onesided output is not possible. + Enabling onesided with real inputs performs a Real-valued fast Fourier + transform (RFFT). When invoked with real or complex valued input, the + default value is False. Values can be True or False. + dft_length + The length of the signal.If greater than the axis dimension, + the signal will be zero-padded up to dft_length. If less than + the axis dimension, only the first dft_length values will be + used as the signal. It’s an optional value. norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing FFT. - Should be a integer greater than 1. + Optional argument, "backward", "ortho" or "forward". Defaults to be + "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by 1/sqrt(n). + "forward" indicates normalization by 1/n. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Optional output array, for writing the result to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The result of the FFT operation. - - Examples - -------- - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, - 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, - 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, - 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, - 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, - 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, - 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, - 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, - 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, - 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) + The Fourier Transform of the input vector.If onesided is False, + the following shape is expected: [batch_idx][signal_dim1][signal_dim2] + …[signal_dimN][2]. If axis=0 and onesided is True, the following shape + is expected: [batch_idx][floor(signal_dim1/2)+1][signal_dim2]…[signal_dimN][2]. + If axis=1 and onesided is True, the following shape is expected: + [batch_idx][signal_dim1][floor(signal_dim2/2)+1]…[signal_dimN][2]. + If axis=N-1 and onesided is True, the following shape is expected: + [batch_idx][signal_dim1][signal_dim2]…[floor(signal_dimN/2)+1][2]. + The signal_dim at the specified axis is equal to the dft_length. """ - return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) + if inverse: + res = ivy.ifft(x, axis, norm=norm, n=dft_length, out=out) + else: + res = ivy.fft(x, axis, norm=norm, n=dft_length, out=out) + + if onesided: + slices = [slice(0, a) for a in res.shape] + slices[axis] = slice(0, ivy.shape(res, as_array=True)[axis] // 2 + 1) + res = res[tuple(slices)] + return res @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back -def fft2( - x: Union[ivy.Array, ivy.NativeArray], - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the 2-dimensional discrete Fourier Transform. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT2. - s - sequence of ints, optional - Shape (length of each transformed axis) of the output (s[0] refers to axis 0, - s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, - if the given shape is smaller than that of the input, the input is cropped. - If it is larger, the input is padded with zeros. if s is not given, the shape - of the input along the axes specified by axes is used. - dim - Axes over which to compute the FFT2. If not given, the last two axes are used. - A repeated index in axes means the transform over that axis is performed - multiple times. A one-element sequence means that a one-dimensional FFT is - performed. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. +@inputs_to_ivy_arrays +def interp(x, xp, fp, left=None, right=None, period=None): + x_arr = ivy.array(x) + fix_later = False + if x_arr.shape == (): + x_arr = ivy.array([x]) + fix_later = True + x = ivy.astype(x_arr, "float64") + xp = ivy.astype(ivy.array(xp), "float64") + fp = ivy.astype(ivy.array(fp), "float64") + ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) + if period is not None: + ivy.utils.assertions.check_equal(period, 0, inverse=True) + period = ivy.abs(period) + x = ivy.remainder(x, period) + xp = ivy.remainder(xp, period) + asort_xp = ivy.argsort(xp) + xp = xp[asort_xp] + fp = fp[asort_xp] + xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) + fp = ivy.concat((fp[-1:], fp, fp[0:1])) - Returns - ------- - ret - The result of the FFT2 operation. + def interp_inner(value): + value = ivy.array(value) + if value < xp[0]: + return left if left is not None else fp[0] + elif value > xp[-1]: + return right if right is not None else fp[-1] + else: + last = None + if xp.shape[0] < 3: + for i in range(xp.shape[0] - 1, -1, -1): + if xp[i] == value: + return fp[i] + elif xp[i] < value: + last = i + else: + first = 0 + last = xp.shape[0] + while first < last: + midpoint = (first + last) // 2 + if xp[midpoint] == value: + already_exists = ivy.argwhere(xp == value) + if already_exists.shape[0] > 0: + return fp[already_exists[-1][0]] + return fp[midpoint] + else: + if value < xp[midpoint]: + last = midpoint - 1 + else: + first = midpoint + 1 + dist = (value - xp[last]) / (xp[last + 1] - xp[last]) + return (fp[last + 1] - fp[last]) * dist + fp[last] - Examples - -------- - >>> x = ivy.array([[0, 0, 0, 0, 0], - ... [1, 1, 1, 1, 1], - ... [2, 2, 2, 2, 2], - ... [3, 3, 3, 3, 3], - ... [4, 4, 4, 4, 4]]) - >>> y = ivy.fft2(x) - >>> print(y) - ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5+17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5-17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ]]) - """ - return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) + ret = ivy.map(interp_inner, unique={"value": x}) + if fix_later: + return ivy.astype(ivy.array(ret[0]), "float64") + else: + return ivy.astype(ivy.array(ret), "float64") -def generate_einsum_equation(dim): - alphabet = "abcdefghijklmnopqrstuvwxyz" - input_indices = alphabet[: dim + 2] - output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] - contraction_indices = ",".join([input_indices, *output_indices]) - output = input_indices[:2] + "".join([output[-1] for output in output_indices]) - einsum_string = contraction_indices + "->" + output - return einsum_string +def _tf_area_dim_scale(index, starting_index, scale, ending_index): + if index < starting_index: + dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index + else: + dim_scale = ending_index - index if index + 1 > ending_index else 1.0 + return dim_scale -def get_interpolate_kernel(mode): - kernel_func = _triangle_kernel - if mode == "bicubic_tensorflow": - kernel_func = lambda inputs: _cubic_kernel(inputs) - elif mode == "lanczos3": - kernel_func = lambda inputs: _lanczos_kernel(3, inputs) - elif mode == "lanczos5": - kernel_func = lambda inputs: _lanczos_kernel(5, inputs) - return kernel_func +def _tf_area_indices(dim_index, scale): + starting_index = dim_index * scale + ending_index = (dim_index + 1) * scale + rounded_indices = ( + int(starting_index), + math.ceil(ending_index), + ) + return starting_index, ending_index, rounded_indices -@handle_exceptions -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def idct( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. +def _tf_area_interpolate(x, size, dims): + ret = ivy.zeros((x.shape[:2] + size)) + scale = ivy.divide(ivy.shape(x)[2:], size) + area = 1.0 / ivy.prod(scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) + sum_data = ivy.zeros( + ( + d_index[1] - d_index[0], + h_index[1] - h_index[0], + w_index[1] - w_index[0], + ) + ) + for d_ind in range(d_index[0], d_index[1]): + scale_z = _tf_area_dim_scale( + d_ind, d_in, scale[0], d_in1 + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale( + h_ind, h_in, scale[1], h_in1 + ) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[2], w_in1 + ) + sum_data[ + d_ind - d_index[0], + h_ind - h_index[0], + w_ind - w_index[0], + ] = ( + ivy.array(ch[d_ind, h_ind, w_ind]) + * scale_x + * scale_y + * scale_z + * area + ) + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) + sum_data = ivy.zeros( + (h_index[1] - h_index[0], w_index[1] - w_index[0]) + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[1], w_in1 + ) + sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( + ivy.array(ch[h_ind, w_ind]) + * scale_x + * scale_y + * area + ) + ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) + else: + for w_dim in range(size[0]): + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) + sum_data = ivy.zeros((w_index[1] - w_index[0],)) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) + sum_data[w_ind - w_index[0]] = ( + ivy.array(ch[w_ind]) * scale_x * area + ) + ret[i, j, w_dim] = ivy.sum(sum_data) + return ret - Parameters - ---------- - x - The input signal. - type - The type of the idct. Must be 1, 2, 3 or 4. - n - The length of the transform. If n is less than the input signal length, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the IDCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". - out - optional output array, for writing the result to. - Returns - ------- - ret - Array containing the transformed input. +def nearest_interpolate(x, dims, size, input_shape, exact): + off = 0.5 if exact else 0 + for d in range(dims): + m = input_shape[d + 2] + n = size[d] + offsets = (ivy.arange(n, dtype="float32") + off) * m / n + offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") + x = ivy.gather(x, offsets, axis=d + 2) + return x - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - Examples - -------- - With :class:`ivy.Array` input: +def _triangle_kernel(x): + return ivy.maximum(0, 1 - ivy.abs(x)) - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.idct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475 , 5.19664812, -1.95411837]) - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], +def _cubic_kernel(x): + out = ((1.5 * x - 2.5) * x) * x + 1.0 + out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) + return ivy.where(x >= 2.0, 0.0, out) - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.idct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) +def _lanczos_kernel(radius, x): + y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) + out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) + return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.idct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) - With one :class:`ivy.Container` input: +def _dim_scale_factor(input_size, output_size, align_corners, scales): + if align_corners: + if output_size > 1: + dim_scale_factor = (input_size - 1) / (output_size - 1) + else: + dim_scale_factor = 0.0 + else: + dim_scale_factor = ( + input_size / (input_size * scales) + if scales is not None + else input_size / output_size + ) + return dim_scale_factor - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.idct(x, type=3, n=None, norm='ortho') - >>> print(y) - { - a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, - -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, - -8.80319249e-08, -4.05617893e-01]), - b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, - -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, - -1.10039906e-08, -5.07022366e-02]) - } - With multiple :class:`ivy.Container` inputs: +def _mitchellcubic_kernel(x): + absx = abs(x) + if absx < 1: + return (7 * absx**3 - 12 * absx**2 + 6) / 6 + elif absx < 2: + return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 + else: + return 0 - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) - >>> print(y) - { - a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, - -16.18951607, 18.06697273, -17.57439613, 11.68861485, - -4.41308832]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } - """ - return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) +def _compute_weight_mat( + input_size, + output_size, + scale, + align_corners, + kernel_fn, + antialias: bool, + dim_scale_factor, +): + inv_scale = 1.0 / scale + kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 + if not align_corners: + sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 + x = ( + ivy.abs( + ivy.expand_dims(sample_f) + - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) + / kernel_scale + ) + else: + sample_f = ivy.arange(output_size) * dim_scale_factor + x = ivy.abs( + ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) / (kernel_scale) + weights = kernel_fn(x) + total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) + weights = ivy.where( + ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), + ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), + 0, + ) + input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 + return ivy.where( + ivy.expand_dims( + ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) + ), + weights, + 0, + ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def ifft( - x: Union[ivy.Array, ivy.NativeArray], - dim: int, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs IFFT. - dim - The dimension along which to take the one dimensional IFFT. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing IFFT. - Should be a integer greater than 1. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. +def _upsample_cubic_convolution1(x, A): + return ((A + 2) * x - (A + 3)) * x * x + 1 - Returns - ------- - ret - The result of the IFFT operation. - Examples - -------- - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, - 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, - 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, - 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, - 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, - 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, - 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, - 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, - 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, - 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) - """ - return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) +def _upsample_cubic_convolution2(x, A): + return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def ifftn( - x: Union[ivy.Array, ivy.NativeArray], - s: Optional[Union[int, Tuple[int, ...]]] = None, - axes: Optional[Union[int, Tuple[int, ...]]] = None, - *, - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the N-dimensional inverse discrete Fourier Transform. +def _upsample_get_cubic_coefficients(t): + A = -0.75 + return ( + _upsample_cubic_convolution2(t + 1.0, A), + _upsample_cubic_convolution1(t, A), + _upsample_cubic_convolution1(1.0 - t, A), + _upsample_cubic_convolution2(2.0 - t, A), + ) - Parameters - ---------- - x - Input array of complex numbers. - s - Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, - `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, - the input is cropped. If larger, input is padded with zeros. If `s` is not - given, shape of input along axes specified by axes is used. - axes - Axes over which to compute the IFFT. If not given, last `len(s)` axes are - used, or all axes if `s` is also not specified. Repeated indices in axes - means inverse transform over that axis is performed multiple times. - norm - Indicates direction of the forward/backward pair of transforms is scaled - and with what normalization factor. "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" - indicates normalization by $\frac{1}{n}$. - out - Optional output array for writing the result to. It must have a shape that - the inputs broadcast to. - Returns - ------- - out - The truncated or zero-padded input, transformed along the axes indicated - by axes, or by a combination of s or x, as explained in the parameters - section above. +def _upsample_cubic_interp1d(coeffs, ts): + coeffs2 = _upsample_get_cubic_coefficients(ts) + return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) - Raises - ------ - ValueError - If `s` and `axes` have different length. - IndexError - If an element of axes is larger than the number of axes of x. - Examples - -------- - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> y = ivy.ifftn(x) - >>> print(y) - ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, - -0.015561 -0.04216015j], - [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, - -0.08371392+0.17252843j], - [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, - 0.05344324+0.07972424j]]) +def _sum_tensors(ts): + return _reduce(ivy.add, ts) - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') - >>> print(b) - ivy.array([[ 0.8344667 +0.98222595j], - [-0.48472244+0.30233797j]]) - """ - return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) +def _upsample_bicubic2d_default( + a, + output_size, + align_corners, + scale_h=None, + scale_w=None, +): + N, C, iH, iW = a.shape + oH, oW = output_size -@handle_exceptions -@handle_nestable -@handle_out_argument -@inputs_to_ivy_arrays -def interp(x, xp, fp, left=None, right=None, period=None): - x_arr = ivy.array(x) - fix_later = False - if x_arr.shape == (): - x_arr = ivy.array([x]) - fix_later = True - x = ivy.astype(x_arr, "float64") - xp = ivy.astype(ivy.array(xp), "float64") - fp = ivy.astype(ivy.array(fp), "float64") - ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) - if period is not None: - ivy.utils.assertions.check_equal(period, 0, inverse=True) - period = ivy.abs(period) - x = ivy.remainder(x, period) - xp = ivy.remainder(xp, period) - asort_xp = ivy.argsort(xp) - xp = xp[asort_xp] - fp = fp[asort_xp] - xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) - fp = ivy.concat((fp[-1:], fp, fp[0:1])) + def compute_scale(in_size, out_size, align_corners, scale=None): + if align_corners: + return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 + else: + return 1 / scale if scale is not None and scale > 0 else in_size / out_size - def interp_inner(value): - value = ivy.array(value) - if value < xp[0]: - return left if left is not None else fp[0] - elif value > xp[-1]: - return right if right is not None else fp[-1] + def compute_source_index(scale, dst_index, align_corners): + if align_corners: + return scale * dst_index else: - last = None - if xp.shape[0] < 3: - for i in range(xp.shape[0] - 1, -1, -1): - if xp[i] == value: - return fp[i] - elif xp[i] < value: - last = i + return scale * (dst_index + 0.5) - 0.5 + + height_scale = compute_scale(iH, oH, align_corners, scale_h) + width_scale = compute_scale(iW, oW, align_corners, scale_w) + + N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) + C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) + out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) + out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) + + real_x = compute_source_index(width_scale, out_x, align_corners) + in_x = ivy.floor(real_x) + t_x = real_x - in_x + ix = ivy.astype(in_x, ivy.int64) + + real_y = compute_source_index(height_scale, out_y, align_corners) + in_y = ivy.floor(real_y) + t_y = real_y - in_y + iy = ivy.astype(in_y, ivy.int64) + + iys_ofs = (iy - 1, iy, iy + 1, iy + 2) + ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) + + def load_bounded(ys, xs): + y_idx = ivy.clip(ys, 0, iH - 1) + x_idx = ivy.clip(xs, 0, iW - 1) + return a[N_idx, C_idx, y_idx, x_idx] + + def get_x_interp(y): + coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) + return _upsample_cubic_interp1d(coeffs_x, t_x) + + coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) + result = _upsample_cubic_interp1d(coeffs_y, t_y) + + return result + + +def area_interpolate(x, dims, size, scale): + ret = ivy.zeros((x.shape[:2] + size)) + inv_scale = ivy.divide(1.0, scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_index = ( + int(d_dim * inv_scale[0]), + math.ceil((d_dim + 1) * inv_scale[0]), + ) + h_index = ( + int(h_dim * inv_scale[1]), + math.ceil((h_dim + 1) * inv_scale[1]), + ) + w_index = ( + int(w_dim * scale[2]), + math.ceil((w_dim + 1) * inv_scale[2]), + ) + scale_z = d_index[1] - d_index[0] + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_z * scale_y * scale_x + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( + ch[ + d_index[0] : d_index[1], + h_index[0] : h_index[1], + w_index[0] : w_index[1], + ] + ) * (1 / area) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_index = ( + int(h_dim * inv_scale[0]), + math.ceil((h_dim + 1) * inv_scale[0]), + ) + w_index = ( + int(w_dim * inv_scale[1]), + math.ceil((w_dim + 1) * inv_scale[1]), + ) + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_y * scale_x + ret[i, j, h_dim, w_dim] = ivy.sum( + ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] + ) * (1 / area) else: - first = 0 - last = xp.shape[0] - while first < last: - midpoint = (first + last) // 2 - if xp[midpoint] == value: - already_exists = ivy.argwhere(xp == value) - if already_exists.shape[0] > 0: - return fp[already_exists[-1][0]] - return fp[midpoint] - else: - if value < xp[midpoint]: - last = midpoint - 1 - else: - first = midpoint + 1 - dist = (value - xp[last]) / (xp[last + 1] - xp[last]) - return (fp[last + 1] - fp[last]) * dist + fp[last] + for w_dim in range(size[0]): + w_index = ( + int(w_dim * inv_scale[0]), + math.ceil((w_dim + 1) * inv_scale[0]), + ) + scale_x = w_index[1] - w_index[0] + ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( + 1 / scale_x + ) + return ret - ret = ivy.map(interp_inner, unique={"value": x}) - if fix_later: - return ivy.astype(ivy.array(ret[0]), "float64") - else: - return ivy.astype(ivy.array(ret), "float64") + +def get_interpolate_kernel(mode): + kernel_func = _triangle_kernel + if mode == "bicubic_tensorflow": + kernel_func = lambda inputs: _cubic_kernel(inputs) + elif mode == "lanczos3": + kernel_func = lambda inputs: _lanczos_kernel(3, inputs) + elif mode == "lanczos5": + kernel_func = lambda inputs: _lanczos_kernel(5, inputs) + return kernel_func + + +def generate_einsum_equation(dim): + alphabet = "abcdefghijklmnopqrstuvwxyz" + input_indices = alphabet[: dim + 2] + output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] + contraction_indices = ",".join([input_indices, *output_indices]) + output = input_indices[:2] + "".join([output[-1] for output in output_indices]) + einsum_string = contraction_indices + "->" + output + return einsum_string + + +def _interpolate_with_kernel( + x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode +): + spatial_dims = [2 + i for i in range(dims)] + equation = generate_einsum_equation(dims) + kernel_func = get_interpolate_kernel(mode) + output_shape = tuple(input_shape[:2]) + size + operands = [] + for i, d in enumerate(spatial_dims): + m = input_shape[d] + n = output_shape[d] + dim_scale_factor = _dim_scale_factor( + m, + n, + align_corners, + scale_factor[i] if scale_factor is not None else None, + ) + w = _compute_weight_mat( + m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor + ).astype(x.dtype) + operands.append(w) + return ivy.einsum(equation, x, *operands) @handle_exceptions @@ -2199,425 +2050,497 @@ def interpolate( return ivy.astype(ret, ivy.dtype(x), out=out) -@handle_backend_invalid +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} + + +def _get_size(scale_factor, size, dims, x_shape): + if scale_factor is not None: + if isinstance(scale_factor, (float, int)): + scale_factor = [scale_factor] * dims + elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: + scale_factor = [scale_factor[0]] * dims + + size = tuple( + [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] + ) + else: + size = (size,) * dims if isinstance(size, int) else tuple(size) + return size + + +def _output_ceil_shape(w, f, p, s): + return math.ceil((w - f + p) / s) + 1 + + +def _padding_ceil_mode(w, f, p, s, return_added_padding=False): + remaining_pixels = (w - f + sum(p)) % s + added_padding = 0 + if s > 1 and remaining_pixels != 0 and f > 1: + input_size = w + sum(p) + # making sure that the remaining pixels are supposed + # to be covered by the window + # they won't be covered if stride is big enough to skip them + if input_size - remaining_pixels - (f - 1) + s > input_size: + return p + output_shape = _output_ceil_shape( + w, + f, + sum(p), + s, + ) + # calculating new padding with ceil_output_shape + new_pad = (output_shape - 1) * s + f - w + # updating pad_list with new padding by adding it to the end + added_padding = new_pad - sum(p) + p = ( + p[0], + p[1] + added_padding, + ) + if return_added_padding: + return p, added_padding + return p + + +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} + + +def _compute_idx(in_size, out_size, device): + out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) + i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) + maxlength = in_size // out_size + 1 + in_size_mod = in_size % out_size + # adaptive = True iff there are kernels with different lengths + adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) + if adaptive: + maxlength += 1 + elif in_size_mod == 0: + maxlength -= 1 + range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) + idx = ivy.expand_dims(i0, axis=-1) + range_max + if adaptive: + maxval = ivy.full_like(idx, fill_value=in_size - 1) + idx = ivy.minimum(idx, maxval) + i1 = ivy.trunc_divide( + (out_range + 1) * in_size + out_size - 1, out_size + ).astype(ivy.int64) + length = i1 - i0 + else: + length = maxlength + return idx, length, range_max, adaptive + + +def _expand_to_dim(x, dim): + for _ in range(dim - len(x.shape)): + x = ivy.expand_dims(x, axis=-1) + return x + + +def _mask(vals, length, range_max, dim, mask_value=0.0): + if isinstance(length, int): + return vals, length + else: + assert dim < 0 + mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) + if dim == -2: + mask = _expand_to_dim(mask, 4) + vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) + length = _expand_to_dim(length, -dim) + return vals, length + + @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool1d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +@inputs_to_ivy_arrays +def adaptive_max_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], +): """ - Compute a 1-D max pool given 3-D input x. + Apply a 2D adaptive maximum pooling over an input signal composed of several input + planes. Parameters ---------- - x - Input image *[batch_size, w, d_in]* if data_format is "NWC". - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) - data_format - "NWC" or "NCW". Defaults to "NWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - The result of the pooling operation. + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + if isinstance(output_size, int): + output_size = (output_size, output_size) - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.max_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 4., 5., 6., 7.]], + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) - [[16., 17., 18., 19.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa - ivy.array([[[ 1., 3.], - [ 5., 7.], - [ 9., 11.]], + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array( + input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device + ) - [[13., 15.], - [17., 19.], - [21., 23.]]]) - """ - return ivy.current_backend(x).max_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, + if not adaptive_h and not adaptive_w: + ret = ivy.max(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret + + vals, length_h = _mask( + vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") + ) + vals, length_w = _mask( + vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") ) + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ivy.maximum(ret, vals[..., i, :, j]) + pooled_output = ret.astype(vals.dtype) + + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output + + +adaptive_max_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool2d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, +@inputs_to_ivy_arrays +def adaptive_avg_pool1d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: int, ) -> ivy.Array: """ - Compute a 2-D max pool given 4-D input x. + Apply a 1D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. Returns ------- - ret - The result of the pooling operation. + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` + """ + squeeze = False + if input.ndim == 2: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 3: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", + ) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + if input.shape[-1] % output_size == 0: + stride = input.shape[-1] // output_size + kernel_size = input.shape[-1] - (output_size - 1) * stride + pooled_output = ivy.avg_pool1d( + input, kernel_size, stride, "VALID", data_format="NCW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 2., 3.], - [ 4., 5.], - [ 4., 5.]]], + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size, input.device + ) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., idxw]) - [[[ 8., 9.], - [10., 11.], - [10., 11.]]]]) + if not adaptive_w: + ret = ivy.mean(vals, axis=-1) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[16., 17.]], + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) - [[22., 23.]]], + ret = None + for i in range(vals.shape[-1]): + if ret is None: + ret = vals[..., i] + else: + ret = ret + vals[..., i] + pooled_output = ret / length_w.astype(ret.dtype) + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output - [[[40., 41.]], - [[46., 47.]]]]) - """ - return ivy.current_backend(x).max_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +adaptive_avg_pool1d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} -@handle_backend_invalid +@handle_exceptions @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool3d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def adaptive_avg_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], ) -> ivy.Array: """ - Compute a 3-D max pool given 5-D input x. + Apply a 2D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - The result of the pooling operation. + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + if isinstance(output_size, int): + output_size = (output_size, output_size) - Examples - -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) - ivy.array([[[[[14., 15.]]]], + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.avg_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output + + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) + if not adaptive_h and not adaptive_w: + ret = ivy.mean(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - [[[[38., 39.]]]]]) - >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) - ivy.array([[[[[14., 15.]]], + vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ret + vals[..., i, :, j] + pooled_output = ret / (length_h * length_w).astype(vals.dtype) + + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output - [[[22., 23.]]]], +adaptive_avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + +def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): + def _pad(arr, pads, pad_value): + out = ivy.astype( + ivy.pad( + arr, + ivy.maximum(0, pads).to_list(), + mode="constant", + constant_values=ivy.to_scalar(pad_value), + ), + arr.dtype, + ) + slices = tuple( + _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) + for (lo, hi), dim in zip(pads, arr.shape) + ) + return out[slices] + if ( + _min(lhs.ndim, len(rhs_shape)) < 2 + or lhs.ndim != len(rhs_shape) + or lhs.shape[1] != rhs_shape[1] + ): + raise ValueError("Dimension mismatch") + if len(window_strides) != len(rhs_shape) - 2: + raise ValueError("Wrong number of strides for spatial dimensions") + if len(pads) != len(rhs_shape) - 2: + raise ValueError("Wrong number of pads for spatial dimensions") - [[[[38., 39.]]], + lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) + in_shape = lhs.shape[2:] + filter_shape = rhs_shape[2:] + dim = len(filter_shape) + out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() + view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] - [[[46., 47.]]]]]) - """ - return ivy.current_backend(x).max_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) + out_shape = [ + (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) + ] + view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] + view = ivy.as_strided(lhs, view_shape, view_strides) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_unpool1d( - x: ivy.Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D max unpooling given the 1-D pooled input x and its indices. + view_axes = list(range(view.ndim)) + sum_axes = view_axes[-dim - 1 :] + rhs_axes = [view.ndim] + sum_axes + out_axes = [0, view.ndim] + list(range(1, dim + 1)) - Parameters - ---------- - x - Pooled input image *[batch_size, w, d_in]*. - indices - Indices obtained from the corresponding max pooling operation. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - out - optional output array, for writing the result to. + return view, view_axes, rhs_axes, out_axes - Returns - ------- - ret - The result of the unpooling operation. - Both the description and the type hints above assume an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. +def _dilate(operand, factors, fill_value): + outspace = list(operand.shape[:2]) + [ + shape + (factors[i] - 1) * (shape - 1) + for i, shape in enumerate(operand.shape[2:]) + ] + out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) + lhs_slices = tuple(_slice(None, None, step) for step in factors) + out[(_slice(None),) * 2 + lhs_slices] = operand + return out - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') - >>> print(pool_result) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') - >>> print(unpool_result) - ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], - [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], +def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): + if padding.upper() == "SAME": + out_shape = [ + math.ceil(in_size / stride) + for in_size, stride in zip(in_shape, window_strides) + ] + pad_sizes = [ + _max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size in zip( + out_shape, window_strides, filter_shape, in_shape + ) + ] + return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] + else: + return [(0, 0)] * len(in_shape) - [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], - [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) - """ - return ivy.current_backend(x).max_unpool1d( - x, indices, kernel, strides, padding, data_format=data_format, out=out - ) +identities = { + "max": -float("inf"), + "min": float("inf"), + "add": 0, + "mul": 1, + "multiply": 1, + "logical_and": True, + "logical_or": False, +} -def nearest_interpolate(x, dims, size, input_shape, exact): - off = 0.5 if exact else 0 - for d in range(dims): - m = input_shape[d + 2] - n = size[d] - offsets = (ivy.arange(n, dtype="float32") + off) * m / n - offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") - x = ivy.gather(x, offsets, axis=d + 2) - return x +def _cast_init(init, dtype): + if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): + if ivy.is_float_dtype(dtype): + info = ivy.finfo(dtype) + else: + info = ivy.iinfo(dtype) + if "float64" not in str(dtype): + init = info.max if init > 0 else info.min + return ivy.array(init, dtype=dtype) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def pool( - x: Union[ivy.Array, ivy.NativeArray], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, - *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Perform an N-D pooling operation. - Parameters - ---------- - x - Input array to pool over. - window_shape - Shape of the pooling window. - pool_type - Type of pooling operation, either 'MAX' or 'AVG'. - strides - Strides of the pooling operation. - padding - Padding type, either 'VALID' or 'SAME'. - data_format - Data format of the input and output data, either 'NCHW' or 'NHWC'. - dilations - Dilation rate of the pooling operation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. +def _get_identity(func, dtype, init): + func_name = func.__name__ + if func_name in identities: + identity = identities[func_name] + return _cast_init(identity, dtype) + return init - Returns - ------- - ret - The result of the pooling operation. - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) - ivy.array([[[[ 8., 9.]], - [[14., 15.]]], - [[[32., 33.]], - [[38., 39.]]]]) - """ - return ivy.current_backend(x).pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ceil_mode=ceil_mode, - out=out, - ) +avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_exceptions @@ -2696,6 +2619,178 @@ def reduce_window( return ret.astype(operand.dtype) +reduce_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def fft2( + x: Union[ivy.Array, ivy.NativeArray], + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the 2-dimensional discrete Fourier Transform. + + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT2. + s + sequence of ints, optional + Shape (length of each transformed axis) of the output (s[0] refers to axis 0, + s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, + if the given shape is smaller than that of the input, the input is cropped. + If it is larger, the input is padded with zeros. if s is not given, the shape + of the input along the axes specified by axes is used. + dim + Axes over which to compute the FFT2. If not given, the last two axes are used. + A repeated index in axes means the transform over that axis is performed + multiple times. A one-element sequence means that a one-dimensional FFT is + performed. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the FFT2 operation. + + Examples + -------- + >>> x = ivy.array([[0, 0, 0, 0, 0], + ... [1, 1, 1, 1, 1], + ... [2, 2, 2, 2, 2], + ... [3, 3, 3, 3, 3], + ... [4, 4, 4, 4, 4]]) + >>> y = ivy.fft2(x) + >>> print(y) + ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5+17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5-17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ]]) + """ + return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) + + +fft2.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def ifftn( + x: Union[ivy.Array, ivy.NativeArray], + s: Optional[Union[int, Tuple[int, ...]]] = None, + axes: Optional[Union[int, Tuple[int, ...]]] = None, + *, + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the N-dimensional inverse discrete Fourier Transform. + + Parameters + ---------- + x + Input array of complex numbers. + s + Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, + `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, + the input is cropped. If larger, input is padded with zeros. If `s` is not + given, shape of input along axes specified by axes is used. + axes + Axes over which to compute the IFFT. If not given, last `len(s)` axes are + used, or all axes if `s` is also not specified. Repeated indices in axes + means inverse transform over that axis is performed multiple times. + norm + Indicates direction of the forward/backward pair of transforms is scaled + and with what normalization factor. "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" + indicates normalization by $\frac{1}{n}$. + out + Optional output array for writing the result to. It must have a shape that + the inputs broadcast to. + + Returns + ------- + out + The truncated or zero-padded input, transformed along the axes indicated + by axes, or by a combination of s or x, as explained in the parameters + section above. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of axes is larger than the number of axes of x. + + Examples + -------- + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> y = ivy.ifftn(x) + >>> print(y) + ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, + -0.015561 -0.04216015j], + [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, + -0.08371392+0.17252843j], + [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, + 0.05344324+0.07972424j]]) + + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') + >>> print(b) + ivy.array([[ 0.8344667 +0.98222595j], + [-0.48472244+0.30233797j]]) + """ + return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2779,87 +2874,3 @@ def rfftn( raise ValueError("s and axes must have the same length.") return ivy.current_backend(x).rfftn(x, s=s, axes=axes, norm=norm, out=out) - - -idct.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -adaptive_max_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -adaptive_avg_pool1d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -adaptive_avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -reduce_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -fft2.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} -_max = builtins.max -_min = builtins.min -_slice = builtins.slice -identities = { - "max": -float("inf"), - "min": float("inf"), - "add": 0, - "mul": 1, - "multiply": 1, - "logical_and": True, - "logical_or": False, -} diff --git a/ivy/functional/ivy/experimental/linear_algebra.py b/ivy/functional/ivy/experimental/linear_algebra.py index 94a5d48422683..249df35212e17 100644 --- a/ivy/functional/ivy/experimental/linear_algebra.py +++ b/ivy/functional/ivy/experimental/linear_algebra.py @@ -17,96 +17,149 @@ ) from ivy.utils.exceptions import handle_exceptions - -# --- Helpers --- # -# --------------- # +# Helpers # +# ------- # def _check_valid_dimension_size(std): ivy.utils.assertions.check_dimensions(std) -def _svd_checks(x, n_eigenvecs=None): +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@handle_array_function +def eigh_tridiagonal( + alpha: Union[ivy.Array, ivy.NativeArray], + beta: Union[ivy.Array, ivy.NativeArray], + /, + *, + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] + ] = None, + tol: Optional[float] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: """ - Run common checks to all of the SVD methods. + Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. Parameters ---------- - matrix : 2D-array - n_eigenvecs : int, optional, default is None - if specified, number of eigen[vectors-values] to return + alpha + A real or complex array of shape (n), the diagonal elements of the + matrix. If alpha is complex, the imaginary part is ignored + (assumed zero) to satisfy the requirement that the matrix be Hermitian. + beta + A real or complex array of shape (n-1), containing the elements of + the first super-diagonal of the matrix. If beta is complex, the first + sub-diagonal of the matrix is assumed to be the conjugate of beta to + satisfy the requirement that the matrix be Hermitian. + eigvals_only + If False, both eigenvalues and corresponding eigenvectors are + computed. If True, only eigenvalues are computed. Default is True. + select + Optional string with values in {'a', 'v', 'i'} (default is 'a') that + determines which eigenvalues to calculate: 'a': all eigenvalues. + 'v': eigenvalues in the interval (min, max] given by select_range. + 'i': eigenvalues with indices min <= i <= max. + select_range + Size 2 tuple or list or array specifying the range of eigenvalues to + compute together with select. If select is 'a', select_range is ignored. + tol + Optional scalar. Ignored when backend is not Tensorflow. The absolute + tolerance to which each eigenvalue is required. An eigenvalue + (or cluster) is considered to have converged if it lies in an interval + of this width. If tol is None (default), the value eps*|T|_2 is used + where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. Returns ------- - n_eigenvecs : int - the number of eigenvectors to solve for - min_dim : int - the minimum dimension of matrix - max_dim : int - the maximum dimension of matrix - """ - # ndims = len(x.shape) - # if ndims != 2: - # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") - - dim_1, dim_2 = ivy.shape(x)[-2:] - min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) + eig_vals + The eigenvalues of the matrix in non-decreasing order. + eig_vectors + If eigvals_only is False the eigenvectors are returned in the second output + argument. - if n_eigenvecs is None: - n_eigenvecs = max_dim + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. - if n_eigenvecs > max_dim: - logging.warning( - f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " - f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." - ) - n_eigenvecs = max_dim + Examples + -------- + With :class:`ivy.Array` input: - return n_eigenvecs, min_dim, max_dim + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + ivy.array([0., 0.38196602, 2.61803389]) + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, select='v', + ... select_range=[0.2,3.0]) + >>> print(y) + ivy.array([0.38196602, 2.61803389]) -# TODO uncommment the code below when these svd -# methods have been added -def _svd_interface( - matrix, - method="truncated_svd", - n_eigenvecs=None, - flip_sign=True, - u_based_flip_sign=True, - non_negative=None, - mask=None, - n_iter_mask_imputation=5, - **kwargs, -): - if method == "truncated_svd": - svd_fun = truncated_svd - # elif method == "symeig_svd": - # svd_fun = symeig_svd - # elif method == "randomized_svd": - # svd_fun = randomized_svd - elif callable(method): - svd_fun = method - else: - raise ValueError("Invalid Choice") + >>> alpha = ivy.array([0., 1., 2., 3.]) + >>> beta = ivy.array([2., 1., 2.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, + ... eigvals_only=False, + ... select='i', + ... select_range=[1,2], + ... tol=1.) + >>> print(y) + (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], + [ 0.06693714, -0.74234426], + [-0.74234426, -0.06693714], + [ 0.56710052, 0.35048741]])) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) - if mask is not None and n_eigenvecs is not None: - for _ in range(n_iter_mask_imputation): - S = S * ivy.eye(U.shape[-1], V.shape[-2]) - matrix = matrix * mask + (U @ S @ V) * (1 - mask) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + With :class:`ivy.Container` input: - if flip_sign: - U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.array([0.,2.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([0., 2., 4.]) + } - if non_negative is not False and non_negative is not None: - U, V = make_svd_non_negative(matrix, U, S, V) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([-0.82842714, 2., 4.82842731]) + } + """ + x = ivy.diag(alpha) + y = ivy.diag(beta, k=1) + z = ivy.diag(beta, k=-1) + w = x + y + z - return U, S, V + eigh_out = ivy.linalg.eigh(w) + eigenvalues = eigh_out.eigenvalues + eigenvectors = eigh_out.eigenvectors + if select == "i": + eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] + eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] + elif select == "v": + condition = ivy.logical_and( + eigenvalues.greater(select_range[0]), + eigenvalues.less_equal(select_range[1]), + ) + eigenvalues = eigenvalues[condition] + eigenvectors = eigenvectors[:, condition] -# --- Main --- # -# ------------ # + if eigvals_only: + return eigenvalues + return eigenvalues, eigenvectors @handle_exceptions @@ -116,19 +169,29 @@ def _svd_interface( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def adjoint( +def diagflat( x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, + offset: int = 0, + padding_value: float = 0, + align: str = "RIGHT_LEFT", + num_rows: int = -1, + num_cols: int = -1, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Compute the complex conjugate transpose of x. + Return a two-dimensional array with the flattened input as a diagonal. Parameters ---------- x - An array with more than one dimension. + Input data, which is flattened and set as the k-th diagonal of the output. + k + Diagonal to set. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -136,18 +199,35 @@ def adjoint( Returns ------- ret - the complex conjugate transpose of the input. + The 2-D output array. - Examples - -------- - >>> x = np.array([[1.-1.j, 2.+2.j], - [3.+3.j, 4.-4.j]]) - >>> x = ivy.array(x) - >>> ivy.adjoint(x) - ivy.array([[1.+1.j, 3.-3.j], - [2.-2.j, 4.+4.j]]) + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.diagflat(x) + ivy.array([[1, 0, 0, 0], + [0, 2, 0, 0], + [0, 0, 3, 0], + [0, 0, 0, 4]]) + + >>> x = ivy.array([1,2]) + >>> ivy.diagflat(x, k=1) + ivy.array([[0, 1, 0], + [0, 0, 2], + [0, 0, 0]]) """ - return current_backend(x).adjoint(x, out=out) + return current_backend(x).diagflat( + x, + offset=offset, + padding_value=padding_value, + align=align, + num_rows=num_rows, + num_cols=num_cols, + out=out, + ) @handle_exceptions @@ -157,22 +237,23 @@ def adjoint( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def cond( - x: Union[ivy.Array, ivy.NativeArray], +def kron( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, - p: Optional[Union[int, float, str]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the condition number of x. + Compute the Kronecker product, a composite array made of blocks of the second array + scaled by the first. Parameters ---------- - x - An array with more than one dimension. - p - The order of the norm of the matrix (see :func:`ivy.norm` for details). + a + First input array. + b + Second input array out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -180,21 +261,16 @@ def cond( Returns ------- ret - the condition number of the input. + Array representing the Kronecker product of the input arrays. Examples -------- - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x) - ivy.array(14.933034) - - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x, p=ivy.inf) - ivy.array(21.0) + >>> a = ivy.array([1,2]) + >>> b = ivy.array([3,4]) + >>> ivy.kron(a, b) + ivy.array([3, 4, 6, 8]) """ - return current_backend(x).cond(x, p=p, out=out) + return current_backend(a, b).kron(a, b, out=out) @handle_exceptions @@ -204,123 +280,41 @@ def cond( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def diagflat( +def matrix_exp( x: Union[ivy.Array, ivy.NativeArray], /, *, - offset: int = 0, - padding_value: float = 0, - align: str = "RIGHT_LEFT", - num_rows: int = -1, - num_cols: int = -1, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: - """ - Return a two-dimensional array with the flattened input as a diagonal. - - Parameters - ---------- - x - Input data, which is flattened and set as the k-th diagonal of the output. - k - Diagonal to set. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The 2-D output array. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.diagflat(x) - ivy.array([[1, 0, 0, 0], - [0, 2, 0, 0], - [0, 0, 3, 0], - [0, 0, 0, 4]]) - - >>> x = ivy.array([1,2]) - >>> ivy.diagflat(x, k=1) - ivy.array([[0, 1, 0], - [0, 0, 2], - [0, 0, 0]]) - """ - return current_backend(x).diagflat( - x, - offset=offset, - padding_value=padding_value, - align=align, - num_rows=num_rows, - num_cols=num_cols, - out=out, - ) - - -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_exceptions -def dot( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], - /, - *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the dot product between two arrays `a` and `b` using the current backend's - implementation. The dot product is defined as the sum of the element-wise product of - the input arrays. + Compute the matrix exponential of a square matrix. Parameters ---------- a - First input array. - b - Second input array. + Square matrix. out - Optional output array. If provided, the output array to store the result. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The dot product of the input arrays. + the matrix exponential of the input. Examples -------- - With :class:`ivy.Array` inputs: - - >>> a = ivy.array([1, 2, 3]) - >>> b = ivy.array([4, 5, 6]) - >>> result = ivy.dot(a, b) - >>> print(result) - ivy.array(32) - - >>> a = ivy.array([[1, 2], [3, 4]]) - >>> b = ivy.array([[5, 6], [7, 8]]) - >>> c = ivy.empty_like(a) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[19, 22], - [43, 50]]) - - >>> a = ivy.array([[1.1, 2.3, -3.6]]) - >>> b = ivy.array([[-4.8], [5.2], [6.1]]) - >>> c = ivy.zeros((1, 1)) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[-15.28]]) + >>> x = ivy.array([[[1., 0.], + [0., 1.]], + [[2., 0.], + [0., 2.]]]) + >>> ivy.matrix_exp(x) + ivy.array([[[2.7183, 1.0000], + [1.0000, 2.7183]], + [[7.3891, 1.0000], + [1.0000, 7.3891]]]) """ - return current_backend(a, b).dot(a, b, out=out) + return current_backend(x).matrix_exp(x, out=out) @handle_exceptions @@ -388,316 +382,259 @@ def eig( @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_array_function -def eigh_tridiagonal( - alpha: Union[ivy.Array, ivy.NativeArray], - beta: Union[ivy.Array, ivy.NativeArray], +@to_native_arrays_and_back +@handle_device_shifting +def eigvals( + x: Union[ivy.Array, ivy.NativeArray], /, - *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] - ] = None, - tol: Optional[float] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: +) -> ivy.Array: """ - Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. + Compute eigenvalues of x. Returns a set of eigenvalues. Parameters ---------- - alpha - A real or complex array of shape (n), the diagonal elements of the - matrix. If alpha is complex, the imaginary part is ignored - (assumed zero) to satisfy the requirement that the matrix be Hermitian. - beta - A real or complex array of shape (n-1), containing the elements of - the first super-diagonal of the matrix. If beta is complex, the first - sub-diagonal of the matrix is assumed to be the conjugate of beta to - satisfy the requirement that the matrix be Hermitian. - eigvals_only - If False, both eigenvalues and corresponding eigenvectors are - computed. If True, only eigenvalues are computed. Default is True. - select - Optional string with values in {'a', 'v', 'i'} (default is 'a') that - determines which eigenvalues to calculate: 'a': all eigenvalues. - 'v': eigenvalues in the interval (min, max] given by select_range. - 'i': eigenvalues with indices min <= i <= max. - select_range - Size 2 tuple or list or array specifying the range of eigenvalues to - compute together with select. If select is 'a', select_range is ignored. - tol - Optional scalar. Ignored when backend is not Tensorflow. The absolute - tolerance to which each eigenvalue is required. An eigenvalue - (or cluster) is considered to have converged if it lies in an interval - of this width. If tol is None (default), the value eps*|T|_2 is used - where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. + x + An array of shape (..., N, N). Returns ------- - eig_vals - The eigenvalues of the matrix in non-decreasing order. - eig_vectors - If eigvals_only is False the eigenvectors are returned in the second output - argument. - - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + w + Not necessarily ordered array(..., N) of eigenvalues in complex type. - Examples - -------- - With :class:`ivy.Array` input: + Functional Examples + ------------------ + With :class:`ivy.Array` inputs: + >>> x = ivy.array([[1,2], [3,4]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array([-0.37228132+0.j, 5.37228132+0.j]) - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - ivy.array([0., 0.38196602, 2.61803389]) + >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array( + [ + [-0.37228132+0.j, 5.37228132+0.j], + [ 0. +0.j, 11. +0.j] + ] + ) + """ + return current_backend(x).eigvals(x) - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, select='v', - ... select_range=[0.2,3.0]) - >>> print(y) - ivy.array([0.38196602, 2.61803389]) - >>> alpha = ivy.array([0., 1., 2., 3.]) - >>> beta = ivy.array([2., 1., 2.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, - ... eigvals_only=False, - ... select='i', - ... select_range=[1,2], - ... tol=1.) - >>> print(y) - (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], - [ 0.06693714, -0.74234426], - [-0.74234426, -0.06693714], - [ 0.56710052, 0.35048741]])) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def adjoint( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the complex conjugate transpose of x. - With :class:`ivy.Container` input: + Parameters + ---------- + x + An array with more than one dimension. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.array([0.,2.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([0., 2., 4.]) - } + Returns + ------- + ret + the complex conjugate transpose of the input. - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([-0.82842714, 2., 4.82842731]) - } + Examples + -------- + >>> x = np.array([[1.-1.j, 2.+2.j], + [3.+3.j, 4.-4.j]]) + >>> x = ivy.array(x) + >>> ivy.adjoint(x) + ivy.array([[1.+1.j, 3.-3.j], + [2.-2.j, 4.+4.j]]) """ - x = ivy.diag(alpha) - y = ivy.diag(beta, k=1) - z = ivy.diag(beta, k=-1) - w = x + y + z - - eigh_out = ivy.linalg.eigh(w) - eigenvalues = eigh_out.eigenvalues - eigenvectors = eigh_out.eigenvectors - - if select == "i": - eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] - eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] - elif select == "v": - condition = ivy.logical_and( - eigenvalues.greater(select_range[0]), - eigenvalues.less_equal(select_range[1]), - ) - eigenvalues = eigenvalues[condition] - eigenvectors = eigenvectors[:, condition] - - if eigvals_only: - return eigenvalues - return eigenvalues, eigenvectors + return current_backend(x).adjoint(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def eigvals( - x: Union[ivy.Array, ivy.NativeArray], +def multi_dot( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], /, + *, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute eigenvalues of x. Returns a set of eigenvalues. + Compute the dot product of two or more matrices in a single function call, while + selecting the fastest evaluation order. Parameters ---------- x - An array of shape (..., N, N). + sequence of matrices to multiply. + out + optional output array, for writing the result to. It must have a valid + shape, i.e. the resulting shape after applying regular matrix multiplication + to the inputs. Returns ------- - w - Not necessarily ordered array(..., N) of eigenvalues in complex type. + ret + dot product of the arrays. - Functional Examples - ------------------ - With :class:`ivy.Array` inputs: - >>> x = ivy.array([[1,2], [3,4]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array([-0.37228132+0.j, 5.37228132+0.j]) + Examples + -------- + With :class:`ivy.Array` input: - >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array( - [ - [-0.37228132+0.j, 5.37228132+0.j], - [ 0. +0.j, 11. +0.j] - ] - ) + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> ivy.multi_dot((A, B, C)) + ivy.array([[ 26, 49], + [ 80, 148]]) + + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> D = ivy.zeros((2, 2)) + >>> ivy.multi_dot((A, B, C), out=D) + >>> print(D) + ivy.array([[ 26, 49], + [ 80, 148]]) """ - return current_backend(x).eigvals(x) + return current_backend(x).multi_dot(x, out=out) -# This function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 +multi_dot.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} -# TODO update svd type hints when other svd methods have been added -# also update the test -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def initialize_tucker( +def cond( x: Union[ivy.Array, ivy.NativeArray], - rank: Sequence[int], - modes: Sequence[int], /, *, - init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", - seed: Optional[int] = None, - svd: Optional[Literal["truncated_svd"]] = "truncated_svd", - non_negative: Optional[bool] = False, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - svd_mask_repeats: Optional[int] = 5, -) -> Tuple[ivy.Array, Sequence[ivy.Array]]: + p: Optional[Union[int, float, str]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Initialize core and factors used in `tucker`. The type of initialization is set - using `init`. If `init == 'random'` then initialize factor matrices using - `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the - `rank` left singular vectors of the `m`th unfolding of the input tensor. + Compute the condition number of x. Parameters ---------- x - input tensor - rank - number of components - modes - modes to consider in the input tensor - seed - Used to create a random seed distribution - when init == 'random' - init - initialization scheme for tucker decomposition. - svd - function to use to compute the SVD - non_negative - if True, non-negative factors are returned - mask - array of booleans with the same shape as ``tensor`` should be 0 where - the values are missing and 1 everywhere else. Note: if tensor is - sparse, then mask should also be sparse with a fill value of 1 (or - True). - svd_mask_repeats - number of iterations for imputing the values in the SVD matrix when - mask is not None + An array with more than one dimension. + p + The order of the norm of the matrix (see :func:`ivy.norm` for details). + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - core - initialized core tensor - factors - list of factors - """ - try: - assert len(x.shape) >= 2 - except ValueError: - raise ValueError( - "expected x to have atleast 2 dimensions but it has only" - f" {len(x.shape)} dimension(s)" - ) - - # Initialisation - if init == "svd": - factors = [] - for index, mode in enumerate(modes): - mask_unfold = None if mask is None else ivy.unfold(mask, mode) - U, _, _ = _svd_interface( - ivy.unfold(x, mode), - n_eigenvecs=rank[index], - method=svd, - non_negative=non_negative, - mask=mask_unfold, - n_iter_mask_imputation=svd_mask_repeats, - # random_state=random_state, - ) - factors.append(U) - - # The initial core approximation is needed here for the masking step - core = multi_mode_dot(x, factors, modes=modes, transpose=True) - - elif init == "random": - core = ( - ivy.random_uniform( - shape=[rank[index] for index in range(len(modes))], - dtype=x.dtype, - seed=seed, - ) - + 0.01 - ) - factors = [ - ivy.random_uniform( - shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed - ) - for index, mode in enumerate(modes) - ] - - else: - (core, factors) = init + ret + the condition number of the input. - if non_negative is True: - factors = [ivy.abs(f) for f in factors] - core = ivy.abs(core) + Examples + -------- + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x) + ivy.array(14.933034) - return (core, factors) + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x, p=ivy.inf) + ivy.array(21.0) + """ + return current_backend(x).cond(x, p=p, out=out) -# The code has been adapated from tensorly.khatri_rao -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_khatri_rao.py#L9 -@handle_nestable +# This code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def khatri_rao( +def kronecker( x: Sequence[Union[ivy.Array, ivy.NativeArray]], - weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - skip_matrix: Optional[Sequence[int]] = None, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + skip_matrix: Optional[int] = None, + reverse: Optional[bool] = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Kronecker product of a list of matrices. + + Parameters + ---------- + x + Sequence of matrices + + skip_matrix + if not None, index of a matrix to skip + + reverse + if True, the order of the matrices is reversed + + Returns + ------- + kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` + where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` + and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` + """ + if skip_matrix is not None: + x = [x[i] for i in range(len(x)) if i != skip_matrix] + + if reverse: + order = -1 + else: + order = 1 + + for i, matrix in enumerate(x[::order]): + if not i: + res = matrix + else: + res = ivy.kron(res, matrix, out=out) + return res + + +# The code has been adapated from tensorly.khatri_rao +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_khatri_rao.py#L9 +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def khatri_rao( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], + weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + skip_matrix: Optional[Sequence[int]] = None, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ @@ -788,468 +725,590 @@ def khatri_rao( return res -@handle_exceptions -@handle_backend_invalid +# The following code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def kron( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def mode_dot( + x: Union[ivy.Array, ivy.NativeArray], /, + matrix_or_vector: Union[ivy.Array, ivy.NativeArray], + mode: int, + transpose: Optional[bool] = False, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kronecker product, a composite array made of blocks of the second array - scaled by the first. + N-mode product of a tensor and a matrix or vector at the specified mode. Parameters ---------- - a - First input array. - b - Second input array + x + tensor of shape ``(i_1, ..., i_k, ..., i_N)`` + matrix_or_vector + 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` + matrix or vectors to which to n-mode multiply the tensor + mode + int in the range(1, N) + transpose + If True, the matrix is transposed. + For complex tensors, the conjugate transpose is used. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + result can broadcast to. Returns ------- - ret - Array representing the Kronecker product of the input arrays. - - Examples - -------- - >>> a = ivy.array([1,2]) - >>> b = ivy.array([3,4]) - >>> ivy.kron(a, b) - ivy.array([3, 4, 6, 8]) - """ - return current_backend(a, b).kron(a, b, out=out) - - -# This code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def kronecker( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], - skip_matrix: Optional[int] = None, - reverse: Optional[bool] = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + ivy.Array + `mode`-mode product of `tensor` by `matrix_or_vector` + * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` + if matrix_or_vector is a matrix + * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` + if matrix_or_vector is a vector """ - Kronecker product of a list of matrices. + # the mode along which to fold might decrease if we take product with a vector + fold_mode = mode + new_shape = list(x.shape) + ndims = len(matrix_or_vector.shape) - Parameters - ---------- - x - Sequence of matrices + if ndims == 2: # Tensor times matrix + # Test for the validity of the operation + dim = 0 if transpose else 1 + if matrix_or_vector.shape[dim] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" + ) - skip_matrix - if not None, index of a matrix to skip + if transpose: + matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) - reverse - if True, the order of the matrices is reversed + new_shape[mode] = matrix_or_vector.shape[0] + vec = False - Returns - ------- - kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` - where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` - and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` - """ - if skip_matrix is not None: - x = [x[i] for i in range(len(x)) if i != skip_matrix] + elif ndims == 1: # Tensor times vector + if matrix_or_vector.shape[0] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[0]} (vector size)" + ) + if len(new_shape) > 1: + new_shape.pop(mode) + else: + new_shape = [] + vec = True - if reverse: - order = -1 else: - order = 1 + raise ValueError( + "Can only take n_mode_product with a vector or a matrix." + f"Provided array of dimension {ndims} not in [1, 2]." + ) - for i, matrix in enumerate(x[::order]): - if not i: - res = matrix - else: - res = ivy.kron(res, matrix, out=out) - return res + res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) + + if vec: # We contracted with a vector, leading to a vector + return ivy.reshape(res, new_shape, out=out) + else: # tensor times vec: refold the unfolding + return ivy.fold(res, fold_mode, new_shape, out=out) -# This function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L65 +# The following code has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def make_svd_non_negative( +def multi_mode_dot( x: Union[ivy.Array, ivy.NativeArray], - U: Union[ivy.Array, ivy.NativeArray], - S: Union[ivy.Array, ivy.NativeArray], - V: Union[ivy.Array, ivy.NativeArray], + mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], /, + modes: Optional[Sequence[int]] = None, + skip: Optional[Sequence[int]] = None, + transpose: Optional[bool] = False, *, - nntype: Optional[Literal["nndsvd", "nndsvda"]] = "nndsvd", -) -> Tuple[ivy.Array, ivy.Array]: - """ - Use NNDSVD method to transform SVD results into a non-negative form. This method - leads to more efficient solving with NNMF [1]. + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + N-mode product of a tensor and several matrices or vectors over several modes. Parameters ---------- x - tensor being decomposed. - U - left singular matrix from SVD. - S - diagonal matrix from SVD. - V - right singular matrix from SVD. - nntype - whether to fill small values with 0.0 (nndsvd), - or the tensor mean (nndsvda, default). + the input tensor - [1]: Boutsidis & Gallopoulos. Pattern Recognition, 41(4): 1350-1362, 2008. - """ - W = ivy.zeros_like(U) - H = ivy.zeros_like(V) + mat_or_vec_list + sequence of matrices or vectors of length ``tensor.ndim`` - # The leading singular triplet is non-negative - # so it can be used as is for initialization. - W[:, 0] = ivy.sqrt(S[0]) * ivy.abs(U[:, 0]) - H[0, :] = ivy.sqrt(S[0]) * ivy.abs(V[0, :]) + skip + None or int, optional, default is None + If not None, index of a matrix to skip. - for j in range(1, len(S)): - a, b = U[:, j], V[j, :] + modes + None or int list, optional, default is None - # extract positive and negative parts of column vectors - a_p, b_p = ivy.where(a < 0.0, 0, a), ivy.where(b < 0.0, 0.0, b) - # a_p, b_p = ivy.clip(a, 0.0), ivy.clip(b, 0.0) - # a_n, b_n = ivy.abs(ivy.clip(a, 0.0)), ivy.abs(ivy.clip(b, 0.0)) - a_n, b_n = ivy.abs(ivy.where(a > 0.0, 0.0, a)), ivy.abs( - ivy.where(b > 0.0, 0.0, b) - ) + transpose + If True, the matrices or vectors in in the list are transposed. + For complex tensors, the conjugate transpose is used. + out + optional output array, for writing the result to. It must have a shape that the + result can broadcast to. - # and their norms - a_p_nrm, b_p_nrm = float(ivy.vector_norm(a_p)), float(ivy.vector_norm(b_p)) - a_n_nrm, b_n_nrm = float(ivy.vector_norm(a_n)), float(ivy.vector_norm(b_n)) + Returns + ------- + ivy.Array + tensor times each matrix or vector in the list at mode `mode` - m_p, m_n = a_p_nrm * b_p_nrm, a_n_nrm * b_n_nrm + Notes + ----- + If no modes are specified, just assumes there is one matrix or vector per mode and returns: + :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa + """ + if modes is None: + modes = range(len(mat_or_vec_list)) - # choose update - if m_p > m_n: - u = a_p / a_p_nrm - v = b_p / b_p_nrm - sigma = m_p + decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor + + res = x + + # Order of mode dots doesn't matter for different modes + # Sorting by mode shouldn't change order for equal modes + factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) + for i, (mat_or_vec_list, mode) in enumerate(factors_modes): + ndims = len(mat_or_vec_list.shape) + if (skip is not None) and (i == skip): + continue + + if transpose and ndims == 2: + res = mode_dot( + res, + ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), + mode - decrement, + ) else: - u = a_n / a_n_nrm - v = b_n / b_n_nrm - sigma = m_n + res = mode_dot(res, mat_or_vec_list, mode - decrement) - lbd = float(ivy.sqrt(S[j] * sigma)) - W[:, j] = lbd * u - H[j, :] = lbd * v + if ndims == 1: + decrement += 1 - # After this point we no longer need H - eps = ivy.finfo(x.dtype).min + if ivy.exists(out): + return ivy.inplace_update(out, res) - if nntype == "nndsvd": - W = ivy.soft_thresholding(W, eps) - H = ivy.soft_thresholding(H, eps) - elif nntype == "nndsvda": - avg = ivy.mean(x) - W = ivy.where(W < eps, ivy.ones(ivy.shape(W)) * avg, W) - H = ivy.where(H < eps, ivy.ones(ivy.shape(H)) * avg, H) - else: - raise ValueError( - f'Invalid nntype parameter: got {nntype} instead of one of ("nndsvd",' - ' "nndsvda")' + return res + + +def _svd_checks(x, n_eigenvecs=None): + """ + Run common checks to all of the SVD methods. + + Parameters + ---------- + matrix : 2D-array + n_eigenvecs : int, optional, default is None + if specified, number of eigen[vectors-values] to return + + Returns + ------- + n_eigenvecs : int + the number of eigenvectors to solve for + min_dim : int + the minimum dimension of matrix + max_dim : int + the maximum dimension of matrix + """ + # ndims = len(x.shape) + # if ndims != 2: + # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") + + dim_1, dim_2 = ivy.shape(x)[-2:] + min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) + + if n_eigenvecs is None: + n_eigenvecs = max_dim + + if n_eigenvecs > max_dim: + logging.warning( + f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " + f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." ) + n_eigenvecs = max_dim - return W, H + return n_eigenvecs, min_dim, max_dim -@handle_exceptions -@handle_backend_invalid +# This function has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def matrix_exp( - x: Union[ivy.Array, ivy.NativeArray], +def svd_flip( + U: Union[ivy.Array, ivy.NativeArray], + V: Union[ivy.Array, ivy.NativeArray], /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + u_based_decision: Optional[bool] = True, +) -> Tuple[ivy.Array, ivy.Array]: """ - Compute the matrix exponential of a square matrix. + Sign correction to ensure deterministic output from SVD. Adjusts the columns of u + and the rows of v such that the loadings in the columns in u that are largest in + absolute value are always positive. This function is borrowed from scikit- + learn/utils/extmath.py. Parameters ---------- - a - Square matrix. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + U + left singular matrix output of SVD + V + right singular matrix output of SVD + u_based_decision + If True, use the columns of u as the basis for sign flipping. + Otherwise, use the rows of v. The choice of which variable to base the + decision on is generally algorithm dependent. Returns ------- - ret - the matrix exponential of the input. - - Examples - -------- - >>> x = ivy.array([[[1., 0.], - [0., 1.]], - [[2., 0.], - [0., 2.]]]) - >>> ivy.matrix_exp(x) - ivy.array([[[2.7183, 1.0000], - [1.0000, 2.7183]], - [[7.3891, 1.0000], - [1.0000, 7.3891]]]) + u_adjusted, v_adjusted : arrays with the same dimensions as the input. """ - return current_backend(x).matrix_exp(x, out=out) + if u_based_decision: + # columns of U, rows of V + max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) + signs = ivy.sign( + ivy.array( + [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], + ) + ) + U = U * signs + if ivy.shape(V)[0] > ivy.shape(U)[1]: + signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) + V = V * signs[: ivy.shape(V)[0]][:, None] + else: + # rows of V, columns of U + max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) + signs = ivy.sign( + ivy.array( + [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], + ) + ) + V = V * signs[:, None] + if ivy.shape(U)[1] > ivy.shape(V)[0]: + signs = ivy.concat( + ( + signs, + ivy.ones( + ivy.shape(U)[1] - ivy.shape(V)[0], + ), + ) + ) + U = U * signs[: ivy.shape(U)[1]] + return U, V -# The following code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 + +# This function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L65 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def mode_dot( +def make_svd_non_negative( x: Union[ivy.Array, ivy.NativeArray], + U: Union[ivy.Array, ivy.NativeArray], + S: Union[ivy.Array, ivy.NativeArray], + V: Union[ivy.Array, ivy.NativeArray], /, - matrix_or_vector: Union[ivy.Array, ivy.NativeArray], - mode: int, - transpose: Optional[bool] = False, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + nntype: Optional[Literal["nndsvd", "nndsvda"]] = "nndsvd", +) -> Tuple[ivy.Array, ivy.Array]: """ - N-mode product of a tensor and a matrix or vector at the specified mode. + Use NNDSVD method to transform SVD results into a non-negative form. This method + leads to more efficient solving with NNMF [1]. Parameters ---------- x - tensor of shape ``(i_1, ..., i_k, ..., i_N)`` - matrix_or_vector - 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` - matrix or vectors to which to n-mode multiply the tensor - mode - int in the range(1, N) - transpose - If True, the matrix is transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. + tensor being decomposed. + U + left singular matrix from SVD. + S + diagonal matrix from SVD. + V + right singular matrix from SVD. + nntype + whether to fill small values with 0.0 (nndsvd), + or the tensor mean (nndsvda, default). - Returns - ------- - ivy.Array - `mode`-mode product of `tensor` by `matrix_or_vector` - * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` - if matrix_or_vector is a matrix - * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` - if matrix_or_vector is a vector + [1]: Boutsidis & Gallopoulos. Pattern Recognition, 41(4): 1350-1362, 2008. """ - # the mode along which to fold might decrease if we take product with a vector - fold_mode = mode - new_shape = list(x.shape) - ndims = len(matrix_or_vector.shape) + W = ivy.zeros_like(U) + H = ivy.zeros_like(V) - if ndims == 2: # Tensor times matrix - # Test for the validity of the operation - dim = 0 if transpose else 1 - if matrix_or_vector.shape[dim] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" - ) + # The leading singular triplet is non-negative + # so it can be used as is for initialization. + W[:, 0] = ivy.sqrt(S[0]) * ivy.abs(U[:, 0]) + H[0, :] = ivy.sqrt(S[0]) * ivy.abs(V[0, :]) - if transpose: - matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) + for j in range(1, len(S)): + a, b = U[:, j], V[j, :] - new_shape[mode] = matrix_or_vector.shape[0] - vec = False + # extract positive and negative parts of column vectors + a_p, b_p = ivy.where(a < 0.0, 0, a), ivy.where(b < 0.0, 0.0, b) + # a_p, b_p = ivy.clip(a, 0.0), ivy.clip(b, 0.0) + # a_n, b_n = ivy.abs(ivy.clip(a, 0.0)), ivy.abs(ivy.clip(b, 0.0)) + a_n, b_n = ivy.abs(ivy.where(a > 0.0, 0.0, a)), ivy.abs( + ivy.where(b > 0.0, 0.0, b) + ) - elif ndims == 1: # Tensor times vector - if matrix_or_vector.shape[0] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[0]} (vector size)" - ) - if len(new_shape) > 1: - new_shape.pop(mode) + # and their norms + a_p_nrm, b_p_nrm = float(ivy.vector_norm(a_p)), float(ivy.vector_norm(b_p)) + a_n_nrm, b_n_nrm = float(ivy.vector_norm(a_n)), float(ivy.vector_norm(b_n)) + + m_p, m_n = a_p_nrm * b_p_nrm, a_n_nrm * b_n_nrm + + # choose update + if m_p > m_n: + u = a_p / a_p_nrm + v = b_p / b_p_nrm + sigma = m_p else: - new_shape = [] - vec = True + u = a_n / a_n_nrm + v = b_n / b_n_nrm + sigma = m_n + + lbd = float(ivy.sqrt(S[j] * sigma)) + W[:, j] = lbd * u + H[j, :] = lbd * v + + # After this point we no longer need H + eps = ivy.finfo(x.dtype).min + if nntype == "nndsvd": + W = ivy.soft_thresholding(W, eps) + H = ivy.soft_thresholding(H, eps) + elif nntype == "nndsvda": + avg = ivy.mean(x) + W = ivy.where(W < eps, ivy.ones(ivy.shape(W)) * avg, W) + H = ivy.where(H < eps, ivy.ones(ivy.shape(H)) * avg, H) else: raise ValueError( - "Can only take n_mode_product with a vector or a matrix." - f"Provided array of dimension {ndims} not in [1, 2]." + f'Invalid nntype parameter: got {nntype} instead of one of ("nndsvd",' + ' "nndsvda")' ) - res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) - - if vec: # We contracted with a vector, leading to a vector - return ivy.reshape(res, new_shape, out=out) - else: # tensor times vec: refold the unfolding - return ivy.fold(res, fold_mode, new_shape, out=out) + return W, H -@handle_exceptions -@handle_backend_invalid +# The following function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def multi_dot( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def truncated_svd( + x: Union[ivy.Array, ivy.NativeArray], /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + compute_uv: bool = True, + n_eigenvecs: Optional[int] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: """ - Compute the dot product of two or more matrices in a single function call, while - selecting the fastest evaluation order. + Compute a truncated SVD on `x` using the standard SVD. Parameters ---------- x - sequence of matrices to multiply. - out - optional output array, for writing the result to. It must have a valid - shape, i.e. the resulting shape after applying regular matrix multiplication - to the inputs. + 2D-array + compute_uv + If ``True`` then left and right singular vectors will be computed and returned + in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be + computed, which can be significantly faster. + n_eigenvecs + if specified, number of eigen[vectors-values] to return + else full matrices will be returned Returns ------- ret - dot product of the arrays. + a namedtuple ``(U, S, Vh)`` + Each returned array must have the same floating-point data type as ``x``. + """ + n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) + full_matrices = True if n_eigenvecs > min_dim else False - Examples - -------- - With :class:`ivy.Array` input: + if compute_uv: + U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) + return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] + else: + S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) + return S[:n_eigenvecs] - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> ivy.multi_dot((A, B, C)) - ivy.array([[ 26, 49], - [ 80, 148]]) - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> D = ivy.zeros((2, 2)) - >>> ivy.multi_dot((A, B, C), out=D) - >>> print(D) - ivy.array([[ 26, 49], - [ 80, 148]]) - """ - return current_backend(x).multi_dot(x, out=out) +# TODO uncommment the code below when these svd +# methods have been added +def _svd_interface( + matrix, + method="truncated_svd", + n_eigenvecs=None, + flip_sign=True, + u_based_flip_sign=True, + non_negative=None, + mask=None, + n_iter_mask_imputation=5, + **kwargs, +): + if method == "truncated_svd": + svd_fun = truncated_svd + # elif method == "symeig_svd": + # svd_fun = symeig_svd + # elif method == "randomized_svd": + # svd_fun = randomized_svd + elif callable(method): + svd_fun = method + else: + raise ValueError("Invalid Choice") + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + if mask is not None and n_eigenvecs is not None: + for _ in range(n_iter_mask_imputation): + S = S * ivy.eye(U.shape[-1], V.shape[-2]) + matrix = matrix * mask + (U @ S @ V) * (1 - mask) + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) -# The following code has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 + if flip_sign: + U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) + + if non_negative is not False and non_negative is not None: + U, V = make_svd_non_negative(matrix, U, S, V) + + return U, S, V + + +# This function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 + + +# TODO update svd type hints when other svd methods have been added +# also update the test @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def multi_mode_dot( +def initialize_tucker( x: Union[ivy.Array, ivy.NativeArray], - mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], + rank: Sequence[int], + modes: Sequence[int], /, - modes: Optional[Sequence[int]] = None, - skip: Optional[Sequence[int]] = None, - transpose: Optional[bool] = False, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - N-mode product of a tensor and several matrices or vectors over several modes. + init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", + seed: Optional[int] = None, + svd: Optional[Literal["truncated_svd"]] = "truncated_svd", + non_negative: Optional[bool] = False, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + svd_mask_repeats: Optional[int] = 5, +) -> Tuple[ivy.Array, Sequence[ivy.Array]]: + """ + Initialize core and factors used in `tucker`. The type of initialization is set + using `init`. If `init == 'random'` then initialize factor matrices using + `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the + `rank` left singular vectors of the `m`th unfolding of the input tensor. Parameters ---------- x - the input tensor - - mat_or_vec_list - sequence of matrices or vectors of length ``tensor.ndim`` - - skip - None or int, optional, default is None - If not None, index of a matrix to skip. - + input tensor + rank + number of components modes - None or int list, optional, default is None - - transpose - If True, the matrices or vectors in in the list are transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. + modes to consider in the input tensor + seed + Used to create a random seed distribution + when init == 'random' + init + initialization scheme for tucker decomposition. + svd + function to use to compute the SVD + non_negative + if True, non-negative factors are returned + mask + array of booleans with the same shape as ``tensor`` should be 0 where + the values are missing and 1 everywhere else. Note: if tensor is + sparse, then mask should also be sparse with a fill value of 1 (or + True). + svd_mask_repeats + number of iterations for imputing the values in the SVD matrix when + mask is not None Returns ------- - ivy.Array - tensor times each matrix or vector in the list at mode `mode` - - Notes - ----- - If no modes are specified, just assumes there is one matrix or vector per mode and returns: - :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa + core + initialized core tensor + factors + list of factors """ - if modes is None: - modes = range(len(mat_or_vec_list)) - - decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor + try: + assert len(x.shape) >= 2 + except ValueError: + raise ValueError( + "expected x to have atleast 2 dimensions but it has only" + f" {len(x.shape)} dimension(s)" + ) - res = x + # Initialisation + if init == "svd": + factors = [] + for index, mode in enumerate(modes): + mask_unfold = None if mask is None else ivy.unfold(mask, mode) + U, _, _ = _svd_interface( + ivy.unfold(x, mode), + n_eigenvecs=rank[index], + method=svd, + non_negative=non_negative, + mask=mask_unfold, + n_iter_mask_imputation=svd_mask_repeats, + # random_state=random_state, + ) + factors.append(U) - # Order of mode dots doesn't matter for different modes - # Sorting by mode shouldn't change order for equal modes - factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) - for i, (mat_or_vec_list, mode) in enumerate(factors_modes): - ndims = len(mat_or_vec_list.shape) - if (skip is not None) and (i == skip): - continue + # The initial core approximation is needed here for the masking step + core = multi_mode_dot(x, factors, modes=modes, transpose=True) - if transpose and ndims == 2: - res = mode_dot( - res, - ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), - mode - decrement, + elif init == "random": + core = ( + ivy.random_uniform( + shape=[rank[index] for index in range(len(modes))], + dtype=x.dtype, + seed=seed, ) - else: - res = mode_dot(res, mat_or_vec_list, mode - decrement) + + 0.01 + ) + factors = [ + ivy.random_uniform( + shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed + ) + for index, mode in enumerate(modes) + ] - if ndims == 1: - decrement += 1 + else: + (core, factors) = init - if ivy.exists(out): - return ivy.inplace_update(out, res) + if non_negative is True: + factors = [ivy.abs(f) for f in factors] + core = ivy.abs(core) - return res + return (core, factors) # This function has been adpated from TensorLy @@ -1403,122 +1462,6 @@ def partial_tucker( return (core, factors) -# This function has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def svd_flip( - U: Union[ivy.Array, ivy.NativeArray], - V: Union[ivy.Array, ivy.NativeArray], - /, - u_based_decision: Optional[bool] = True, -) -> Tuple[ivy.Array, ivy.Array]: - """ - Sign correction to ensure deterministic output from SVD. Adjusts the columns of u - and the rows of v such that the loadings in the columns in u that are largest in - absolute value are always positive. This function is borrowed from scikit- - learn/utils/extmath.py. - - Parameters - ---------- - U - left singular matrix output of SVD - V - right singular matrix output of SVD - u_based_decision - If True, use the columns of u as the basis for sign flipping. - Otherwise, use the rows of v. The choice of which variable to base the - decision on is generally algorithm dependent. - - Returns - ------- - u_adjusted, v_adjusted : arrays with the same dimensions as the input. - """ - if u_based_decision: - # columns of U, rows of V - max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) - signs = ivy.sign( - ivy.array( - [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], - ) - ) - U = U * signs - if ivy.shape(V)[0] > ivy.shape(U)[1]: - signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) - V = V * signs[: ivy.shape(V)[0]][:, None] - else: - # rows of V, columns of U - max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) - signs = ivy.sign( - ivy.array( - [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], - ) - ) - V = V * signs[:, None] - if ivy.shape(U)[1] > ivy.shape(V)[0]: - signs = ivy.concat( - ( - signs, - ivy.ones( - ivy.shape(U)[1] - ivy.shape(V)[0], - ), - ) - ) - U = U * signs[: ivy.shape(U)[1]] - - return U, V - - -# The following function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def truncated_svd( - x: Union[ivy.Array, ivy.NativeArray], - /, - compute_uv: bool = True, - n_eigenvecs: Optional[int] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: - """ - Compute a truncated SVD on `x` using the standard SVD. - - Parameters - ---------- - x - 2D-array - compute_uv - If ``True`` then left and right singular vectors will be computed and returned - in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be - computed, which can be significantly faster. - n_eigenvecs - if specified, number of eigen[vectors-values] to return - else full matrices will be returned - - Returns - ------- - ret - a namedtuple ``(U, S, Vh)`` - Each returned array must have the same floating-point data type as ``x``. - """ - n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) - full_matrices = True if n_eigenvecs > min_dim else False - - if compute_uv: - U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) - return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] - else: - S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) - return S[:n_eigenvecs] - - @handle_nestable @handle_exceptions @handle_array_like_without_promotion @@ -1667,7 +1610,59 @@ def tucker( return ivy.TuckerTensor((core, factors)) -multi_dot.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_exceptions +def dot( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the dot product between two arrays `a` and `b` using the current backend's + implementation. The dot product is defined as the sum of the element-wise product of + the input arrays. + + Parameters + ---------- + a + First input array. + b + Second input array. + out + Optional output array. If provided, the output array to store the result. + + Returns + ------- + ret + The dot product of the input arrays. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> a = ivy.array([1, 2, 3]) + >>> b = ivy.array([4, 5, 6]) + >>> result = ivy.dot(a, b) + >>> print(result) + ivy.array(32) + + >>> a = ivy.array([[1, 2], [3, 4]]) + >>> b = ivy.array([[5, 6], [7, 8]]) + >>> c = ivy.empty_like(a) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[19, 22], + [43, 50]]) + + >>> a = ivy.array([[1.1, 2.3, -3.6]]) + >>> b = ivy.array([[-4.8], [5.2], [6.1]]) + >>> c = ivy.zeros((1, 1)) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[-15.28]]) + """ + return current_backend(a, b).dot(a, b, out=out) diff --git a/ivy/functional/ivy/experimental/losses.py b/ivy/functional/ivy/experimental/losses.py index f873fb6e6d1dd..7b9f55d010c5a 100644 --- a/ivy/functional/ivy/experimental/losses.py +++ b/ivy/functional/ivy/experimental/losses.py @@ -12,67 +12,86 @@ from ivy.utils.exceptions import handle_exceptions +# log_poisson_loss @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def huber_loss( +def log_poisson_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - delta: Optional[float] = 1.0, - reduction: Optional[str] = "mean", + compute_full_loss: bool = False, + axis: int = -1, + reduction: str = "none", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Huber loss (smooth L1 loss) between true and predicted values. + Compute the log-likelihood loss between the prediction and the target under the + assumption that the target has a Poisson distribution. Caveat: By default, this is + not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect + for optimization, but does not play well with relative loss comparisons. To compute + an approximation of the log factorial term, specify ``compute_full_loss=True`` to + enable Stirling's Approximation. Parameters ---------- - true: array_like - The true (ground truth) values. - pred : array_like - The predicted values by the model. - delta : float, optional - The threshold parameter that determines the point where the loss transitions fro - -m - squared error to absolute error. Default is 1.0. - reduction : str, optional - The type of reduction to apply to the loss. Possible values are "mean" (default) - and "sum". - out : array_like, optional - Optional output array, for writing the result to. It must have a shape + true + input array containing true labels. + pred + input array containing Predicted labels. + compute_full_loss + whether to compute the full loss. If false, a constant term is dropped + in favor of more efficient optimization. Default: ``False``. + axis + the axis along which to compute the log-likelihood loss. If axis is ``-1``, + the log-likelihood loss will be computed along the last dimension. + Default: ``-1``. + reduction + ``'none'``: No reduction will be applied to the output. + ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'``. + out + optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret : array_like - The Huber loss between the true and predicted values. + ret + The binary log-likelihood loss between the given distributions. + Examples -------- - >>> true = ivy.array([2, 4, 7, 1]) - >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) - >>> huber_loss(true, pred, delta=1.0) - ivy.array([0.125, 0.125, 0.5 , 0.125]) - - >>> huber_loss(true, pred, delta=2.0) - ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.log_poisson_loss(x, z)) + ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) - >>> huber_loss(true, pred, delta=0.5) - ivy.array([0.25 , 0.25 , 0. , 0.125]) + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) + ivy.array(1.1573164) """ - abs_diff = ivy.abs(true - pred) - quadratic_loss = 0.5 * (abs_diff**2) - linear_loss = delta * (abs_diff - 0.5 * delta) - loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) + try: + assert true.shape == pred.shape + except ValueError: + raise ValueError( + "`pred` and `true` must have the same shape, received " + f"({pred.shape} vs {true.shape})." + ) + loss = ivy.exp(pred) - pred * true + if compute_full_loss: + stirling_approx = ( + (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) + ) + cond = ivy.logical_and(true >= 0.0, true <= 1.0) + loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, out=out) + return ivy.sum(loss, axis=axis, out=out) elif reduction == "mean": - return ivy.mean(loss, out=out) + return ivy.mean(loss, axis=axis, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss @@ -135,86 +154,67 @@ def l1_loss( return ivy.inplace_update(out, loss) if out is not None else loss -# log_poisson_loss @handle_exceptions @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def log_poisson_loss( +def huber_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - compute_full_loss: bool = False, - axis: int = -1, - reduction: str = "none", + delta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the log-likelihood loss between the prediction and the target under the - assumption that the target has a Poisson distribution. Caveat: By default, this is - not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect - for optimization, but does not play well with relative loss comparisons. To compute - an approximation of the log factorial term, specify ``compute_full_loss=True`` to - enable Stirling's Approximation. + Compute the Huber loss (smooth L1 loss) between true and predicted values. Parameters ---------- - true - input array containing true labels. - pred - input array containing Predicted labels. - compute_full_loss - whether to compute the full loss. If false, a constant term is dropped - in favor of more efficient optimization. Default: ``False``. - axis - the axis along which to compute the log-likelihood loss. If axis is ``-1``, - the log-likelihood loss will be computed along the last dimension. - Default: ``-1``. - reduction - ``'none'``: No reduction will be applied to the output. - ``'mean'``: The output will be averaged. - ``'sum'``: The output will be summed. Default: ``'none'``. - out - optional output array, for writing the result to. It must have a shape + true: array_like + The true (ground truth) values. + pred : array_like + The predicted values by the model. + delta : float, optional + The threshold parameter that determines the point where the loss transitions fro + -m + squared error to absolute error. Default is 1.0. + reduction : str, optional + The type of reduction to apply to the loss. Possible values are "mean" (default) + and "sum". + out : array_like, optional + Optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret - The binary log-likelihood loss between the given distributions. - + ret : array_like + The Huber loss between the true and predicted values. Examples -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.log_poisson_loss(x, z)) - ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) + >>> true = ivy.array([2, 4, 7, 1]) + >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) + >>> huber_loss(true, pred, delta=1.0) + ivy.array([0.125, 0.125, 0.5 , 0.125]) - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) - ivy.array(1.1573164) + >>> huber_loss(true, pred, delta=2.0) + ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + + >>> huber_loss(true, pred, delta=0.5) + ivy.array([0.25 , 0.25 , 0. , 0.125]) """ - try: - assert true.shape == pred.shape - except ValueError: - raise ValueError( - "`pred` and `true` must have the same shape, received " - f"({pred.shape} vs {true.shape})." - ) + abs_diff = ivy.abs(true - pred) + quadratic_loss = 0.5 * (abs_diff**2) + linear_loss = delta * (abs_diff - 0.5 * delta) + loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) - loss = ivy.exp(pred) - pred * true - if compute_full_loss: - stirling_approx = ( - (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) - ) - cond = ivy.logical_and(true >= 0.0, true <= 1.0) - loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, axis=axis, out=out) + return ivy.sum(loss, out=out) elif reduction == "mean": - return ivy.mean(loss, axis=axis, out=out) + return ivy.mean(loss, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index c3a21d4a63e19..46b3e98ff1e70 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -33,91 +33,667 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # +@handle_exceptions +@handle_nestable +@handle_partial_mixed_function +@handle_array_like_without_promotion +@handle_view +@inputs_to_ivy_arrays +@handle_array_function +def flatten( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + start_dim: Optional[int] = 0, + end_dim: Optional[int] = -1, + order: Optional[str] = "C", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Flattens input by reshaping it into a one-dimensional tensor. If start_dim or + end_dim are passed, only dimensions starting with start_dim and ending with end_dim + are flattened. The order of elements in input is unchanged. + Parameters + ---------- + x + input array to flatten. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + start_dim + first dim to flatten. If not set, defaults to 0. + end_dim + last dim to flatten. If not set, defaults to -1. + order + Read the elements of the input container using this index order, + and place the elements into the reshaped array using this index order. + ‘C’ means to read / write the elements using C-like index order, + with the last axis index changing fastest, back to the first axis index + changing slowest. + ‘F’ means to read / write the elements using Fortran-like index order, with + the first index changing fastest, and the last index changing slowest. + Note that the ‘C’ and ‘F’ options take no account of the memory layout + of the underlying array, and only refer to the order of indexing. + Default order is 'C' + out + optional output array, for writing the result to. -def _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, -): - ivy.utils.assertions.check_true( - callable(mode) - or mode - in [ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - message="the provided mode is not supported", - ) - _check_tuple_arg(pad_width, "pad_width") - if mode not in ["dilated"]: - ivy.utils.assertions.check_true( - all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), - message="the pad_widths must be greater or equal to zero", + Returns + ------- + ret + the flattened array over the specified dimensions. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x) + ivy.array([1, 2, 3, 4]) + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x, order='F') + ivy.array([1, 3, 2, 4]) + + >>> x = ivy.array( + [[[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12]], + + [[ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]]], + + + [[[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17]], + + [[ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]]], + + + [[[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1]], + + [[19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]]] + ) + >>> ivy.flatten(x, start_dim = 1, end_dim = 2) + ivy.array( + [[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12], + [ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]], + + [[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17], + [ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]], + + [[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1], + [19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]])) + """ + if x.shape == (): + x = ivy.reshape(x, (1, -1))[0, :] + if start_dim == end_dim: + return ivy.inplace_update(out, x) if ivy.exists(out) else x + if start_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" ) - if mode in ["maximum", "mean", "median", "minimum"]: - _check_tuple_arg(stat_length, "stat_length") - ivy.utils.assertions.check_true( - all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), - message="the stat lengths must be greater than zero", + if end_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" ) - elif mode == "constant": - _check_tuple_arg(constant_values, "constant_values", force_integer=False) - elif mode == "linear_ramp": - _check_tuple_arg(end_values, "end_values", force_integer=False) - ivy.utils.assertions.check_true( - reflect_type in ["even", "odd"], - message="the provided reflect_type is not supported", - ) + if start_dim < 0: + start_dim = len(x.shape) + start_dim + if end_dim < 0: + end_dim = len(x.shape) + end_dim + c = 1 + for i in range(start_dim, end_dim + 1): + c *= x.shape[i] + lst = [c] + if start_dim != 0: + for i in range(0, start_dim): + lst.insert(i, x.shape[i]) + for i in range(end_dim + 1, len(x.shape)): + lst.insert(i, x.shape[i]) + return ivy.reshape(x, tuple(lst), order=order, out=out) -def _check_bounds(shape0, shape1, strides1, itemsize): - numel0 = math.prod(shape0) - ndim1 = len(shape1) - return ( - sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize - <= numel0 * itemsize - ) +flatten.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} -def _check_tuple_arg(arg, name, force_integer=True): - is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype - flag_assert = False - if isinstance(arg, (tuple, list)): - for nested in arg: - if isinstance(nested, (tuple, list)): - for sub_nested in nested: - if not is_scalar(sub_nested): - flag_assert = True - break - elif not is_scalar(nested): - flag_assert = True - elif not is_scalar(arg): - flag_assert = True - if flag_assert: - if not force_integer: - raise ivy.utils.exceptions.IvyException( - name + " should be scalar, tuple of scalars or tuple of scalar tuples" - ) - else: - raise ivy.utils.exceptions.IvyException( - name + " should be int, tuple of ints or tuple of int tuples" - ) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def moveaxis( + a: Union[ivy.Array, ivy.NativeArray], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Move axes of an array to new positions.. + + Parameters + ---------- + a + The array whose axes should be reordered. + source + Original positions of the axes to move. These must be unique. + destination + Destination positions for each of the original axes. + These must also be unique. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array with moved axes. This array is a view of the input array. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.zeros((3, 4, 5)) + >>> ivy.moveaxis(x, 0, -1).shape + (4, 5, 3) + >>> ivy.moveaxis(x, -1, 0).shape + (5, 3, 4) + """ + return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def heaviside( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Heaviside step function for each element in x1. + + Parameters + ---------- + x1 + input array. + x2 + values to use where x1 is zero. + out + optional output array, for writing the result to. + + Returns + ------- + ret + output array with element-wise Heaviside step function of x1. + This is a scalar if both x1 and x2 are scalars. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([0.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0.0000, 0.5000, 1.0000]) + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([1.2, -2.0, 3.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0., -2., 1.]) + """ + return ivy.current_backend().heaviside(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def flipud( + m: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Flip array in the up/down direction. Flip the entries in each column in the up/down + direction. Rows are preserved, but appear in a different order than before. + + Parameters + ---------- + m + The array to be flipped. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array corresponding to input array with elements + order reversed along axis 0. + + Examples + -------- + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.flipud(m) + ivy.array([[ 0., 0., 3.], + [ 0., 2., 0.], + [ 1., 0., 0.]]) + """ + return ivy.current_backend().flipud(m, copy=copy, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def vstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> ivy.Array: + """ + Stack arrays in sequence vertically (row wise). + + Parameters + ---------- + arrays + Sequence of arrays to be stacked. + + Returns + ------- + ret + The array formed by stacking the given arrays. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.vstack((x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4]]) + >>> ivy.vstack((x, y, x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [1, 2, 3], + [2, 3, 4]]) + + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.vstack(y)) + ivy.array([[5, 6], + [7, 8]]) + """ + return ivy.current_backend().vstack(arrays, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def hstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> ivy.Array: + """ + Stack arrays in sequence horizotally (column wise). + + Parameters + ---------- + arrays + Sequence of arrays to be stacked. + + Returns + ------- + ret + The array formed by stacking the given arrays. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.hstack((x, y)) + ivy.array([1, 2, 3, 2, 3, 4]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0, 0, 0]) + >>> ivy.hstack((x, y, x)) + ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.hstack(y)) + ivy.array([[5, 6, 7, 8]]) + """ + return ivy.current_backend().hstack(arrays, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def rot90( + m: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is + from the first towards the second axis. + + Parameters + ---------- + m + Input array of two or more dimensions. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + k + Number of times the array is rotated by 90 degrees. + axes + The array is rotated in the plane defined by the axes. Axes must be + different. + out + optional output container, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + A rotated view of m. + + Examples + -------- + With :code:`ivy.Array` input: + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) + With :code:`ivy.NativeArray` input: + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.native_array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) + """ + return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def top_k( + x: Union[ivy.Array, ivy.NativeArray], + k: int, + /, + *, + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[tuple] = None, +) -> Tuple[ivy.Array, ivy.NativeArray]: + """ + Return the `k` largest elements of the given input array along a given axis. + + Parameters + ---------- + x + The array to compute top_k for. + k + Number of top elements to retun must not exceed the array size. + axis + The axis along which we must return the top elements default value is 1. + largest + If largest is set to False we return k smallest elements of the array. + sorted + If sorted is set to True we return the elements in sorted order. + out: + Optional output tuple, for writing the result to. Must have two arrays inside, + with a shape that the returned tuple broadcast to. + + Returns + ------- + ret + A named tuple with values and indices of top k elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 2) + >>> print(y) + top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) + + >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) + >>> y = ivy.top_k(x, 2, axis=1, largest=False) + >>> print(y) + top_k(values=ivy.array([[-2., 0.], + [-8., -1.]]), indices=ivy.array([[0, 3], + [0, 2]])) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 3) + >>> print(y) + top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) + >>> y = x.top_k(2) + >>> print(y) + [{ + a: ivy.array([2, -1]), + b: ivy.array([5., 4.]) + }, { + a: ivy.array([1, 0]), + b: ivy.array([1, 0]) + }] + """ + return current_backend(x).top_k( + x, k, axis=axis, largest=largest, sorted=sorted, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fliplr( + m: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Flip array in the left/right direction. Flip the entries in each column in the + left/right direction. Columns are preserved, but appear in a different order than + before. + + Parameters + ---------- + m + The array to be flipped. Must be at least 2-D. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array corresponding to input array with elements + order reversed along axis 1. + + Examples + -------- + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.fliplr(m) + ivy.array([[0, 0, 1], + [0, 2, 0], + [3, 0, 0]]) + """ + return ivy.current_backend().fliplr(m, copy=copy, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def i0( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Bessel i0 function of x element-wise. + + Parameters + ---------- + x + Array input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array with the modified Bessel function + evaluated at each of the elements of x. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> ivy.i0(x) + ivy.array([1.26606588, 2.2795853 , 4.88079259]) + """ + return ivy.current_backend(x).i0(x, out=out) + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +def _set_pad_area(padded, axis, width_pair, value_pair): + if width_pair[0] > 0: + left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) + padded[left_slice] = value_pair[0] + if width_pair[1] > 0: + right_slice = _slice_at_axis( + slice(padded.shape[axis] - width_pair[1], None), axis + ) + padded[right_slice] = value_pair[1] + return padded def _get_edges(padded, axis, width_pair): @@ -182,82 +758,6 @@ def _get_stats(padded, axis, width_pair, length_pair, stat_func): return left_stat, right_stat -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = ivy.full(new_shape, padding_value, dtype=operand.dtype) - src_indices = ivy.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array - - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) - - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = ivy.constant_pad(padded, pad_width, value=padding_value) - return padded - - -def _interleave(a, b, axis): - assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 - a_pad = [(0, 0, 0)] * a.ndim - b_pad = [(0, 0, 0)] * b.ndim - a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) - b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) - a = _interior_pad(a, 0.0, a_pad) - b = _interior_pad(b, 0.0, b_pad) - return ivy.add(a, b) - - -def _pad_simple(array, pad_width, fill_value=None): - new_shape = tuple( - [left + size + right for size, (left, right) in zip(array.shape, pad_width)] - ) - padded = ivy.zeros(new_shape, dtype=array.dtype) - if fill_value is not None: - padded = ivy.ones_like(padded) * fill_value - original_area_slice = tuple( - [ - slice(left, left + size) - for size, (left, right) in zip(array.shape, pad_width) - ] - ) - padded[original_area_slice] = array - return padded, original_area_slice - - -def _set_pad_area(padded, axis, width_pair, value_pair): - if width_pair[0] > 0: - left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) - padded[left_slice] = value_pair[0] - if width_pair[1] > 0: - right_slice = _slice_at_axis( - slice(padded.shape[axis] - width_pair[1], None), axis - ) - padded[right_slice] = value_pair[1] - return padded - - def _set_reflect_both(padded, axis, width_pair, method, include_edge=False): left_pad, right_pad = width_pair old_length = padded.shape[axis] - right_pad - left_pad @@ -330,489 +830,455 @@ def _set_wrap_both(padded, axis, width_pair): pad_area = _slice_at_axis(slice(-right_pad, -right_pad + period), axis) new_right_pad = right_pad - period else: - pad_area = _slice_at_axis(slice(-right_pad, None), axis) - padded[pad_area] = left_chunk - return new_left_pad, new_right_pad, padded - - -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides - - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] - - -def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): - if axis >= 0: - slices = [slice(None)] * axis + [slice(start, stop, stride)] - else: - slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) - return x[tuple(slices)] - - -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - -def _to_dilated(x, n): - if ivy.isscalar(x): - return ((x, x, x),) * n - elif len(x) == 3 and ivy.isscalar(x[0]): - return ((x[0], x[1], x[2]),) * n - elif len(x) != n: - ivy.utils.assertions.check_equal( - ivy.asarray(list(x)).shape, - (n, 3), - message=( - "tuple argument should contain " - "ndim groups where ndim is the number of " - "the input's dimensions" - ), - as_array=False, - ) - return x - - -def _to_pairs(x, n): - if ivy.isscalar(x): - return ((x, x),) * n - elif len(x) == 2 and ivy.isscalar(x[0]): - return ((x[0], x[1]),) * n - elif len(x) != n: - ivy.utils.assertions.check_equal( - ivy.asarray(list(x)).shape, - (n, 2), - message=( - "tuple argument should contain " - "ndim pairs where ndim is the number of " - "the input's dimensions" - ), - as_array=False, - ) - return x - - -# --- Main --- # -# ------------ # - - -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@inputs_to_native_shapes -def as_strided( - x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - strides: Sequence[int], - /, -) -> ivy.Array: - """ - Create a copy of the input array with the given shape and strides. - - Parameters - ---------- - x - Input Array. - shape - The shape of the new array. - strides - The strides of the new array (specified in bytes). - - Returns - ------- - ret - Output Array - - Examples - -------- - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> ivy.as_strided(x, (4, 3), (8, 8)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [3, 4, 5], - [4, 5, 6]]) - """ - itemsize = x.itemsize - if not _check_bounds(x.shape, shape, strides, itemsize): - raise ivy.exceptions.IvyException("attempted unsafe memory access") - if any(strides[i] % itemsize != 0 for i in range(len(strides))): - raise ivy.exceptions.IvyException("strides must be multiple of itemsize") - - src = memoryview(ivy.to_numpy(x)).cast("b") - - src_ind = ivy.inner( - ivy.indices(shape).reshape((len(shape), -1)).T, - ivy.array(strides), - ) - src_ind = ivy.expand_dims(src_ind, axis=-1) - src_ind = src_ind + ivy.arange(itemsize) - src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() - - temp_list = [src[i] for i in src_ind] - temp_array = ivy.asarray(temp_list, dtype=ivy.int8) - result = bytearray(temp_array.to_numpy()) - - return ivy.reshape( - ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), - shape, - ) - - -@handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def associative_scan( - x: Union[ivy.Array, ivy.NativeArray], - fn: Callable, - /, - *, - reverse: bool = False, - axis: int = 0, -) -> ivy.Array: - """ - Perform an associative scan over the given array. - - Parameters - ---------- - x - The array to scan over. - fn - The associative function to apply. - reverse - Whether to scan in reverse with respect to the given axis. - axis - The axis to scan over. - - Returns - ------- - ret - The result of the scan. - """ - elems = [x] + pad_area = _slice_at_axis(slice(-right_pad, None), axis) + padded[pad_area] = left_chunk + return new_left_pad, new_right_pad, padded - if reverse: - elems = [ivy.flip(elem, axis=[axis]) for elem in elems] - def _combine(a, b): - a = a[0] - b = b[0] - if a.shape[axis] == 0: - return [a] - c = fn(a, b) - return [c] +def _pad_simple(array, pad_width, fill_value=None): + new_shape = tuple( + [left + size + right for size, (left, right) in zip(array.shape, pad_width)] + ) + padded = ivy.zeros(new_shape, dtype=array.dtype) + if fill_value is not None: + padded = ivy.ones_like(padded) * fill_value + original_area_slice = tuple( + [ + slice(left, left + size) + for size, (left, right) in zip(array.shape, pad_width) + ] + ) + padded[original_area_slice] = array + return padded, original_area_slice - def _scan(elems): - num_elems = elems[0].shape[axis] - if num_elems < 2: - return elems +def _to_pairs(x, n): + if ivy.isscalar(x): + return ((x, x),) * n + elif len(x) == 2 and ivy.isscalar(x[0]): + return ((x[0], x[1]),) * n + elif len(x) != n: + ivy.utils.assertions.check_equal( + ivy.asarray(list(x)).shape, + (n, 2), + message=( + "tuple argument should contain " + "ndim pairs where ndim is the number of " + "the input's dimensions" + ), + as_array=False, + ) + return x - reduced_elems = _combine( - [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], - [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], + +def _to_dilated(x, n): + if ivy.isscalar(x): + return ((x, x, x),) * n + elif len(x) == 3 and ivy.isscalar(x[0]): + return ((x[0], x[1], x[2]),) * n + elif len(x) != n: + ivy.utils.assertions.check_equal( + ivy.asarray(list(x)).shape, + (n, 3), + message=( + "tuple argument should contain " + "ndim groups where ndim is the number of " + "the input's dimensions" + ), + as_array=False, ) + return x - odd_elems = _scan(reduced_elems) - if num_elems % 2 == 0: - even_elems = _combine( - [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], +def _check_tuple_arg(arg, name, force_integer=True): + is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype + flag_assert = False + if isinstance(arg, (tuple, list)): + for nested in arg: + if isinstance(nested, (tuple, list)): + for sub_nested in nested: + if not is_scalar(sub_nested): + flag_assert = True + break + elif not is_scalar(nested): + flag_assert = True + elif not is_scalar(arg): + flag_assert = True + if flag_assert: + if not force_integer: + raise ivy.utils.exceptions.IvyException( + name + " should be scalar, tuple of scalars or tuple of scalar tuples" ) else: - even_elems = _combine( - odd_elems, - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + raise ivy.utils.exceptions.IvyException( + name + " should be int, tuple of ints or tuple of int tuples" ) - even_elems = [ - ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) - for (elem, result) in zip(elems, even_elems) - ] - return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) - - scans = _scan(elems) - if reverse: - scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] - return ivy.reshape(ivy.asarray(scans), elems[0].shape) +def _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, +): + ivy.utils.assertions.check_true( + callable(mode) + or mode + in [ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + message="the provided mode is not supported", + ) + _check_tuple_arg(pad_width, "pad_width") + if mode not in ["dilated"]: + ivy.utils.assertions.check_true( + all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), + message="the pad_widths must be greater or equal to zero", + ) + if mode in ["maximum", "mean", "median", "minimum"]: + _check_tuple_arg(stat_length, "stat_length") + ivy.utils.assertions.check_true( + all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), + message="the stat lengths must be greater than zero", + ) + elif mode == "constant": + _check_tuple_arg(constant_values, "constant_values", force_integer=False) + elif mode == "linear_ramp": + _check_tuple_arg(end_values, "end_values", force_integer=False) + ivy.utils.assertions.check_true( + reflect_type in ["even", "odd"], + message="the provided reflect_type is not supported", + ) -@handle_backend_invalid +@handle_exceptions @handle_nestable @handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_device_shifting -def atleast_1d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], - copy: Optional[bool] = None, -) -> List[ivy.Array]: +@inputs_to_ivy_arrays +@handle_array_function +def pad( + input: Union[ivy.Array, ivy.NativeArray], + pad_width: Union[Iterable[Tuple[int]], int], + /, + *, + mode: Union[ + Literal[ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + Callable, + ] = "constant", + stat_length: Union[Iterable[Tuple[int]], int] = 1, + constant_values: Union[Iterable[Tuple[Number]], Number] = 0, + end_values: Union[Iterable[Tuple[Number]], Number] = 0, + reflect_type: Literal["even", "odd"] = "even", + **kwargs: Optional[Any], +) -> ivy.Array: """ - Convert inputs to arrays with at least one dimension. Scalar inputs are converted to - 1-dimensional arrays, whilst higher-dimensional inputs are preserved. + Pad an array. Parameters ---------- - arys - One or more input arrays. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + input + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths + for each axis. + - ((before, after),) yields same before and after pad for each axis. + - pad (integer) is shortcut for before = after = pad width for all axes. + mode + One of the following string values or a user-supplied function. + - "constant": Pads with a constant value. + - "edge": Pads with the input's edge values. + - "linear_ramp": Pads with the linear ramp between end_value + and the input's edge value. + - "maximum": Pads with the maximum value of all or part of the vector + along each axis. + - "mean": Pads with the mean value of all or part of the vector along + each axis. + - "median": Pads with the median value of all or part of the vector + along each axis. + - "minimum": Pads with the minimum value of all or part of the vector + along each axis. + - "reflect": Pads with the reflection mirrored on the first and last + values of the vector along each axis. + - "symmetric": Pads with the reflection of the vector mirrored along + the edge of the input. + - "wrap": Pads with the wrap of the vector along the axis. + The first values are used to pad the end and the end values are used + to pad the beginning. + - "empty": Pads with undefined values. + - : Pads with a user-defined padding function. The padding + function should modify a rank 1 array following the signature + `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: + - `vector` is a rank 1 array already padded with zeros. Padded + values are `vector[:iaxis_pad_width[0]]` and + `vector[-iaxis_pad_width[1]:]`. + - `iaxis_pad_width` is a 2-tuple of ints, where + `iaxis_pad_width[0]` represents the number of values padded at + the beginning of `vector` and `iaxis_pad_width[1]` represents + the number of values padded at the end of `vector`. + - `iaxis` is the axis currently being calculated. + - `kwargs` is a dict of keyword arguments the function requires. + stat_length + Used in "maximum", "mean", "median", and "minimum". Number of values at edge + of each axis used to calculate the statistic value. + - ((before_1, after_1), … (before_N, after_N)) yields unique statistic + lengths for each axis. + - ((before, after),) yields same before and after statistic lengths for + each axis. + - stat_length (integer) is a shortcut for before = after = stat_length + length for all axes. + - None uses the entire axis. + constant_values + Used in "constant". The values to set the padded values for each axis. + - ((before_1, after_1), ... (before_N, after_N)) yields unique pad + constants for each axis. + - ((before, after),) yields same before and after constants for each axis. + - constant (integer) is a shortcut for before = after = constant for + all axes. + end_values + Used in "linear_ramp". The values used for the ending value of the linear_ramp + and that will form the edge of the padded array. + - ((before_1, after_1), ... (before_N, after_N)) yields unique end values + for each axis. + - ((before, after),) yields same before and after end values for each axis + - end (integer) is a shortcut for before = after = end for all axes. + reflect_type + Used in "reflect", and "symmetric". The "even" style is the default with an + unaltered reflection around the edge value. For the "odd" style, the extended + part of the array is created by subtracting the reflected values from two + times the edge value. Returns ------- ret - An array, or list of arrays, each with atleast 1D. - Copies are made only if necessary. - - Examples - -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_1d(ary1) - ivy.array([5]) - >>> ary2 = ivy.array([[3,4]]) - >>> ivy.atleast_1d(ary2) - ivy.array([[3, 4]]) - >>> ivy.atleast_1d(6,7,8) - [ivy.array([6]), ivy.array([7]), ivy.array([8])] - """ - return ivy.current_backend().atleast_1d(*arys, copy=copy) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_device_shifting -def atleast_2d( - *arys: Union[ivy.Array, ivy.NativeArray], - copy: Optional[bool] = None, -) -> List[ivy.Array]: - """ - Convert inputs to arrays with at least two dimension. Scalar inputs are converted to - 2-dimensional arrays, whilst higher-dimensional inputs are preserved. + Padded array of the same rank as the input but with shape increased according + to pad_width. - Parameters - ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have two or more - dimensions are preserved. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - Returns - ------- - ret - An array, or list of arrays, each with atleast 2D. - Copies are made only if necessary. + Both the description and the type hints above assume an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_2d(ary1) - ivy.array([[5]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_2d(ary2) - ivy.array([[[3, 4]]]) - >>> ivy.atleast_2d(6,7,8) - [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] - """ - return ivy.current_backend().atleast_2d(*arys, copy=copy) - + With :class:`ivy.Array` input: -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_device_shifting -def atleast_3d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], - copy: Optional[bool] = None, -) -> List[ivy.Array]: - """ - Convert inputs to arrays with at least three dimension. Scalar inputs are converted - to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 3, 0, 0], + [0, 0, 4, 5, 6, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) - Parameters - ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have three or more - dimensions are preserved. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="reflect") + >>> print(y) + ivy.array([[6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1], + [6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1]]) - Returns - ------- - ret - An array, or list of arrays, each with a.ndim >= 3. Copies - are avoided where possible, and views with three or more - dimensions are returned. For example, a 1-D array of shape - (N,) becomes a view of shape (1, N, 1), and a 2-D array of - shape (M, N) becomes a view of shape (M, N, 1). + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="symmetric") + >>> print(y) + ivy.array([[2, 1, 1, 2, 3, 3, 2], + [2, 1, 1, 2, 3, 3, 2], + [5, 4, 4, 5, 6, 6, 5], + [5, 4, 4, 5, 6, 6, 5]]) - Examples - -------- - >>> ary1 = ivy.array([5,6]) - >>> ivy.atleast_3d(ary1) - ivy.array([[[5], - [6]]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_3d(ary2) - ivy.array([[[3, 4]]]) - >>> ary3 = ivy.array([[3,4],[9,10]]) - >>> ivy.atleast_3d(6,7,ary3) - [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], - [ 4]], + With :class:`ivy.NativeArray` input: - [[ 9], - [10]]])] - """ - return ivy.current_backend().atleast_3d(*arys, copy=copy) + >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) + >>> print(y) + ivy.array([[7, 7, 7, 7, 7, 7, 7], + [7, 7, 1, 2, 3, 7, 7], + [7, 7, 4, 5, 6, 7, 7], + [7, 7, 7, 7, 7, 7, 7]]) + With :class:`ivy.Container` input: -@handle_exceptions -@inputs_to_native_shapes -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) + >>> padding = (1, 1) + >>> y = ivy.pad(x, padding, mode="constant") + >>> print(y) + { + a: ivy.array([0, 0, 1, 2, 0]), + b: ivy.array([0, 4, 5, 6, 0]) + } """ - Broadcasts shapes. - - Parameters - ---------- - shapes - The shapes to broadcast. - - Returns - ------- - ret - The broadcasted shape. + _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, + ) + if mode == "dilated": + pad_width = _to_dilated(pad_width, input.ndim) + if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: + constant_values = ivy.asarray(constant_values, dtype=input.dtype) + return _interior_pad(input, constant_values, pad_width) + pad_width = _to_pairs(pad_width, len(input.shape)) + if callable(mode): + func = mode + padded, _ = _pad_simple(input, pad_width, fill_value=0) + for axis in range(padded.ndim): + padded = ivy.moveaxis(padded, axis, -1) + inds = ivy.ndindex(padded.shape[:-1]) + for ind in inds: + padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) + return padded + padded, original_area_slice = _pad_simple(input, pad_width) + axes = range(padded.ndim) + stat_functions = { + "maximum": ivy.max, + "minimum": ivy.min, + "mean": ivy.mean, + "median": ivy.median, + } + if mode == "constant": + constant_values = _to_pairs(constant_values, padded.ndim) + constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) + for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): + padded = _set_pad_area(padded, axis, width_pair, value_pair) + elif mode == "empty": + pass + elif mode == "edge": + for axis, width_pair in zip(axes, pad_width): + edge_pair = _get_edges(padded, axis, width_pair) + padded = _set_pad_area(padded, axis, width_pair, edge_pair) + elif mode == "linear_ramp": + end_values = _to_pairs(end_values, padded.ndim) + for axis, width_pair, value_pair in zip(axes, pad_width, end_values): + ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) + padded = _set_pad_area(padded, axis, width_pair, ramp_pair) + elif mode in stat_functions: + func = stat_functions[mode] + stat_length = _to_pairs(stat_length, padded.ndim) + if mode == "median": + ivy.utils.assertions.check_true( + ivy.is_float_dtype(input), + message="median interpolation is only supported for floats", + ) + for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): + stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) + padded = _set_pad_area(padded, axis, width_pair, stat_pair) + elif mode in {"reflect", "symmetric"}: + include_edge = True if mode == "symmetric" else False + for axis, (left_index, right_index) in zip(axes, pad_width): + if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): + edge_pair = _get_edges(padded, axis, (left_index, right_index)) + padded = _set_pad_area( + padded, axis, (left_index, right_index), edge_pair + ) + continue + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_reflect_both( + padded, axis, (left_index, right_index), reflect_type, include_edge + ) + elif mode == "wrap": + for axis, (left_index, right_index) in zip(axes, pad_width): + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_wrap_both( + padded, axis, (left_index, right_index) + ) + return padded - Examples - -------- - >>> x = [(3, 3), (3, 1)] - >>> print(ivy.broadcast_shapes(*x)) - (3, 3) - >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) - (3, 3) - """ - return ivy.current_backend().broadcast_shapes(*shapes) +pad.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument +@handle_view @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def choose( - arr: Union[ivy.Array, ivy.NativeArray], - choices: Union[ivy.Array, ivy.NativeArray], +def vsplit( + ary: Union[ivy.Array, ivy.NativeArray], + indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], /, *, - out: None = None, - mode: Union[str, None] = None, -) -> ivy.Array: + copy: Optional[bool] = None, +) -> List[ivy.Array]: """ - Take values from the input array by matching 1d index and data slices. + Split an array vertically into multiple sub-arrays. Parameters ---------- - arr - The source array. - choices - The indices of the values to extract. - out - The output array. - mode - One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices - will be handled. + ary + Array input. + indices_or_sections + If indices_or_sections is an integer n, the array is split into n + equal sections, provided that n must be a divisor of the split axis. + If indices_or_sections is a sequence of ints or 1-D array, + then input is split at each of the indices. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - The returned array has the same shape as `indices`. + input array split vertically. Examples -------- - >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], - [20, 21, 22, 23], [30, 31, 32, 33]]) - >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) - ivy.array([20, 1, 12, 3]) - """ - return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def concat_from_sequence( - input_sequence: Union[ - Tuple[Union[ivy.Array, ivy.NativeArray]], - List[Union[ivy.Array, ivy.NativeArray]], - ], - /, - *, - new_axis: int = 0, - axis: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Concatenate a sequence of arrays along a new or an existing axis. - - Parameters - ---------- - input_sequence - A sequence of arrays. - new_axis - Insert and concatenate on a new axis or not, - default 0 means do not insert new axis. - new_axis = 0: concatenate - new_axis = 1: stack - axis - axis along which the arrays will be concatenated. - - out - optional output array, for writing the result to. - - Returns - ------- - ret - Output Array + >>> ary = ivy.array( + [[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]] + ) + >>> ivy.vsplit(ary, 2) + [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) """ - return current_backend(input_sequence).concat_from_sequence( - input_sequence, new_axis=new_axis, axis=axis, out=out - ) + return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) @handle_exceptions @@ -872,449 +1338,241 @@ def dsplit( @handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def dstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Stack arrays in sequence depth wise (along third axis). - - Parameters - ---------- - arrays - Sequence of arrays to be stacked. - - Returns - ------- - ret - The array formed by stacking the given arrays. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2], - [2, 3], - [3, 4]]]) - >>> x = ivy.array([[1], [2], [3]]) - >>> y = ivy.array([[2], [3], [4]]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2]], - [[2, 3]], - [[3, 4]]]) - """ - return ivy.current_backend().dstack(arrays, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_device_shifting -def expand( - x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape], - /, - *, - copy: Optional[bool] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Broadcast the input Array following the given shape and the broadcast rule. - - Parameters - ---------- - x - Array input. - shape - A 1-D Array indicates the shape you want to expand to, - following the broadcast rule. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Output Array - """ - return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_array_function -def fill_diagonal( - a: Union[ivy.Array, ivy.NativeArray], - v: Union[int, float], - /, - *, - wrap: bool = False, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Fill the main diagonal of the given array of any dimensionality.. - - Parameters - ---------- - a - Array at least 2D. - v - The value to write on the diagonal. - wrap - The diagonal ‘wrapped’ after N columns for tall matrices. - - Returns - ------- - ret - Array with the diagonal filled. - """ - return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) - - -@handle_exceptions -@handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_view -@inputs_to_ivy_arrays -@handle_array_function -def flatten( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_1d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], copy: Optional[bool] = None, - start_dim: Optional[int] = 0, - end_dim: Optional[int] = -1, - order: Optional[str] = "C", - out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> List[ivy.Array]: """ - Flattens input by reshaping it into a one-dimensional tensor. If start_dim or - end_dim are passed, only dimensions starting with start_dim and ending with end_dim - are flattened. The order of elements in input is unchanged. + Convert inputs to arrays with at least one dimension. Scalar inputs are converted to + 1-dimensional arrays, whilst higher-dimensional inputs are preserved. Parameters ---------- - x - input array to flatten. + arys + One or more input arrays. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. - start_dim - first dim to flatten. If not set, defaults to 0. - end_dim - last dim to flatten. If not set, defaults to -1. - order - Read the elements of the input container using this index order, - and place the elements into the reshaped array using this index order. - ‘C’ means to read / write the elements using C-like index order, - with the last axis index changing fastest, back to the first axis index - changing slowest. - ‘F’ means to read / write the elements using Fortran-like index order, with - the first index changing fastest, and the last index changing slowest. - Note that the ‘C’ and ‘F’ options take no account of the memory layout - of the underlying array, and only refer to the order of indexing. - Default order is 'C' - out - optional output array, for writing the result to. Returns ------- ret - the flattened array over the specified dimensions. + An array, or list of arrays, each with atleast 1D. + Copies are made only if necessary. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x) - ivy.array([1, 2, 3, 4]) - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x, order='F') - ivy.array([1, 3, 2, 4]) - - >>> x = ivy.array( - [[[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12]], - - [[ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]]], - - - [[[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17]], - - [[ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]]], - - - [[[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1]], - - [[19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]]] - ) - >>> ivy.flatten(x, start_dim = 1, end_dim = 2) - ivy.array( - [[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12], - [ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]], - - [[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17], - [ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]], - - [[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1], - [19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]])) + >>> ary1 = ivy.array(5) + >>> ivy.atleast_1d(ary1) + ivy.array([5]) + >>> ary2 = ivy.array([[3,4]]) + >>> ivy.atleast_1d(ary2) + ivy.array([[3, 4]]) + >>> ivy.atleast_1d(6,7,8) + [ivy.array([6]), ivy.array([7]), ivy.array([8])] """ - if x.shape == (): - x = ivy.reshape(x, (1, -1))[0, :] - if start_dim == end_dim: - return ivy.inplace_update(out, x) if ivy.exists(out) else x - if start_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" - ) - if end_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" - ) - if start_dim < 0: - start_dim = len(x.shape) + start_dim - if end_dim < 0: - end_dim = len(x.shape) + end_dim - c = 1 - for i in range(start_dim, end_dim + 1): - c *= x.shape[i] - lst = [c] - if start_dim != 0: - for i in range(0, start_dim): - lst.insert(i, x.shape[i]) - for i in range(end_dim + 1, len(x.shape)): - lst.insert(i, x.shape[i]) - return ivy.reshape(x, tuple(lst), order=order, out=out) + return ivy.current_backend().atleast_1d(*arys, copy=copy) @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_view @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def fliplr( - m: Union[ivy.Array, ivy.NativeArray], +def dstack( + arrays: Sequence[ivy.Array], /, *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Flip array in the left/right direction. Flip the entries in each column in the - left/right direction. Columns are preserved, but appear in a different order than - before. + Stack arrays in sequence depth wise (along third axis). Parameters ---------- - m - The array to be flipped. Must be at least 2-D. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. + arrays + Sequence of arrays to be stacked. Returns ------- ret - Array corresponding to input array with elements - order reversed along axis 1. + The array formed by stacking the given arrays. Examples -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.fliplr(m) - ivy.array([[0, 0, 1], - [0, 2, 0], - [3, 0, 0]]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2], + [2, 3], + [3, 4]]]) + >>> x = ivy.array([[1], [2], [3]]) + >>> y = ivy.array([[2], [3], [4]]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2]], + [[2, 3]], + [[3, 4]]]) """ - return ivy.current_backend().fliplr(m, copy=copy, out=out) + return ivy.current_backend().dstack(arrays, out=out) @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_view -@handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def flipud( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, +def atleast_2d( + *arys: Union[ivy.Array, ivy.NativeArray], copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +) -> List[ivy.Array]: """ - Flip array in the up/down direction. Flip the entries in each column in the up/down - direction. Rows are preserved, but appear in a different order than before. + Convert inputs to arrays with at least two dimension. Scalar inputs are converted to + 2-dimensional arrays, whilst higher-dimensional inputs are preserved. Parameters ---------- - m - The array to be flipped. + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have two or more + dimensions are preserved. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. Returns ------- ret - Array corresponding to input array with elements - order reversed along axis 0. + An array, or list of arrays, each with atleast 2D. + Copies are made only if necessary. Examples -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.flipud(m) - ivy.array([[ 0., 0., 3.], - [ 0., 2., 0.], - [ 1., 0., 0.]]) + >>> ary1 = ivy.array(5) + >>> ivy.atleast_2d(ary1) + ivy.array([[5]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_2d(ary2) + ivy.array([[[3, 4]]]) + >>> ivy.atleast_2d(6,7,8) + [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] """ - return ivy.current_backend().flipud(m, copy=copy, out=out) + return ivy.current_backend().atleast_2d(*arys, copy=copy) +@handle_backend_invalid @handle_nestable -@handle_exceptions @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_view +@to_native_arrays_and_back @handle_device_shifting -def fold( - x: Union[ivy.Array, ivy.NativeArray], - /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def atleast_3d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], + copy: Optional[bool] = None, +) -> List[ivy.Array]: """ - Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, - refolds the n-mode unfolded tensor into the original tensor of the specified shape. + Convert inputs to arrays with at least three dimension. Scalar inputs are converted + to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. Parameters - ---------- - input - unfolded tensor of shape ``(shape[mode], -1)`` - mode - the mode of the unfolding - shape - shape of the original tensor before unfolding - out - optional output array, for writing the result to. + ---------- + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have three or more + dimensions are preserved. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - folded_tensor of shape `shape` + An array, or list of arrays, each with a.ndim >= 3. Copies + are avoided where possible, and views with three or more + dimensions are returned. For example, a 1-D array of shape + (N,) becomes a view of shape (1, N, 1), and a 2-D array of + shape (M, N) becomes a view of shape (M, N, 1). + + Examples + -------- + >>> ary1 = ivy.array([5,6]) + >>> ivy.atleast_3d(ary1) + ivy.array([[[5], + [6]]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_3d(ary2) + ivy.array([[[3, 4]]]) + >>> ary3 = ivy.array([[3,4],[9,10]]) + >>> ivy.atleast_3d(6,7,ary3) + [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], + [ 4]], + + [[ 9], + [10]]])] """ - full_shape = list(shape) - mode_dim = full_shape.pop(mode) - full_shape.insert(0, mode_dim) - return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) + return ivy.current_backend().atleast_3d(*arys, copy=copy) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def heaviside( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def take_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + axis: int, /, *, + mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Heaviside step function for each element in x1. + Take values from the input array by matching 1d index and data slices. Parameters ---------- - x1 - input array. - x2 - values to use where x1 is zero. + arr + The source array. + indices + The indices of the values to extract. + axis + The axis over which to select values. + If axis is None, arr is treated as a flattened 1D array. + mode + One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices + will be handled. out - optional output array, for writing the result to. + The output array. Returns ------- ret - output array with element-wise Heaviside step function of x1. - This is a scalar if both x1 and x2 are scalars. + The returned array has the same shape as `indices`. Examples -------- - With :class:`ivy.Array` input: - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([0.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0.0000, 0.5000, 1.0000]) - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([1.2, -2.0, 3.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0., -2., 1.]) + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> y = ivy.take_along_axis(arr, indices, 1) + >>> print(y) + ivy.array([[4, 3, 3], [1, 1, 1]]) """ - return ivy.current_backend().heaviside(x1, x2, out=out) + return ivy.current_backend(arr).take_along_axis( + arr, indices, axis, mode=mode, out=out + ) @handle_exceptions @@ -1376,595 +1634,524 @@ def hsplit( return ivy.current_backend(ary).hsplit(ary, indices_or_sections, copy=copy) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def hstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: +@handle_exceptions +@inputs_to_native_shapes +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: """ - Stack arrays in sequence horizotally (column wise). + Broadcasts shapes. Parameters ---------- - arrays - Sequence of arrays to be stacked. + shapes + The shapes to broadcast. Returns ------- ret - The array formed by stacking the given arrays. + The broadcasted shape. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.hstack((x, y)) - ivy.array([1, 2, 3, 2, 3, 4]) - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0, 0, 0]) - >>> ivy.hstack((x, y, x)) - ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.hstack(y)) - ivy.array([[5, 6, 7, 8]]) + >>> x = [(3, 3), (3, 1)] + >>> print(ivy.broadcast_shapes(*x)) + (3, 3) + + >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) + (3, 3) """ - return ivy.current_backend().hstack(arrays, out=out) + return ivy.current_backend().broadcast_shapes(*shapes) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion +@handle_view @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_device_shifting -def i0( +def expand( x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape], /, *, + copy: Optional[bool] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Bessel i0 function of x element-wise. + Broadcast the input Array following the given shape and the broadcast rule. Parameters ---------- x Array input. + shape + A 1-D Array indicates the shape you want to expand to, + following the broadcast rule. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. out optional output array, for writing the result to. Returns ------- ret - Array with the modified Bessel function - evaluated at each of the elements of x. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.i0(x) - ivy.array([1.26606588, 2.2795853 , 4.88079259]) + Output Array """ - return ivy.current_backend(x).i0(x, out=out) + return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) -@handle_nestable @handle_exceptions +@handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def matricize( - x: Union[ivy.Array, ivy.NativeArray], +def put_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + values: Union[ivy.Array, ivy.NativeArray], + axis: int, /, - row_modes: Sequence[int], - column_modes: Optional[Sequence[int]] = None, *, + mode: str = "raise", out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> None: """ - Matricizes the given tensor. + Put values into the input array by matching 1d index and data slices along a + specified axis. Parameters ---------- - x - the input tensor - row_modes - modes to use as row of the matrix (in the desired order) - column_modes - modes to use as column of the matrix, in the desired order - if None, the modes not in `row_modes` will be used in ascending order - out - optional output array, for writing the result to. + arr : array_like + The input array to modify. + indices : array_like + The indices of the values to put into `arr`. + values : array_like + The values to put into `arr`. + axis : int + The axis over which to put the `values`. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will be handled. + The following modes are available: - ret + - 'raise': a `ValueError` is raised when an index is out of bounds. + - 'wrap': the index is wrapped around to the corresponding index + at the other end of the axis. + - 'clip': the index is clipped to the closest in-bounds index. + out : ndarray, optional + Output array in which to place the result. + If not specified, a new array is created. + + Returns ------- - ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) + None + + Examples + -------- + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) + >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') + >>> print(arr) + ivy.array([[3, 7, 5], + [6, 4, 1]]) """ - ndims = len(x.shape) - row_indices = list(row_modes) + if out is None: + out = ivy.zeros_like(arr) - if column_modes: - column_indices = list(column_modes) - else: - column_indices = [i for i in range(ndims) if i not in row_indices] - if sorted(column_indices + row_indices) != list(range(ndims)): - msg = ( - "If you provide both column and row modes for the matricization then" - " column_modes + row_modes must contain all the modes of the tensor." - f" Yet, got row_modes={row_modes} and column_modes={column_modes}." - ) - raise ValueError(msg) + indices = ivy.expand_dims(indices, axis=axis) + values = ivy.expand_dims(values, axis=axis) + + stacked = ivy.concat((arr, values), axis=axis) + + sorted_indices = ivy.argsort(indices, axis=axis) + sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), + sorted_stacked, + arr, + ) + + if mode == "clip": + indices = ivy.clip(indices, 0, arr.shape[axis] - 1) + elif mode == "wrap": + indices = ivy.mod(indices, arr.shape[axis]) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values + ) + + ivy.assign(out, arr) - row_size, column_size = 1, 1 - row_size = int(ivy.prod([x.shape[i] for i in row_indices])) - column_size = int(ivy.prod([x.shape[i] for i in column_indices])) - return ivy.reshape( - ivy.permute_dims(x, row_indices + column_indices), - (row_size, column_size), - out=out, +def _check_bounds(shape0, shape1, strides1, itemsize): + numel0 = math.prod(shape0) + ndim1 = len(shape1) + return ( + sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize + <= numel0 * itemsize ) -@handle_backend_invalid +@handle_exceptions @handle_nestable @handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def moveaxis( - a: Union[ivy.Array, ivy.NativeArray], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], +@inputs_to_ivy_arrays +@inputs_to_native_shapes +def as_strided( + x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + strides: Sequence[int], /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +) -> ivy.Array: """ - Move axes of an array to new positions.. + Create a copy of the input array with the given shape and strides. Parameters ---------- - a - The array whose axes should be reordered. - source - Original positions of the axes to move. These must be unique. - destination - Destination positions for each of the original axes. - These must also be unique. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. + x + Input Array. + shape + The shape of the new array. + strides + The strides of the new array (specified in bytes). Returns ------- ret - Array with moved axes. This array is a view of the input array. + Output Array Examples -------- - With :class:`ivy.Array` input: - >>> x = ivy.zeros((3, 4, 5)) - >>> ivy.moveaxis(x, 0, -1).shape - (4, 5, 3) - >>> ivy.moveaxis(x, -1, 0).shape - (5, 3, 4) + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> ivy.as_strided(x, (4, 3), (8, 8)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [3, 4, 5], + [4, 5, 6]]) """ - return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) + itemsize = x.itemsize + if not _check_bounds(x.shape, shape, strides, itemsize): + raise ivy.exceptions.IvyException("attempted unsafe memory access") + if any(strides[i] % itemsize != 0 for i in range(len(strides))): + raise ivy.exceptions.IvyException("strides must be multiple of itemsize") + + src = memoryview(ivy.to_numpy(x)).cast("b") + + src_ind = ivy.inner( + ivy.indices(shape).reshape((len(shape), -1)).T, + ivy.array(strides), + ) + src_ind = ivy.expand_dims(src_ind, axis=-1) + src_ind = src_ind + ivy.arange(itemsize) + src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() + + temp_list = [src[i] for i in src_ind] + temp_array = ivy.asarray(temp_list, dtype=ivy.int8) + result = bytearray(temp_array.to_numpy()) + + return ivy.reshape( + ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), + shape, + ) + + +as_strided.unsupported_dtypes = ("bfloat16",) +as_strided.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def pad( - input: Union[ivy.Array, ivy.NativeArray], - pad_width: Union[Iterable[Tuple[int]], int], +@handle_device_shifting +def concat_from_sequence( + input_sequence: Union[ + Tuple[Union[ivy.Array, ivy.NativeArray]], + List[Union[ivy.Array, ivy.NativeArray]], + ], /, *, - mode: Union[ - Literal[ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - Callable, - ] = "constant", - stat_length: Union[Iterable[Tuple[int]], int] = 1, - constant_values: Union[Iterable[Tuple[Number]], Number] = 0, - end_values: Union[Iterable[Tuple[Number]], Number] = 0, - reflect_type: Literal["even", "odd"] = "even", - **kwargs: Optional[Any], + new_axis: int = 0, + axis: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Pad an array. + Concatenate a sequence of arrays along a new or an existing axis. Parameters ---------- - input - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths - for each axis. - - ((before, after),) yields same before and after pad for each axis. - - pad (integer) is shortcut for before = after = pad width for all axes. - mode - One of the following string values or a user-supplied function. - - "constant": Pads with a constant value. - - "edge": Pads with the input's edge values. - - "linear_ramp": Pads with the linear ramp between end_value - and the input's edge value. - - "maximum": Pads with the maximum value of all or part of the vector - along each axis. - - "mean": Pads with the mean value of all or part of the vector along - each axis. - - "median": Pads with the median value of all or part of the vector - along each axis. - - "minimum": Pads with the minimum value of all or part of the vector - along each axis. - - "reflect": Pads with the reflection mirrored on the first and last - values of the vector along each axis. - - "symmetric": Pads with the reflection of the vector mirrored along - the edge of the input. - - "wrap": Pads with the wrap of the vector along the axis. - The first values are used to pad the end and the end values are used - to pad the beginning. - - "empty": Pads with undefined values. - - : Pads with a user-defined padding function. The padding - function should modify a rank 1 array following the signature - `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: - - `vector` is a rank 1 array already padded with zeros. Padded - values are `vector[:iaxis_pad_width[0]]` and - `vector[-iaxis_pad_width[1]:]`. - - `iaxis_pad_width` is a 2-tuple of ints, where - `iaxis_pad_width[0]` represents the number of values padded at - the beginning of `vector` and `iaxis_pad_width[1]` represents - the number of values padded at the end of `vector`. - - `iaxis` is the axis currently being calculated. - - `kwargs` is a dict of keyword arguments the function requires. - stat_length - Used in "maximum", "mean", "median", and "minimum". Number of values at edge - of each axis used to calculate the statistic value. - - ((before_1, after_1), … (before_N, after_N)) yields unique statistic - lengths for each axis. - - ((before, after),) yields same before and after statistic lengths for - each axis. - - stat_length (integer) is a shortcut for before = after = stat_length - length for all axes. - - None uses the entire axis. - constant_values - Used in "constant". The values to set the padded values for each axis. - - ((before_1, after_1), ... (before_N, after_N)) yields unique pad - constants for each axis. - - ((before, after),) yields same before and after constants for each axis. - - constant (integer) is a shortcut for before = after = constant for - all axes. - end_values - Used in "linear_ramp". The values used for the ending value of the linear_ramp - and that will form the edge of the padded array. - - ((before_1, after_1), ... (before_N, after_N)) yields unique end values - for each axis. - - ((before, after),) yields same before and after end values for each axis - - end (integer) is a shortcut for before = after = end for all axes. - reflect_type - Used in "reflect", and "symmetric". The "even" style is the default with an - unaltered reflection around the edge value. For the "odd" style, the extended - part of the array is created by subtracting the reflected values from two - times the edge value. + input_sequence + A sequence of arrays. + new_axis + Insert and concatenate on a new axis or not, + default 0 means do not insert new axis. + new_axis = 0: concatenate + new_axis = 1: stack + axis + axis along which the arrays will be concatenated. + + out + optional output array, for writing the result to. Returns ------- ret - Padded array of the same rank as the input but with shape increased according - to pad_width. - - - Both the description and the type hints above assume an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: + Output Array + """ + return current_backend(input_sequence).concat_from_sequence( + input_sequence, new_axis=new_axis, axis=axis, out=out + ) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 3, 0, 0], - [0, 0, 4, 5, 6, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="reflect") - >>> print(y) - ivy.array([[6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1], - [6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1]]) +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="symmetric") - >>> print(y) - ivy.array([[2, 1, 1, 2, 3, 3, 2], - [2, 1, 1, 2, 3, 3, 2], - [5, 4, 4, 5, 6, 6, 5], - [5, 4, 4, 5, 6, 6, 5]]) + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] - With :class:`ivy.NativeArray` input: - >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) - >>> print(y) - ivy.array([[7, 7, 7, 7, 7, 7, 7], - [7, 7, 1, 2, 3, 7, 7], - [7, 7, 4, 5, 6, 7, 7], - [7, 7, 7, 7, 7, 7, 7]]) +def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): + if axis >= 0: + slices = [slice(None)] * axis + [slice(start, stop, stride)] + else: + slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) + return x[tuple(slices)] - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) - >>> padding = (1, 1) - >>> y = ivy.pad(x, padding, mode="constant") - >>> print(y) - { - a: ivy.array([0, 0, 1, 2, 0]), - b: ivy.array([0, 4, 5, 6, 0]) - } - """ - _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, - ) - if mode == "dilated": - pad_width = _to_dilated(pad_width, input.ndim) - if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: - constant_values = ivy.asarray(constant_values, dtype=input.dtype) - return _interior_pad(input, constant_values, pad_width) - pad_width = _to_pairs(pad_width, len(input.shape)) - if callable(mode): - func = mode - padded, _ = _pad_simple(input, pad_width, fill_value=0) - for axis in range(padded.ndim): - padded = ivy.moveaxis(padded, axis, -1) - inds = ivy.ndindex(padded.shape[:-1]) - for ind in inds: - padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) - return padded - padded, original_area_slice = _pad_simple(input, pad_width) - axes = range(padded.ndim) - stat_functions = { - "maximum": ivy.max, - "minimum": ivy.min, - "mean": ivy.mean, - "median": ivy.median, - } - if mode == "constant": - constant_values = _to_pairs(constant_values, padded.ndim) - constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) - for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): - padded = _set_pad_area(padded, axis, width_pair, value_pair) - elif mode == "empty": - pass - elif mode == "edge": - for axis, width_pair in zip(axes, pad_width): - edge_pair = _get_edges(padded, axis, width_pair) - padded = _set_pad_area(padded, axis, width_pair, edge_pair) - elif mode == "linear_ramp": - end_values = _to_pairs(end_values, padded.ndim) - for axis, width_pair, value_pair in zip(axes, pad_width, end_values): - ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) - padded = _set_pad_area(padded, axis, width_pair, ramp_pair) - elif mode in stat_functions: - func = stat_functions[mode] - stat_length = _to_pairs(stat_length, padded.ndim) - if mode == "median": - ivy.utils.assertions.check_true( - ivy.is_float_dtype(input), - message="median interpolation is only supported for floats", +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = ivy.full( + new_shape, padding_value, dtype=operand.dtype ) - for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): - stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) - padded = _set_pad_area(padded, axis, width_pair, stat_pair) - elif mode in {"reflect", "symmetric"}: - include_edge = True if mode == "symmetric" else False - for axis, (left_index, right_index) in zip(axes, pad_width): - if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): - edge_pair = _get_edges(padded, axis, (left_index, right_index)) - padded = _set_pad_area( - padded, axis, (left_index, right_index), edge_pair - ) - continue - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_reflect_both( - padded, axis, (left_index, right_index), reflect_type, include_edge - ) - elif mode == "wrap": - for axis, (left_index, right_index) in zip(axes, pad_width): - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_wrap_both( - padded, axis, (left_index, right_index) - ) + src_indices = ivy.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) + + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = ivy.constant_pad(padded, pad_width, value=padding_value) return padded -@handle_nestable +def _interleave(a, b, axis): + assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 + a_pad = [(0, 0, 0)] * a.ndim + b_pad = [(0, 0, 0)] * b.ndim + a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) + b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) + a = _interior_pad(a, 0.0, a_pad) + b = _interior_pad(b, 0.0, b_pad) + return ivy.add(a, b) + + @handle_exceptions -@handle_array_like_without_promotion +@handle_nestable @inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def partial_fold( +def associative_scan( x: Union[ivy.Array, ivy.NativeArray], + fn: Callable, /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, *, - out: Optional[ivy.Array] = None, + reverse: bool = False, + axis: int = 0, ) -> ivy.Array: """ - Re-folds a partially unfolded tensor. + Perform an associative scan over the given array. Parameters ---------- x - a partially unfolded tensor - mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions left untouched at the beginning - out - optional output array, for writing the result to. + The array to scan over. + fn + The associative function to apply. + reverse + Whether to scan in reverse with respect to the given axis. + axis + The axis to scan over. Returns ------- ret - partially re-folded tensor + The result of the scan. """ - transposed_shape = list(shape) - mode_dim = transposed_shape.pop(skip_begin + mode) - transposed_shape.insert(skip_begin, mode_dim) - return ivy.moveaxis( - ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out - ) + elems = [x] + + if reverse: + elems = [ivy.flip(elem, axis=[axis]) for elem in elems] + + def _combine(a, b): + a = a[0] + b = b[0] + if a.shape[axis] == 0: + return [a] + c = fn(a, b) + return [c] + + def _scan(elems): + num_elems = elems[0].shape[axis] + + if num_elems < 2: + return elems + + reduced_elems = _combine( + [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], + [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], + ) + + odd_elems = _scan(reduced_elems) + + if num_elems % 2 == 0: + even_elems = _combine( + [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + else: + even_elems = _combine( + odd_elems, + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + even_elems = [ + ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) + for (elem, result) in zip(elems, even_elems) + ] + return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) + + scans = _scan(elems) + + if reverse: + scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] + + return ivy.reshape(ivy.asarray(scans), elems[0].shape) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def partial_tensor_to_vec( +def unique_consecutive( x: Union[ivy.Array, ivy.NativeArray], /, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + axis: Optional[int] = None, +) -> Tuple[ + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], +]: """ - Partial vectorization of a tensor while ignoring the specified dimension at the - beginning and the end. + Eliminates all but the first element from every consecutive group of equivalent + elements in ``x``. Parameters ---------- x - tensor to partially vectorise - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - out - optional output array, for writing the result to. + input array. + + axis + the axis to apply unique on. If None, unique is applied on flattened ``x``. Returns ------- ret - partially vectorised tensor with the - `skip_begin` first and `skip_end` last dimensions untouched + a namedtuple ``(output, inverse_indices, counts)`` whose + - first element has the field name ``output`` and is an array + containing ``x`` with its equivalent consecutive elements eliminated. + - second element has the field name ``inverse_indices`` and is an + array containing the indices of ``output`` that reconstruct ``x``. + - third element has the field name ``counts`` and is an array + containing the number of occurrences for each unique value or array in ``x``. + + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) + >>> ivy..unique_consecutive(x) + Results(values=ivy.array([1, 2, 3, 1, 2]), + inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), + counts=ivy.array([2, 2, 1, 2, 1])) """ - return partial_unfold( - x, - mode=0, - skip_begin=skip_begin, - skip_end=skip_end, - ravel_tensors=True, - out=out, - ) + return ivy.current_backend(x).unique_consecutive(x, axis=axis) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def partial_unfold( - x: Union[ivy.Array, ivy.NativeArray], +def fill_diagonal( + a: Union[ivy.Array, ivy.NativeArray], + v: Union[int, float], /, - mode: Optional[int] = 0, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, - ravel_tensors: Optional[bool] = False, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + wrap: bool = False, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Partial unfolding of a tensor while ignoring the specified number of dimensions at - the beginning and the end. For instance, if the first dimension of the tensor is the - number of samples, to unfold each sample, set skip_begin=1. This would, for each i - in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. + Fill the main diagonal of the given array of any dimensionality.. Parameters ---------- - x - tensor of shape n_samples x n_1 x n_2 x ... x n_i - mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - ravel_tensors - if True, the unfolded tensors are also flattened - out - optional output array, for writing the result to. + a + Array at least 2D. + v + The value to write on the diagonal. + wrap + The diagonal ‘wrapped’ after N columns for tall matrices. Returns ------- ret - partially unfolded tensor + Array with the diagonal filled. """ - if ravel_tensors: - new_shape = [-1] - else: - new_shape = [x.shape[mode + skip_begin], -1] - - if skip_begin: - new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape - - if skip_end: - new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] - - return ivy.reshape( - ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out - ) + return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) @handle_nestable @@ -1973,205 +2160,71 @@ def partial_unfold( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def partial_vec_to_tensor( +def unfold( x: Union[ivy.Array, ivy.NativeArray], /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, + mode: Optional[int] = 0, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Refolds a partially vectorised tensor into a full one. + Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. Parameters ---------- x - a partially vectorised tensor - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions to leave untouched at the beginning + input tensor to be unfolded + mode + indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` out optional output array, for writing the result to. Returns ------- ret - full tensor + unfolded_tensor of shape ``(tensor.shape[mode], -1)`` """ - return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) + return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) -@handle_exceptions @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -def put_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - values: Union[ivy.Array, ivy.NativeArray], - axis: int, - /, - *, - mode: str = "raise", - out: Optional[ivy.Array] = None, -) -> None: - """ - Put values into the input array by matching 1d index and data slices along a - specified axis. - - Parameters - ---------- - arr : array_like - The input array to modify. - indices : array_like - The indices of the values to put into `arr`. - values : array_like - The values to put into `arr`. - axis : int - The axis over which to put the `values`. - mode : {'raise', 'wrap', 'clip'}, optional - Specifies how out-of-bounds indices will be handled. - The following modes are available: - - - 'raise': a `ValueError` is raised when an index is out of bounds. - - 'wrap': the index is wrapped around to the corresponding index - at the other end of the axis. - - 'clip': the index is clipped to the closest in-bounds index. - out : ndarray, optional - Output array in which to place the result. - If not specified, a new array is created. - - Returns - ------- - None - - Examples - -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) - >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') - >>> print(arr) - ivy.array([[3, 7, 5], - [6, 4, 1]]) - """ - if out is None: - out = ivy.zeros_like(arr) - - indices = ivy.expand_dims(indices, axis=axis) - values = ivy.expand_dims(values, axis=axis) - - stacked = ivy.concat((arr, values), axis=axis) - - sorted_indices = ivy.argsort(indices, axis=axis) - sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), - sorted_stacked, - arr, - ) - - if mode == "clip": - indices = ivy.clip(indices, 0, arr.shape[axis] - 1) - elif mode == "wrap": - indices = ivy.mod(indices, arr.shape[axis]) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values - ) - - ivy.assign(out, arr) - - @handle_exceptions -@handle_backend_invalid -@handle_nestable @handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def rot90( - m: Union[ivy.Array, ivy.NativeArray], +def fold( + x: Union[ivy.Array, ivy.NativeArray], /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is - from the first towards the second axis. + Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, + refolds the n-mode unfolded tensor into the original tensor of the specified shape. Parameters ---------- - m - Input array of two or more dimensions. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - k - Number of times the array is rotated by 90 degrees. - axes - The array is rotated in the plane defined by the axes. Axes must be - different. + input + unfolded tensor of shape ``(shape[mode], -1)`` + mode + the mode of the unfolding + shape + shape of the original tensor before unfolding out - optional output container, for writing the result to. It must have a shape - that the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - A rotated view of m. - - Examples - -------- - With :code:`ivy.Array` input: - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) - With :code:`ivy.NativeArray` input: - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.native_array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) + folded_tensor of shape `shape` """ - return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) + full_shape = list(shape) + mode_dim = full_shape.pop(mode) + full_shape.insert(0, mode_dim) + return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) @handle_nestable @@ -2180,190 +2233,144 @@ def rot90( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def soft_thresholding( +def partial_unfold( x: Union[ivy.Array, ivy.NativeArray], /, - threshold: Union[float, ivy.Array, ivy.NativeArray], + mode: Optional[int] = 0, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, + ravel_tensors: Optional[bool] = False, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Soft-thresholding operator. - - sign(tensor) * max[abs(tensor) - threshold, 0] + Partial unfolding of a tensor while ignoring the specified number of dimensions at + the beginning and the end. For instance, if the first dimension of the tensor is the + number of samples, to unfold each sample, set skip_begin=1. This would, for each i + in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. Parameters ---------- x - input array - threshold - float or array with shape tensor.shape - * If float the threshold is applied to the whole tensor - * If array, one threshold is applied per elements, 0 values are ignored + tensor of shape n_samples x n_1 x n_2 x ... x n_i + mode + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + ravel_tensors + if True, the unfolded tensors are also flattened out optional output array, for writing the result to. Returns ------- - ivy.Array - thresholded tensor on which the operator has been applied - - Examples - -------- - Basic shrinkage - - >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) - >>> soft_thresholding(x, 1.1) - array([[ 0. , -0.9, 0.4], - [-2.9, 1.9, 0. ]]) - + ret + partially unfolded tensor + """ + if ravel_tensors: + new_shape = [-1] + else: + new_shape = [x.shape[mode + skip_begin], -1] - Example with missing values + if skip_begin: + new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape - >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) - >>> soft_thresholding(x, mask*1.1) - array([[ 1. , -2. , 0.4], - [-2.9, 3. , 0. ]]) - """ - res = ivy.abs(x) - threshold - res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) + if skip_end: + new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + return ivy.reshape( + ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def take_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - axis: int, +def partial_fold( + x: Union[ivy.Array, ivy.NativeArray], /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, - mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Take values from the input array by matching 1d index and data slices. + Re-folds a partially unfolded tensor. Parameters ---------- - arr - The source array. - indices - The indices of the values to extract. - axis - The axis over which to select values. - If axis is None, arr is treated as a flattened 1D array. + x + a partially unfolded tensor mode - One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices - will be handled. + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions left untouched at the beginning out - The output array. + optional output array, for writing the result to. Returns ------- ret - The returned array has the same shape as `indices`. - - Examples - -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> y = ivy.take_along_axis(arr, indices, 1) - >>> print(y) - ivy.array([[4, 3, 3], [1, 1, 1]]) + partially re-folded tensor """ - return ivy.current_backend(arr).take_along_axis( - arr, indices, axis, mode=mode, out=out + transposed_shape = list(shape) + mode_dim = transposed_shape.pop(skip_begin + mode) + transposed_shape.insert(skip_begin, mode_dim) + return ivy.moveaxis( + ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def top_k( +def partial_tensor_to_vec( x: Union[ivy.Array, ivy.NativeArray], - k: int, /, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[tuple] = None, -) -> Tuple[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Return the `k` largest elements of the given input array along a given axis. + Partial vectorization of a tensor while ignoring the specified dimension at the + beginning and the end. Parameters ---------- x - The array to compute top_k for. - k - Number of top elements to retun must not exceed the array size. - axis - The axis along which we must return the top elements default value is 1. - largest - If largest is set to False we return k smallest elements of the array. - sorted - If sorted is set to True we return the elements in sorted order. - out: - Optional output tuple, for writing the result to. Must have two arrays inside, - with a shape that the returned tuple broadcast to. + tensor to partially vectorise + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + out + optional output array, for writing the result to. Returns ------- ret - A named tuple with values and indices of top k elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 2) - >>> print(y) - top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) - - >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) - >>> y = ivy.top_k(x, 2, axis=1, largest=False) - >>> print(y) - top_k(values=ivy.array([[-2., 0.], - [-8., -1.]]), indices=ivy.array([[0, 3], - [0, 2]])) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 3) - >>> print(y) - top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) - >>> y = x.top_k(2) - >>> print(y) - [{ - a: ivy.array([2, -1]), - b: ivy.array([5., 4.]) - }, { - a: ivy.array([1, 0]), - b: ivy.array([1, 0]) - }] + partially vectorised tensor with the + `skip_begin` first and `skip_end` last dimensions untouched """ - return current_backend(x).top_k( - x, k, axis=axis, largest=largest, sorted=sorted, out=out + return partial_unfold( + x, + mode=0, + skip_begin=skip_begin, + skip_end=skip_end, + ravel_tensors=True, + out=out, ) @@ -2373,209 +2380,200 @@ def top_k( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unfold( +def partial_vec_to_tensor( x: Union[ivy.Array, ivy.NativeArray], /, - mode: Optional[int] = 0, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. + Refolds a partially vectorised tensor into a full one. Parameters ---------- x - input tensor to be unfolded - mode - indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` + a partially vectorised tensor + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions to leave untouched at the beginning out optional output array, for writing the result to. Returns ------- ret - unfolded_tensor of shape ``(tensor.shape[mode], -1)`` + full tensor """ - return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) + return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unique_consecutive( +def matricize( x: Union[ivy.Array, ivy.NativeArray], /, + row_modes: Sequence[int], + column_modes: Optional[Sequence[int]] = None, *, - axis: Optional[int] = None, -) -> Tuple[ - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], -]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Eliminates all but the first element from every consecutive group of equivalent - elements in ``x``. + Matricizes the given tensor. Parameters ---------- x - input array. - - axis - the axis to apply unique on. If None, unique is applied on flattened ``x``. + the input tensor + row_modes + modes to use as row of the matrix (in the desired order) + column_modes + modes to use as column of the matrix, in the desired order + if None, the modes not in `row_modes` will be used in ascending order + out + optional output array, for writing the result to. - Returns - ------- ret - a namedtuple ``(output, inverse_indices, counts)`` whose - - first element has the field name ``output`` and is an array - containing ``x`` with its equivalent consecutive elements eliminated. - - second element has the field name ``inverse_indices`` and is an - array containing the indices of ``output`` that reconstruct ``x``. - - third element has the field name ``counts`` and is an array - containing the number of occurrences for each unique value or array in ``x``. + ------- + ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) + """ + ndims = len(x.shape) + row_indices = list(row_modes) + + if column_modes: + column_indices = list(column_modes) + else: + column_indices = [i for i in range(ndims) if i not in row_indices] + if sorted(column_indices + row_indices) != list(range(ndims)): + msg = ( + "If you provide both column and row modes for the matricization then" + " column_modes + row_modes must contain all the modes of the tensor." + f" Yet, got row_modes={row_modes} and column_modes={column_modes}." + ) + raise ValueError(msg) + row_size, column_size = 1, 1 + row_size = int(ivy.prod([x.shape[i] for i in row_indices])) + column_size = int(ivy.prod([x.shape[i] for i in column_indices])) - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) - >>> ivy..unique_consecutive(x) - Results(values=ivy.array([1, 2, 3, 1, 2]), - inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), - counts=ivy.array([2, 2, 1, 2, 1])) - """ - return ivy.current_backend(x).unique_consecutive(x, axis=axis) + return ivy.reshape( + ivy.permute_dims(x, row_indices + column_indices), + (row_size, column_size), + out=out, + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def vsplit( - ary: Union[ivy.Array, ivy.NativeArray], - indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], +def soft_thresholding( + x: Union[ivy.Array, ivy.NativeArray], /, + threshold: Union[float, ivy.Array, ivy.NativeArray], *, - copy: Optional[bool] = None, -) -> List[ivy.Array]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Split an array vertically into multiple sub-arrays. + Soft-thresholding operator. + + sign(tensor) * max[abs(tensor) - threshold, 0] Parameters ---------- - ary - Array input. - indices_or_sections - If indices_or_sections is an integer n, the array is split into n - equal sections, provided that n must be a divisor of the split axis. - If indices_or_sections is a sequence of ints or 1-D array, - then input is split at each of the indices. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + x + input array + threshold + float or array with shape tensor.shape + * If float the threshold is applied to the whole tensor + * If array, one threshold is applied per elements, 0 values are ignored + out + optional output array, for writing the result to. Returns ------- - ret - input array split vertically. + ivy.Array + thresholded tensor on which the operator has been applied Examples -------- - >>> ary = ivy.array( - [[[0., 1.], - [2., 3.]], - [[4., 5.], - [6., 7.]]] - ) - >>> ivy.vsplit(ary, 2) - [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) + Basic shrinkage + + >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) + >>> soft_thresholding(x, 1.1) + array([[ 0. , -0.9, 0.4], + [-2.9, 1.9, 0. ]]) + + + Example with missing values + + >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) + >>> soft_thresholding(x, mask*1.1) + array([[ 1. , -2. , 0.4], + [-2.9, 3. , 0. ]]) """ - return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) + res = ivy.abs(x) - threshold + res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) + + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def vstack( - arrays: Sequence[ivy.Array], +def choose( + arr: Union[ivy.Array, ivy.NativeArray], + choices: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: None = None, + mode: Union[str, None] = None, ) -> ivy.Array: """ - Stack arrays in sequence vertically (row wise). + Take values from the input array by matching 1d index and data slices. Parameters ---------- - arrays - Sequence of arrays to be stacked. + arr + The source array. + choices + The indices of the values to extract. + out + The output array. + mode + One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices + will be handled. Returns ------- ret - The array formed by stacking the given arrays. + The returned array has the same shape as `indices`. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.vstack((x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4]]) - >>> ivy.vstack((x, y, x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [1, 2, 3], - [2, 3, 4]]) - - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.vstack(y)) - ivy.array([[5, 6], - [7, 8]]) + >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], + [20, 21, 22, 23], [30, 31, 32, 33]]) + >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) + ivy.array([20, 1, 12, 3]) """ - return ivy.current_backend().vstack(arrays, out=out) - - -flatten.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -pad.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -as_strided.unsupported_dtypes = ("bfloat16",) -as_strided.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) diff --git a/ivy/functional/ivy/experimental/norms.py b/ivy/functional/ivy/experimental/norms.py index d37b4d6c0e954..4d59a2aaa6907 100644 --- a/ivy/functional/ivy/experimental/norms.py +++ b/ivy/functional/ivy/experimental/norms.py @@ -17,35 +17,89 @@ from ivy.utils.exceptions import handle_exceptions -batch_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -instance_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -group_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l1_normalize( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """Normalize the input array along the given axis to have L1 norm equal to + 1. + + Parameters + ---------- + x + Input array. + axis + Axis or axes along which to normalize. If ``None``, + the whole array is normalized. + out + Optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + + Returns + ------- + ret + The normalized array. + + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l1_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.33333334, 1.33333337], + [1.28571439, 2.28571439]]) + """ + return current_backend(x).l1_normalize(x, axis=axis, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l2_normalize( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """Normalize the input array along the given axis to have L2 norm equal to + 1. + + Parameters + ---------- + x + Input array. + axis + Axis along which to normalize. If ``None``, the whole array is normalized. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The normalized array. + + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l2_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.44721359, 0.89442718], + [0.60000002, 0.80000001]]) + """ + return current_backend(x).l2_normalize(x, axis=axis, out=out) @handle_exceptions @@ -158,80 +212,16 @@ def batch_norm( return xnormalized, runningmean, runningvariance -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def group_norm( - x: Union[ivy.NativeArray, ivy.Array], - num_groups: int = 1, - /, - *, - offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Apply group normalization to the input array and returns the normalized input. - - Parameters - ---------- - x - Input array of default shape (N, *S, C), where N is the batch dimension, - *S corresponds to any number of spatial dimensions and - C corresponds to the channel dimension. - num_groups - number of groups to separate the channels into - offset - An offset array of size C. If present, will be added - to the normalized input. - scale - A scale array of size C. If present, the scale is - applied to the normalized input. - eps - A small float number to avoid dividing by 0. - data_format - The ordering of the dimensions in the input, one of "NSC" or "NCS", - where N is the batch dimension, S represents any number of spatial - dimensions and C is the channel dimension. Default is "NSC". - out - optional output arrays, for writing the result to. - - Returns - ------- - ret - The normalized array. - """ - xdims = ivy.get_num_dims(x) - if data_format == "NSC": - x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) - N = x.shape[0] - C = x.shape[1] - S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 - assert C % num_groups == 0 - x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) - mean = ivy.mean(x_, axis=(2, 3), keepdims=True) - var = ivy.var(x_, axis=(2, 3), keepdims=True) - x_normalized = (x_ - mean) / ivy.sqrt(var + eps) - x_normalized = ivy.reshape(x_normalized, x.shape) - - if ivy.exists(scale): - scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized * scale - - if ivy.exists(offset): - offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized + offset - - if data_format == "NSC": - x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) - - if ivy.exists(out): - x_normalized = ivy.inplace_update(out, x_normalized) - return x_normalized +batch_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} @handle_exceptions @@ -354,89 +344,103 @@ def instance_norm( return (xnormalized, runningmean, runningvariance) +instance_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} + + @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l1_normalize( - x: Union[ivy.Array, ivy.NativeArray], +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def group_norm( + x: Union[ivy.NativeArray, ivy.Array], + num_groups: int = 1, /, *, - axis: Optional[Union[int, Tuple[int, ...]]] = None, + offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", out: Optional[ivy.Array] = None, ) -> ivy.Array: - """Normalize the input array along the given axis to have L1 norm equal to - 1. + """ + Apply group normalization to the input array and returns the normalized input. Parameters ---------- x - Input array. - axis - Axis or axes along which to normalize. If ``None``, - the whole array is normalized. + Input array of default shape (N, *S, C), where N is the batch dimension, + *S corresponds to any number of spatial dimensions and + C corresponds to the channel dimension. + num_groups + number of groups to separate the channels into + offset + An offset array of size C. If present, will be added + to the normalized input. + scale + A scale array of size C. If present, the scale is + applied to the normalized input. + eps + A small float number to avoid dividing by 0. + data_format + The ordering of the dimensions in the input, one of "NSC" or "NCS", + where N is the batch dimension, S represents any number of spatial + dimensions and C is the channel dimension. Default is "NSC". out - Optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + optional output arrays, for writing the result to. Returns ------- ret The normalized array. - - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l1_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.33333334, 1.33333337], - [1.28571439, 2.28571439]]) """ - return current_backend(x).l1_normalize(x, axis=axis, out=out) + xdims = ivy.get_num_dims(x) + if data_format == "NSC": + x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) + N = x.shape[0] + C = x.shape[1] + S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 + assert C % num_groups == 0 + x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) + mean = ivy.mean(x_, axis=(2, 3), keepdims=True) + var = ivy.var(x_, axis=(2, 3), keepdims=True) + x_normalized = (x_ - mean) / ivy.sqrt(var + eps) + x_normalized = ivy.reshape(x_normalized, x.shape) + if ivy.exists(scale): + scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized * scale -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l2_normalize( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """Normalize the input array along the given axis to have L2 norm equal to - 1. + if ivy.exists(offset): + offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized + offset - Parameters - ---------- - x - Input array. - axis - Axis along which to normalize. If ``None``, the whole array is normalized. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + if data_format == "NSC": + x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) - Returns - ------- - ret - The normalized array. + if ivy.exists(out): + x_normalized = ivy.inplace_update(out, x_normalized) + return x_normalized - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l2_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.44721359, 0.89442718], - [0.60000002, 0.80000001]]) - """ - return current_backend(x).l2_normalize(x, axis=axis, out=out) + +group_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_exceptions diff --git a/ivy/functional/ivy/experimental/random.py b/ivy/functional/ivy/experimental/random.py index ff5a5d32f56bb..70586db0f3ccf 100644 --- a/ivy/functional/ivy/experimental/random.py +++ b/ivy/functional/ivy/experimental/random.py @@ -14,65 +14,69 @@ from ivy.utils.exceptions import handle_exceptions +# dirichlet @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back -@infer_dtype @handle_device_shifting -@infer_device -def bernoulli( - probs: Union[float, ivy.Array, ivy.NativeArray], +def dirichlet( + alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], + /, *, - logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from Bernoulli distrubution paramterized by probs or logits (but not - both) + Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- + distributed random variable can be seen as a multivariate generalization of a Beta + distribution. The Dirichlet distribution is a conjugate prior of a multinomial + distribution in Bayesian inference. Parameters ---------- - logits - An N-D Array representing the log-odds of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution where the probability of an event is sigmoid - (logits). Only one of logits or probs should be passed in. - probs - An N-D Array representing the probability of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution. Only one of logits or probs should be passed in - shape - If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. - (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). + alpha + Sequence of floats of length k + size + optional int or tuple of ints, Output shape. If the given shape is, + e.g., (m, n), then m * n * k samples are drawn. Default is None, + in which case a vector of length k is returned. dtype output array data type. If ``dtype`` is ``None``, the output array data type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Drawn samples from the Bernoulli distribution + The drawn samples, of shape (size, k). + + Functional Examples + ------------------- + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha) + ivy.array([0.10598304, 0.21537054, 0.67864642]) + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha, size = (2,3)) + ivy.array([[[0.48006698, 0.07472073, 0.44521229], + [0.55479872, 0.05426367, 0.39093761], + [0.19531053, 0.51675832, 0.28793114]], + + [[0.12315625, 0.29823365, 0.5786101 ], + [0.15564976, 0.50542368, 0.33892656], + [0.1325352 , 0.44439589, 0.42306891]]]) """ - return ivy.current_backend(probs).bernoulli( - probs, - logits=logits, - shape=shape, - device=device, + return ivy.current_backend().dirichlet( + alpha, + size=size, dtype=dtype, seed=seed, out=out, @@ -134,75 +138,6 @@ def beta( ) -# dirichlet -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def dirichlet( - alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], - /, - *, - size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- - distributed random variable can be seen as a multivariate generalization of a Beta - distribution. The Dirichlet distribution is a conjugate prior of a multinomial - distribution in Bayesian inference. - - Parameters - ---------- - alpha - Sequence of floats of length k - size - optional int or tuple of ints, Output shape. If the given shape is, - e.g., (m, n), then m * n * k samples are drawn. Default is None, - in which case a vector of length k is returned. - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. - - Returns - ------- - ret - The drawn samples, of shape (size, k). - - Functional Examples - ------------------- - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha) - ivy.array([0.10598304, 0.21537054, 0.67864642]) - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha, size = (2,3)) - ivy.array([[[0.48006698, 0.07472073, 0.44521229], - [0.55479872, 0.05426367, 0.39093761], - [0.19531053, 0.51675832, 0.28793114]], - - [[0.12315625, 0.29823365, 0.5786101 ], - [0.15564976, 0.50542368, 0.33892656], - [0.1325352 , 0.44439589, 0.42306891]]]) - """ - return ivy.current_backend().dirichlet( - alpha, - size=size, - dtype=dtype, - seed=seed, - out=out, - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -325,3 +260,68 @@ def poisson( fill_value=fill_value, out=out, ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +@infer_device +def bernoulli( + probs: Union[float, ivy.Array, ivy.NativeArray], + *, + logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Draws samples from Bernoulli distrubution paramterized by probs or logits (but not + both) + + Parameters + ---------- + logits + An N-D Array representing the log-odds of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution where the probability of an event is sigmoid + (logits). Only one of logits or probs should be passed in. + probs + An N-D Array representing the probability of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution. Only one of logits or probs should be passed in + shape + If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. + (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Drawn samples from the Bernoulli distribution + """ + return ivy.current_backend(probs).bernoulli( + probs, + logits=logits, + shape=shape, + device=device, + dtype=dtype, + seed=seed, + out=out, + ) diff --git a/ivy/functional/ivy/experimental/sparse_array.py b/ivy/functional/ivy/experimental/sparse_array.py index 785eccaa6b3bd..46e6eb64fa428 100644 --- a/ivy/functional/ivy/experimental/sparse_array.py +++ b/ivy/functional/ivy/experimental/sparse_array.py @@ -4,440 +4,270 @@ from ivy.utils.exceptions import handle_exceptions -class SparseArray: - def __init__( - self, - data=None, - *, - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format=None, - ): - if _is_data_not_indices_values_and_shape( - data, - coo_indices, - crow_indices, - col_indices, - ccol_indices, - row_indices, - values, - dense_shape, - ): - self._init_data(data) - elif _is_valid_format( - coo_indices, - crow_indices, - col_indices, - ccol_indices, - row_indices, - values, - dense_shape, - format=format, - ): - format = format.lower() - - if format == "coo": - self._init_coo_components(coo_indices, values, dense_shape, format) - elif format == "csr" or format == "bsr": - self._init_compressed_row_components( - crow_indices, col_indices, values, dense_shape, format - ) - else: - print(format) - self._init_compressed_column_components( - ccol_indices, row_indices, values, dense_shape, format - ) - - else: - print( - format, - ccol_indices, - row_indices, - values, - dense_shape, - crow_indices, - col_indices, - values, - ) - - raise ivy.utils.exceptions.IvyException( - "specify all coo components (coo_indices, values and " - " dense_shape), all csr components (crow_indices, " - "col_indices, values and dense_shape), all csc components " - "(ccol_indices, row_indices, values and dense_shape). all " - "bsc components (ccol_indices, row_indices, values and " - "dense_shape), or all bsr components (crow_indices, " - "col_indices, values and dense_shape)." - ) - - def _init_data(self, data): - if ivy.is_ivy_sparse_array(data): - self._data = data.data - self._coo_indices = data.coo_indices - self._crow_indices = data.crow_indices - self._col_indices = data.col_indices - self._ccol_indices = data.ccol_indices - self._row_indices = data.row_indices - self._values = data.values - self._dense_shape = data.dense_shape - self._format = data.format.lower() - else: - ivy.utils.assertions.check_true( - ivy.is_native_sparse_array(data), message="not a native sparse array" - ) - self._data = data - self._native_sparse_array_to_indices_values_and_shape() - - def _native_sparse_array_to_indices_values_and_shape(self): - indices, values, shape = ivy.native_sparse_array_to_indices_values_and_shape( - self._data - ) - - if "coo_indices" in indices: - self._coo_indices = ivy.array(indices["coo_indices"], dtype="int64") - self._crow_indices = None - self._col_indices = None - self._ccol_indices = None - self._row_indices = None - - elif "crow_indices" in indices and "col_indices" in indices: - self._crow_indices = ivy.array(indices["crow_indices"], dtype="int64") - self._col_indices = ivy.array(indices["col_indices"], dtype="int64") - self._coo_indices = None - self._ccol_indices = None - self._row_indices = None - - else: - self._ccol_indices = ivy.array(indices["ccol_indices"], dtype="int64") - self._row_indices = ivy.array(indices["row_indices"], dtype="int64") - self._coo_indices = None - self._crow_indices = None - self._col_indices = None - - self._values = ivy.array(values) - self._dense_shape = ivy.Shape(shape) - self._format = self._data.format.lower() - - def _init_coo_components(self, coo_indices, values, shape, format): - coo_indices = ivy.array(coo_indices, dtype="int64") - values = ivy.array(values) - shape = ivy.Shape(shape) - self._data = ivy.native_sparse_array( - coo_indices=coo_indices, values=values, dense_shape=shape, format=format - ) - self._coo_indices = coo_indices - self._values = values - self._dense_shape = shape - self._format = format - self._crow_indices = None - self._col_indices = None - self._ccol_indices = None - self._row_indices = None - - def _init_compressed_row_components( - self, crow_indices, col_indices, values, shape, format - ): - crow_indices = ivy.array(crow_indices, dtype="int64") - col_indices = ivy.array(col_indices, dtype="int64") - values = ivy.array(values) - shape = ivy.Shape(shape) - self._data = ivy.native_sparse_array( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=shape, - format=format, - ) - self._crow_indices = crow_indices - self._col_indices = col_indices - self._values = values - self._dense_shape = shape - self._format = format - self._coo_indices = None - self._ccol_indices = None - self._row_indices = None - - def _init_compressed_column_components( - self, ccol_indices, row_indices, values, shape, format - ): - ccol_indices = ivy.array(ccol_indices, dtype="int64") - row_indices = ivy.array(row_indices, dtype="int64") - values = ivy.array(values) - shape = ivy.Shape(shape) - self._data = ivy.native_sparse_array( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=shape, - format=format, - ) - self._ccol_indices = ccol_indices - self._row_indices = row_indices - self._values = values - self._dense_shape = shape - self._format = format - self._coo_indices = None - self._crow_indices = None - self._col_indices = None - - # Properties # - # -----------# - - @property - def data(self): - return self._data - - @property - def coo_indices(self): - return self._coo_indices - - @property - def crow_indices(self): - return self._crow_indices - - @property - def col_indices(self): - return self._col_indices - - @property - def ccol_indices(self): - return self._ccol_indices - - @property - def row_indices(self): - return self._row_indices - - @property - def values(self): - return self._values - - @property - def dense_shape(self): - return self._dense_shape - - @property - def format(self): - return self._format - - # Setters # - # --------# - - @data.setter - def data(self, data): - self._init_data(data) - - @coo_indices.setter - def coo_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - _verify_coo_components( - indices=indices, values=self._values, dense_shape=self._dense_shape +# helpers +def _verify_coo_components(indices=None, values=None, dense_shape=None): + ivy.utils.assertions.check_all_or_any_fn( + indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message="indices, values and dense_shape must all be specified", + ) + # coordinates style (COO), must be shaped (x, y) + ivy.utils.assertions.check_equal( + len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.to_ivy_shape(dense_shape)), + ivy.shape(indices)[0], + message="shape and indices shape do not match", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(values)[0], + ivy.shape(indices)[1], + message="values and indices do not match", + as_array=False, + ) + for i in range(ivy.shape(indices)[0]): + ivy.utils.assertions.check_less( + indices[i], + ivy.to_ivy_shape(dense_shape)[i], + message="indices is larger than shape", ) - self._coo_indices = indices - - @crow_indices.setter - def crow_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._crow_indices = indices - - @col_indices.setter - def col_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._col_indices = indices - @ccol_indices.setter - def ccol_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._ccol_indices = indices - @row_indices.setter - def row_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._row_indices = indices +def _verify_common_row_format_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" +): + ivy.utils.assertions.check_all_or_any_fn( + crow_indices, + col_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "crow_indices, col_indices, values and dense_shape must all be specified." + ), + ) - @values.setter - def values(self, values): - values = ivy.array(values) - _verify_coo_components( - indices=self._coo_indices, values=values, dense_shape=self._dense_shape - ) - self._values = values + ivy.utils.assertions.check_equal( + len(ivy.shape(crow_indices)), + 1, + message="crow_indices must be 1D.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(col_indices)), + 1, + message="col_indices must be 1D.", + as_array=False, + ) - @dense_shape.setter - def dense_shape(self, dense_shape): - dense_shape = ivy.Shape(dense_shape) - _verify_coo_components( - indices=self._coo_indices, values=self._values, dense_shape=dense_shape - ) - self._dense_shape = dense_shape + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", + as_array=False, + ) - @format.setter - def format(self, format): - self._format = format + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + crow_indices[-1], + message="size of col_indices does not match with last element of crow_indices", + ) - # Instance Methods # - # ---------------- # + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + ivy.shape(values)[0], + message="values and col_indices do not match", + as_array=False, + ) - def _coo_to_dense_coordinates(self): - all_coordinates = [] - for i in range(self._values.shape[0]): - coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) - coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) - all_coordinates.append(coordinate.to_list()) - return all_coordinates + # index in crow_indices must not exceed length of col_indices + ivy.utils.assertions.check_less( + crow_indices, + ivy.shape(col_indices)[0], + allow_equal=True, + message="index in crow_indices does not match the number of col_indices", + ) - def _csr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._col_indices.to_list() - all_cols = self._crow_indices.to_list() - for row in range(total_rows): - cols = all_rows[all_cols[row] : all_cols[row + 1]] - for col in cols: - all_coordinates.append([row, col]) - return all_coordinates - def _csc_to_dense_coordinates(self): - # CSC sparse array - all_coordinates = [] - total_rows = self._dense_shape[1] - all_cols = self._row_indices.to_list() - all_rows = self._ccol_indices.to_list() - for col in range(total_rows): - rows = all_cols[all_rows[col] : all_rows[col + 1]] - for row in rows: - all_coordinates.append([row, col]) - return all_coordinates +def _verify_csr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="csr", + ) - def _bsr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._crow_indices.to_list() - all_cols = self._col_indices.to_list() + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False + ) + # number of intervals must be equal to x in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False + ) - nblockrows, nblockcols = self._values.shape[-2:] + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1], + message="index in col_indices does not match shape", + ) - for row in range(total_rows // nblockrows): - cols = all_cols[all_rows[row] : all_rows[row + 1]] - for col in cols: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates - def _bsc_to_dense_coordinates(self): - all_coordinates = [] - total_cols = self._dense_shape[1] - all_rows = self._row_indices.to_list() - all_cols = self._ccol_indices.to_list() +def _verify_bsr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="bsr", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="The number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="The number of cols of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False + ) + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1] // ncolblocks, + message="index in col_indices does not match shape", + ) - nblockrows, nblockcols = self._values.shape[-2:] - for col in range(total_cols // nblockcols): - rows = all_rows[all_cols[col] : all_cols[col + 1]] - for row in rows: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates +def _verify_common_column_format_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" +): + ivy.utils.assertions.check_all_or_any_fn( + ccol_indices, + row_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "ccol_indices, row_indices, values and dense_shape must all be specified" + ), + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(ccol_indices)), + 1, + message="ccol_indices must be 1D", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False + ) - def to_dense_array(self, *, native=False): - if self._format == "coo": - all_coordinates = self._coo_to_dense_coordinates() - elif self._format == "csr": - all_coordinates = self._csr_to_dense_coordinates() - elif self._format == "csc": - all_coordinates = self._csc_to_dense_coordinates() - elif self._format == "bsc": - all_coordinates = self._bsc_to_dense_coordinates() - else: - all_coordinates = self._bsr_to_dense_coordinates() + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(row_indices)[0], + ivy.shape(values)[0], + message="values and row_indices do not match", + as_array=False, + ) + # index in ccol_indices must not exceed length of row_indices + ivy.utils.assertions.check_less( + ccol_indices, + ivy.shape(row_indices)[0], + allow_equal=True, + message="index in ccol_indices does not match the number of row_indices", + ) - # make dense array - ret = ivy.scatter_nd( - ivy.array(all_coordinates), - ivy.flatten(self._values), - ivy.array(self._dense_shape), - ) - return ret.to_native() if native else ret +def _verify_csc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="csc", + ) -class NativeSparseArray: - pass + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0], + message="index in row_indices does not match shape", + ) -# --- Helpers --- # -# --------------- # +def _verify_bsc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="bsc", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="number of cols of array must be divisible by that of block.", + as_array=False, + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0] // nrowblocks, + message="index in row_indices does not match shape", + ) def _is_data_not_indices_values_and_shape( @@ -474,324 +304,486 @@ def _is_data_not_indices_values_and_shape( return False -def _is_valid_format( - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format="coo", -): - valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] +def _is_valid_format( + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format="coo", +): + valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] + + if not isinstance(format, str) or not format.lower() in valid_formats: + return False + + if format.endswith("o"): + # format is coo + return ( + ivy.exists(coo_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and crow_indices is None + and col_indices is None + and ccol_indices is None + and row_indices is None + ) + + if format.endswith("r"): + # format is either csr or bsr + return ( + ivy.exists(crow_indices) + and ivy.exists(col_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and ccol_indices is None + and row_indices is None + ) + # format is either csc or bsc + return ( + ivy.exists(ccol_indices) + and ivy.exists(row_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and crow_indices is None + and col_indices is None + ) + + +class SparseArray: + def __init__( + self, + data=None, + *, + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format=None, + ): + if _is_data_not_indices_values_and_shape( + data, + coo_indices, + crow_indices, + col_indices, + ccol_indices, + row_indices, + values, + dense_shape, + ): + self._init_data(data) + elif _is_valid_format( + coo_indices, + crow_indices, + col_indices, + ccol_indices, + row_indices, + values, + dense_shape, + format=format, + ): + format = format.lower() + + if format == "coo": + self._init_coo_components(coo_indices, values, dense_shape, format) + elif format == "csr" or format == "bsr": + self._init_compressed_row_components( + crow_indices, col_indices, values, dense_shape, format + ) + else: + print(format) + self._init_compressed_column_components( + ccol_indices, row_indices, values, dense_shape, format + ) + + else: + print( + format, + ccol_indices, + row_indices, + values, + dense_shape, + crow_indices, + col_indices, + values, + ) + + raise ivy.utils.exceptions.IvyException( + "specify all coo components (coo_indices, values and " + " dense_shape), all csr components (crow_indices, " + "col_indices, values and dense_shape), all csc components " + "(ccol_indices, row_indices, values and dense_shape). all " + "bsc components (ccol_indices, row_indices, values and " + "dense_shape), or all bsr components (crow_indices, " + "col_indices, values and dense_shape)." + ) + + def _init_data(self, data): + if ivy.is_ivy_sparse_array(data): + self._data = data.data + self._coo_indices = data.coo_indices + self._crow_indices = data.crow_indices + self._col_indices = data.col_indices + self._ccol_indices = data.ccol_indices + self._row_indices = data.row_indices + self._values = data.values + self._dense_shape = data.dense_shape + self._format = data.format.lower() + else: + ivy.utils.assertions.check_true( + ivy.is_native_sparse_array(data), message="not a native sparse array" + ) + self._data = data + self._native_sparse_array_to_indices_values_and_shape() + + def _native_sparse_array_to_indices_values_and_shape(self): + indices, values, shape = ivy.native_sparse_array_to_indices_values_and_shape( + self._data + ) + + if "coo_indices" in indices: + self._coo_indices = ivy.array(indices["coo_indices"], dtype="int64") + self._crow_indices = None + self._col_indices = None + self._ccol_indices = None + self._row_indices = None + + elif "crow_indices" in indices and "col_indices" in indices: + self._crow_indices = ivy.array(indices["crow_indices"], dtype="int64") + self._col_indices = ivy.array(indices["col_indices"], dtype="int64") + self._coo_indices = None + self._ccol_indices = None + self._row_indices = None + + else: + self._ccol_indices = ivy.array(indices["ccol_indices"], dtype="int64") + self._row_indices = ivy.array(indices["row_indices"], dtype="int64") + self._coo_indices = None + self._crow_indices = None + self._col_indices = None + + self._values = ivy.array(values) + self._dense_shape = ivy.Shape(shape) + self._format = self._data.format.lower() + + def _init_coo_components(self, coo_indices, values, shape, format): + coo_indices = ivy.array(coo_indices, dtype="int64") + values = ivy.array(values) + shape = ivy.Shape(shape) + self._data = ivy.native_sparse_array( + coo_indices=coo_indices, values=values, dense_shape=shape, format=format + ) + self._coo_indices = coo_indices + self._values = values + self._dense_shape = shape + self._format = format + self._crow_indices = None + self._col_indices = None + self._ccol_indices = None + self._row_indices = None + + def _init_compressed_row_components( + self, crow_indices, col_indices, values, shape, format + ): + crow_indices = ivy.array(crow_indices, dtype="int64") + col_indices = ivy.array(col_indices, dtype="int64") + values = ivy.array(values) + shape = ivy.Shape(shape) + self._data = ivy.native_sparse_array( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=shape, + format=format, + ) + self._crow_indices = crow_indices + self._col_indices = col_indices + self._values = values + self._dense_shape = shape + self._format = format + self._coo_indices = None + self._ccol_indices = None + self._row_indices = None + + def _init_compressed_column_components( + self, ccol_indices, row_indices, values, shape, format + ): + ccol_indices = ivy.array(ccol_indices, dtype="int64") + row_indices = ivy.array(row_indices, dtype="int64") + values = ivy.array(values) + shape = ivy.Shape(shape) + self._data = ivy.native_sparse_array( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=shape, + format=format, + ) + self._ccol_indices = ccol_indices + self._row_indices = row_indices + self._values = values + self._dense_shape = shape + self._format = format + self._coo_indices = None + self._crow_indices = None + self._col_indices = None + + # Properties # + # -----------# + + @property + def data(self): + return self._data - if not isinstance(format, str) or not format.lower() in valid_formats: - return False + @property + def coo_indices(self): + return self._coo_indices - if format.endswith("o"): - # format is coo - return ( - ivy.exists(coo_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and crow_indices is None - and col_indices is None - and ccol_indices is None - and row_indices is None - ) + @property + def crow_indices(self): + return self._crow_indices - if format.endswith("r"): - # format is either csr or bsr - return ( - ivy.exists(crow_indices) - and ivy.exists(col_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and ccol_indices is None - and row_indices is None - ) - # format is either csc or bsc - return ( - ivy.exists(ccol_indices) - and ivy.exists(row_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and crow_indices is None - and col_indices is None - ) + @property + def col_indices(self): + return self._col_indices + @property + def ccol_indices(self): + return self._ccol_indices -def _verify_bsc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="bsc", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="number of cols of array must be divisible by that of block.", - as_array=False, - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0] // nrowblocks, - message="index in row_indices does not match shape", - ) + @property + def row_indices(self): + return self._row_indices + @property + def values(self): + return self._values -def _verify_bsr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="bsr", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="The number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="The number of cols of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False - ) - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1] // ncolblocks, - message="index in col_indices does not match shape", - ) + @property + def dense_shape(self): + return self._dense_shape + @property + def format(self): + return self._format -def _verify_common_column_format_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" -): - ivy.utils.assertions.check_all_or_any_fn( - ccol_indices, - row_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "ccol_indices, row_indices, values and dense_shape must all be specified" - ), - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(ccol_indices)), - 1, - message="ccol_indices must be 1D", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False - ) + # Setters # + # --------# - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(row_indices)[0], - ivy.shape(values)[0], - message="values and row_indices do not match", - as_array=False, - ) - # index in ccol_indices must not exceed length of row_indices - ivy.utils.assertions.check_less( - ccol_indices, - ivy.shape(row_indices)[0], - allow_equal=True, - message="index in ccol_indices does not match the number of row_indices", - ) + @data.setter + def data(self, data): + self._init_data(data) + + @coo_indices.setter + def coo_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + _verify_coo_components( + indices=indices, values=self._values, dense_shape=self._dense_shape + ) + self._coo_indices = indices + + @crow_indices.setter + def crow_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._crow_indices = indices + + @col_indices.setter + def col_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._col_indices = indices + + @ccol_indices.setter + def ccol_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._ccol_indices = indices + @row_indices.setter + def row_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._row_indices = indices -def _verify_common_row_format_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" -): - ivy.utils.assertions.check_all_or_any_fn( - crow_indices, - col_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "crow_indices, col_indices, values and dense_shape must all be specified." - ), - ) + @values.setter + def values(self, values): + values = ivy.array(values) + _verify_coo_components( + indices=self._coo_indices, values=values, dense_shape=self._dense_shape + ) + self._values = values - ivy.utils.assertions.check_equal( - len(ivy.shape(crow_indices)), - 1, - message="crow_indices must be 1D.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(col_indices)), - 1, - message="col_indices must be 1D.", - as_array=False, - ) + @dense_shape.setter + def dense_shape(self, dense_shape): + dense_shape = ivy.Shape(dense_shape) + _verify_coo_components( + indices=self._coo_indices, values=self._values, dense_shape=dense_shape + ) + self._dense_shape = dense_shape - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", - as_array=False, - ) + @format.setter + def format(self, format): + self._format = format - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - crow_indices[-1], - message="size of col_indices does not match with last element of crow_indices", - ) + # Instance Methods # + # ---------------- # - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - ivy.shape(values)[0], - message="values and col_indices do not match", - as_array=False, - ) + def _coo_to_dense_coordinates(self): + all_coordinates = [] + for i in range(self._values.shape[0]): + coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) + coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) + all_coordinates.append(coordinate.to_list()) + return all_coordinates - # index in crow_indices must not exceed length of col_indices - ivy.utils.assertions.check_less( - crow_indices, - ivy.shape(col_indices)[0], - allow_equal=True, - message="index in crow_indices does not match the number of col_indices", - ) + def _csr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._col_indices.to_list() + all_cols = self._crow_indices.to_list() + for row in range(total_rows): + cols = all_rows[all_cols[row] : all_cols[row + 1]] + for col in cols: + all_coordinates.append([row, col]) + return all_coordinates + def _csc_to_dense_coordinates(self): + # CSC sparse array + all_coordinates = [] + total_rows = self._dense_shape[1] + all_cols = self._row_indices.to_list() + all_rows = self._ccol_indices.to_list() + for col in range(total_rows): + rows = all_cols[all_rows[col] : all_rows[col + 1]] + for row in rows: + all_coordinates.append([row, col]) + return all_coordinates -# helpers -def _verify_coo_components(indices=None, values=None, dense_shape=None): - ivy.utils.assertions.check_all_or_any_fn( - indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message="indices, values and dense_shape must all be specified", - ) - # coordinates style (COO), must be shaped (x, y) - ivy.utils.assertions.check_equal( - len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.to_ivy_shape(dense_shape)), - ivy.shape(indices)[0], - message="shape and indices shape do not match", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(values)[0], - ivy.shape(indices)[1], - message="values and indices do not match", - as_array=False, - ) - for i in range(ivy.shape(indices)[0]): - ivy.utils.assertions.check_less( - indices[i], - ivy.to_ivy_shape(dense_shape)[i], - message="indices is larger than shape", - ) + def _bsr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._crow_indices.to_list() + all_cols = self._col_indices.to_list() + nblockrows, nblockcols = self._values.shape[-2:] -def _verify_csc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="csc", - ) + for row in range(total_rows // nblockrows): + cols = all_cols[all_rows[row] : all_rows[row + 1]] + for col in cols: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0], - message="index in row_indices does not match shape", - ) + def _bsc_to_dense_coordinates(self): + all_coordinates = [] + total_cols = self._dense_shape[1] + all_rows = self._row_indices.to_list() + all_cols = self._ccol_indices.to_list() + nblockrows, nblockcols = self._values.shape[-2:] -def _verify_csr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="csr", - ) + for col in range(total_cols // nblockcols): + rows = all_rows[all_cols[col] : all_cols[col + 1]] + for row in rows: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False - ) - # number of intervals must be equal to x in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False - ) + def to_dense_array(self, *, native=False): + if self._format == "coo": + all_coordinates = self._coo_to_dense_coordinates() + elif self._format == "csr": + all_coordinates = self._csr_to_dense_coordinates() + elif self._format == "csc": + all_coordinates = self._csc_to_dense_coordinates() + elif self._format == "bsc": + all_coordinates = self._bsc_to_dense_coordinates() + else: + all_coordinates = self._bsr_to_dense_coordinates() - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1], - message="index in col_indices does not match shape", - ) + # make dense array + ret = ivy.scatter_nd( + ivy.array(all_coordinates), + ivy.flatten(self._values), + ivy.array(self._dense_shape), + ) + return ret.to_native() if native else ret -# --- Main --- # -# ------------ # +class NativeSparseArray: + pass def is_ivy_sparse_array(x): diff --git a/ivy/functional/ivy/experimental/statistical.py b/ivy/functional/ivy/experimental/statistical.py index e9392dcff411a..1349cc5752326 100644 --- a/ivy/functional/ivy/experimental/statistical.py +++ b/ivy/functional/ivy/experimental/statistical.py @@ -13,47 +13,141 @@ from ivy.utils.exceptions import handle_exceptions +# TODO: Make bins optional by offering an automatic bins creation like numpy. +# Make density argument work in tensorflow +# Bins as str is not defined (check Numpy implementation). +# Permit multiple axis. +# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def bincount( - x: ivy.Array, +def histogram( + a: Union[ivy.Array, ivy.NativeArray], /, *, - weights: Optional[ivy.Array] = None, - minlength: int = 0, + bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + density: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Count the number of occurrences of each value in an integer array. + Compute the histogram of the array ``a``. + + .. note:: + Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), + ..., I_{K-1} = [c_{K-1}, cK]. Parameters ---------- - self - Input array. + a + input array. + bins + if ``bins`` is an int, it defines the number of equal-width bins in the given + range. + if ``bins`` is an array, it defines a monotonically increasing array of bin + edges, including the rightmost edge, allowing for non-uniform bin widths. + axis + dimension along which maximum values must be computed. By default, the maximum + value must be computed over the entire array. Default: ``None``. + extend_lower_interval + if True, extend the lowest interval I0 to (-inf, c1]. + extend_upper_interval + ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). + dtype + the output type. + range + the lower and upper range of the bins. The first element of the range must be + less than or equal to the second. weights - An optional input array. - minlength - A minimum number of bins for the output array. + each value in ``a`` only contributes its associated weight towards the bin count + (instead of 1). Must be of the same shape as a. + density + if True, the result is the value of the probability density function at the + bin, normalized such that the integral over the range of bins is 1. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The bincount of the array elements. + a tuple containing the values of the histogram and the bin edges. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) - >>> a.bincount(a) - 3.0 - >>> a.bincount(a, axis=0) - array([6.5, 2. , 2.5]) + With :class:`ivy.Array` input: + + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) + >>> z = ivy.histogram(x, bins=y) + >>> print(z) + ivy.array([1., 0., 1., 1.]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [4.4, 5.5, .6]]) + >>> bins = 4 + >>> range = (0., 5.) + >>> dtype = ivy.int32 + >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) + >>> print(y) + ivy.array([2, 1, 1, 1]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> axis = 1 + >>> extend_lower_interval = True + >>> extend_upper_interval = True + >>> dtype = ivy.float32 + >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) + >>> z = ivy.histogram( + ... x, + ... bins=y, + ... axis=axis, + ... extend_lower_interval=extend_lower_interval, + ... extend_upper_interval=extend_upper_interval, + ... dtype=dtype, + ... weights=weights) + >>> print(z) + ivy.array([[0., 3.], + [1., 0.], + [1., 0.], + [1., 0.], + [0., 0.]]) + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> dtype = ivy.int32 + >>> z = ivy.histogram(x, bins=y, dtype=dtype) + >>> print(z) + { + a: ivy.array([1, 1, 1, 0, 0]), + b: ivy.array([0, 0, 0, 1, 2]) + } """ - return ivy.current_backend(x).bincount( - x, weights=weights, minlength=minlength, out=out + return ivy.current_backend(a).histogram( + a, + bins=bins, + axis=axis, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + range=range, + weights=weights, + density=density, + out=out, ) @@ -63,533 +157,276 @@ def bincount( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def corrcoef( - x: ivy.Array, +def median( + input: ivy.Array, /, *, - y: Optional[ivy.Array] = None, - rowvar: bool = True, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: - return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) + """ + Compute the median along the specified axis. + + Parameters + ---------- + input + Input array. + axis + Axis or axes along which the medians are computed. The default is to compute + the median along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. + out + optional output array, for writing the result to. + + Returns + ------- + ret + The median of the array elements. + + Functional Examples + ------------------- + >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> ivy.median(a) + 3.5 + >>> ivy.median(a, axis=0) + ivy.array([6.5, 4.5, 2.5]) + """ + return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -def cov( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray] = None, +@infer_dtype +@handle_device_shifting +def nanmean( + a: ivy.Array, /, *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[ivy.Array] = None, - aweights: Optional[ivy.Array] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the covariance of matrix x1, or variables x1 and x2. + Compute the mean of all non-NaN elements along the specified dimensions. Parameters ---------- - x1 - a 1D or 2D input array, with a numeric data type. - x2 - optional second 1D or 2D input array, with a numeric data type. - Must have the same shape as ``self``. - rowVar - optional variable where each row of input is interpreted as a variable - (default = True). If set to False, each column is instead interpreted - as a variable. - bias - optional variable for normalizing input (default = False) by (N - 1) where - N is the number of given observations. If set to True, then normalization - is instead by N. Can be overridden by keyword ``ddof``. - ddof - optional variable to override ``bias`` (default = None). ddof=1 will return - the unbiased estimate, even with fweights and aweights given. ddof=0 will - return the simple average. - fweights - optional 1D array of integer frequency weights; the number of times each - observation vector should be repeated. - aweights - optional 1D array of observation vector weights. These relative weights are - typically large for observations considered "important" and smaller for - observations considered less "important". If ddof=0 is specified, the array - of weights can be used to assign probabilities to observation vectors. + a + Input array. + axis + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of sub-classes + of ndarray. If the sub-classes methods does not implement keepdims any + exceptions will be raised. dtype - optional variable to set data-type of the result. By default, data-type - will have at least ``numpy.float64`` precision. + The desired data type of returned tensor. Default is None. out - optional output array, for writing the result to. It must have a shape that - the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array containing the covariance matrix of an input matrix, or the - covariance matrix of two variables. The returned array must have a - floating-point data type determined by Type Promotion Rules and must be - a square matrix of shape (N, N), where N is the number of variables in the - input(s). + The nanmean of the array elements. - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([[1, 2, 3], - ... [4, 5, 6]]) - >>> y = x[0].cov(x[1]) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) - >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) - >>> z = ivy.Container.static_cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) - >>> z = ivy.cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With :class:`ivy.Array` input and rowVar flag set to False (True by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], rowVar=False) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input and bias flag set to True (False by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], bias=True) - >>> print(y) - ivy.array([[0.66666667, 0.66666667], - [0.66666667, 0.66666667]]) - - With :class:`ivy.Array` input with both fweights and aweights given: - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> fw = ivy.array([1,2,3]) - >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) - >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) - >>> print(y) - ivy.array([[0.48447205, 0.48447205], - [0.48447205, 0.48447205]]) - """ - return ivy.current_backend(x1).cov( - x1, - x2, - rowVar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - dtype=dtype, - ) + Functional Examples + ------------------- + >>> a = ivy.array([[1, ivy.nan], [3, 4]]) + >>> ivy.nanmean(a) + 2.6666666666666665 + >>> ivy.nanmean(a, axis=0) + ivy.array([2., 4.]) + """ + return ivy.current_backend(a).nanmean( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummax( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def quantile( + a: ivy.Array, + q: Union[ivy.Array, float], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: bool = False, + interpolation: str = "linear", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a tuple containing the cumulative maximum of elements of input along the - given axis and index location of each maximum value found along the given axis. + Compute the q-th quantile of the data along the specified axis. Parameters ---------- - x + a Input array. + q + Quantile or sequence of quantiles to compute, which must be + between 0 and 1 inclusive. axis - Axis along which the cumulative maximum is computed. Default is ``0``. - exclusive - Whether to perform cummax exclusively. Default is ``False``. - reverse - Whether to perform the cummax from last to first element in the selected - axis. Default is ``False`` (from first to last element) + Axis or axes along which the quantiles are computed. The default + is to compute the quantile(s) along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original array a. + interpolation + {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. + Default value: 'linear'. + This specifies the interpolation method to use when the desired quantile lies + between two data points i < j: + - linear: i + (j - i) * fraction, where fraction is the fractional part of the + index surrounded by i and j. + - lower: i. + - higher: j. + - nearest: i or j, whichever is nearest. + - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with + integer dtypes. + - nearest_jax: provides jax-like computation for interpolation='nearest'. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cummax at each - original array elements along the specified axis. + A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis + is None, a rank(q) array. The first rank(q) dimensions index quantiles for + different values of q. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) - >>> y = ivy.cummax(x, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), - ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = ivy.array(0.5) + >>> ivy.quantile(a, q) + ivy.array(3.5) - >>> x = ivy.array([ 14, 15, 49, -24, -39]) - >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = 0.5 + >>> ivy.quantile(a, q) + ivy.array(3.5) - >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) - >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) - >>> print(x) - ivy.array([[0, 0, 0, 0], - [0, 1, 1, 1]]) + >>> ivy.quantile(a, q, axis=0) + ivy.array([6.5, 4.5, 2.5]) - >>> x = ivy.array([[-36, 83, -81], - ... [ 23, 29, 63], - ... [-83, 85, 2], - ... [ 31, 25, -86], - ... [-10, -52, 0], - ... [ 22, 38, 55], - ... [ 33, 54, -16]]) - >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) - >>> print(y) - (ivy.array([[ 0, 0, 83], - [ 0, 23, 29], - [ 0, 0, 85], - [ 0, 31, 31], - [ 0, 0, 0], - [ 0, 22, 38], - [ 0, 33, 54]]), ivy.array([[0, 0, 2], - [0, 1, 2], - [0, 0, 2], - [0, 1, 1], - [0, 0, 0], - [0, 1, 2], - [0, 1, 2]])) + >>> ivy.quantile(a, q, axis=1) + ivy.array([7., 2.]) - >>> x = ivy.array([73, 15, 47]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) - >>> print(y) - (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) + >>> ivy.quantile(a, q, axis=1, keepdims=True) + ivy.array([[7.],[2.]]) - >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) - >>> print(y) - (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) + >>> a = ivy.array([1., 2., 3., 4.]) + >>> q = ivy.array([0.3, 0.7]) + >>> ivy.quantile(a, q, interpolation='lower') + ivy.array([1., 3.]) """ - return ivy.current_backend(x).cummax( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + return ivy.current_backend(a).quantile( + a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummin( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def corrcoef( + x: ivy.Array, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + y: Optional[ivy.Array] = None, + rowvar: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Return the cumulative minimum of the elements along a given axis. - - Parameters - ---------- - x - Input array. - axis - Axis along which the cumulative minimum is computed. Default is ``0``. - reverse - Whether to perform the cummin from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Array which holds the result of applying cummin at each - original array elements along the specified axis. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cummin(x) - >>> print(y) - ivy.array([1, 1, 1, 0]) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cummin(x, axis=0, reverse=True, out=y) - >>> print(y) - ivy.array([[1., 3., 0.], - [1., 3., 0.]]) - - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[ 2, 4, 5], - [ 3, 5, 5], - [ 1, 3, 10]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cummin(x, axis= 0) - >>> print(y) - { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) - } - - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cummin(x,axis=1,reverse=True, out=y) - >>> print(y) - { - a: ivy.array([[1., 3., 4.]]), - b: ivy.array([[3., 5., 8.], - [5., 5., 5.]]), - c: ivy.array([[1., 1., 1.], - [3., 6., 9.], - [0., 2., 3.]]) - } - - >>> x = ivy.Container(a=ivy.array([[0],[5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cummin(x,axis=0,out=x) - >>> print(x) - { - a: ivy.array([[0], - [0]]), - b: ivy.array([[6, 8, 7], - [4, 2, 3]]), - c: ivy.array([[1, 2], - [1, 2], - [1, 2]]) - } - """ - return ivy.current_backend(x).cummin( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out - ) + return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) -# TODO: Make bins optional by offering an automatic bins creation like numpy. -# Make density argument work in tensorflow -# Bins as str is not defined (check Numpy implementation). -# Permit multiple axis. -# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def histogram( - a: Union[ivy.Array, ivy.NativeArray], +def nanmedian( + input: ivy.Array, /, *, - bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - density: Optional[bool] = False, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the histogram of the array ``a``. - - .. note:: - Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), - ..., I_{K-1} = [c_{K-1}, cK]. + ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the + function, and so the docstring for ivy.nanmedian also applies to this method with + minimal changes. Parameters ---------- - a - input array. - bins - if ``bins`` is an int, it defines the number of equal-width bins in the given - range. - if ``bins`` is an array, it defines a monotonically increasing array of bin - edges, including the rightmost edge, allowing for non-uniform bin widths. + self + Input array. axis - dimension along which maximum values must be computed. By default, the maximum - value must be computed over the entire array. Default: ``None``. - extend_lower_interval - if True, extend the lowest interval I0 to (-inf, c1]. - extend_upper_interval - ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). - dtype - the output type. - range - the lower and upper range of the bins. The first element of the range must be - less than or equal to the second. - weights - each value in ``a`` only contributes its associated weight towards the bin count - (instead of 1). Must be of the same shape as a. - density - if True, the result is the value of the probability density function at the - bin, normalized such that the integral over the range of bins is 1. + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of + sub-classes of ndarray. If the sub-classes methods does not implement + keepdims any exceptions will be raised. + overwrite_input + If True, then allow use of memory of input array a for calculations. + The input array will be modified by the call to median. This will + save memory when you do not need to preserve the contents of the input array. + Treat the input as undefined, but it will probably be fully or partially sorted. + Default is False. If overwrite_input is True and a is not already an ndarray, + an error will be raised. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - a tuple containing the values of the histogram and the bin edges. + A new array holding the result. If the input contains integers - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) - >>> z = ivy.histogram(x, bins=y) - >>> print(z) - ivy.array([1., 0., 1., 1.]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [4.4, 5.5, .6]]) - >>> bins = 4 - >>> range = (0., 5.) - >>> dtype = ivy.int32 - >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) - >>> print(y) - ivy.array([2, 1, 1, 1]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> axis = 1 - >>> extend_lower_interval = True - >>> extend_upper_interval = True - >>> dtype = ivy.float32 - >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) - >>> z = ivy.histogram( - ... x, - ... bins=y, - ... axis=axis, - ... extend_lower_interval=extend_lower_interval, - ... extend_upper_interval=extend_upper_interval, - ... dtype=dtype, - ... weights=weights) - >>> print(z) - ivy.array([[0., 3.], - [1., 0.], - [1., 0.], - [1., 0.], - [0., 0.]]) - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> dtype = ivy.int32 - >>> z = ivy.histogram(x, bins=y, dtype=dtype) - >>> print(z) - { - a: ivy.array([1, 1, 1, 0, 0]), - b: ivy.array([0, 0, 0, 1, 2]) - } - """ - return ivy.current_backend(a).histogram( - a, - bins=bins, - axis=axis, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - range=range, - weights=weights, - density=density, - out=out, + >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) + >>> ivy.nanmedian(x) + ivy.array(23.) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), + b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) + >>> ivy.nanmedian(x) + { + a: ivy.array(3.), + b: ivy.array(23.) + } + """ + return ivy.current_backend().nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out ) @@ -599,39 +436,42 @@ def histogram( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def igamma( - a: Union[ivy.Array, ivy.NativeArray], +def bincount( + x: ivy.Array, /, *, - x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + weights: Optional[ivy.Array] = None, + minlength: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the regularized lower gamma function of ``a`` and ``x``. + Count the number of occurrences of each value in an integer array. Parameters ---------- self Input array. - x - An additional input array. - `x` has the same type as `a`. - out - optional output array, for writing the result to. + weights + An optional input array. + minlength + A minimum number of bins for the output array. Returns ------- ret - The lower incomplete gamma function of the array elements. + The bincount of the array elements. Examples -------- - >>> a = ivy.array([2.5]) - >>> x = ivy.array([1.7, 1.2]) - >>> a.igamma(x) - ivy.array([0.3614, 0.2085]) + >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) + >>> a.bincount(a) + 3.0 + >>> a.bincount(a, axis=0) + array([6.5, 2. , 2.5]) """ - return ivy.current_backend().igamma(a, x=x, out=out) + return ivy.current_backend(x).bincount( + x, weights=weights, minlength=minlength, out=out + ) @handle_exceptions @@ -640,257 +480,417 @@ def igamma( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def median( - input: ivy.Array, +def igamma( + a: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[ivy.Array] = None, + x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Compute the median along the specified axis. + Compute the regularized lower gamma function of ``a`` and ``x``. Parameters ---------- - input + self Input array. - axis - Axis or axes along which the medians are computed. The default is to compute - the median along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. + x + An additional input array. + `x` has the same type as `a`. out optional output array, for writing the result to. Returns ------- ret - The median of the array elements. + The lower incomplete gamma function of the array elements. - Functional Examples - ------------------- - >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> ivy.median(a) - 3.5 - >>> ivy.median(a, axis=0) - ivy.array([6.5, 4.5, 2.5]) + Examples + -------- + >>> a = ivy.array([2.5]) + >>> x = ivy.array([1.7, 1.2]) + >>> a.igamma(x) + ivy.array([0.3614, 0.2085]) """ - return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) + return ivy.current_backend().igamma(a, x=x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_out_argument +@handle_array_like_without_promotion @to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nanmean( - a: ivy.Array, +def cov( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray] = None, /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[ivy.Array] = None, + aweights: Optional[ivy.Array] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the mean of all non-NaN elements along the specified dimensions. + Compute the covariance of matrix x1, or variables x1 and x2. Parameters ---------- - a - Input array. - axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of sub-classes - of ndarray. If the sub-classes methods does not implement keepdims any - exceptions will be raised. + x1 + a 1D or 2D input array, with a numeric data type. + x2 + optional second 1D or 2D input array, with a numeric data type. + Must have the same shape as ``self``. + rowVar + optional variable where each row of input is interpreted as a variable + (default = True). If set to False, each column is instead interpreted + as a variable. + bias + optional variable for normalizing input (default = False) by (N - 1) where + N is the number of given observations. If set to True, then normalization + is instead by N. Can be overridden by keyword ``ddof``. + ddof + optional variable to override ``bias`` (default = None). ddof=1 will return + the unbiased estimate, even with fweights and aweights given. ddof=0 will + return the simple average. + fweights + optional 1D array of integer frequency weights; the number of times each + observation vector should be repeated. + aweights + optional 1D array of observation vector weights. These relative weights are + typically large for observations considered "important" and smaller for + observations considered less "important". If ddof=0 is specified, the array + of weights can be used to assign probabilities to observation vectors. dtype - The desired data type of returned tensor. Default is None. + optional variable to set data-type of the result. By default, data-type + will have at least ``numpy.float64`` precision. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that + the inputs broadcast to. Returns ------- ret - The nanmean of the array elements. + an array containing the covariance matrix of an input matrix, or the + covariance matrix of two variables. The returned array must have a + floating-point data type determined by Type Promotion Rules and must be + a square matrix of shape (N, N), where N is the number of variables in the + input(s). - Functional Examples - ------------------- - >>> a = ivy.array([[1, ivy.nan], [3, 4]]) - >>> ivy.nanmean(a) - 2.6666666666666665 - >>> ivy.nanmean(a, axis=0) - ivy.array([2., 4.]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([[1, 2, 3], + ... [4, 5, 6]]) + >>> y = x[0].cov(x[1]) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Container` inputs: + >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) + >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) + >>> z = ivy.Container.static_cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) + >>> z = ivy.cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With :class:`ivy.Array` input and rowVar flag set to False (True by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], rowVar=False) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Array` input and bias flag set to True (False by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], bias=True) + >>> print(y) + ivy.array([[0.66666667, 0.66666667], + [0.66666667, 0.66666667]]) + + With :class:`ivy.Array` input with both fweights and aweights given: + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> fw = ivy.array([1,2,3]) + >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) + >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) + >>> print(y) + ivy.array([[0.48447205, 0.48447205], + [0.48447205, 0.48447205]]) """ - return ivy.current_backend(a).nanmean( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + return ivy.current_backend(x1).cov( + x1, + x2, + rowVar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, + dtype=dtype, ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def nanmedian( - input: ivy.Array, +@handle_array_function +def cummax( + x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the - function, and so the docstring for ivy.nanmedian also applies to this method with - minimal changes. + Return a tuple containing the cumulative maximum of elements of input along the + given axis and index location of each maximum value found along the given axis. Parameters ---------- - self + x Input array. axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of - sub-classes of ndarray. If the sub-classes methods does not implement - keepdims any exceptions will be raised. - overwrite_input - If True, then allow use of memory of input array a for calculations. - The input array will be modified by the call to median. This will - save memory when you do not need to preserve the contents of the input array. - Treat the input as undefined, but it will probably be fully or partially sorted. - Default is False. If overwrite_input is True and a is not already an ndarray, - an error will be raised. + Axis along which the cumulative maximum is computed. Default is ``0``. + exclusive + Whether to perform cummax exclusively. Default is ``False``. + reverse + Whether to perform the cummax from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A new array holding the result. If the input contains integers - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Array which holds the result of applying cummax at each + original array elements along the specified axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) - >>> ivy.nanmedian(x) - ivy.array(23.) - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), - b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) - >>> ivy.nanmedian(x) - { - a: ivy.array(3.), - b: ivy.array(23.) - } + >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) + >>> y = ivy.cummax(x, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), + ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) + + >>> x = ivy.array([ 14, 15, 49, -24, -39]) + >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) + + >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) + >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) + >>> print(x) + ivy.array([[0, 0, 0, 0], + [0, 1, 1, 1]]) + + >>> x = ivy.array([[-36, 83, -81], + ... [ 23, 29, 63], + ... [-83, 85, 2], + ... [ 31, 25, -86], + ... [-10, -52, 0], + ... [ 22, 38, 55], + ... [ 33, 54, -16]]) + >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) + >>> print(y) + (ivy.array([[ 0, 0, 83], + [ 0, 23, 29], + [ 0, 0, 85], + [ 0, 31, 31], + [ 0, 0, 0], + [ 0, 22, 38], + [ 0, 33, 54]]), ivy.array([[0, 0, 2], + [0, 1, 2], + [0, 0, 2], + [0, 1, 1], + [0, 0, 0], + [0, 1, 2], + [0, 1, 2]])) + + >>> x = ivy.array([73, 15, 47]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) + >>> print(y) + (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) + + >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) + >>> print(y) + (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) """ - return ivy.current_backend().nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + return ivy.current_backend(x).cummax( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def quantile( - a: ivy.Array, - q: Union[ivy.Array, float], +@handle_array_function +def cummin( + x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: bool = False, - interpolation: str = "linear", + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the q-th quantile of the data along the specified axis. + Return the cumulative minimum of the elements along a given axis. Parameters ---------- - a + x Input array. - q - Quantile or sequence of quantiles to compute, which must be - between 0 and 1 inclusive. axis - Axis or axes along which the quantiles are computed. The default - is to compute the quantile(s) along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original array a. - interpolation - {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. - Default value: 'linear'. - This specifies the interpolation method to use when the desired quantile lies - between two data points i < j: - - linear: i + (j - i) * fraction, where fraction is the fractional part of the - index surrounded by i and j. - - lower: i. - - higher: j. - - nearest: i or j, whichever is nearest. - - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with - integer dtypes. - - nearest_jax: provides jax-like computation for interpolation='nearest'. + Axis along which the cumulative minimum is computed. Default is ``0``. + reverse + Whether to perform the cummin from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis - is None, a rank(q) array. The first rank(q) dimensions index quantiles for - different values of q. + Array which holds the result of applying cummin at each + original array elements along the specified axis. Examples -------- - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = ivy.array(0.5) - >>> ivy.quantile(a, q) - ivy.array(3.5) + With :class:`ivy.Array` input: - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = 0.5 - >>> ivy.quantile(a, q) - ivy.array(3.5) + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cummin(x) + >>> print(y) + ivy.array([1, 1, 1, 0]) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cummin(x, axis=0, reverse=True, out=y) + >>> print(y) + ivy.array([[1., 3., 0.], + [1., 3., 0.]]) - >>> ivy.quantile(a, q, axis=0) - ivy.array([6.5, 4.5, 2.5]) + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[ 2, 4, 5], + [ 3, 5, 5], + [ 1, 3, 10]]) - >>> ivy.quantile(a, q, axis=1) - ivy.array([7., 2.]) + With :class:`ivy.Container` input: - >>> ivy.quantile(a, q, axis=1, keepdims=True) - ivy.array([[7.],[2.]]) + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cummin(x, axis= 0) + >>> print(y) + { + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) + } - >>> a = ivy.array([1., 2., 3., 4.]) - >>> q = ivy.array([0.3, 0.7]) - >>> ivy.quantile(a, q, interpolation='lower') - ivy.array([1., 3.]) + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cummin(x,axis=1,reverse=True, out=y) + >>> print(y) + { + a: ivy.array([[1., 3., 4.]]), + b: ivy.array([[3., 5., 8.], + [5., 5., 5.]]), + c: ivy.array([[1., 1., 1.], + [3., 6., 9.], + [0., 2., 3.]]) + } + + >>> x = ivy.Container(a=ivy.array([[0],[5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cummin(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [0]]), + b: ivy.array([[6, 8, 7], + [4, 2, 3]]), + c: ivy.array([[1, 2], + [1, 2], + [1, 2]]) + } """ - return ivy.current_backend(a).quantile( - a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out + return ivy.current_backend(x).cummin( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out ) diff --git a/ivy/functional/ivy/general.py b/ivy/functional/ivy/general.py index 4d0aa36bfef8a..1267e86d08269 100644 --- a/ivy/functional/ivy/general.py +++ b/ivy/functional/ivy/general.py @@ -46,18 +46,23 @@ FN_CACHE = dict() INF = float("inf") -array_mode_stack = list() -exception_trace_mode_stack = list() -inplace_mode_stack = list() -min_base_stack = list() -min_denominator_stack = list() -nestable_mode_stack = list() + precise_mode_stack = list() queue_timeout_stack = list() +array_mode_stack = list() shape_array_mode_stack = list() +nestable_mode_stack = list() +exception_trace_mode_stack = list() +inplace_mode_stack = list() +trace_mode_dict = dict() +trace_mode_dict["frontend"] = "ivy/functional/frontends" +trace_mode_dict["ivy"] = "ivy/" +trace_mode_dict["full"] = "" +trace_mode_dict["none"] = "" show_func_wrapper_trace_mode_stack = list() +min_denominator_stack = list() +min_base_stack = list() tmp_dir_stack = list() -trace_mode_dict = dict() # Extra # @@ -83,6 +88,78 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self +ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True + + +@handle_exceptions +def set_precise_mode(mode: bool) -> None: + """ + Set the mode of whether to use a promotion table that avoids any precision loss or a + compute effecient table that avoids most wider-than-necessary promotions. + + Parameter + --------- + mode + boolean whether to use high precision promtion table + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.set_precise_mode(True) + >>> ivy.precise_mode + True + """ + global precise_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + precise_mode_stack.append(mode) + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) + + +@handle_exceptions +def unset_precise_mode() -> None: + """ + Reset the mode of whether to use a promotion table that avoids any precision loss or + a compute effecient table that avoids most wider-than-necessary promotions. + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.unset_precise_mode() + >>> ivy.precise_mode + True + """ + global precise_mode_stack + if precise_mode_stack: + precise_mode_stack.pop(-1) + mode = precise_mode_stack[-1] if precise_mode_stack else True + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) + + +def _update_promotion_table(precise): + """Update the current datatype promotion table.""" + if precise: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.precise_extra_promotion_table, + } + + else: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.extra_promotion_table, + } + + class ArrayMode: """Array Mode Context Manager.""" @@ -102,371 +179,426 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self -# --- Helpers --- # -# --------------- # - +def get_referrers_recursive( + item, depth=0, max_depth=None, seen_set=None, local_set=None +): + """ + Summary. -def _all_dnd_combinations(): - all_comb = {} - for device in ivy.all_devices: - all_comb[device] = ivy.all_dtypes - return all_comb + Parameters + ---------- + item + depth + (Default value = 0) + max_depth + (Default value = None) + seen_set + (Default value = None) + local_set + (Default value = None`) + """ + seen_set = ivy.default(seen_set, set()) + local_set = ivy.default(local_set, set()) + ret_cont = ivy.Container( + repr=str(item).replace(" ", ""), + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + referrers = [ + ref + for ref in gc.get_referrers(item) + if not ( + isinstance(ref, dict) + and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) + ) + ] + local_set.add(str(id(referrers))) + for ref in referrers: + ref_id = str(id(ref)) + if ref_id in local_set or hasattr(ref, "cell_contents"): + continue + seen = ref_id in seen_set + seen_set.add(ref_id) + refs_rec = lambda: get_referrers_recursive( + ref, depth + 1, max_depth, seen_set, local_set + ) + this_repr = "tracked" if seen else str(ref).replace(" ", "") + if not seen and (not max_depth or depth < max_depth): + val = ivy.Container( + repr=this_repr, + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + refs = refs_rec() + for k, v in refs.items(): + val[k] = v + else: + val = this_repr + ret_cont[str(ref_id)] = val + return ret_cont -def _broadcast_to(input, target_shape): - input = ivy.squeeze(input) - if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): - return ivy.reshape(input, target_shape) - else: - input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input - new_dims = () - i_i = len(input.shape) - 1 - for i_t in range(len(target_shape) - 1, -1, -1): - if len(input.shape) + len(new_dims) >= len(target_shape): - break - if i_i < 0 or target_shape[i_t] != input.shape[i_i]: - new_dims += (i_t,) - else: - i_i -= 1 - input = ivy.expand_dims(input, axis=new_dims) - return ivy.broadcast_to(input, target_shape) +@handle_exceptions +@handle_backend_invalid +def is_native_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False +) -> bool: + """ + Determine whether the input x is an :class:`ivy.NativeArray` instance. -def _dnd_dict_difference(a, b): - res = a - for device in list(a): - if device in b: - difference = set.difference(set(a[device]), set(b[device])) - if difference: - res[device] = tuple(difference) - else: - del res[device] - return res + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. + Returns + ------- + ret + Boolean, whether or not x is an :class:`ivy.NativeArray`. -def _dnd_dict_intersection(a, b): - res = {} - for device in a: - if device in b: - intersection = set.intersection(set(a[device]), set(b[device])) - if intersection: - res[device] = tuple(intersection) - return res + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_native_array(x) + False + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_native_array(x, exclusive=True) + True + """ + try: + return current_backend(x).is_native_array(x, exclusive=exclusive) + except ValueError: + return False -def _dnd_dict_union(a, b): - res = {} - for device in set(list(a) + list(b)): - u1 = set(a.get(device, ())) - u2 = set(b.get(device, ())) - res[device] = tuple(set.union(u1, u2)) - return res +@handle_exceptions +@handle_backend_invalid +def is_ivy_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False +) -> bool: + """ + Determine whether the input x is a valid Ivy Array. + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. -def _get_devices_and_dtypes(fn, recurse=False, complement=True): - supported_devices = ivy.function_supported_devices(fn, recurse=recurse) - supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) + Returns + ------- + ret + Boolean, whether or not x is a valid Ivy Array. - if hasattr(fn, "partial_mixed_handler"): - supported_devices = supported_devices["primary"] - supported_dtypes = supported_dtypes["primary"] + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_ivy_array(x) + True - supported = {} - # Generate a base supported set from other attributes - for device in supported_devices: - supported[device] = supported_dtypes + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_ivy_array(x, exclusive=True) + False + """ + return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported - backend = ivy.current_backend_str() +@handle_exceptions +@handle_backend_invalid +def is_array(x: Any, /, *, exclusive: bool = False) -> bool: + """ + Determine whether the input x is either an Ivy Array or a Native Array. - # Their values are formatted like either - # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype.__get__() + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. - if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): - fn_supported_dnd = fn_supported_dnd.get(backend, supported) + Returns + ------- + ret + Boolean, whether or not x is an array. - ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) - # dict intersection - supported = _dnd_dict_intersection(supported, fn_supported_dnd) + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> print(ivy.is_array(x)) + True - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> print(ivy.is_array(x, exclusive=True)) + True - if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) + >>> x = [2, 3] + >>> print(ivy.is_array(x)) + False + """ + return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( + x, exclusive=exclusive + ) - ivy.utils.assertions.check_isinstance( - list(fn_unsupported_dnd.values())[0], tuple - ) - # dict difference - supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - if complement: - # dict difference - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported +@handle_exceptions +def is_ivy_container(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Container. + Parameters + ---------- + x + The input to check -def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: - fn_unsupported_dnd = {} - fn_supported_dnd = {} - backend = ivy.current_backend_str() - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): - fn_supported_dnd = fn_supported_dnd.get(backend, {}) + Returns + ------- + ret + Boolean, whether or not x is an ivy container. - ivy.utils.assertions.check_false( - fn_unsupported_dnd and fn_supported_dnd, - "unsupported_device_and_dtype and supported_device_and_dtype cannot" - " both be defined for the same function", - ) + Examples + -------- + >>> x = ivy.Container() + >>> print(ivy.is_ivy_container(x)) + True - us = "unsupported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") + >>> x = [2, 3] + >>> print(ivy.is_ivy_container(x)) + False + """ + return isinstance(x, ivy.Container) - ss = "supported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") - return True +ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True -def _numel(shape): - shape = tuple(shape) - return ivy.prod(shape).to_scalar() if shape != () else 1 +@handle_exceptions +def set_array_mode(mode: bool) -> None: + """ + Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs + back to ivy.Array. + It Stops the conversion of ivy.NativeArray to ivy.Array in the + case when it is set to False. -def _parse_query(query, x_shape): - query = (query,) if not isinstance(query, tuple) else query - query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions - # array containing all of x's flat indices - x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - # use numpy's __getitem__ to get the queried indices - x_idxs = ivy.array(x_.to_numpy()[query_]) - target_shape = x_idxs.shape + >>> ivy.set_array_mode(True) + >>> ivy.array_mode + True + """ + global array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + array_mode_stack.append(mode) + ivy.__setattr__("array_mode", mode, True) - if 0 in x_idxs.shape or 0 in x_shape: - return None, target_shape - # convert the flat indices to multi-D indices - x_idxs = ivy.unravel_index(x_idxs, x_shape) +@handle_exceptions +def unset_array_mode() -> None: + """ + Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back + to ivy.Array to the previous state. - # stack the multi-D indices to bring them to gather_nd/scatter_nd format - x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - return x_idxs, target_shape + >>> ivy.unset_shape_array_mode() + >>> ivy.array_mode + True + """ + global array_mode_stack + if array_mode_stack: + array_mode_stack.pop(-1) + mode = array_mode_stack[-1] if array_mode_stack else True + ivy.__setattr__("array_mode", mode, True) -def _update_promotion_table(precise): - """Update the current datatype promotion table.""" - if precise: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.precise_extra_promotion_table, - } +ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True - else: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.extra_promotion_table, - } +@handle_exceptions +def set_nestable_mode(mode: bool) -> None: + """ + Set the mode of whether to check if function inputs are ivy.Container. -def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): - attr_list = () - if hasattr(fn, other_attr_name): - attr_list = getattr(fn, other_attr_name) - if isinstance(attr_list, dict): - attr_list = attr_list.get(backend, ()) - ivy.utils.assertions.check_false( - dnd_dict and attr_list, - f"Cannot specify both {first_attr_name} and {other_attr_name} " - "cannot both be defined for the same function", - ) + Parameter + --------- + mode + boolean whether to check if function inputs are ivy.Container + Examples + -------- + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False -# --- Main --- # -# ------------ # + >>> ivy.set_nestable_mode(True) + >>> ivy.nestable_mode + True + """ + global nestable_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + nestable_mode_stack.append(mode) + ivy.__setattr__("nestable_mode", mode, True) @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def all_equal( - *xs: Iterable[Any], equality_matrix: bool = False -) -> Union[bool, ivy.Array, ivy.NativeArray]: +def unset_nestable_mode() -> None: """ - Determine whether the inputs are all equal. - - Parameters - ---------- - xs - inputs to compare. - equality_matrix - Whether to return a matrix of equalities comparing each input with every other. - Default is ``False``. - - Returns - ------- - ret - Boolean, whether or not the inputs are equal, or matrix array of booleans if - equality_matrix=True is set. + Reset the mode of whether to check if function inputs are ivy.Container to the + previous state. Examples -------- - With :class:`ivy.Array` inputs: + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False - >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> y = ivy.all_equal(x1, x2) - >>> print(y) + >>> ivy.unset_nestable_mode() + >>> ivy.nestable_mode True + """ + global nestable_mode_stack + if nestable_mode_stack: + nestable_mode_stack.pop(-1) + mode = nestable_mode_stack[-1] if nestable_mode_stack else True + ivy.__setattr__("nestable_mode", mode, True) - >>> x1 = ivy.array([0, 0]) - >>> x2 = ivy.array([0, 0]) - >>> x3 = ivy.array([1, 0]) - >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) - >>> print(y) - ivy.array([[ True, True, False], - [ True, True, False], - [False, False, True]]) - With one :class:`ivy.Container` inputs: +ivy.exception_trace_mode = ( + exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" +) - >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), - ... b=ivy.array([0, 0, -1, 1, 0])) - >>> x2 = ivy.array([0, 0, -1, 1, 0]) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) - { - a: True, - b: True - } - With multiple :class:`ivy.Container` inputs: +@handle_exceptions +def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: + """ + Set the mode of whether to show frontend-truncated exception stack traces, ivy- + truncated exception stack traces or full exception stack traces. - >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, 0, 1])) - >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, -1, -1])) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) - { - a: True, - b: False - } + Parameter + --------- + mode + str exeption trace mode, one of `ivy`, `full` or `frontend` + + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' + + >>> ivy.set_exception_trace_mode("full") + >>> ivy.exception_trace_mode + 'full' """ - equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b - if equality_matrix: - num_arrays = len(xs) - mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] - for i, xa in enumerate(xs): - for j_, xb in enumerate(xs[i:]): - j = j_ + i - res = equality_fn(xa, xb) - if ivy.is_native_array(res): - # noinspection PyTypeChecker - res = ivy.to_scalar(res) - # noinspection PyTypeChecker - mat[i][j] = res - # noinspection PyTypeChecker - mat[j][i] = res - return ivy.array(mat) - x0 = xs[0] - for x in xs[1:]: - if not equality_fn(x0, x): - return False - return True + global exception_trace_mode_stack + trace_modes = list(trace_mode_dict.keys()) + ivy.utils.assertions.check_elem_in_list( + mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) + ) + exception_trace_mode_stack.append(mode) + ivy.__setattr__("exception_trace_mode", mode, True) @handle_exceptions -def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): +def unset_exception_trace_mode() -> None: """ - Return the index and `inspect.Parameter` representation of the specified argument. - In the form of a dict with keys "idx" and "param". + Reset the trace mode to the previously set mode. - Parameters - ---------- - fn - The function to retrieve the argument information for - name - The name of the argument - idx - the index of the argument in the inputs + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' - Returns - ------- - ret - a `dict` containing the idx, and the `inspect.Parameter` for the argument, - which itself contains the parameter name, type, and other helpful information. + >>> ivy.unset_exception_trace_mode() + >>> ivy.exception_trace_mode + 'full' """ - ivy.utils.assertions.check_all_or_any_fn( - name, - idx, - fn=ivy.exists, - type="any", - limit=[1], - message="exactly one of the keyword arguments name or idx must be provided", - as_array=False, - ) - params = inspect.signature(fn).parameters - if ivy.exists(name): - return {"idx": list(params).index(name), "param": params[name]} - return {"idx": idx, "param": list(params.values())[idx]} + global exception_trace_mode_stack + if exception_trace_mode_stack: + exception_trace_mode_stack.pop(-1) + mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" + ivy.__setattr__("exception_trace_mode", mode, True) + + +ivy.show_func_wrapper_trace_mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True +) @handle_exceptions -def arg_names(receiver): +def set_show_func_wrapper_trace_mode(mode: bool) -> None: """ - Get the expected keyword arguments for a function or class constructor. + Set the mode of whether to show the full stack trace with function wrapping traces. - Parameters - ---------- - receiver - Function or class constructor + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions - Returns - ------- - ret - List containing the keyword arguments' names for a function or class constructor + Examples + -------- + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False + + >>> ivy.set_show_func_wrapper_trace_mode(True) + >>> ivy.show_func_wrapper_trace_mode + True + """ + global show_func_wrapper_trace_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + show_func_wrapper_trace_mode_stack.append(mode) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + + +@handle_exceptions +def unset_show_func_wrapper_trace_mode() -> None: + """ + Reset the mode of whether to show the full stack trace with function wrapping + traces. Examples -------- - >>> x = ivy.arg_names(ivy.tan) - >>> print(x) - ['x', 'out'] + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False - >>> x = ivy.arg_names(ivy.optimizers.Adam) - >>> print(x) - ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', - 'stop_gradients', 'compile_on_next_step', 'device'] + >>> ivy.unset_show_func_wrapper_trace_mode() + >>> ivy.show_func_wrapper_trace_mode + True """ - return list(inspect.signature(receiver).parameters.keys()) + global show_func_wrapper_trace_mode_stack + if show_func_wrapper_trace_mode_stack: + show_func_wrapper_trace_mode_stack.pop(-1) + mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True + ) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) @handle_exceptions @@ -521,254 +653,347 @@ def array_equal( @handle_exceptions @handle_nestable -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: +def all_equal( + *xs: Iterable[Any], equality_matrix: bool = False +) -> Union[bool, ivy.Array, ivy.NativeArray]: """ - Assert that inplace operations are supported for x. + Determine whether the inputs are all equal. Parameters ---------- - x - Input variable or array to check for inplace support for. + xs + inputs to compare. + equality_matrix + Whether to return a matrix of equalities comparing each input with every other. + Default is ``False``. Returns ------- ret - True if supports, raises IvyBackendException otherwise - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Boolean, whether or not the inputs are equal, or matrix array of booleans if + equality_matrix=True is set. Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: + With :class:`ivy.Array` inputs: - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) + >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> y = ivy.all_equal(x1, x2) + >>> print(y) True - With :class:`ivy.Array` input and default backend set as `torch`: - - >>> ivy.set_backend("torch") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) - True + >>> x1 = ivy.array([0, 0]) + >>> x2 = ivy.array([0, 0]) + >>> x3 = ivy.array([1, 0]) + >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) + >>> print(y) + ivy.array([[ True, True, False], + [ True, True, False], + [False, False, True]]) - With :class:`ivy.Container` input and default backend set as `numpy`: + With one :class:`ivy.Container` inputs: - >>> ivy.set_backend("numpy") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) + >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), + ... b=ivy.array([0, 0, -1, 1, 0])) + >>> x2 = ivy.array([0, 0, -1, 1, 0]) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) { a: True, b: True } - With :class:`ivy.Container` input and default backend set as `torch`: + With multiple :class:`ivy.Container` inputs: - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) + >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, 0, 1])) + >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, -1, -1])) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) { a: True, - b: True + b: False } """ - ivy.utils.assertions.check_true( - ivy.supports_inplace_updates(x), - "Inplace operations are not supported {} types with {} backend".format( - type(x), ivy.current_backend_str() - ), - ) + equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b + if equality_matrix: + num_arrays = len(xs) + mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] + for i, xa in enumerate(xs): + for j_, xb in enumerate(xs[i:]): + j = j_ + i + res = equality_fn(xa, xb) + if ivy.is_native_array(res): + # noinspection PyTypeChecker + res = ivy.to_scalar(res) + # noinspection PyTypeChecker + mat[i][j] = res + # noinspection PyTypeChecker + mat[j][i] = res + return ivy.array(mat) + x0 = xs[0] + for x in xs[1:]: + if not equality_fn(x0, x): + return False return True @handle_exceptions -def cache_fn(func: Callable) -> Callable: +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_native_arrays +@handle_array_function +@handle_device_shifting +def to_numpy( + x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True +) -> np.ndarray: """ - Cache function outputs. - - A decorator to wrap a function, such that computed outputs are cached to avoid - recalculating them later. + Convert an array into a numpy array. Parameters ---------- - func - The function to wrap, whose output should be cached for later. + x + input array + copy + whether to copy the array to a new address or not. + Default is ``True``. Returns ------- ret - The newly cache wrapped function. + a numpy array copying all the element of the array ``x``. Examples -------- - With positional arguments only: + With :class:`ivy.Array` inputs: - >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 - >>> cached_sum = ivy.cache_fn(my_sum) - >>> print(cached_sum(3, 5)) - 8 + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_numpy(x, copy=True) + >>> print(y) + [-1 0 1] - With keyword arguments: + >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) + >>> y = ivy.to_numpy(x, copy=True) + >>> print(y) + [[-1 0 1] + [-1 0 1] + [ 1 0 -1]] - >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc - >>> cached_line_eq = ivy.cache_fn(line_eq) - >>> print(cached_line_eq(3, itc=5, slp=2)) - 11 + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([-1, 0, 1], dtype=int32) + } + + >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), + ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([[-1., 0., 1.], + [-1., 0., 1.], + [1., 0., -1.]], dtype=float32), + b: array([[-1, 0, 0], + [1, 0, 1], + [1, 1, 1]], dtype=int32) + } """ - global FN_CACHE - if func not in FN_CACHE: - FN_CACHE[func] = dict() + return current_backend(x).to_numpy(x, copy=copy) - @wraps(func) - def cached_fn(*args, **kwargs): - key = "".join( - [str(i) + ", " for i in args] - + [" kw, "] - + [str(i) + ", " for i in sorted(kwargs.items())] - ) - cache = FN_CACHE[func] - if key in cache: - return cache[key] - ret = func(*args, **kwargs) - cache[key] = ret - return ret - return cached_fn +@handle_exceptions +@handle_nestable +def isscalar(x: Any, /) -> bool: + return np.isscalar(x) @handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@handle_array_like_without_promotion +@inputs_to_native_arrays @handle_array_function -def clip_matrix_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, - /, - *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +@handle_device_shifting +def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: """ - Clips (limits) the matrix norm of an array. + Convert an array with a single element into a scalar. Parameters ---------- x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array with a single element. Returns ------- ret - An array with the matrix norm downscaled to the max norm if needed. + a scalar copying the element of the array ``x``. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Functional Examples ------------------- With :class:`ivy.Array` input: - >>> x = ivy.array([[0., 1., 2.]]) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = ivy.array([3]) + >>> y = ivy.to_scalar(x) >>> print(y) - ivy.array([[0. , 0.894, 1.79 ]]) + 3 - >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) - >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) - >>> print(y) - ivy.array([[ 0.0353, -0.424 , 1.31 ], - [ 0. , 2.58 , -0.176 ]]) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> x = ivy.array([[[5., 4.], [-2., 6.]], - ... [[3., 7.], [0., -5.]]]) - >>> y = ivy.empty((2, 2, 2)) - >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) + >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) + >>> y = ivy.to_scalar(x) >>> print(y) - ivy.array([[[ 0.339, 0.271], - [-0.135, 0.406]], - [[ 0.168, 0.391], - [ 0. , -0.279]]]) - - >>> x = ivy.array([[0., 1.], - ... [2., 3.]]) - >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) - >>> print(x) - ivy.array([[0., 1.], - [2., 3.]]) - - With :class:`ivy.Container` input: + { + a: -1, + b: 3 + } - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), - ... b=ivy.array([[3., 4., 5.]])) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), + ... c=ivy.array([-1])) + >>> y = ivy.to_scalar(x) >>> print(y) { - a: ivy.array([[0., 0.894, 1.79]]), - b: ivy.array([[0.849, 1.13, 1.41]]) + a: 1, + b: 0, + c: -1 } """ - norms = ivy.matrix_norm(x, ord=p, keepdims=True) - ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) - return ivy.multiply(ratios, x, out=out) + return current_backend(x).to_scalar(x) @handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@handle_array_like_without_promotion +@inputs_to_native_arrays @handle_array_function -def clip_vector_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, - /, - *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +@handle_device_shifting +def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: """ - Clips (limits) the vector p-norm of an array. + Create a (possibly nested) list from input array. Parameters ---------- x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Input array. Returns ------- ret - An array with the vector norm downscaled to the max norm if needed. - - Functional Examples - ------------------ + A list representation of the input array ``x``. + Examples + -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.clip_vector_norm(x, 2.0) + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_list(x) >>> print(y) - ivy.array([0. , 0.894, 1.79 ]) + [-1, 0, 1] - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) + >>> x = ivy.array([[ 1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.to_list(x) >>> print(y) - ivy.array([ 0.417, -0.583, 2. ]) + [[1.100000023841858,2.200000047683716,3.299999952316284], + [-4.400000095367432,-5.5,-6.599999904632568]] + + >>> x = ivy.array([[[-1, 0, 1], + ... [ 1, 0, -1]], + ... [[ 1, -1, 0], + ... [ 1, 0, -1]]]) + >>> y = ivy.to_list(x) + >>> print(y) + [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] + + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_list(x) + >>> print(y) + { + a: [-1, 0, 1] + } + + >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], + ... [-1, 0, 1], + ... [1, 0, -1]])) + >>> y = ivy.to_list(x) + >>> print(y) + { + a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] + } + + >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], + ... [[1, -1, 0],[1, 0, -1]]])) + >>> y = ivy.to_list(x) + >>> print(y) + { + a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] + } + """ + return current_backend(x).to_list(x) + + +@handle_exceptions +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +def clip_vector_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, + /, + *, + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Clips (limits) the vector p-norm of an array. + + Parameters + ---------- + x + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + + Returns + ------- + ret + An array with the vector norm downscaled to the max norm if needed. + + Functional Examples + ------------------ + + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.clip_vector_norm(x, 2.0) + >>> print(y) + ivy.array([0. , 0.894, 1.79 ]) + + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) + >>> print(y) + ivy.array([ 0.417, -0.583, 2. ]) >>> x = ivy.array([[[0., 0.], [1., 3.], [2., 6.]], ... [[3., 9.], [4., 12.], [5., 15.]]]) @@ -812,139 +1037,85 @@ def clip_vector_norm( @handle_exceptions -def container_types(): - """ - Summary. - - Returns - ------- - ret - a key-value structure, and exposes public methods .keys(), .values() and - items(). - """ - # noinspection PyBroadException - try: - return current_backend().container_types() - except ValueError: - return [] - - -@handle_exceptions -def current_backend_str() -> Union[str, None]: - """ - Return framework string. - - Returns - ------- - ret - The framework string. - """ - fw = current_backend() - if not backend_stack: - return "" - return fw.current_backend_str() - - -@handle_exceptions -def default( - x: Any, +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +def clip_matrix_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, /, - default_val: Any, *, - catch_exceptions: bool = False, - rev: bool = False, - with_callable: bool = False, -) -> Any: + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Return x provided it exists (is not None), else returns default value. + Clips (limits) the matrix norm of an array. Parameters ---------- x - Input which may or may not exist (be None). - default_val - The default value. - catch_exceptions - Whether to catch exceptions from callable x. - Default is ``False``. - rev - Whether to reverse the input x and default_val. - Default is ``False``. - with_callable - Whether either of the arguments might be callable functions. - Default is ``False``. + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - x if x exists (is not None), else default. + An array with the matrix norm downscaled to the max norm if needed. Functional Examples - ------------------ - With :code:`Any` input: - - >>> x = None - >>> y = ivy.default(x, "default_string") - >>> print(y) - default_string - - >>> x = "" - >>> y = ivy.default(x, "default_string") - >>> print(y) + ------------------- + With :class:`ivy.Array` input: - >>> x = ivy.array([4, 5, 6]) - >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) + >>> x = ivy.array([[0., 1., 2.]]) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) - ivy.array([1, 2, 3]) + ivy.array([[0. , 0.894, 1.79 ]]) - >>> x = lambda: ivy.array([1, 2, 3]) - >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) + >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) + >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) >>> print(y) - ivy.array([1, 2, 3]) + ivy.array([[ 0.0353, -0.424 , 1.31 ], + [ 0. , 2.58 , -0.176 ]]) - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) + >>> x = ivy.array([[[5., 4.], [-2., 6.]], + ... [[3., 7.], [0., -5.]]]) + >>> y = ivy.empty((2, 2, 2)) + >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) >>> print(y) - ivy.array([1, 2, 3]) + ivy.array([[[ 0.339, 0.271], + [-0.135, 0.406]], + [[ 0.168, 0.391], + [ 0. , -0.279]]]) - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) + >>> x = ivy.array([[0., 1.], + ... [2., 3.]]) + >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) + >>> print(x) + ivy.array([[0., 1.], + [2., 3.]]) - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) + With :class:`ivy.Container` input: - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True, rev=True) + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), + ... b=ivy.array([[3., 4., 5.]])) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) - ivy.array([1, 2, 3]) + { + a: ivy.array([[0., 0.894, 1.79]]), + b: ivy.array([[0.849, 1.13, 1.41]]) + } """ - with_callable = catch_exceptions or with_callable - if rev: - tmp = x - x = default_val - default_val = tmp - if with_callable: - x_callable = callable(x) - default_callable = callable(default_val) - else: - x_callable = False - default_callable = False - if catch_exceptions: - # noinspection PyBroadException - try: - x = x() if x_callable else x - except Exception: - return default_val() if default_callable else default_val - else: - x = x() if x_callable else x - return x if exists(x) else default_val() if default_callable else default_val + norms = ivy.matrix_norm(x, ord=p, keepdims=True) + ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) + return ivy.multiply(ratios, x, out=out) @handle_exceptions @@ -952,267 +1123,234 @@ def default( @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def einops_rearrange( +def fourier_encode( x: Union[ivy.Array, ivy.NativeArray], - pattern: str, + max_freq: Union[float, ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: - """ - Perform einops rearrange operation on input array x. + num_bands: int = 4, + linear: bool = False, + concat: bool = True, + flatten: bool = False, +) -> Union[ivy.Array, ivy.NativeArray, Tuple]: + """ + Pad an array with fourier encodings. Parameters ---------- x - Input array to be re-arranged. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array to encode. + max_freq + The maximum frequency of the encoding. + num_bands + The number of frequency bands for the encoding. + Default is 4. + linear + Whether to space the frequency bands linearly as opposed to geometrically. + Default is ``False``. + concat + Whether to concatenate the position, sin and cos values, or return seperately. + Default is ``True``. + flatten + Whether to flatten the position dimension into the batch dimension. + Default is False. Returns ------- ret - New array with einops.rearrange having been applied. + New array with the final dimension expanded, and the encodings stored in this + channel. Examples -------- - With :class:`ivy.Array` instance method: - - >>> x = ivy.array([[1, 2, 3], - ... [-4, -5, -6]]) - >>> y = x.einops_rearrange("height width -> width height") - >>> print(y) - ivy.array([[ 1, -4], - [ 2, -5], - [ 3, -6]]) + >>> x = ivy.array([1,2,3]) + >>> y = 1.5 + >>> z = ivy.fourier_encode(x,y) + >>> print(z) + ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00]]) - >>> x = ivy.array([[[ 1, 2, 3], - ... [ 4, 5, 6]], - ... [[ 7, 8, 9], - ... [10, 11, 12]]]) - >>> y = x.einops_rearrange("c h w -> c (h w)") - >>> print(y) - ivy.array([[ 1, 2, 3, 4, 5, 6], - [ 7, 8, 9, 10, 11, 12]]) - >>> x = ivy.array([[1, 2, 3, 4, 5, 6], - ... [7, 8, 9, 10, 11, 12]]) - >>> y = ivy.zeros((4,3)) - >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) - >>> print(y) - ivy.array([[ 1, 2, 3], - [ 4, 5, 6], - [ 7, 8, 9], - [10, 11, 12]]) + >>> x = ivy.array([3,10]) + >>> y = 2.5 + >>> z = ivy.fourier_encode(x, y, num_bands=3) + >>> print(z) + ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, + -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], + [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, + 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) + """ + x_in = x + dim = x.shape[-1] + x = ivy.expand_dims(x, axis=-1) + orig_x = x + if linear: + scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) + else: + if ivy.backend == "torch" and isinstance(max_freq, float): + scales = ivy.logspace( + 0.0, + ivy.log(ivy.array(max_freq / 2)) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + else: + scales = ivy.logspace( + 0.0, + ivy.log(max_freq / 2) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + scales = ivy.astype(scales, ivy.dtype(x)) + scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] + x = x * scales * math.pi + sin_x = ivy.sin(x) + cos_x = ivy.cos(x) + if flatten: + orig_x = x_in + sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) + cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) + if concat: + return ivy.concat([orig_x, sin_x, cos_x], axis=-1) + return sin_x, cos_x - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> y = ivy.einops_rearrange(x, 'a b -> b a') - >>> print(y) - { - a: ivy.array([[-4.46999979, 3.66000009], - [0.93000001, 24.29000092], - [-3.33999991, 3.6400001]]), - b: ivy.array([[4.96000004, 4.36000013], - [1.51999998, 13.96000004], - [-10.67000008, 0.30000001]]) - } +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def value_is_nan( + x: Union[ivy.Array, ivy.NativeArray, Number], + /, + *, + include_infs: bool = True, +) -> bool: + """ + Determine whether the single valued array or scalar is of nan type. - With varying pattern: + Parameters + ---------- + x + The input to check Input array. + include_infs + Whether to include infs and -infs in the check. + Default is ``True``. - Suppose we have a set of 32 images in "h w c" format (height-width-channel) - and concatenate images along height (vertical axis), 960 = 32 * 30 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') - >>> print(x.shape) - (960, 40, 3) + Returns + ------- + ret + Boolean as to whether the input value is a nan or not. - # Concatenate images along horizontal axis, 1280 = 32 * 40 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') - >>> print(x.shape) - (30, 1280, 3) + Examples + -------- + >>> x = ivy.array([451]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False - # Reorder axes to "b c h w" format for deep learning - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') - >>> print(x.shape) - (32, 3, 30, 40) + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + True - # Flatten each image into a vector, 3600 = 30 * 40 * 3 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') - >>> print(x.shape) - (32, 3600) + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + False - # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), - # 128 = 32 * 2 * 2 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', - ... h1=2, w1=2) - >>> print(x.shape) - (128, 15, 20, 3) + >>> x = ivy.array([float('nan')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + True - # Space-to-depth operation - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, - ... w1=2) - >>> print(x.shape) - (32, 15, 20, 12) + >>> x = ivy.array([0]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False """ - ret = einops.rearrange(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x + if not x_scalar == x: + return True + if include_infs and (x_scalar == INF or x_scalar == -INF): + return True + return False @handle_exceptions @handle_nestable @handle_array_like_without_promotion -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -def einops_reduce( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - reduction: Union[str, Callable], - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: +def has_nans( + x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True +) -> bool: """ - Perform einops reduce operation on input array x. + Determine whether the array contains any nans, as well as infs or -infs if + specified. Parameters ---------- x - Input array to be reduced. - pattern - Reduction pattern. - reduction - One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array. + include_infs + Whether to include ``+infinity`` and ``-infinity`` in the check. + Default is ``True``. Returns ------- ret - New array with einops.reduce having been applied. + Boolean as to whether the array contains nans. - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]) - >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') - >>> print(reduced) - ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') - >>> print(reduced) - { - a: ivy.array([-2.29333329, 10.53000069]), - b: ivy.array([-1.39666676, 6.20666695]) - } - """ - ret = einops.reduce(x, pattern, reduction, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def einops_repeat( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: - """ - Perform einops repeat operation on input array x. - - Parameters - ---------- - x - Input array to be repeated. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - New array with einops.repeat having been applied. + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + False - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + >>> x = ivy.array([float('nan'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True - Examples - -------- - With :class:`ivy.Array` input: + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True - >>> x = ivy.array([1, 2, 3, 4]) - >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) - >>> print(repeated) - ivy.array([[1, 2, 3, 4], - [1, 2, 3, 4]]) + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x, include_infs=False) + >>> print(y) + False With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[4,5], - ... [1, 3]]), - ... b=ivy.array([[9, 10], - ... [4, 2]])) - >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) - >>> print(repeated) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.has_nans(x) + >>> print(y) { - a: ivy.array([[4, 5, 4, 5], - [1, 3, 1, 3]]), - b: ivy.array([[9, 10, 9, 10], - [4, 2, 4, 2]]) + a: False, + b: False } """ - ret = einops.repeat(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) @handle_exceptions @@ -1290,1254 +1428,986 @@ def exists(x: Any) -> bool: @handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def fourier_encode( - x: Union[ivy.Array, ivy.NativeArray], - max_freq: Union[float, ivy.Array, ivy.NativeArray], +def default( + x: Any, /, + default_val: Any, *, - num_bands: int = 4, - linear: bool = False, - concat: bool = True, - flatten: bool = False, -) -> Union[ivy.Array, ivy.NativeArray, Tuple]: + catch_exceptions: bool = False, + rev: bool = False, + with_callable: bool = False, +) -> Any: """ - Pad an array with fourier encodings. + Return x provided it exists (is not None), else returns default value. Parameters ---------- x - Input array to encode. - max_freq - The maximum frequency of the encoding. - num_bands - The number of frequency bands for the encoding. - Default is 4. - linear - Whether to space the frequency bands linearly as opposed to geometrically. + Input which may or may not exist (be None). + default_val + The default value. + catch_exceptions + Whether to catch exceptions from callable x. + Default is ``False``. + rev + Whether to reverse the input x and default_val. + Default is ``False``. + with_callable + Whether either of the arguments might be callable functions. Default is ``False``. - concat - Whether to concatenate the position, sin and cos values, or return seperately. - Default is ``True``. - flatten - Whether to flatten the position dimension into the batch dimension. - Default is False. Returns ------- ret - New array with the final dimension expanded, and the encodings stored in this - channel. + x if x exists (is not None), else default. - Examples - -------- - >>> x = ivy.array([1,2,3]) - >>> y = 1.5 - >>> z = ivy.fourier_encode(x,y) - >>> print(z) - ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00]]) + Functional Examples + ------------------ + With :code:`Any` input: + >>> x = None + >>> y = ivy.default(x, "default_string") + >>> print(y) + default_string - >>> x = ivy.array([3,10]) - >>> y = 2.5 - >>> z = ivy.fourier_encode(x, y, num_bands=3) - >>> print(z) - ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, - -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], - [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, - 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) - """ - x_in = x - dim = x.shape[-1] - x = ivy.expand_dims(x, axis=-1) - orig_x = x - if linear: - scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) - else: - if ivy.backend == "torch" and isinstance(max_freq, float): - scales = ivy.logspace( - 0.0, - ivy.log(ivy.array(max_freq / 2)) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - else: - scales = ivy.logspace( - 0.0, - ivy.log(max_freq / 2) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - scales = ivy.astype(scales, ivy.dtype(x)) - scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] - x = x * scales * math.pi - sin_x = ivy.sin(x) - cos_x = ivy.cos(x) - if flatten: - orig_x = x_in - sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) - cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) - if concat: - return ivy.concat([orig_x, sin_x, cos_x], axis=-1) - return sin_x, cos_x + >>> x = "" + >>> y = ivy.default(x, "default_string") + >>> print(y) -@handle_exceptions -@handle_nestable -def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: - """ - Return the supported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the supported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + >>> x = ivy.array([4, 5, 6]) + >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) + >>> print(y) + ivy.array([1, 2, 3]) - Parameters - ---------- - fn - The function to check for the supported device and dtype attribute - recurse - Whether to recurse into used ivy functions. - Default is ``True``. + >>> x = lambda: ivy.array([1, 2, 3]) + >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) + >>> print(y) + ivy.array([1, 2, 3]) - Returns - ------- - ret - Tuple or dict containing the supported devices and dtypes of the function - """ - ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", - ) + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) + >>> print(y) + ivy.array([1, 2, 3]) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=False), - } + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) + >>> print(y) + ivy.array([1, 2, 3]) + + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True) + >>> print(y) + ivy.array([1, 2, 3]) + + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True, rev=True) + >>> print(y) + ivy.array([1, 2, 3]) + """ + with_callable = catch_exceptions or with_callable + if rev: + tmp = x + x = default_val + default_val = tmp + if with_callable: + x_callable = callable(x) + default_callable = callable(default_val) else: - supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) - if recurse: - supported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - supported_devices_dtypes, - _dnd_dict_intersection, - function_supported_devices_and_dtypes, - wrapper=lambda x: x, - ) + x_callable = False + default_callable = False + if catch_exceptions: + # noinspection PyBroadException + try: + x = x() if x_callable else x + except Exception: + return default_val() if default_callable else default_val + else: + x = x() if x_callable else x + return x if exists(x) else default_val() if default_callable else default_val - return supported_devices_dtypes + +@handle_exceptions +def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: + """ + Return the input shape in ivy.Shape form. + + Parameters + ---------- + shape + The input to be converted + + Returns + ------- + ret + the input in ivy.Shape form + """ + if isinstance(shape, ivy.Shape): + return shape + return ivy.Shape(shape) @handle_exceptions -@handle_nestable -def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: +def to_native_shape( + shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] +) -> ivy.NativeShape: """ - Return the unsupported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the unsupported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + Return the input shape in its native backend framework form. Parameters ---------- - fn - The function to check for the unsupported device and dtype attribute - recurse - Whether to recurse into used ivy functions. - Default is ``True``. + shape + The input to be converted Returns ------- - ret - Tuple or dict containing the unsupported devices and dtypes of the function + ret + the input in its native framework form """ + native_shape_type = (ivy.NativeShape,) + if ivy.current_backend_str() == "torch": + native_shape_type += (tuple,) + if len(backend_stack) != 0 and isinstance(shape, native_shape_type): + return shape + ivy.utils.assertions.check_isinstance( + shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) + ) + if isinstance(shape, int): + shape = (shape,) + elif isinstance(shape, list): + shape = tuple(shape) + elif is_array(shape): + shape = ivy.to_numpy(shape).tolist() + elif isinstance(shape, ivy.Shape): + shape = shape.shape + ivy.utils.assertions.check_all( + [isinstance(v, int) for v in shape if not is_array(v)], + "shape must take integers only", + as_array=False, + ) ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", + not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=True), - } - else: - unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) - if recurse: - unsupported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - unsupported_devices_dtypes, - _dnd_dict_union, - function_unsupported_devices_and_dtypes, - wrapper=lambda x: x, - ) - return unsupported_devices_dtypes + return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def gather( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: """ - Gather slices from params at axis according to indices. + Try and return the function, otherwise return None if an exception was raised during + function execution. Parameters ---------- - params - The array from which to gather values. - indices - The array which indicates the indices that will be gathered along - the specified axis. - axis - optional int, the axis from which to gather from. - Default is ``-1``. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional array, for writing the result to. It must have a shape - that the inputs broadcast to. + fn + Function to try and call and return. + args + list of arguments. + kwargs + dictionay of keyword arguments Returns ------- - ret - New array with the values gathered at the specified indices along the - specified axis. - - - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Either the function itself or None if an exception was raised + during function execution. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.array([1, 2]) - >>> print(ivy.gather(x, y)) - ivy.array([1., 2.]) + with a function that is executed without any exception: - >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) - >>> y = ivy.array([[0, 1],[1, 2]]) - >>> z = ivy.zeros((2, 2, 2)) - >>> ivy.gather(x, y, out=z) - >>> print(z) - ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> z = ivy.try_else_none(ivy.add, x, y) + >>> print(z.__name__) + add - >>> x = ivy.array([[[0., 1.], [2., 3.]], - ... [[8., 9.], [10., 11.]]]) - >>> y = ivy.array([[0, 1]]) - >>> z = ivy.zeros((1, 2, 2, 2)) - >>> ivy.gather(x, y, axis=0, out=z) - >>> print(z) - ivy.array( - [[[[ 0., 1.], - [ 2., 3.]], - [[ 8., 9.], - [10., 11.]]]]) + with a function that is executed with an exception: - >>> x = ivy.array([[0, 10, 20, 0, 0], - ... [0, 0, 0, 30, 40], - ... [0, 10, 0, 0, 40]]) - >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) - >>> z = ivy.gather(x, y, batch_dims=1) + >>> x = ivy.array([1, 2, 3]) + >>> y = 'hemant' + >>> z = ivy.try_else_none(ivy.add,x, y) >>> print(z) - ivy.array([[10, 20], [30, 40],[10, 40]]) + None + """ + try: + _ = fn(*args, **kwargs) + return fn + except Exception: + return None - With :class:`ivy.Container` input: - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.Container(a = ivy.array([0, 1]), - ... b = ivy.array([1, 2])) - >>> print(ivy.gather(x, y)) - { - a: ivy.array([0., 1.]), - b: ivy.array([5., 6.]) - } +@handle_exceptions +def arg_names(receiver): + """ + Get the expected keyword arguments for a function or class constructor. - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + Parameters + ---------- + receiver + Function or class constructor - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.array([0, 1]) - >>> print(ivy.gather(x, y)) - { - a: ivy.array([0., 1.]), - b: ivy.array([4., 5.]) - } + Returns + ------- + ret + List containing the keyword arguments' names for a function or class constructor + + Examples + -------- + >>> x = ivy.arg_names(ivy.tan) + >>> print(x) + ['x', 'out'] + + >>> x = ivy.arg_names(ivy.optimizers.Adam) + >>> print(x) + ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', + 'stop_gradients', 'compile_on_next_step', 'device'] """ - return current_backend(params, indices).gather( - params, indices, axis=axis, batch_dims=batch_dims, out=out - ) + return list(inspect.signature(receiver).parameters.keys()) @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def gather_nd( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - batch_dims: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def match_kwargs( + kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False +) -> Union[List[Dict], Dict]: """ - Gather slices from params into a array with shape specified by indices. + Match keyword arguments to either class or function receivers. Parameters ---------- - params - The array from which to gather values. - indices - Index array. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + kwargs + Keyword arguments to match. + receivers + Functions and/or classes to match the keyword arguments to. + allow_duplicates + Whether to allow one keyword argument to be used for multiple receivers. + Default is ``False``. Returns ------- ret - New array of given shape, with the values gathered at the indices. + Sequence of keyword arguments split as best as possible. Examples -------- - With :class:`ivy.Array` input: + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) + >>> print(x) + [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) - ivy.array(1.) + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) + >>> print(x) + [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] + """ + split_kwargs = list() + for receiver in receivers: + expected_kwargs = arg_names(receiver) + found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} + if not allow_duplicates: + for k in found_kwargs.keys(): + del kwargs[k] + split_kwargs.append(found_kwargs) + if len(split_kwargs) == 1: + return split_kwargs[0] + return split_kwargs - >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) - >>> y = ivy.array([[0],[1],[1]], dtype='int32') - >>> z = ivy.gather_nd(x,y,batch_dims=1) - ivy.array([0., 3., 5.]) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: +@handle_exceptions +def cache_fn(func: Callable) -> Callable: + """ + Cache function outputs. - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) - { - a: ivy.array(1.), - b: ivy.array(5.) - } + A decorator to wrap a function, such that computed outputs are cached to avoid + recalculating them later. - With :class:`ivy.Container` input: + Parameters + ---------- + func + The function to wrap, whose output should be cached for later. - >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), - ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) - >>> y = ivy.Container(a=ivy.array([1,0]), - ... b=ivy.array([0])) - >>> print(ivy.gather_nd(x, y)) - { - a: ivy.array(30.), - b: ivy.array([0., 100., 200.]) - } + Returns + ------- + ret + The newly cache wrapped function. + + Examples + -------- + With positional arguments only: + + >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 + >>> cached_sum = ivy.cache_fn(my_sum) + >>> print(cached_sum(3, 5)) + 8 + + With keyword arguments: + + >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc + >>> cached_line_eq = ivy.cache_fn(line_eq) + >>> print(cached_line_eq(3, itc=5, slp=2)) + 11 """ - res = current_backend(params, indices).gather_nd( - params, indices, batch_dims=batch_dims - ) - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + global FN_CACHE + if func not in FN_CACHE: + FN_CACHE[func] = dict() + + @wraps(func) + def cached_fn(*args, **kwargs): + key = "".join( + [str(i) + ", " for i in args] + + [" kw, "] + + [str(i) + ", " for i in sorted(kwargs.items())] + ) + cache = FN_CACHE[func] + if key in cache: + return cache[key] + ret = func(*args, **kwargs) + cache[key] = ret + return ret + + return cached_fn @handle_exceptions -def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: +def current_backend_str() -> Union[str, None]: """ - Get all arrays which are currently alive. + Return framework string. Returns ------- ret - All arrays which are alive. - - Examples - -------- - >>> ivy.get_all_arrays_in_memory() - [] - >>> x = ivy.get_all_arrays_in_memory() - >>> x - [] - >>> y = ivy.array([0, 1, 2]) - >>> x - [ivy.array([0, 1, 2])] + The framework string. """ - all_arrays = list() - for obj in gc.get_objects(): - try: - if ivy.current_backend_str() in ["", "numpy"]: - if ivy.is_ivy_array(obj): - all_arrays.append(obj) - else: - if ivy.is_native_array(obj): - all_arrays.append(obj) - - except Exception: - pass - return all_arrays + fw = current_backend() + if not backend_stack: + return "" + return fw.current_backend_str() +@handle_exceptions @handle_nestable -@handle_partial_mixed_function -@handle_view_indexing +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def get_item( +def einops_rearrange( x: Union[ivy.Array, ivy.NativeArray], + pattern: str, /, - query: Union[ivy.Array, ivy.NativeArray, Tuple], *, - copy: Optional[bool] = None, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], ) -> ivy.Array: """ - Gather slices from x according to query array, identical to x[query]. + Perform einops rearrange operation on input array x. Parameters ---------- x - array, the array from which to gather values. - query - array, index array, integer indices or boolean mask. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + Input array to be re-arranged. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - New array with the values gathered at the specified indices. + New array with einops.rearrange having been applied. - Functional Examples - ------------------- + Examples + -------- + With :class:`ivy.Array` instance method: - >>> x = ivy.array([0, -1, 20]) - >>> query = ivy.array([0, 1]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 0, -1]) + >>> x = ivy.array([[1, 2, 3], + ... [-4, -5, -6]]) + >>> y = x.einops_rearrange("height width -> width height") + >>> print(y) + ivy.array([[ 1, -4], + [ 2, -5], + [ 3, -6]]) - >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) - >>> query = ivy.array([[True, False], [False, False], [True, True]]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 4, -2, -10]) + >>> x = ivy.array([[[ 1, 2, 3], + ... [ 4, 5, 6]], + ... [[ 7, 8, 9], + ... [10, 11, 12]]]) + >>> y = x.einops_rearrange("c h w -> c (h w)") + >>> print(y) + ivy.array([[ 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12]]) + + >>> x = ivy.array([[1, 2, 3, 4, 5, 6], + ... [7, 8, 9, 10, 11, 12]]) + >>> y = ivy.zeros((4,3)) + >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) + >>> print(y) + ivy.array([[ 1, 2, 3], + [ 4, 5, 6], + [ 7, 8, 9], + [10, 11, 12]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> y = ivy.einops_rearrange(x, 'a b -> b a') + >>> print(y) + { + a: ivy.array([[-4.46999979, 3.66000009], + [0.93000001, 24.29000092], + [-3.33999991, 3.6400001]]), + b: ivy.array([[4.96000004, 4.36000013], + [1.51999998, 13.96000004], + [-10.67000008, 0.30000001]]) + } + + With varying pattern: + + Suppose we have a set of 32 images in "h w c" format (height-width-channel) + and concatenate images along height (vertical axis), 960 = 32 * 30 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') + >>> print(x.shape) + (960, 40, 3) + + # Concatenate images along horizontal axis, 1280 = 32 * 40 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') + >>> print(x.shape) + (30, 1280, 3) + + # Reorder axes to "b c h w" format for deep learning + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') + >>> print(x.shape) + (32, 3, 30, 40) + + # Flatten each image into a vector, 3600 = 30 * 40 * 3 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') + >>> print(x.shape) + (32, 3600) + + # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), + # 128 = 32 * 2 * 2 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', + ... h1=2, w1=2) + >>> print(x.shape) + (128, 15, 20, 3) + + # Space-to-depth operation + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, + ... w1=2) + >>> print(x.shape) + (32, 15, 20, 12) """ - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return ivy.array([], shape=(0,), dtype=x.dtype) - return ivy.expand_dims(x, axis=0) - query = ivy.nonzero(query, as_tuple=False) - ret = ivy.gather_nd(x, query) - else: - indices, target_shape = _parse_query(query, x.shape) - if indices is None: - return ivy.empty(target_shape, dtype=x.dtype) - ret = ivy.gather_nd(x, indices) - ret = ivy.reshape(ret, target_shape) + ret = einops.rearrange(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) return ret -@handle_backend_invalid +@handle_exceptions @handle_nestable @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function -@handle_device_shifting -def get_num_dims( - x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False -) -> int: +def einops_reduce( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + reduction: Union[str, Callable], + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Return the number of dimensions of the array x. + Perform einops reduce operation on input array x. Parameters ---------- x - Input array to infer the number of dimensions for. - as_array - Whether to return the shape as a array, default False. + Input array to be reduced. + pattern + Reduction pattern. + reduction + One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Shape of the array + New array with einops.reduce having been applied. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) - >>> b = ivy.get_num_dims(a, as_array=False) - >>> print(b) - 3 + >>> x = ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]) + >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') + >>> print(reduced) + ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) With :class:`ivy.Container` input: - >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) - >>> print(ivy.get_num_dims(a)) - { - b: 2 - } - - >>> b = ivy.get_num_dims(a, as_array=True) - >>> print(b) + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') + >>> print(reduced) { - b: ivy.array(2) + a: ivy.array([-2.29333329, 10.53000069]), + b: ivy.array([-1.39666676, 6.20666695]) } """ - return current_backend(x).get_num_dims(x, as_array=as_array) - - -def get_referrers_recursive( - item, depth=0, max_depth=None, seen_set=None, local_set=None -): - """ - Summary. + ret = einops.reduce(x, pattern, reduction, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret - Parameters - ---------- - item - depth - (Default value = 0) - max_depth - (Default value = None) - seen_set - (Default value = None) - local_set - (Default value = None`) - """ - seen_set = ivy.default(seen_set, set()) - local_set = ivy.default(local_set, set()) - ret_cont = ivy.Container( - repr=str(item).replace(" ", ""), - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - referrers = [ - ref - for ref in gc.get_referrers(item) - if not ( - isinstance(ref, dict) - and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) - ) - ] - local_set.add(str(id(referrers))) - for ref in referrers: - ref_id = str(id(ref)) - if ref_id in local_set or hasattr(ref, "cell_contents"): - continue - seen = ref_id in seen_set - seen_set.add(ref_id) - refs_rec = lambda: get_referrers_recursive( - ref, depth + 1, max_depth, seen_set, local_set - ) - this_repr = "tracked" if seen else str(ref).replace(" ", "") - if not seen and (not max_depth or depth < max_depth): - val = ivy.Container( - repr=this_repr, - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - refs = refs_rec() - for k, v in refs.items(): - val[k] = v - else: - val = this_repr - ret_cont[str(ref_id)] = val - return ret_cont +# IMPORTANT: assign attribute directly to function instead of wrapper here +einops_reduce.unsupported_dtypes = { + "torch": ("float16",), + "tensorflow": ("complex",), + "paddle": ("complex", "uint8", "int8", "int16", "float16"), +} @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def has_nans( - x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True -) -> bool: +def einops_repeat( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Determine whether the array contains any nans, as well as infs or -infs if - specified. + Perform einops repeat operation on input array x. Parameters ---------- x - Input array. - include_infs - Whether to include ``+infinity`` and ``-infinity`` in the check. - Default is ``True``. - - Returns + Input array to be repeated. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns ------- ret - Boolean as to whether the array contains nans. - + New array with einops.repeat having been applied. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - False - - >>> x = ivy.array([float('nan'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True - - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True - - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x, include_infs=False) - >>> print(y) - False + >>> x = ivy.array([1, 2, 3, 4]) + >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) + >>> print(repeated) + ivy.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.has_nans(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[4,5], + ... [1, 3]]), + ... b=ivy.array([[9, 10], + ... [4, 2]])) + >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) + >>> print(repeated) { - a: False, - b: False + a: ivy.array([[4, 5, 4, 5], + [1, 3, 1, 3]]), + b: ivy.array([[9, 10, 9, 10], + [4, 2, 4, 2]]) } """ - return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) - + ret = einops.repeat(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret -@handle_exceptions -def inplace_arrays_supported() -> bool: - """ - Determine whether inplace arrays are supported for the current backend framework. - Returns - ------- - ret - Boolean, whether or not inplace arrays are supported. - """ - return current_backend().inplace_arrays_supported() +ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 @handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def inplace_decrement( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: +def set_min_denominator(val: float) -> None: """ - Perform in-place decrement for the input array. + Set the global minimum denominator used by ivy for numerically stable division. Parameters ---------- - x - The input array to be decremented by the defined value. val - The value of decrement. - - Returns - ------- - ret - The array following the in-place decrement. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The value to set the global minimum denominator to. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.min_denominator + >>> print(x) + 1e-12 - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_decrement(x, 1.25) + >>> ivy.set_min_denominator(1e-13) + >>> y = ivy.min_denominator >>> print(y) - ivy.array([[ 4.05, 5.75, -1.25], - [ 5.55, 6.75, 2.65], - [-1.25, 8.75, 5.05]]) + 1e-13 + """ + global min_denominator_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_denominator_stack.append(val) + ivy.__setattr__("min_denominator", val, True) - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) - >>> y = ivy.inplace_decrement(x, 1.5) - >>> print(y) - { - a: ivy.array([-1., -6.5, 28.5]), - b: ivy.array([-1.5, -26.5, 48.5]) - } +@handle_exceptions +def unset_min_denominator() -> None: + """ + Reset the global minimum denominator used by ivy for numerically stable division to + the previous value. - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_decrement(x, y) - >>> print(z) - { - a: ivy.array([0., 0., 0.]), - b: ivy.array([0., 0., 0.]) - } + Examples + -------- + >>> ivy.set_min_denominator(1e-10) + >>> y = ivy.min_denominator + >>> print(y) + 1e-10 - >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) - >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) - >>> z = ivy.inplace_decrement(x, y) - >>> print(z) - { - a: ivy.array([1., 1.5, 3.]), - b: ivy.array([0., 50., 3.5]) - } + >>> ivy.unset_min_denominator() + >>> ivy.min_denominator + 1e-12 """ - return current_backend(x).inplace_decrement(x, val) + global min_denominator_stack + if min_denominator_stack: + min_denominator_stack.pop(-1) + val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 + ivy.__setattr__("min_denominator", val, True) + + +ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 @handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def inplace_increment( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: +def set_min_base(val: float) -> None: """ - Perform in-place increment for the input array. + Set the global minimum base used by ivy for numerically stable power raising. Parameters ---------- - x - The input array to be incremented by the defined value. val - The value of increment. - - Returns - ------- - ret - The array following the in-place increment. + The new value to set the minimum base to. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.min_base + >>> print(x) + 1e-05 - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_increment(x, 3.) + >>> ivy.set_min_base(1e-04) + >>> y = ivy.min_base >>> print(y) - ivy.array([[ 8.3, 10., 3.], - [ 9.8, 11., 6.9], - [ 3., 13., 9.3]]) + 1e-04 + """ + global min_base_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_base_stack.append(val) + ivy.__setattr__("min_base", val, True) - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.inplace_increment(x, 2.5) - >>> print(y) - { - a: ivy.array([2.5, 17.5, 32.5]), - b: ivy.array([2.5, 27.5, 52.5]) - } +@handle_exceptions +def unset_min_base() -> None: + """ + Reset the global minimum base used by ivy for numerically stable power raising to + the previous value. + Examples + -------- + >>> ivy.set_min_base(1e-07) + >>> y = ivy.min_base + >>> print(y) + 1e-07 - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_increment(x, y) - >>> print(z) - { - a: ivy.array([0., 30., 60.]), - b: ivy.array([0., 50., 100.]) - } + >>> ivy.unset_min_base() + >>> ivy.min_base + 1e-05 """ - return current_backend(x).inplace_increment(x, val) + global min_base_stack + if min_base_stack: + min_base_stack.pop(-1) + val = min_base_stack[-1] if min_base_stack else 1e-05 + ivy.__setattr__("min_base", val, True) @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -# @handle_device_shifting -def inplace_update( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], +def stable_divide( + numerator: Union[Number, ivy.Array, ivy.NativeArray], + denominator: Union[Number, ivy.Array, ivy.NativeArray], /, *, - ensure_in_backend: bool = False, - keep_input_dtype: bool = False, -) -> ivy.Array: + min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, +) -> Union[Number, ivy.Array]: """ - Perform in-place update for the input array. - - This will always be performed on ivy.Array instances pass in the input, and will - also be performed on the native array classes in the backend when the backend - supports this. If the backend does not natively support inplace updates, and x is an - ivy.NativeArray instance, then an - exception will be thrown. + Divide the numerator by the denominator, with min denominator added to the + denominator for numerical stability. Parameters ---------- - x - The variable to update. - val - The array to update the variable with. - ensure_in_backend - Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. - In cases where it should be, backends which do not natively support inplace - updates will raise an exception. - keep_input_dtype - Whether or not to preserve `x` data type after the update, otherwise `val` - data type will be applied. Defaults to False. + numerator + The numerator of the division. + denominator + The denominator of the division. + min_denominator + The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) + by default. Returns ------- ret - The array following the in-place update. - - Raises - ------ - IvyException - If backend set doesn't natively support inplace updates and ensure_in_backend is - True, above exception will be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the arguments. + The new item following the numerically stable division. Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: + With :code:`int` input: - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0]) - >>> ivy.inplace_update(x, y) + >>> x = ivy.stable_divide(1, 2) >>> print(x) - ivy.array([0]) - - With :class:`ivy.Array` input and default backend set as `numpy`: + 0.49999999999975 - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) - >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) - >>> ivy.inplace_update(x, y, keep_input_dtype=True) + >>> x = ivy.stable_divide(1, 4, min_denominator=1) >>> print(x) - ivy.array([0., 0., 0.]) + 0.2 - With :class:`ivy.Container` instances:, and backend set as `torch`: + With float input: - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) - >>> ivy.inplace_update(x, y) + >>> x = ivy.stable_divide(5.0, 3.33) >>> print(x) - { - a: ivy.array([1, 1]), - b: ivy.array([2, 2]) - } + 1.5015015015010504 - With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend - set as `torch`: + With :code:`complex` input: - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.array([1, 2]) - >>> ivy.inplace_update(x, y) + >>> x = ivy.stable_divide(1+1j, 1-1j) >>> print(x) - { - a: ivy.array([1, 2]), - b: ivy.array([1, 2]) - } - """ - return current_backend(x).inplace_update( - x, - val, - ensure_in_backend=ensure_in_backend, - keep_input_dtype=keep_input_dtype, - ) - + (5.000444502911705e-13+0.9999999999995j) -@handle_exceptions -def inplace_variables_supported() -> bool: - """ - Determine whether inplace variables are supported for the current backend framework. + With :class:`ivy.Array` input: - Returns - ------- - ret - Boolean, whether or not inplace variables are supported. - """ - return current_backend().inplace_variables_supported() + >>> x = ivy.asarray([[10., 20., 30.], + ... [40., 50., 60.]]) + >>> y = ivy.stable_divide(x, 10.) + >>> print(y) + ivy.array([[1., 2., 3.], + [4., 5., 6.]]) -@handle_exceptions -@handle_backend_invalid -def is_array(x: Any, /, *, exclusive: bool = False) -> bool: - """ - Determine whether the input x is either an Ivy Array or a Native Array. + >>> x = ivy.asarray([1,2,3]) + >>> y = np.array((1., 3., 5.)) + >>> z = ivy.stable_divide(x, y) + >>> print(z) + ivy.array([1. , 0.667, 0.6 ]) - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + >>> x = ivy.asarray([1., 2., 4.]) + >>> y = ivy.asarray([1., 0.5, 0.25]) + >>> z = ivy.asarray([0.01, 0.02, 0.03]) + >>> w = ivy.stable_divide(x, y, min_denominator=z) + >>> print(w) + ivy.array([ 0.99, 3.85, 14.3 ]) - Returns - ------- - ret - Boolean, whether or not x is an array. + With :class:`ivy.Container` input: - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> print(ivy.is_array(x)) - True + >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) + >>> y = ivy.stable_divide(x, 0.5) + >>> print(y) + { + a: ivy.array([20., 30.]), + b: ivy.array([40., 50.]) + } - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> print(ivy.is_array(x, exclusive=True)) - True - >>> x = [2, 3] - >>> print(ivy.is_array(x)) - False + >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) + >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) + >>> z = ivy.stable_divide(x, y) + >>> print(z) + { + a: ivy.array([2., 0.8]), + b: ivy.array([0.857, 10.]) + } """ - return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( - x, exclusive=exclusive - ) + return numerator / (denominator + default(min_denominator, ivy.min_denominator)) @handle_exceptions -@handle_backend_invalid -def is_ivy_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False -) -> bool: +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +def stable_pow( + base: Union[Number, ivy.Array, ivy.NativeArray], + exponent: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + min_base: float = None, +) -> Any: """ - Determine whether the input x is a valid Ivy Array. + Raise the base by the power, with ivy.min_base added to the base when exponent > 1 + for numerical stability. Parameters ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + base + The base number. + exponent + The exponent number. + min_base + The minimum base to use, use global ivy.min_base by default. Returns ------- ret - Boolean, whether or not x is a valid Ivy Array. + The new item following the numerically stable power. Examples -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_ivy_array(x) - True + With :code:`int` input: - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_ivy_array(x, exclusive=True) - False - """ - return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) + >>> x = ivy.stable_pow(2, 2) + >>> print(x) + ivy.array(4.00004) + >>> x = ivy.stable_pow(2, 2, min_base=2) + >>> print(x) + ivy.array(16) -@handle_exceptions -def is_ivy_container(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Container. + With float input: - Parameters - ---------- - x - The input to check + >>> x = ivy.stable_pow(4.0, .5) + >>> print(x) + ivy.array(2.00000262) - Returns - ------- - ret - Boolean, whether or not x is an ivy container. + With :code:`complex` input: - Examples - -------- - >>> x = ivy.Container() - >>> print(ivy.is_ivy_container(x)) - True + >>> x = ivy.stable_pow(3+4j, 2j) + >>> print(x) + ivy.array(-0.15605032-0.01208451j) - >>> x = [2, 3] - >>> print(ivy.is_ivy_container(x)) - False - """ - return isinstance(x, ivy.Container) + With :class:`ivy.Array` input: + >>> x = ivy.asarray([[2, 4], + ... [6, 8]]) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) + ivy.array([[ 4.00004, 16.00008], + [36.00012, 64.00016]]) -def is_ivy_nested_array(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Nested Array. + >>> x = ivy.asarray([2, 4, 6]) + >>> y = ivy.asarray([2, 3, 4]) + >>> z = ivy.stable_pow(x, y) + >>> print(z) + ivy.array([ 4.00004, 64.00048, 1296.00864]) - Parameters - ---------- - x - The input to check - Returns - ------- - ret - Boolean, whether or not x is an ivy nested array. - """ - return isinstance(x, ivy.NestedArray) - - -@handle_exceptions -@handle_backend_invalid -def is_native_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False -) -> bool: - """ - Determine whether the input x is an :class:`ivy.NativeArray` instance. - - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. - - Returns - ------- - ret - Boolean, whether or not x is an :class:`ivy.NativeArray`. - - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_native_array(x) - False - - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_native_array(x, exclusive=True) - True - """ - try: - return current_backend(x).is_native_array(x, exclusive=exclusive) - except ValueError: - return False - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@to_native_arrays_and_back -@handle_device_shifting -def isin( - elements: Union[ivy.Array, ivy.NativeArray], - test_elements: Union[ivy.Array, ivy.NativeArray], - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> ivy.Array: - """ - Test if each element of elements is in test_elements. - - Parameters - ---------- - elements - input array - test_elements - values against which to test for each input element - assume_unique - If True, assumes both elements and test_elements contain unique elements, - which can speed up the calculation. Default value is False. - invert - If True, inverts the boolean return array, resulting in True values for - elements not in test_elements. Default value is False. - - Returns - ------- - ret - output a boolean array of the same shape as elements that is True for elements - in test_elements and False otherwise. + With :class:`ivy.Container` input: - Examples - -------- - >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y) - ivy.array([[False, False, False], [ True, True, True]]) + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) + { + a: ivy.array([4.00004, 16.00008]), + b: ivy.array([36.00012, 64.00016]) + } - >>> x = ivy.array([3, 2, 1, 0]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y, invert=True) - ivy.array([False, False, False, True]) + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) + >>> z = ivy.stable_pow(x, y) + >>> print(z) + { + a: ivy.array([2.00001, 64.00048]), + b: ivy.array([1296.00864, 32768.2048]) + } """ - return ivy.current_backend(elements, test_elements).isin( - elements, test_elements, assume_unique=assume_unique, invert=invert + return_dtype = ivy.promote_types( + ivy.default_dtype(item=base), + ivy.default_dtype(item=default(min_base, ivy.min_base)), ) + return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) + ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) + return ret.astype(return_dtype) -@handle_exceptions -@handle_nestable -def isscalar(x: Any, /) -> bool: - return np.isscalar(x) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def itemsize( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> int: - """ - Return the size of the input array's elements. - - Parameters - ---------- - x - The input array. - - Returns - ------- - ret - An integer specifying the element size in bytes. - - Examples - -------- - >>> x = ivy.array([1,2,3], dtype=ivy.float64) - >>> ivy.itemsize(x) - 8 - - >>> x = ivy.array([1,2,3], dtype=ivy.complex128) - >>> ivy.itemsize(x) - 16 - """ - return ivy.current_backend(x).itemsize(x) +stable_pow.unsupported_dtypes = ("bfloat16",) @handle_exceptions -def match_kwargs( - kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False -) -> Union[List[Dict], Dict]: +def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: """ - Match keyword arguments to either class or function receivers. - - Parameters - ---------- - kwargs - Keyword arguments to match. - receivers - Functions and/or classes to match the keyword arguments to. - allow_duplicates - Whether to allow one keyword argument to be used for multiple receivers. - Default is ``False``. + Get all arrays which are currently alive. Returns ------- ret - Sequence of keyword arguments split as best as possible. + All arrays which are alive. Examples -------- - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) - >>> print(x) - [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] - - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) - >>> print(x) - [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] - """ - split_kwargs = list() - for receiver in receivers: - expected_kwargs = arg_names(receiver) - found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} - if not allow_duplicates: - for k in found_kwargs.keys(): - del kwargs[k] - split_kwargs.append(found_kwargs) - if len(split_kwargs) == 1: - return split_kwargs[0] - return split_kwargs - - -@handle_exceptions -@handle_nestable -@handle_array_function -def multiprocessing(context: Optional[str] = None): + >>> ivy.get_all_arrays_in_memory() + [] + >>> x = ivy.get_all_arrays_in_memory() + >>> x + [] + >>> y = ivy.array([0, 1, 2]) + >>> x + [ivy.array([0, 1, 2])] """ - Return backend-specific multiprocessing module. - - Parameters - ---------- - context - The context of the multiprocessing, either fork, forkserver or spawn. - Default is ``None``. + all_arrays = list() + for obj in gc.get_objects(): + try: + if ivy.current_backend_str() in ["", "numpy"]: + if ivy.is_ivy_array(obj): + all_arrays.append(obj) + else: + if ivy.is_native_array(obj): + all_arrays.append(obj) - Returns - ------- - ret - Multiprocessing module - """ - return current_backend().multiprocessing(context) + except Exception: + pass + return all_arrays @handle_exceptions @@ -2577,286 +2447,365 @@ def print_all_arrays_in_memory(): print(type(arr), arr.shape) +ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 + + @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def scatter_flat( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], - /, - *, - size: Optional[int] = None, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def set_queue_timeout(timeout: float): """ - Scatter flat updates into a new flat array according to flat indices. + Set a timeout value (in seconds) for the global queue. + + Set the global queue timeout value (in seconds) Default value without this function + being called is 15 seconds. Parameters ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - size - The size of the result. Default is `None`, in which case tensor - argument out must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - New array of given shape, with the values scattered at the indices. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + timeout + The timeout when waiting for containers to arrive from the queues. + To be set in seconds. Examples -------- - With :class:`ivy.Array` input: - >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) - >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) - >>> ivy.scatter_flat(indices, updates, out=out) - >>> print(out) - ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) - - - With :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) - ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) - - - With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) - { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) - } - + >>> x = ivy.set_queue_timeout(10) + >>> x = ivy.queue_timeout + >>> print(x) + 10.0 - With :class:`ivy.Container` input: - >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), - ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) - { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) - } + >>> ivy.set_queue_timeout(30) + >>> y = ivy.queue_timeout + >>> print(y) + 30 """ - return current_backend(indices).scatter_flat( - indices, updates, size=size, reduction=reduction, out=out - ) + global queue_timeout_stack + ivy.utils.assertions.check_isinstance(timeout, (int, float)) + queue_timeout_stack.append(timeout) + ivy.__setattr__("queue_timeout", timeout, True) @handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def scatter_nd( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], - /, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - *, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def unset_queue_timeout() -> None: """ - Scatter updates into a new array according to indices. - - Parameters - ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - shape - The shape of the result. Default is ``None``, in which case tensor - argument must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - New array of given shape, with the values scattered at the indices. + Reset the global queue timeout value (in seconds) to the previous state. Examples -------- - scatter values into an empty array, With :class:`ivy.Array` input: - - >>> indices = ivy.array([[4], [3], [1], [7]]) - >>> updates = ivy.array([9, 10, 11, 12]) - >>> shape = ivy.array([8]) - >>> scatter = ivy.scatter_nd(indices, updates, shape) - >>> print(scatter) - ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) - - With scatter into an empty array, With :class:`ivy.Container` input: + >>> ivy.set_queue_timeout(10.0) + >>> y = ivy.queue_timeout + >>> print(y) + 10.0 - >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), - ... b=ivy.array([[5],[1],[2]])) - >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), - ... b=ivy.array([20, 30, 40])) - >>> shape = ivy.Container(a=ivy.array([10]), - ... b = ivy.array([10])) - >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') - >>> print(z) - { - a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), - b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) - } + >>> ivy.unset_queue_timeout() + >>> ivy.queue_timeout + 15.0 + """ + global queue_timeout_stack + if queue_timeout_stack: + queue_timeout_stack.pop(-1) + timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 + ivy.__setattr__("queue_timeout", timeout, True) - With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> indices = ivy.array([[4],[3],[1]]) - >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), - ... b=ivy.array([200, 300, 400])) - >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), - ... b=ivy.array([10, 20, 30, 40, 50])) - >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) - >>> print(z) - { - a: ivy.array([1, 30, 3, 20, 10]), - b: ivy.array([10, 400, 30, 300, 200]) - } - """ - return current_backend(indices).scatter_nd( - indices, updates, shape=shape, reduction=reduction, out=out - ) +ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" @handle_exceptions -def set_array_mode(mode: bool) -> None: +def set_tmp_dir(tmp_dr: str) -> None: """ - Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs - back to ivy.Array. - - It Stops the conversion of ivy.NativeArray to ivy.Array in the - case when it is set to False. + Set the directory for saving temporary files. - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions + Parameters + ---------- + tmp_dr + The new directory for saving temporary files Examples -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False + >>> x = ivy.tmp_dir + >>> print(x) + /tmp - >>> ivy.set_array_mode(True) - >>> ivy.array_mode - True + >>> ivy.set_tmp_dir("/my_tmp") + >>> y = ivy.tmp_dir + >>> print(y) + /my_tmp """ - global array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - array_mode_stack.append(mode) - ivy.__setattr__("array_mode", mode, True) + global tmp_dir_stack + ivy.utils.assertions.check_isinstance(tmp_dr, str) + tmp_dir_stack.append(tmp_dr) + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions -def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: +def unset_tmp_dir() -> None: """ - Set the mode of whether to show frontend-truncated exception stack traces, ivy- - truncated exception stack traces or full exception stack traces. - - Parameter - --------- - mode - str exeption trace mode, one of `ivy`, `full` or `frontend` + Reset the directory for saving temporary files to the previous value. Examples -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' + >>> ivy.set_tmp_dir("/my_dir") + >>> y = ivy.tmp_dir + >>> print(y) + /my_dir - >>> ivy.set_exception_trace_mode("full") - >>> ivy.exception_trace_mode - 'full' + >>> ivy.unset_tmp_dir() + >>> ivy.tmp_dir + /tmp """ - global exception_trace_mode_stack - trace_modes = list(trace_mode_dict.keys()) - ivy.utils.assertions.check_elem_in_list( - mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) - ) - exception_trace_mode_stack.append(mode) - ivy.__setattr__("exception_trace_mode", mode, True) + global tmp_dir_stack + if tmp_dir_stack: + tmp_dir_stack.pop(-1) + tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions -def set_inplace_mode(mode: str = "lenient") -> None: +def container_types(): """ - Set the memory management behavior for in-place updates in Ivy. - - By default, Ivy creates new arrays in the backend for in-place updates. - However, this behavior can be controlled by the user - using the 'inplace_mode' parameter. - - Parameters - ---------- - mode : str - The mode for memory management during in-place updates. - - 'lenient': (Default) In this mode, new arrays will be created during - in-place updates to avoid breaking existing code. - This is the default behavior. - - 'strict': In this mode, an error will be raised if the - 'inplace_update' function is called - in a backend that doesn't support inplace updates natively. + Summary. Returns ------- - None + ret + a key-value structure, and exposes public methods .keys(), .values() and + items(). + """ + # noinspection PyBroadException + try: + return current_backend().container_types() + except ValueError: + return [] + + +@handle_exceptions +def inplace_arrays_supported() -> bool: + """ + Determine whether inplace arrays are supported for the current backend framework. + + Returns + ------- + ret + Boolean, whether or not inplace arrays are supported. + """ + return current_backend().inplace_arrays_supported() + + +@handle_exceptions +def inplace_variables_supported() -> bool: + """ + Determine whether inplace variables are supported for the current backend framework. + + Returns + ------- + ret + Boolean, whether or not inplace variables are supported. + """ + return current_backend().inplace_variables_supported() + + +@handle_exceptions +@handle_nestable +@inputs_to_native_arrays +@handle_array_function +def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: + """ + Return if in-place operations are supported for x's data type. + + Determine whether in-place operations are supported for x's data type, by the + current backend framework setting. + + Parameters + ---------- + x + Input variable for whose data type we check whether the current backend + framework supports in-place operations. + + Returns + ------- + ret + Value depends on whether in-place operations are supported for + data type of x. + + Raises + ------ + IvyException + If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will + be raised. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - >>> set_inplace_mode('lenient') - >>> ivy.inplace_mode - 'lenient' + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) + True - Note - ---- - Enabling strict mode can help users have more control over memory management - but may lead to errors if the backend doesn't support inplace updates natively. + With :class:`ivy.Container` input and backend set as `torch`: + + >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) + { + a: True, + b: True + } + + With `ivy.Array` input and backend set as "tensorflow": + + >>> x = ivy.array([1., 4.2, 2.2]) + >>> ret = x.supports_inplace_updates() + >>> print(ret) + False """ - global inplace_mode_stack - inplace_modes = ["lenient", "strict"] - ivy.utils.assertions.check_elem_in_list( - mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" + if _is_variable(x): + return ivy.inplace_variables_supported() + elif ivy.is_native_array(x): + return ivy.inplace_arrays_supported() + raise ivy.utils.exceptions.IvyException( + "Input x must be either a variable or an array." ) - inplace_mode_stack.append(mode) - ivy.__setattr__("inplace_mode", mode, True) + + +@handle_exceptions +@handle_nestable +@inputs_to_native_arrays +@handle_array_function +def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: + """ + Assert that inplace operations are supported for x. + + Parameters + ---------- + x + Input variable or array to check for inplace support for. + + Returns + ------- + ret + True if supports, raises IvyBackendException otherwise + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. + + Examples + -------- + With :class:`ivy.Array` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) + True + + With :class:`ivy.Array` input and default backend set as `torch`: + + >>> ivy.set_backend("torch") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) + True + + With :class:`ivy.Container` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) + { + a: True, + b: True + } + + With :class:`ivy.Container` input and default backend set as `torch`: + + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) + { + a: True, + b: True + } + """ + ivy.utils.assertions.check_true( + ivy.supports_inplace_updates(x), + "Inplace operations are not supported {} types with {} backend".format( + type(x), ivy.current_backend_str() + ), + ) + return True + + +@handle_nestable +@handle_partial_mixed_function +@handle_view_indexing +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def get_item( + x: Union[ivy.Array, ivy.NativeArray], + /, + query: Union[ivy.Array, ivy.NativeArray, Tuple], + *, + copy: Optional[bool] = None, +) -> ivy.Array: + """ + Gather slices from x according to query array, identical to x[query]. + + Parameters + ---------- + x + array, the array from which to gather values. + query + array, index array, integer indices or boolean mask. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + + Returns + ------- + ret + New array with the values gathered at the specified indices. + + Functional Examples + ------------------- + + >>> x = ivy.array([0, -1, 20]) + >>> query = ivy.array([0, 1]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 0, -1]) + + >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) + >>> query = ivy.array([[True, False], [False, False], [True, True]]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 4, -2, -10]) + """ + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return ivy.array([], shape=(0,), dtype=x.dtype) + return ivy.expand_dims(x, axis=0) + query = ivy.nonzero(query, as_tuple=False) + ret = ivy.gather_nd(x, query) + else: + indices, target_shape = _parse_query(query, x.shape) + if indices is None: + return ivy.empty(target_shape, dtype=x.dtype) + ret = ivy.gather_nd(x, indices) + ret = ivy.reshape(ret, target_shape) + return ret + + +get_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} @handle_nestable @@ -2930,593 +2879,758 @@ def set_item( return ret -@handle_exceptions -@handle_array_function -def set_min_base(val: float) -> None: - """ - Set the global minimum base used by ivy for numerically stable power raising. - - Parameters - ---------- - val - The new value to set the minimum base to. +set_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} - Examples - -------- - >>> x = ivy.min_base - >>> print(x) - 1e-05 - >>> ivy.set_min_base(1e-04) - >>> y = ivy.min_base - >>> print(y) - 1e-04 - """ - global min_base_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_base_stack.append(val) - ivy.__setattr__("min_base", val, True) +def _parse_query(query, x_shape): + query = (query,) if not isinstance(query, tuple) else query + query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) + + # array containing all of x's flat indices + x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) + + # use numpy's __getitem__ to get the queried indices + x_idxs = ivy.array(x_.to_numpy()[query_]) + target_shape = x_idxs.shape + + if 0 in x_idxs.shape or 0 in x_shape: + return None, target_shape + + # convert the flat indices to multi-D indices + x_idxs = ivy.unravel_index(x_idxs, x_shape) + + # stack the multi-D indices to bring them to gather_nd/scatter_nd format + x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) + + return x_idxs, target_shape + + +def _numel(shape): + shape = tuple(shape) + return ivy.prod(shape).to_scalar() if shape != () else 1 + + +def _broadcast_to(input, target_shape): + input = ivy.squeeze(input) + if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): + return ivy.reshape(input, target_shape) + else: + input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input + new_dims = () + i_i = len(input.shape) - 1 + for i_t in range(len(target_shape) - 1, -1, -1): + if len(input.shape) + len(new_dims) >= len(target_shape): + break + if i_i < 0 or target_shape[i_t] != input.shape[i_i]: + new_dims += (i_t,) + else: + i_i -= 1 + input = ivy.expand_dims(input, axis=new_dims) + return ivy.broadcast_to(input, target_shape) @handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays @handle_array_function -def set_min_denominator(val: float) -> None: +# @handle_device_shifting +def inplace_update( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], + /, + *, + ensure_in_backend: bool = False, + keep_input_dtype: bool = False, +) -> ivy.Array: """ - Set the global minimum denominator used by ivy for numerically stable division. + Perform in-place update for the input array. + + This will always be performed on ivy.Array instances pass in the input, and will + also be performed on the native array classes in the backend when the backend + supports this. If the backend does not natively support inplace updates, and x is an + ivy.NativeArray instance, then an + exception will be thrown. Parameters ---------- + x + The variable to update. val - The value to set the global minimum denominator to. + The array to update the variable with. + ensure_in_backend + Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. + In cases where it should be, backends which do not natively support inplace + updates will raise an exception. + keep_input_dtype + Whether or not to preserve `x` data type after the update, otherwise `val` + data type will be applied. Defaults to False. + + Returns + ------- + ret + The array following the in-place update. + + Raises + ------ + IvyException + If backend set doesn't natively support inplace updates and ensure_in_backend is + True, above exception will be raised. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the arguments. Examples -------- - >>> x = ivy.min_denominator - >>> print(x) - 1e-12 - - >>> ivy.set_min_denominator(1e-13) - >>> y = ivy.min_denominator - >>> print(y) - 1e-13 - """ - global min_denominator_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_denominator_stack.append(val) - ivy.__setattr__("min_denominator", val, True) + With :class:`ivy.Array` input and default backend set as `numpy`: + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0]) + >>> ivy.inplace_update(x, y) + >>> print(x) + ivy.array([0]) -@handle_exceptions -def set_nestable_mode(mode: bool) -> None: - """ - Set the mode of whether to check if function inputs are ivy.Container. + With :class:`ivy.Array` input and default backend set as `numpy`: - Parameter - --------- - mode - boolean whether to check if function inputs are ivy.Container + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) + >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) + >>> ivy.inplace_update(x, y, keep_input_dtype=True) + >>> print(x) + ivy.array([0., 0., 0.]) - Examples - -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False + With :class:`ivy.Container` instances:, and backend set as `torch`: - >>> ivy.set_nestable_mode(True) - >>> ivy.nestable_mode - True - """ - global nestable_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - nestable_mode_stack.append(mode) - ivy.__setattr__("nestable_mode", mode, True) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) + >>> ivy.inplace_update(x, y) + >>> print(x) + { + a: ivy.array([1, 1]), + b: ivy.array([2, 2]) + } + With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend + set as `torch`: -@handle_exceptions -def set_precise_mode(mode: bool) -> None: + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.array([1, 2]) + >>> ivy.inplace_update(x, y) + >>> print(x) + { + a: ivy.array([1, 2]), + b: ivy.array([1, 2]) + } """ - Set the mode of whether to use a promotion table that avoids any precision loss or a - compute effecient table that avoids most wider-than-necessary promotions. + return current_backend(x).inplace_update( + x, + val, + ensure_in_backend=ensure_in_backend, + keep_input_dtype=keep_input_dtype, + ) - Parameter - --------- - mode - boolean whether to use high precision promtion table - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False +inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} - >>> ivy.set_precise_mode(True) - >>> ivy.precise_mode - True - """ - global precise_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - precise_mode_stack.append(mode) - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) +ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" @handle_exceptions -@handle_array_function -def set_queue_timeout(timeout: float): +def set_inplace_mode(mode: str = "lenient") -> None: """ - Set a timeout value (in seconds) for the global queue. + Set the memory management behavior for in-place updates in Ivy. - Set the global queue timeout value (in seconds) Default value without this function - being called is 15 seconds. + By default, Ivy creates new arrays in the backend for in-place updates. + However, this behavior can be controlled by the user + using the 'inplace_mode' parameter. Parameters ---------- - timeout - The timeout when waiting for containers to arrive from the queues. - To be set in seconds. + mode : str + The mode for memory management during in-place updates. + - 'lenient': (Default) In this mode, new arrays will be created during + in-place updates to avoid breaking existing code. + This is the default behavior. + - 'strict': In this mode, an error will be raised if the + 'inplace_update' function is called + in a backend that doesn't support inplace updates natively. + + Returns + ------- + None Examples -------- - >>> x = ivy.set_queue_timeout(10) - >>> x = ivy.queue_timeout - >>> print(x) - 10.0 + >>> set_inplace_mode('lenient') + >>> ivy.inplace_mode + 'lenient' - >>> ivy.set_queue_timeout(30) - >>> y = ivy.queue_timeout - >>> print(y) - 30 - """ - global queue_timeout_stack - ivy.utils.assertions.check_isinstance(timeout, (int, float)) - queue_timeout_stack.append(timeout) - ivy.__setattr__("queue_timeout", timeout, True) - - -@handle_exceptions -def set_shape_array_mode(mode: bool) -> None: - """ - Set the mode of returning shape as ivy.Array to the given mode instance. - - Parameter - --------- - mode - boolean whether to return shape as ivy.Array - - Examples - -------- - >>> ivy.set_shape_array_mode(False) - >>> ivy.shape_array_mode - False + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True + Note + ---- + Enabling strict mode can help users have more control over memory management + but may lead to errors if the backend doesn't support inplace updates natively. """ - global shape_array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - shape_array_mode_stack.append(mode) - ivy.__setattr__("shape_array_mode", mode, True) + global inplace_mode_stack + inplace_modes = ["lenient", "strict"] + ivy.utils.assertions.check_elem_in_list( + mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" + ) + inplace_mode_stack.append(mode) + ivy.__setattr__("inplace_mode", mode, True) @handle_exceptions -def set_show_func_wrapper_trace_mode(mode: bool) -> None: +def unset_inplace_mode() -> None: """ - Set the mode of whether to show the full stack trace with function wrapping traces. - - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions + Reset the memory management behavior for in-place updates in Ivy to the previous + state. Examples -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' - >>> ivy.set_show_func_wrapper_trace_mode(True) - >>> ivy.show_func_wrapper_trace_mode - True + >>> unset_inplace_mode() + >>> ivy.inplace_mode + 'lenient' """ - global show_func_wrapper_trace_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - show_func_wrapper_trace_mode_stack.append(mode) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + global inplace_mode_stack + if inplace_mode_stack: + inplace_mode_stack.pop(-1) + mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" + ivy.__setattr__("inplace_mode", mode, True) @handle_exceptions -def set_tmp_dir(tmp_dr: str) -> None: +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def inplace_decrement( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Set the directory for saving temporary files. + Perform in-place decrement for the input array. Parameters ---------- - tmp_dr - The new directory for saving temporary files + x + The input array to be decremented by the defined value. + val + The value of decrement. + + Returns + ------- + ret + The array following the in-place decrement. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.tmp_dir - >>> print(x) - /tmp + With :class:`ivy.Array` input: - >>> ivy.set_tmp_dir("/my_tmp") - >>> y = ivy.tmp_dir + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_decrement(x, 1.25) >>> print(y) - /my_tmp + ivy.array([[ 4.05, 5.75, -1.25], + [ 5.55, 6.75, 2.65], + [-1.25, 8.75, 5.05]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) + >>> y = ivy.inplace_decrement(x, 1.5) + >>> print(y) + { + a: ivy.array([-1., -6.5, 28.5]), + b: ivy.array([-1.5, -26.5, 48.5]) + } + + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([0., 0., 0.]), + b: ivy.array([0., 0., 0.]) + } + + >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) + >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([1., 1.5, 3.]), + b: ivy.array([0., 50., 3.5]) + } """ - global tmp_dir_stack - ivy.utils.assertions.check_isinstance(tmp_dr, str) - tmp_dir_stack.append(tmp_dr) - ivy.__setattr__("tmp_dir", tmp_dr, True) + return current_backend(x).inplace_decrement(x, val) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays -@outputs_to_ivy_shapes -@outputs_to_ivy_arrays +@inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def shape( +def inplace_increment( x: Union[ivy.Array, ivy.NativeArray], - /, - *, - as_array: bool = False, -) -> Union[ivy.Shape, ivy.NativeShape]: + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Return the shape of the array ``x``. + Perform in-place increment for the input array. Parameters ---------- x - Input array to infer the shape of. - as_array - Whether to return the shape as an array. - Default is False. + The input array to be incremented by the defined value. + val + The value of increment. Returns ------- ret - Shape of the array ``x``. + The array following the in-place increment. Examples -------- - >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) - >>> y = ivy.shape(x) - >>> z = ivy.shape(x, as_array = True) + With :class:`ivy.Array` input: + + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_increment(x, 3.) >>> print(y) - (2, 3) + ivy.array([[ 8.3, 10., 3.], + [ 9.8, 11., 6.9], + [ 3., 13., 9.3]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.inplace_increment(x, 2.5) + >>> print(y) + { + a: ivy.array([2.5, 17.5, 32.5]), + b: ivy.array([2.5, 27.5, 52.5]) + } + + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_increment(x, y) >>> print(z) - ivy.array([2, 3]) + { + a: ivy.array([0., 30., 60.]), + b: ivy.array([0., 50., 100.]) + } """ - return current_backend(x).shape(x, as_array=as_array) + return current_backend(x).inplace_increment(x, val) @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function -def stable_divide( - numerator: Union[Number, ivy.Array, ivy.NativeArray], - denominator: Union[Number, ivy.Array, ivy.NativeArray], +@handle_device_shifting +def scatter_flat( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], /, *, - min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, -) -> Union[Number, ivy.Array]: + size: Optional[int] = None, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Divide the numerator by the denominator, with min denominator added to the - denominator for numerical stability. + Scatter flat updates into a new flat array according to flat indices. Parameters ---------- - numerator - The numerator of the division. - denominator - The denominator of the division. - min_denominator - The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) - by default. + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + size + The size of the result. Default is `None`, in which case tensor + argument out must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new item following the numerically stable division. + New array of given shape, with the values scattered at the indices. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - With :code:`int` input: - - >>> x = ivy.stable_divide(1, 2) - >>> print(x) - 0.49999999999975 + With :class:`ivy.Array` input: + >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) + >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) + >>> ivy.scatter_flat(indices, updates, out=out) + >>> print(out) + ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) - >>> x = ivy.stable_divide(1, 4, min_denominator=1) - >>> print(x) - 0.2 - - With float input: - - >>> x = ivy.stable_divide(5.0, 3.33) - >>> print(x) - 1.5015015015010504 - - With :code:`complex` input: - - >>> x = ivy.stable_divide(1+1j, 1-1j) - >>> print(x) - (5.000444502911705e-13+0.9999999999995j) With :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) - >>> x = ivy.asarray([[10., 20., 30.], - ... [40., 50., 60.]]) - >>> y = ivy.stable_divide(x, 10.) - >>> print(y) - ivy.array([[1., 2., 3.], - [4., 5., 6.]]) - - - >>> x = ivy.asarray([1,2,3]) - >>> y = np.array((1., 3., 5.)) - >>> z = ivy.stable_divide(x, y) - >>> print(z) - ivy.array([1. , 0.667, 0.6 ]) - - >>> x = ivy.asarray([1., 2., 4.]) - >>> y = ivy.asarray([1., 0.5, 0.25]) - >>> z = ivy.asarray([0.01, 0.02, 0.03]) - >>> w = ivy.stable_divide(x, y, min_denominator=z) - >>> print(w) - ivy.array([ 0.99, 3.85, 14.3 ]) - - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) - >>> y = ivy.stable_divide(x, 0.5) - >>> print(y) + With :class:`ivy.Container` and :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) { - a: ivy.array([20., 30.]), - b: ivy.array([40., 50.]) + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) } - >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) - >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) - >>> z = ivy.stable_divide(x, y) - >>> print(z) + With :class:`ivy.Container` input: + >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), + ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) { - a: ivy.array([2., 0.8]), - b: ivy.array([0.857, 10.]) + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) } """ - return numerator / (denominator + default(min_denominator, ivy.min_denominator)) + return current_backend(indices).scatter_flat( + indices, updates, size=size, reduction=reduction, out=out + ) @handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@inputs_to_native_shapes +@to_native_arrays_and_back @handle_array_function -def stable_pow( - base: Union[Number, ivy.Array, ivy.NativeArray], - exponent: Union[Number, ivy.Array, ivy.NativeArray], +@handle_device_shifting +def scatter_nd( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], /, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, *, - min_base: float = None, -) -> Any: + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Raise the base by the power, with ivy.min_base added to the base when exponent > 1 - for numerical stability. + Scatter updates into a new array according to indices. Parameters ---------- - base - The base number. - exponent - The exponent number. - min_base - The minimum base to use, use global ivy.min_base by default. + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + shape + The shape of the result. Default is ``None``, in which case tensor + argument must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new item following the numerically stable power. + New array of given shape, with the values scattered at the indices. Examples -------- - With :code:`int` input: - - >>> x = ivy.stable_pow(2, 2) - >>> print(x) - ivy.array(4.00004) - - >>> x = ivy.stable_pow(2, 2, min_base=2) - >>> print(x) - ivy.array(16) - - With float input: - - >>> x = ivy.stable_pow(4.0, .5) - >>> print(x) - ivy.array(2.00000262) - - With :code:`complex` input: - - >>> x = ivy.stable_pow(3+4j, 2j) - >>> print(x) - ivy.array(-0.15605032-0.01208451j) + scatter values into an empty array, With :class:`ivy.Array` input: - With :class:`ivy.Array` input: + >>> indices = ivy.array([[4], [3], [1], [7]]) + >>> updates = ivy.array([9, 10, 11, 12]) + >>> shape = ivy.array([8]) + >>> scatter = ivy.scatter_nd(indices, updates, shape) + >>> print(scatter) + ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) - >>> x = ivy.asarray([[2, 4], - ... [6, 8]]) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) - ivy.array([[ 4.00004, 16.00008], - [36.00012, 64.00016]]) + With scatter into an empty array, With :class:`ivy.Container` input: - >>> x = ivy.asarray([2, 4, 6]) - >>> y = ivy.asarray([2, 3, 4]) - >>> z = ivy.stable_pow(x, y) + >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), + ... b=ivy.array([[5],[1],[2]])) + >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), + ... b=ivy.array([20, 30, 40])) + >>> shape = ivy.Container(a=ivy.array([10]), + ... b = ivy.array([10])) + >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') >>> print(z) - ivy.array([ 4.00004, 64.00048, 1296.00864]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) { - a: ivy.array([4.00004, 16.00008]), - b: ivy.array([36.00012, 64.00016]) + a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), + b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) } - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) - >>> z = ivy.stable_pow(x, y) + With :class:`ivy.Container` and :class:`ivy.Array` input: + + >>> indices = ivy.array([[4],[3],[1]]) + >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), + ... b=ivy.array([200, 300, 400])) + >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), + ... b=ivy.array([10, 20, 30, 40, 50])) + >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) >>> print(z) { - a: ivy.array([2.00001, 64.00048]), - b: ivy.array([1296.00864, 32768.2048]) + a: ivy.array([1, 30, 3, 20, 10]), + b: ivy.array([10, 400, 30, 300, 200]) } """ - return_dtype = ivy.promote_types( - ivy.default_dtype(item=base), - ivy.default_dtype(item=default(min_base, ivy.min_base)), + return current_backend(indices).scatter_nd( + indices, updates, shape=shape, reduction=reduction, out=out ) - return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) - ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) - return ret.astype(return_dtype) @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def strides( - x: Union[ivy.Array, ivy.NativeArray], +def gather( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], /, -) -> Tuple[int]: + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Return the input array's strides across each dimension. + Gather slices from params at axis according to indices. Parameters ---------- - x - The input array. + params + The array from which to gather values. + indices + The array which indicates the indices that will be gathered along + the specified axis. + axis + optional int, the axis from which to gather from. + Default is ``-1``. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional array, for writing the result to. It must have a shape + that the inputs broadcast to. Returns ------- ret - A tuple containing the strides. + New array with the values gathered at the specified indices along the + specified axis. + + + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) - >>> ivy.strides(x) - (4, 8) - """ - if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): - return ivy.to_numpy(x).strides - # if x is an ivy array with a base, - # convert it to a numpy array with the same base: - ret = ivy.to_numpy(x.base) - ivy_numpy = ivy.with_backend("numpy") - for fn, args, kwargs, index in x._manipulation_stack: - ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) - ret = ret[index] if ivy.exists(index) else ret - return ret.to_native().strides + With :class:`ivy.Array` input: + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.array([1, 2]) + >>> print(ivy.gather(x, y)) + ivy.array([1., 2.]) -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -@handle_array_function -def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: + >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) + >>> y = ivy.array([[0, 1],[1, 2]]) + >>> z = ivy.zeros((2, 2, 2)) + >>> ivy.gather(x, y, out=z) + >>> print(z) + ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) + + >>> x = ivy.array([[[0., 1.], [2., 3.]], + ... [[8., 9.], [10., 11.]]]) + >>> y = ivy.array([[0, 1]]) + >>> z = ivy.zeros((1, 2, 2, 2)) + >>> ivy.gather(x, y, axis=0, out=z) + >>> print(z) + ivy.array( + [[[[ 0., 1.], + [ 2., 3.]], + [[ 8., 9.], + [10., 11.]]]]) + + >>> x = ivy.array([[0, 10, 20, 0, 0], + ... [0, 0, 0, 30, 40], + ... [0, 10, 0, 0, 40]]) + >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) + >>> z = ivy.gather(x, y, batch_dims=1) + >>> print(z) + ivy.array([[10, 20], [30, 40],[10, 40]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.Container(a = ivy.array([0, 1]), + ... b = ivy.array([1, 2])) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([5., 6.]) + } + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.array([0, 1]) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([4., 5.]) + } """ - Return if in-place operations are supported for x's data type. + return current_backend(params, indices).gather( + params, indices, axis=axis, batch_dims=batch_dims, out=out + ) - Determine whether in-place operations are supported for x's data type, by the - current backend framework setting. + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def gather_nd( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + batch_dims: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Gather slices from params into a array with shape specified by indices. Parameters ---------- - x - Input variable for whose data type we check whether the current backend - framework supports in-place operations. + params + The array from which to gather values. + indices + Index array. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Value depends on whether in-place operations are supported for - data type of x. - - Raises - ------ - IvyException - If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will - be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + New array of given shape, with the values gathered at the indices. Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: + With :class:`ivy.Array` input: - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) - True + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) + ivy.array(1.) - With :class:`ivy.Container` input and backend set as `torch`: + >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) + >>> y = ivy.array([[0],[1],[1]], dtype='int32') + >>> z = ivy.gather_nd(x,y,batch_dims=1) + ivy.array([0., 3., 5.]) - >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) { - a: True, - b: True + a: ivy.array(1.), + b: ivy.array(5.) } - With `ivy.Array` input and backend set as "tensorflow": + With :class:`ivy.Container` input: - >>> x = ivy.array([1., 4.2, 2.2]) - >>> ret = x.supports_inplace_updates() - >>> print(ret) - False + >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), + ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) + >>> y = ivy.Container(a=ivy.array([1,0]), + ... b=ivy.array([0])) + >>> print(ivy.gather_nd(x, y)) + { + a: ivy.array(30.), + b: ivy.array([0., 100., 200.]) + } """ - if _is_variable(x): - return ivy.inplace_variables_supported() - elif ivy.is_native_array(x): - return ivy.inplace_arrays_supported() - raise ivy.utils.exceptions.IvyException( - "Input x must be either a variable or an array." + res = current_backend(params, indices).gather_nd( + params, indices, batch_dims=batch_dims ) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res @handle_exceptions -def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: +@handle_nestable +@handle_array_function +def multiprocessing(context: Optional[str] = None): """ - Return the input shape in ivy.Shape form. + Return backend-specific multiprocessing module. Parameters ---------- - shape - The input to be converted + context + The context of the multiprocessing, either fork, forkserver or spawn. + Default is ``None``. Returns ------- - ret - the input in ivy.Shape form + ret + Multiprocessing module """ - if isinstance(shape, ivy.Shape): - return shape - return ivy.Shape(shape) + return current_backend().multiprocessing(context) @handle_exceptions @@ -3524,608 +3638,422 @@ def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: @handle_nestable @handle_array_like_without_promotion @inputs_to_native_arrays +@outputs_to_ivy_shapes +@outputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: +def shape( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + as_array: bool = False, +) -> Union[ivy.Shape, ivy.NativeShape]: """ - Create a (possibly nested) list from input array. + Return the shape of the array ``x``. Parameters ---------- x - Input array. + Input array to infer the shape of. + as_array + Whether to return the shape as an array. + Default is False. Returns ------- ret - A list representation of the input array ``x``. + Shape of the array ``x``. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_list(x) - >>> print(y) - [-1, 0, 1] - - >>> x = ivy.array([[ 1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.to_list(x) - >>> print(y) - [[1.100000023841858,2.200000047683716,3.299999952316284], - [-4.400000095367432,-5.5,-6.599999904632568]] - - >>> x = ivy.array([[[-1, 0, 1], - ... [ 1, 0, -1]], - ... [[ 1, -1, 0], - ... [ 1, 0, -1]]]) - >>> y = ivy.to_list(x) + >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) + >>> y = ivy.shape(x) + >>> z = ivy.shape(x, as_array = True) >>> print(y) - [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + (2, 3) - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [-1, 0, 1] - } + >>> print(z) + ivy.array([2, 3]) + """ + return current_backend(x).shape(x, as_array=as_array) - >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], - ... [-1, 0, 1], - ... [1, 0, -1]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] - } - >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], - ... [[1, -1, 0],[1, 0, -1]]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - } - """ - return current_backend(x).to_list(x) +ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False @handle_exceptions -def to_native_shape( - shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] -) -> ivy.NativeShape: +def set_shape_array_mode(mode: bool) -> None: """ - Return the input shape in its native backend framework form. + Set the mode of returning shape as ivy.Array to the given mode instance. - Parameters - ---------- - shape - The input to be converted + Parameter + --------- + mode + boolean whether to return shape as ivy.Array - Returns - ------- - ret - the input in its native framework form + Examples + -------- + >>> ivy.set_shape_array_mode(False) + >>> ivy.shape_array_mode + False + + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True """ - native_shape_type = (ivy.NativeShape,) - if ivy.current_backend_str() == "torch": - native_shape_type += (tuple,) - if len(backend_stack) != 0 and isinstance(shape, native_shape_type): - return shape - ivy.utils.assertions.check_isinstance( - shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) - ) - if isinstance(shape, int): - shape = (shape,) - elif isinstance(shape, list): - shape = tuple(shape) - elif is_array(shape): - shape = ivy.to_numpy(shape).tolist() - elif isinstance(shape, ivy.Shape): - shape = shape.shape - ivy.utils.assertions.check_all( - [isinstance(v, int) for v in shape if not is_array(v)], - "shape must take integers only", - as_array=False, - ) - ivy.utils.assertions.check_true( - not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" - ) - return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) + global shape_array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + shape_array_mode_stack.append(mode) + ivy.__setattr__("shape_array_mode", mode, True) @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays -@handle_array_function -@handle_device_shifting -def to_numpy( - x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True -) -> np.ndarray: +def unset_shape_array_mode() -> None: """ - Convert an array into a numpy array. - - Parameters - ---------- - x - input array - copy - whether to copy the array to a new address or not. - Default is ``True``. - - Returns - ------- - ret - a numpy array copying all the element of the array ``x``. + Reset the mode of returning shape as ivy.Array to the previous state. Examples -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [-1 0 1] - - >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [[-1 0 1] - [-1 0 1] - [ 1 0 -1]] - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([-1, 0, 1], dtype=int32) - } + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True - >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), - ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([[-1., 0., 1.], - [-1., 0., 1.], - [1., 0., -1.]], dtype=float32), - b: array([[-1, 0, 0], - [1, 0, 1], - [1, 1, 1]], dtype=int32) - } + >>> ivy.unset_shape_array_mode() + >>> ivy.shape_array_mode + False """ - return current_backend(x).to_numpy(x, copy=copy) + global shape_array_mode_stack + if shape_array_mode_stack: + shape_array_mode_stack.pop(-1) + mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False + ivy.__setattr__("shape_array_mode", mode, True) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: +def get_num_dims( + x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False +) -> int: """ - Convert an array with a single element into a scalar. + Return the number of dimensions of the array x. Parameters ---------- x - Input array with a single element. + Input array to infer the number of dimensions for. + as_array + Whether to return the shape as a array, default False. Returns ------- ret - a scalar copying the element of the array ``x``. + Shape of the array Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - Functional Examples - ------------------- - + Examples + -------- With :class:`ivy.Array` input: - >>> x = ivy.array([3]) - >>> y = ivy.to_scalar(x) - >>> print(y) + >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) + >>> b = ivy.get_num_dims(a, as_array=False) + >>> print(b) 3 - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) - >>> y = ivy.to_scalar(x) - >>> print(y) + >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) + >>> print(ivy.get_num_dims(a)) { - a: -1, - b: 3 + b: 2 } - >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), - ... c=ivy.array([-1])) - >>> y = ivy.to_scalar(x) - >>> print(y) + >>> b = ivy.get_num_dims(a, as_array=True) + >>> print(b) { - a: 1, - b: 0, - c: -1 + b: ivy.array(2) } """ - return current_backend(x).to_scalar(x) + return current_backend(x).get_num_dims(x, as_array=as_array) @handle_exceptions -@handle_nestable -def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: +def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): """ - Try and return the function, otherwise return None if an exception was raised during - function execution. + Return the index and `inspect.Parameter` representation of the specified argument. + In the form of a dict with keys "idx" and "param". Parameters ---------- fn - Function to try and call and return. - args - list of arguments. - kwargs - dictionay of keyword arguments + The function to retrieve the argument information for + name + The name of the argument + idx + the index of the argument in the inputs Returns ------- - Either the function itself or None if an exception was raised - during function execution. - - Examples - -------- - with a function that is executed without any exception: + ret + a `dict` containing the idx, and the `inspect.Parameter` for the argument, + which itself contains the parameter name, type, and other helpful information. + """ + ivy.utils.assertions.check_all_or_any_fn( + name, + idx, + fn=ivy.exists, + type="any", + limit=[1], + message="exactly one of the keyword arguments name or idx must be provided", + as_array=False, + ) + params = inspect.signature(fn).parameters + if ivy.exists(name): + return {"idx": list(params).index(name), "param": params[name]} + return {"idx": idx, "param": list(params.values())[idx]} - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> z = ivy.try_else_none(ivy.add, x, y) - >>> print(z.__name__) - add - with a function that is executed with an exception: +def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): + attr_list = () + if hasattr(fn, other_attr_name): + attr_list = getattr(fn, other_attr_name) + if isinstance(attr_list, dict): + attr_list = attr_list.get(backend, ()) + ivy.utils.assertions.check_false( + dnd_dict and attr_list, + f"Cannot specify both {first_attr_name} and {other_attr_name} " + "cannot both be defined for the same function", + ) - >>> x = ivy.array([1, 2, 3]) - >>> y = 'hemant' - >>> z = ivy.try_else_none(ivy.add,x, y) - >>> print(z) - None - """ - try: - _ = fn(*args, **kwargs) - return fn - except Exception: - return None +def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: + fn_unsupported_dnd = {} + fn_supported_dnd = {} + backend = ivy.current_backend_str() + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): + fn_supported_dnd = fn_supported_dnd.get(backend, {}) -@handle_exceptions -def unset_array_mode() -> None: - """ - Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back - to ivy.Array to the previous state. + ivy.utils.assertions.check_false( + fn_unsupported_dnd and fn_supported_dnd, + "unsupported_device_and_dtype and supported_device_and_dtype cannot" + " both be defined for the same function", + ) - Examples - -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False - - >>> ivy.unset_shape_array_mode() - >>> ivy.array_mode - True - """ - global array_mode_stack - if array_mode_stack: - array_mode_stack.pop(-1) - mode = array_mode_stack[-1] if array_mode_stack else True - ivy.__setattr__("array_mode", mode, True) - - -@handle_exceptions -def unset_exception_trace_mode() -> None: - """ - Reset the trace mode to the previously set mode. - - Examples - -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' - - >>> ivy.unset_exception_trace_mode() - >>> ivy.exception_trace_mode - 'full' - """ - global exception_trace_mode_stack - if exception_trace_mode_stack: - exception_trace_mode_stack.pop(-1) - mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" - ivy.__setattr__("exception_trace_mode", mode, True) - - -@handle_exceptions -def unset_inplace_mode() -> None: - """ - Reset the memory management behavior for in-place updates in Ivy to the previous - state. - - Examples - -------- - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' - - >>> unset_inplace_mode() - >>> ivy.inplace_mode - 'lenient' - """ - global inplace_mode_stack - if inplace_mode_stack: - inplace_mode_stack.pop(-1) - mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" - ivy.__setattr__("inplace_mode", mode, True) + us = "unsupported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") + ss = "supported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") -@handle_exceptions -def unset_min_base() -> None: - """ - Reset the global minimum base used by ivy for numerically stable power raising to - the previous value. + return True - Examples - -------- - >>> ivy.set_min_base(1e-07) - >>> y = ivy.min_base - >>> print(y) - 1e-07 - >>> ivy.unset_min_base() - >>> ivy.min_base - 1e-05 - """ - global min_base_stack - if min_base_stack: - min_base_stack.pop(-1) - val = min_base_stack[-1] if min_base_stack else 1e-05 - ivy.__setattr__("min_base", val, True) +def _all_dnd_combinations(): + all_comb = {} + for device in ivy.all_devices: + all_comb[device] = ivy.all_dtypes + return all_comb -@handle_exceptions -def unset_min_denominator() -> None: - """ - Reset the global minimum denominator used by ivy for numerically stable division to - the previous value. +def _dnd_dict_intersection(a, b): + res = {} + for device in a: + if device in b: + intersection = set.intersection(set(a[device]), set(b[device])) + if intersection: + res[device] = tuple(intersection) + return res - Examples - -------- - >>> ivy.set_min_denominator(1e-10) - >>> y = ivy.min_denominator - >>> print(y) - 1e-10 - >>> ivy.unset_min_denominator() - >>> ivy.min_denominator - 1e-12 - """ - global min_denominator_stack - if min_denominator_stack: - min_denominator_stack.pop(-1) - val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 - ivy.__setattr__("min_denominator", val, True) +def _dnd_dict_difference(a, b): + res = a + for device in list(a): + if device in b: + difference = set.difference(set(a[device]), set(b[device])) + if difference: + res[device] = tuple(difference) + else: + del res[device] + return res -@handle_exceptions -def unset_nestable_mode() -> None: - """ - Reset the mode of whether to check if function inputs are ivy.Container to the - previous state. +def _dnd_dict_union(a, b): + res = {} + for device in set(list(a) + list(b)): + u1 = set(a.get(device, ())) + u2 = set(b.get(device, ())) + res[device] = tuple(set.union(u1, u2)) - Examples - -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False + return res - >>> ivy.unset_nestable_mode() - >>> ivy.nestable_mode - True - """ - global nestable_mode_stack - if nestable_mode_stack: - nestable_mode_stack.pop(-1) - mode = nestable_mode_stack[-1] if nestable_mode_stack else True - ivy.__setattr__("nestable_mode", mode, True) +def _get_devices_and_dtypes(fn, recurse=False, complement=True): + supported_devices = ivy.function_supported_devices(fn, recurse=recurse) + supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) -@handle_exceptions -def unset_precise_mode() -> None: - """ - Reset the mode of whether to use a promotion table that avoids any precision loss or - a compute effecient table that avoids most wider-than-necessary promotions. + if hasattr(fn, "partial_mixed_handler"): + supported_devices = supported_devices["primary"] + supported_dtypes = supported_dtypes["primary"] - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False + supported = {} + # Generate a base supported set from other attributes + for device in supported_devices: + supported[device] = supported_dtypes - >>> ivy.unset_precise_mode() - >>> ivy.precise_mode - True - """ - global precise_mode_stack - if precise_mode_stack: - precise_mode_stack.pop(-1) - mode = precise_mode_stack[-1] if precise_mode_stack else True - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported + backend = ivy.current_backend_str() -@handle_exceptions -def unset_queue_timeout() -> None: - """ - Reset the global queue timeout value (in seconds) to the previous state. + # Their values are formatted like either + # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype.__get__() - Examples - -------- - >>> ivy.set_queue_timeout(10.0) - >>> y = ivy.queue_timeout - >>> print(y) - 10.0 + if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): + fn_supported_dnd = fn_supported_dnd.get(backend, supported) - >>> ivy.unset_queue_timeout() - >>> ivy.queue_timeout - 15.0 - """ - global queue_timeout_stack - if queue_timeout_stack: - queue_timeout_stack.pop(-1) - timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 - ivy.__setattr__("queue_timeout", timeout, True) + ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) + # dict intersection + supported = _dnd_dict_intersection(supported, fn_supported_dnd) + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() -@handle_exceptions -def unset_shape_array_mode() -> None: - """ - Reset the mode of returning shape as ivy.Array to the previous state. + if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) - Examples - -------- - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True + ivy.utils.assertions.check_isinstance( + list(fn_unsupported_dnd.values())[0], tuple + ) + # dict difference + supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - >>> ivy.unset_shape_array_mode() - >>> ivy.shape_array_mode - False - """ - global shape_array_mode_stack - if shape_array_mode_stack: - shape_array_mode_stack.pop(-1) - mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False - ivy.__setattr__("shape_array_mode", mode, True) + if complement: + # dict difference + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported @handle_exceptions -def unset_show_func_wrapper_trace_mode() -> None: - """ - Reset the mode of whether to show the full stack trace with function wrapping - traces. - - Examples - -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False - - >>> ivy.unset_show_func_wrapper_trace_mode() - >>> ivy.show_func_wrapper_trace_mode - True +@handle_nestable +def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: """ - global show_func_wrapper_trace_mode_stack - if show_func_wrapper_trace_mode_stack: - show_func_wrapper_trace_mode_stack.pop(-1) - mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True - ) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + Return the supported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the supported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. + Parameters + ---------- + fn + The function to check for the supported device and dtype attribute + recurse + Whether to recurse into used ivy functions. + Default is ``True``. -@handle_exceptions -def unset_tmp_dir() -> None: + Returns + ------- + ret + Tuple or dict containing the supported devices and dtypes of the function """ - Reset the directory for saving temporary files to the previous value. + ivy.utils.assertions.check_true( + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", + ) - Examples - -------- - >>> ivy.set_tmp_dir("/my_dir") - >>> y = ivy.tmp_dir - >>> print(y) - /my_dir + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=False), + } + else: + supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) + if recurse: + supported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + supported_devices_dtypes, + _dnd_dict_intersection, + function_supported_devices_and_dtypes, + wrapper=lambda x: x, + ) - >>> ivy.unset_tmp_dir() - >>> ivy.tmp_dir - /tmp - """ - global tmp_dir_stack - if tmp_dir_stack: - tmp_dir_stack.pop(-1) - tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" - ivy.__setattr__("tmp_dir", tmp_dr, True) + return supported_devices_dtypes @handle_exceptions @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def value_is_nan( - x: Union[ivy.Array, ivy.NativeArray, Number], - /, - *, - include_infs: bool = True, -) -> bool: +def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: """ - Determine whether the single valued array or scalar is of nan type. + Return the unsupported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the unsupported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. Parameters ---------- - x - The input to check Input array. - include_infs - Whether to include infs and -infs in the check. + fn + The function to check for the unsupported device and dtype attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. Returns ------- ret - Boolean as to whether the input value is a nan or not. - - Examples - -------- - >>> x = ivy.array([451]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False - - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - True - - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - False - - >>> x = ivy.array([float('nan')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - True - - >>> x = ivy.array([0]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False + Tuple or dict containing the unsupported devices and dtypes of the function """ - x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x - if not x_scalar == x: - return True - if include_infs and (x_scalar == INF or x_scalar == -INF): - return True - return False + ivy.utils.assertions.check_true( + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=True), + } + else: + unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) + if recurse: + unsupported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + unsupported_devices_dtypes, + _dnd_dict_union, + function_unsupported_devices_and_dtypes, + wrapper=lambda x: x, + ) + return unsupported_devices_dtypes @handle_exceptions @@ -4187,48 +4115,142 @@ def vmap( return current_backend().vmap(func, in_axes, out_axes) -trace_mode_dict["frontend"] = "ivy/functional/frontends" -trace_mode_dict["ivy"] = "ivy/" -trace_mode_dict["full"] = "" -trace_mode_dict["none"] = "" -ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True -ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True -ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True -ivy.exception_trace_mode = ( - exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" -) -ivy.show_func_wrapper_trace_mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True -) -# IMPORTANT: assign attribute directly to function instead of wrapper here -einops_reduce.unsupported_dtypes = { - "torch": ("float16",), - "tensorflow": ("complex",), - "paddle": ("complex", "uint8", "int8", "int16", "float16"), -} -ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 -ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 -stable_pow.unsupported_dtypes = ("bfloat16",) -ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 -ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" -get_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -set_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} -inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} -ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" -ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +@handle_device_shifting +def isin( + elements: Union[ivy.Array, ivy.NativeArray], + test_elements: Union[ivy.Array, ivy.NativeArray], + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> ivy.Array: + """ + Test if each element of elements is in test_elements. + + Parameters + ---------- + elements + input array + test_elements + values against which to test for each input element + assume_unique + If True, assumes both elements and test_elements contain unique elements, + which can speed up the calculation. Default value is False. + invert + If True, inverts the boolean return array, resulting in True values for + elements not in test_elements. Default value is False. + + Returns + ------- + ret + output a boolean array of the same shape as elements that is True for elements + in test_elements and False otherwise. + + Examples + -------- + >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y) + ivy.array([[False, False, False], [ True, True, True]]) + + >>> x = ivy.array([3, 2, 1, 0]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y, invert=True) + ivy.array([False, False, False, True]) + """ + return ivy.current_backend(elements, test_elements).isin( + elements, test_elements, assume_unique=assume_unique, invert=invert + ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def itemsize( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> int: + """ + Return the size of the input array's elements. + + Parameters + ---------- + x + The input array. + + Returns + ------- + ret + An integer specifying the element size in bytes. + + Examples + -------- + >>> x = ivy.array([1,2,3], dtype=ivy.float64) + >>> ivy.itemsize(x) + 8 + + >>> x = ivy.array([1,2,3], dtype=ivy.complex128) + >>> ivy.itemsize(x) + 16 + """ + return ivy.current_backend(x).itemsize(x) + + +@handle_exceptions +@handle_nestable +@handle_device_shifting +def strides( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> Tuple[int]: + """ + Return the input array's strides across each dimension. + + Parameters + ---------- + x + The input array. + + Returns + ------- + ret + A tuple containing the strides. + + Examples + -------- + >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) + >>> ivy.strides(x) + (4, 8) + """ + if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): + return ivy.to_numpy(x).strides + # if x is an ivy array with a base, + # convert it to a numpy array with the same base: + ret = ivy.to_numpy(x.base) + ivy_numpy = ivy.with_backend("numpy") + for fn, args, kwargs, index in x._manipulation_stack: + ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) + ret = ret[index] if ivy.exists(index) else ret + return ret.to_native().strides + + +def is_ivy_nested_array(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Nested Array. + + Parameters + ---------- + x + The input to check + Returns + ------- + ret + Boolean, whether or not x is an ivy nested array. + """ + return isinstance(x, ivy.NestedArray) diff --git a/ivy/functional/ivy/gradients.py b/ivy/functional/ivy/gradients.py index 72e3961ffc7e9..fc0c9b2a1f411 100644 --- a/ivy/functional/ivy/gradients.py +++ b/ivy/functional/ivy/gradients.py @@ -22,8 +22,18 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # +# Helpers # +# ------- # + + +def _get_duplicate_index_chains(xs): + """Generate a list of duplicate index chains for a given nested structure.""" + duplicate_index_chains = () + if isinstance(xs, ivy.Container): + duplicate_index_chains = xs.cont_duplicate_array_keychains() + elif isinstance(xs, (list, tuple, dict)): + duplicate_index_chains = ivy.duplicate_array_index_chains(xs) + return duplicate_index_chains def _arrays_to_float_variables(xs, xs_grad_idxs=None): @@ -50,89 +60,6 @@ def inner_fn(x): return ivy.nested_map(xs, map_fn, include_derived=True, shallow=False) -def _get_duplicate_index_chains(xs): - """Generate a list of duplicate index chains for a given nested structure.""" - duplicate_index_chains = () - if isinstance(xs, ivy.Container): - duplicate_index_chains = xs.cont_duplicate_array_keychains() - elif isinstance(xs, (list, tuple, dict)): - duplicate_index_chains = ivy.duplicate_array_index_chains(xs) - return duplicate_index_chains - - -def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): - """Extract all relevant results from the output nested structure of a function.""" - - def map_fn(x_): - if ivy.is_array(x_): - x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ - if create_var: - x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ - if len(x_.shape) == 0: - return ivy.to_native(x_) - if reshape: - if x_.size == 1: - if reshape: - return ivy.to_native(ivy.reshape(x_, [])) - return ivy.to_native(x_) - else: - return ivy.to_ivy(x_) - else: - return ivy.to_native(x_) - return x_ - - if ivy.is_array(x): - return [], map_fn(x) - - x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) - arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) - if _check_if_empty(arr_idxs): - return arr_idxs, [] - else: - if idxs is not None: - arr_idxs = [ - arr_idx - for arr_idx in arr_idxs - if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) - ] - arr_values = ivy.multi_index_nest(x, arr_idxs) - arr_idxs = _idxs_to_str(arr_idxs) - return arr_idxs, arr_values - - -def _get_native_y(y): - """Convert all outputs to native arrays.""" - array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) - y_final = [] - if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: - y_final = ivy.multi_index_nest(y, array_idxs) - return y_final - - -def _get_required_float_variables(xs, xs_grad_idxs): - """ - Convert all required arrays to float variables for gradient calculation. - - Also, returns a list of duplicate index chains for the nested - structure. - """ - if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: - xs_grad_idxs = None - duplicate_index_chains = _get_duplicate_index_chains(xs) - xs = _to_ivy(xs) - xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) - xs = _set_duplicates(xs, duplicate_index_chains) - xs_required = _get_required_native_variables(xs, xs_grad_idxs) - required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) - return ( - xs, - xs_grad_idxs, - xs_required, - required_duplicate_index_chains, - duplicate_index_chains, - ) - - def _get_required_native_variables(xs, xs_grad_idxs): """Extract all required native variables from a nested structure.""" # To make sure that only the required arrays are converted to native arrays @@ -175,46 +102,68 @@ def map_fn(x): return xs -def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): - """Get the relevant outputs from the function return value.""" - if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ - [0] - ]: - ret_grad_idxs = None - ret_idxs, ret_values = _get_native_variables_and_indices( - func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape +def _get_required_float_variables(xs, xs_grad_idxs): + """ + Convert all required arrays to float variables for gradient calculation. + + Also, returns a list of duplicate index chains for the nested + structure. + """ + if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: + xs_grad_idxs = None + duplicate_index_chains = _get_duplicate_index_chains(xs) + xs = _to_ivy(xs) + xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) + xs = _set_duplicates(xs, duplicate_index_chains) + xs_required = _get_required_native_variables(xs, xs_grad_idxs) + required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) + return ( + xs, + xs_grad_idxs, + xs_required, + required_duplicate_index_chains, + duplicate_index_chains, ) - if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): - return func_ret, {} - if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: - y = ret_values[0] - else: - y = ret_values - return ret_grad_idxs, y, ret_idxs -def _is_variable(x, exclusive=False, to_ignore=None) -> bool: - x = ivy.to_native(x, nested=True, to_ignore=to_ignore) - return ivy.nested_map( - x, - lambda x: current_backend(x).is_variable(x, exclusive=exclusive), - include_derived=True, - shallow=False, - to_ignore=to_ignore, - ) +def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): + """Extract all relevant results from the output nested structure of a function.""" + def map_fn(x_): + if ivy.is_array(x_): + x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ + if create_var: + x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ + if len(x_.shape) == 0: + return ivy.to_native(x_) + if reshape: + if x_.size == 1: + if reshape: + return ivy.to_native(ivy.reshape(x_, [])) + return ivy.to_native(x_) + else: + return ivy.to_ivy(x_) + else: + return ivy.to_native(x_) + return x_ -def _process_func_ret_and_grads(func_ret, grads, retain_grads): - """ - Stop gradients propagation. + if ivy.is_array(x): + return [], map_fn(x) - Set the gradients of non-finite values to zero, and stopping - gradient propagation of the function results. - """ - grads = _non_finite_to_zero(grads) - func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) - grads = _to_ivy(grads) - return func_ret, grads + x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) + arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) + if _check_if_empty(arr_idxs): + return arr_idxs, [] + else: + if idxs is not None: + arr_idxs = [ + arr_idx + for arr_idx in arr_idxs + if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) + ] + arr_values = ivy.multi_index_nest(x, arr_idxs) + arr_idxs = _idxs_to_str(arr_idxs) + return arr_idxs, arr_values def _set_duplicates(xs, duplicate_index_chains): @@ -243,6 +192,33 @@ def _set_duplicates(xs, duplicate_index_chains): return xs +def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): + """Get the relevant outputs from the function return value.""" + if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ + [0] + ]: + ret_grad_idxs = None + ret_idxs, ret_values = _get_native_variables_and_indices( + func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape + ) + if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): + return func_ret, {} + if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: + y = ret_values[0] + else: + y = ret_values + return ret_grad_idxs, y, ret_idxs + + +def _get_native_y(y): + """Convert all outputs to native arrays.""" + array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) + y_final = [] + if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: + y_final = ivy.multi_index_nest(y, array_idxs) + return y_final + + def _stop_grad_and_index(func_ret, retain_grads, grads): """Stop gradient propagation of the function results.""" if not retain_grads: @@ -256,6 +232,46 @@ def _stop_grad_and_index(func_ret, retain_grads, grads): return func_ret, grads +def _process_func_ret_and_grads(func_ret, grads, retain_grads): + """ + Stop gradients propagation. + + Set the gradients of non-finite values to zero, and stopping + gradient propagation of the function results. + """ + grads = _non_finite_to_zero(grads) + func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) + grads = _to_ivy(grads) + return func_ret, grads + + +_check_if_empty = ( + lambda idxs: not isinstance(idxs, list) + or np.asarray(idxs, dtype="object").size == 0 +) + + +_idxs_to_str = lambda idxs: [ + "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) +] + + +_to_ivy = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) + + +_non_finite_to_zero = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) + + # Private Variable Helpers # # -------------------------# @@ -268,12 +284,23 @@ def _variable(x): return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) -def _variable_data( - x: Union[ivy.Array, ivy.NativeArray] -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Get the contents of the input. - +def _is_variable(x, exclusive=False, to_ignore=None) -> bool: + x = ivy.to_native(x, nested=True, to_ignore=to_ignore) + return ivy.nested_map( + x, + lambda x: current_backend(x).is_variable(x, exclusive=exclusive), + include_derived=True, + shallow=False, + to_ignore=to_ignore, + ) + + +def _variable_data( + x: Union[ivy.Array, ivy.NativeArray] +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Get the contents of the input. + Parameters ---------- x @@ -291,340 +318,101 @@ def _variable_data( return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) -# --- Main --- # -# ------------ # - - -# Optimizer Steps # - - @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def adam_step( - dcdw: Union[ivy.Array, ivy.NativeArray], - mw: Union[ivy.Array, ivy.NativeArray], - vw: Union[ivy.Array, ivy.NativeArray], - step: Union[int, float], +@handle_device_shifting +def stop_gradient( + x: Union[ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, + preserve_type: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Compute adam step delta, given the derivatives of some cost c with respect to - weights ws, using ADAM update. `[reference] - - `_ + Stop gradient computation. Parameters ---------- - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - mw - running average of the gradients - vw - running average of second moments of the gradients - step - training step - beta1 - gradient forgetting factor (Default value = 0.9) - beta2 - second moment of gradient forgetting factor (Default value = 0.999) - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7) + x + Array for which to stop the gradient. + preserve_type + Whether to preserve gradient computation on ivy.Array instances. Default is + True. out - optional output array, for writing the effective grad of adam_step to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The adam step delta. + The same array x, but with no gradient information. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` inputs: - >>> dcdw = ivy.array([1, 2, 3]) - >>> mw = ivy.ones(3) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) - >>> print(adam_step_delta) - (ivy.array([0.2020105 , 0.22187898, 0.24144873]), - ivy.array([0.99999998, 1.09999998, 1.19999998]), - ivy.array([1.00000001, 1.00300001, 1.00800001])) - - >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) - >>> mw = ivy.zeros((2,3)) - >>> vw = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.86 - >>> beta2 = 0.95 - >>> epsilon = 1e-6 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - (ivy.array([[ 1., 1., -1.], - [ 1., 1., 1.]]), - ivy.array([[ 0.14, 0.56, -0.42], - [ 0.28, 0.42, 0.07]]), - ivy.array([[0.05 , 0.8 , 0.45 ], - [0.2 , 0.45 , 0.0125]])) + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.stop_gradient(x, preserve_type=True) + >>> print(y) + ivy.array([1., 2., 3.]) - >>> dcdw = ivy.array([0.1, -0.7, 2]) - >>> mw = ivy.ones(1) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3.6) - >>> out = ivy.zeros_like(dcdw) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) - >>> print(out) - ivy.array([0.17294501, 0.15770318, 0.20863818]) + >>> x = ivy.zeros((2, 3)) + >>> ivy.stop_gradient(x, preserve_type=False, out=x) + >>> print(x) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) - With one :class:`ivy.Container` input: + With one :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.array([1., 4., 9.]) - >>> vw = ivy.array([0.,]) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), - b: ivy.array([2.02, 4.82, 8.17]) - }, { - a: ivy.array([0.87, 3.61, 8.09]), - b: ivy.array([1.26, 4., 8.48]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.stop_gradient(x, preserve_type=False) + >>> print(y) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> vw = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([0., 0.626, 0.626]), - b: ivy.array([0.626, 0.626, 0.626]) - }, { - a: ivy.array([0., 0.13, 0.26]), - b: ivy.array([0.39, 0.52, 0.65]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> ivy.stop_gradient(x, preserve_type=True, out=x) + >>> print(x) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } """ - step = float(step) - mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) - dcdw_sqrd = dcdw**2 - vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) - vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 - beta1_pow = beta1**step - beta2_pow = beta2**step - alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) - return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw + return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) + + +# AutoGrad # @handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def adam_update( - w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, +@handle_device_shifting +def execute_with_gradients( + func, + xs: Union[ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, - stop_gradients: bool = True, - out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: + retain_grads: bool = False, + xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], + ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], +) -> Tuple[ivy.Array, ivy.Array]: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, using ADAM update. `[reference] - - `_ - - Parameters - ---------- - w - Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. - out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. - - Returns - ------- - ret - The new function weights ws_new, and also new mw and vw, following the adam - updates. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = 1 - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(updated_weights) - (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), - ivy.array([0.05, 0.02, 0.01]), - ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) - - >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) - >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = 2 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, beta2=beta2, - ... epsilon=epsilon, out=out, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ( - ivy.array([[0.92558873, 1.92558754, 2.92558718], - [3.92558694, 1.92558682, 3.92558861], - [5.92558861, 3.92558694, 1.92558718]]), - ivy.array([[0.01, 0.02, 0.03], - [0.04, 0.05, 0.01], - [0.01, 0.05, 0.03]]), - ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], - [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], - [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) - ) - - With one :class:`ivy.Container` input: - - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([0.5, 0.2, 0.4]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(0.01) - >>> step = 2 - >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(updated_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) - - With multiple :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), - ... b=ivy.array([0.3,0.2,0.2])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = 3 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> stop_gradients = False - >>> lr = ivy.array(0.001) - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ({ - a: ivy.array([0.99936122, 1.99936116, 2.99936128]), - b: ivy.array([3.99936128, 4.99936104, 5.99936104]) - }, { - a: ivy.array([0.01, 0.03, 0.03]), - b: ivy.array([0.03, 0.02, 0.02]) - }, { - a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), - b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) - }) - """ - effective_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - return ( - ivy.optimizer_update( - w, effective_grads, lr, stop_gradients=stop_gradients, out=out - ), - mw, - vw, - ) - - -# AutoGrad # - - -@handle_exceptions -@handle_device_shifting -def execute_with_gradients( - func, - xs: Union[ivy.Array, ivy.NativeArray], - /, - *, - retain_grads: bool = False, - xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], - ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], -) -> Tuple[ivy.Array, ivy.Array]: - """ - Call function func with input of xs variables, and return the function result - func_ret and the gradients of each output variable w.r.t each input variable, + Call function func with input of xs variables, and return the function result + func_ret and the gradients of each output variable w.r.t each input variable, Parameters ---------- @@ -695,6 +483,77 @@ def execute_with_gradients( ) +execute_with_gradients.computes_gradients = True + + +@handle_exceptions +def value_and_grad(func: Callable) -> Callable: + """ + Create a function that evaluates both func and the gradient of func. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + + Returns + ------- + ret + A function that returns both func and the gradient of func. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> grad_fn = ivy.value_and_grad(func) + >>> value_grad = grad_fn(x) + >>> print(value_grad) + (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], + [0.93333334, 0.43333334, 2.0666666 ]])) + """ + return current_backend(None).value_and_grad(func) + + +value_and_grad.computes_gradients = True + + +@handle_exceptions +def jac(func: Callable) -> Callable: + """ + Call function func, and return func's Jacobian partial derivatives. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + + Returns + ------- + ret + the Jacobian function + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> jac_fn = ivy.jac(func) + >>> jacobian = jac_fn(x) + >>> print(jacobian) + ivy.array([[1.53 , 0.7 , 1.67 ], + ... [0.933, 0.433, 2.07 ]]) + """ + return current_backend(None).jac(func) + + +jac.computes_gradients = True + + @handle_exceptions def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: """ @@ -726,297 +585,380 @@ def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: return current_backend(None).grad(func, argnums=argnums) +grad.computes_gradients = True + + +# Optimizer Steps # + + @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def gradient_descent_update( - w: Union[ivy.Array, ivy.NativeArray], +def adam_step( dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], + mw: Union[ivy.Array, ivy.NativeArray], + vw: Union[ivy.Array, ivy.NativeArray], + step: Union[int, float], /, *, - stop_gradients: bool = True, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws]. + Compute adam step delta, given the derivatives of some cost c with respect to + weights ws, using ADAM update. `[reference] + + `_ Parameters ---------- - w - Weights of the function to be updated. dcdw Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. + mw + running average of the gradients + vw + running average of second moments of the gradients + step + training step + beta1 + gradient forgetting factor (Default value = 0.9) + beta2 + second moment of gradient forgetting factor (Default value = 0.999) + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7) out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the effective grad of adam_step to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The new weights, following the gradient descent updates. + The adam step delta. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([[1., 2, 3], - ... [4, 6, 1], - ... [1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1], - ... [0.3, 0.6, 0.4], - ... [0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) - >>> print(new_weights) - ivy.array([[ 0.95, 1.98, 2.99], - ... [ 3.97, 5.94, 0.96], - ... [ 0.96, -0.07, 6.98]]) + >>> dcdw = ivy.array([1, 2, 3]) + >>> mw = ivy.ones(3) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) + >>> print(adam_step_delta) + (ivy.array([0.2020105 , 0.22187898, 0.24144873]), + ivy.array([0.99999998, 1.09999998, 1.19999998]), + ivy.array([1.00000001, 1.00300001, 1.00800001])) - >>> w = ivy.array([1., 2., 3.]) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> out = ivy.zeros_like(w) - >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) + >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) + >>> mw = ivy.zeros((2,3)) + >>> vw = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.86 + >>> beta2 = 0.95 + >>> epsilon = 1e-6 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + (ivy.array([[ 1., 1., -1.], + [ 1., 1., 1.]]), + ivy.array([[ 0.14, 0.56, -0.42], + [ 0.28, 0.42, 0.07]]), + ivy.array([[0.05 , 0.8 , 0.45 ], + [0.2 , 0.45 , 0.0125]])) + + >>> dcdw = ivy.array([0.1, -0.7, 2]) + >>> mw = ivy.ones(1) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3.6) + >>> out = ivy.zeros_like(dcdw) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) >>> print(out) - ivy.array([0.85, 1.94, 2.97]) + ivy.array([0.17294501, 0.15770318, 0.20863818]) - With one :class:`ivy.Container` inputs: + With one :class:`ivy.Container` input: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([3.33, 5.66, 1.95]) - } + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.array([1., 4., 9.]) + >>> vw = ivy.array([0.,]) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), + b: ivy.array([2.02, 4.82, 8.17]) + }, { + a: ivy.array([0.87, 3.61, 8.09]), + b: ivy.array([1.26, 4., 8.48]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), - ... b=ivy.array([2., 3.42, 1.69])) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([2.88, 4.69, 1.47]) - } - """ - return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) - - -@handle_exceptions -def jac(func: Callable) -> Callable: + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> vw = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([0., 0.626, 0.626]), + b: ivy.array([0.626, 0.626, 0.626]) + }, { + a: ivy.array([0., 0.13, 0.26]), + b: ivy.array([0.39, 0.52, 0.65]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) """ - Call function func, and return func's Jacobian partial derivatives. + step = float(step) + mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) + dcdw_sqrd = dcdw**2 + vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) + vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 + beta1_pow = beta1**step + beta2_pow = beta2**step + alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) + return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - Returns - ------- - ret - the Jacobian function +adam_step.out_index = 0 - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> jac_fn = ivy.jac(func) - >>> jacobian = jac_fn(x) - >>> print(jacobian) - ivy.array([[1.53 , 0.7 , 1.67 ], - ... [0.933, 0.433, 2.07 ]]) - """ - return current_backend(None).jac(func) +# Optimizer Updates # @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def lamb_update( +def optimizer_update( w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], + effective_grad: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, - max_trust_ratio: Union[int, float] = 10, - decay_lambda: float = 0, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws], by applying LAMB method. + Update weights ws of some function, given the true or effective derivatives of some + cost c with respect to ws, [dc/dw for w in ws]. Parameters ---------- w Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + effective_grad + Effective gradients of the cost c with respect to the weights ws, + [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). - max_trust_ratio - The maximum value for the trust ratio. (Default value = 10) - decay_lambda - The factor used for weight decay. (Default value = 0). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the LAMB updates. + The new function weights ws_new, following the optimizer updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(new_weights) - (ivy.array([0.784, 1.78 , 2.78 ]), - ... ivy.array([0.05, 0.02, 0.01]), - ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... out=None, stop_gradients=True) + >>> print(ws_new) + ivy.array([1., 2., 3.]) + + >>> w = ivy.array([[1., 2.], [4., 5.]]) >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, out=out, - ... stop_gradients=stop_gradients) + >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) + >>> lr = ivy.array([3e-4, 1e-2]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) >>> print(out) - ivy.array([[ 0.639, 1.64 , 2.64 ], - ... [ 3.64 , 5.64 , 0.639], - ... [ 0.639, -0.361, 6.64 ]]) + ivy.array([[0.999, 1.95], + [4., 4.92]]) - With one :class:`ivy.Container` inputs: + >>> w = ivy.array([1., 2., 3.]) + >>> out = ivy.zeros_like(w) + >>> effective_grad = ivy.array([4., 5., 6.]) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False, out=out) + >>> print(out) + ivy.array([0.999, 2. , 3. ]) - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([3., 4., 5.]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(1.) - >>> step = ivy.array([2]) - >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(new_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) + With one :class:`ivy.Container` input: + + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.array([0., 0., 0.]) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), - ... b=ivy.array([3.,4.,2.])) - >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), - ... b=ivy.array([0.6,0.4,0.7])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) + >>> print(w) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 - >>> stop_gradients = True - >>> lr = ivy.array(0.5) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, - ... stop_gradients=stop_gradients) + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + """ + deltas = effective_grad * lr + w = ivy.subtract(w, deltas, out=out) + if stop_gradients: + return ivy.stop_gradient(w, preserve_type=True, out=out) + return w + + +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def gradient_descent_update( + w: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], + lr: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + stop_gradients: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws]. + + Parameters + ---------- + w + Weights of the function to be updated. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The new weights, following the gradient descent updates. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> w = ivy.array([[1., 2, 3], + ... [4, 6, 1], + ... [1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1], + ... [0.3, 0.6, 0.4], + ... [0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) >>> print(new_weights) - ({ - a: ivy.array([-0.708, 1.29, 3.29]), - b: ivy.array([1.45, 2.45, 0.445]) - }, { - a: ivy.array([0.02, 0.03, 0.06]), - b: ivy.array([0.06, 0.04, 0.07]) - }, { - a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), - b: ivy.array([0.00036, 0.00016, 0.00049]) - }) + ivy.array([[ 0.95, 1.98, 2.99], + ... [ 3.97, 5.94, 0.96], + ... [ 0.96, -0.07, 6.98]]) + + >>> w = ivy.array([1., 2., 3.]) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> out = ivy.zeros_like(w) + >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) + >>> print(out) + ivy.array([0.85, 1.94, 2.97]) + + With one :class:`ivy.Container` inputs: + + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([3.33, 5.66, 1.95]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), + ... b=ivy.array([2., 3.42, 1.69])) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([2.88, 4.69, 1.47]) + } """ - r1 = ivy.vector_norm(w) - eff_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - if decay_lambda > 0: - r2 = ivy.vector_norm(eff_grads + decay_lambda * w) - else: - r2 = ivy.vector_norm(eff_grads) - r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) - lr = r * lr - return ( - ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), - mw, - vw, - ) + return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) @handle_exceptions @@ -1070,264 +1012,338 @@ def lars_update( ) -# Optimizer Updates # - - @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def optimizer_update( +def adam_update( w: Union[ivy.Array, ivy.NativeArray], - effective_grad: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Update weights ws of some function, given the true or effective derivatives of some - cost c with respect to ws, [dc/dw for w in ws]. + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, using ADAM update. `[reference] + + `_ Parameters ---------- w Weights of the function to be updated. - effective_grad - Effective gradients of the cost c with respect to the weights ws, - [dc/dw for w in ws]. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the optimizer updates. + The new function weights ws_new, and also new mw and vw, following the adam + updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - ivy.array([1., 2., 3.]) - - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... out=None, stop_gradients=True) - >>> print(ws_new) - ivy.array([1., 2., 3.]) + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = 1 + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(updated_weights) + (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), + ivy.array([0.05, 0.02, 0.01]), + ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) - >>> w = ivy.array([[1., 2.], [4., 5.]]) - >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) - >>> lr = ivy.array([3e-4, 1e-2]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) - >>> print(out) - ivy.array([[0.999, 1.95], - [4., 4.92]]) - - >>> w = ivy.array([1., 2., 3.]) + >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) + >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = 2 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([4., 5., 6.]) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False, out=out) - >>> print(out) - ivy.array([0.999, 2. , 3. ]) + >>> stop_gradients = True + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, beta2=beta2, + ... epsilon=epsilon, out=out, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ( + ivy.array([[0.92558873, 1.92558754, 2.92558718], + [3.92558694, 1.92558682, 3.92558861], + [5.92558861, 3.92558694, 1.92558718]]), + ivy.array([[0.01, 0.02, 0.03], + [0.04, 0.05, 0.01], + [0.01, 0.05, 0.03]]), + ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], + [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], + [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) + ) With one :class:`ivy.Container` input: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.array([0., 0., 0.]) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([0.5, 0.2, 0.4]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(0.01) + >>> step = 2 + >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(updated_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) - >>> print(w) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), + ... b=ivy.array([0.3,0.2,0.2])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = 3 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> stop_gradients = False + >>> lr = ivy.array(0.001) + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ({ + a: ivy.array([0.99936122, 1.99936116, 2.99936128]), + b: ivy.array([3.99936128, 4.99936104, 5.99936104]) + }, { + a: ivy.array([0.01, 0.03, 0.03]), + b: ivy.array([0.03, 0.02, 0.02]) + }, { + a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), + b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) + }) """ - deltas = effective_grad * lr - w = ivy.subtract(w, deltas, out=out) - if stop_gradients: - return ivy.stop_gradient(w, preserve_type=True, out=out) - return w + effective_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + return ( + ivy.optimizer_update( + w, effective_grads, lr, stop_gradients=stop_gradients, out=out + ), + mw, + vw, + ) + + +adam_update.out_index = 0 @handle_exceptions -@handle_backend_invalid -@handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def stop_gradient( - x: Union[ivy.Array, ivy.NativeArray], +def lamb_update( + w: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], + lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, - preserve_type: bool = True, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, + max_trust_ratio: Union[int, float] = 10, + decay_lambda: float = 0, + stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Stop gradient computation. + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws], by applying LAMB method. Parameters ---------- - x - Array for which to stop the gradient. - preserve_type - Whether to preserve gradient computation on ivy.Array instances. Default is - True. + w + Weights of the function to be updated. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). + max_trust_ratio + The maximum value for the trust ratio. (Default value = 10) + decay_lambda + The factor used for weight decay. (Default value = 0). + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The same array x, but with no gradient information. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The new function weights ws_new, following the LAMB updates. Examples -------- With :class:`ivy.Array` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.stop_gradient(x, preserve_type=True) - >>> print(y) - ivy.array([1., 2., 3.]) + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(new_weights) + (ivy.array([0.784, 1.78 , 2.78 ]), + ... ivy.array([0.05, 0.02, 0.01]), + ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) - >>> x = ivy.zeros((2, 3)) - >>> ivy.stop_gradient(x, preserve_type=False, out=x) - >>> print(x) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) + >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 + >>> out = ivy.zeros_like(w) + >>> stop_gradients = True + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, out=out, + ... stop_gradients=stop_gradients) + >>> print(out) + ivy.array([[ 0.639, 1.64 , 2.64 ], + ... [ 3.64 , 5.64 , 0.639], + ... [ 0.639, -0.361, 6.64 ]]) With one :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.stop_gradient(x, preserve_type=False) - >>> print(y) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([3., 4., 5.]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(1.) + >>> step = ivy.array([2]) + >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(new_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> ivy.stop_gradient(x, preserve_type=True, out=x) - >>> print(x) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - """ - return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) - - -@handle_exceptions -def value_and_grad(func: Callable) -> Callable: - """ - Create a function that evaluates both func and the gradient of func. - - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - - Returns - ------- - ret - A function that returns both func and the gradient of func. - - Examples - -------- - With :class:`ivy.Array` input: + >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), + ... b=ivy.array([3.,4.,2.])) + >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), + ... b=ivy.array([0.6,0.4,0.7])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> grad_fn = ivy.value_and_grad(func) - >>> value_grad = grad_fn(x) - >>> print(value_grad) - (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], - [0.93333334, 0.43333334, 2.0666666 ]])) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 + >>> stop_gradients = True + >>> lr = ivy.array(0.5) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, + ... stop_gradients=stop_gradients) + >>> print(new_weights) + ({ + a: ivy.array([-0.708, 1.29, 3.29]), + b: ivy.array([1.45, 2.45, 0.445]) + }, { + a: ivy.array([0.02, 0.03, 0.06]), + b: ivy.array([0.06, 0.04, 0.07]) + }, { + a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), + b: ivy.array([0.00036, 0.00016, 0.00049]) + }) """ - return current_backend(None).value_and_grad(func) + r1 = ivy.vector_norm(w) + eff_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + if decay_lambda > 0: + r2 = ivy.vector_norm(eff_grads + decay_lambda * w) + else: + r2 = ivy.vector_norm(eff_grads) + r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) + lr = r * lr + return ( + ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), + mw, + vw, + ) -execute_with_gradients.computes_gradients = True -value_and_grad.computes_gradients = True -jac.computes_gradients = True -grad.computes_gradients = True -adam_step.out_index = 0 -adam_update.out_index = 0 lamb_update.out_index = 0 -_check_if_empty = ( - lambda idxs: not isinstance(idxs, list) - or np.asarray(idxs, dtype="object").size == 0 -) -_idxs_to_str = lambda idxs: [ - "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) -] -_to_ivy = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) -_non_finite_to_zero = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) diff --git a/ivy/functional/ivy/layers.py b/ivy/functional/ivy/layers.py index 1a239557e8b53..c8096d4d1fe26 100644 --- a/ivy/functional/ivy/layers.py +++ b/ivy/functional/ivy/layers.py @@ -20,111 +20,6 @@ ) from ivy.utils.exceptions import handle_exceptions - -# --- Helpers --- # -# --------------- # - - -def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): - kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) - if padding == "SAME": - dim_size = dim_size * stride_size - else: - dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) - return dim_size - - -def _depth_max_pooling_helper( - x_shape, kernel, strides, dims, data_format="channel_last" -): - # Determine depth pooling. - # We assume that the kernel and the data have the same data_format. - depth_pooling = False - CHANNEL_LAST = "channel_last" - channel_idx = -1 if data_format == CHANNEL_LAST else 1 - if len(kernel) == dims + 2: - spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] - if kernel[channel_idx] != 1: - depth_pooling = True - if any(i != 1 for i in spatial_kernel): - raise NotImplementedError( - "MaxPooling supports exactly one of pooling across" - " depth or pooling across width/height." - ) - if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to equal the depth" - " stride" - ) - if x_shape[channel_idx] % kernel[channel_idx] != 0: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to evenly divide" - " the input depth" - ) - kernel = [kernel[channel_idx], *[1] * (dims - 1)] - strides = [strides[channel_idx], *[1] * (dims - 1)] - else: - kernel = spatial_kernel - if len(strides) == dims + 2: - strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] - return kernel, strides, depth_pooling - - -def _get_num_padded_values(i, p, n, k, s): - """ - Get number of padded values in a specific window. - - Parameters - ---------- - i window index - p total amount of padding - n input size - k kernel size - s stride - - Returns - ------- - number of padded values in a particular window represented by i - """ - current_index = s * i - left_padding = p // 2 - return max(0, left_padding - current_index) + max( - 0, current_index + k - n - left_padding - ) - - -def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): - if dims == 1: - if data_format == "channel_first": - return "NCW" - else: - return "NWC" - if dims == 2: - if data_format == "channel_first": - return "NCHW" - else: - return "NHWC" - elif dims == 3: - if data_format == "channel_first": - return "NCDHW" - else: - return "NDHWC" - - -# Helpers # - - -def _handle_padding(x, strides, filters, padding): - if isinstance(padding, str) and padding.upper() == "SAME": - if x % strides == 0: - pad = max(filters - strides, 0) - else: - pad = max(filters - (x % strides), 0) - else: - pad = 0 - return pad - - # Extra # # ------# @@ -177,134 +72,34 @@ def _in_projection( ) -def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): - if isinstance(kernel, int): - kernel = (kernel,) * dims - elif len(kernel) == 1: - kernel = (kernel[0],) * dims - elif (len(kernel) != dims) and (len(kernel) != dims + 2): - raise ValueError( - "The kernel should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) - - if isinstance(strides, int): - strides = (strides,) * dims - elif len(strides) == 1: - strides = (strides[0],) * dims - elif (len(strides) != dims) and (len(strides) != dims + 2): - raise ValueError( - "The stride should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) - - if isinstance(padding, int): - padding = [(padding,) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == 1: - padding = [(padding[0],) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == dims: - padding = [(padding[i],) * 2 for i in range(dims)] - elif isinstance(padding, list) and len(padding) == dims: - if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): - raise ValueError("Explicit padding must be a list of tuple of two integers") - if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: - raise ValueError( - f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" - ) - - if isinstance(dilation, int): - dilation = (dilation,) * dims - elif len(dilation) == 1: - dilation = (dilation[0],) * dims - elif len(dilation) != dims: - raise ValueError( - f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" - ) - if min(dilation) < 1: - raise ValueError("All values of `dilation` must be positive") - - # Other errors - if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: - raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") - assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" - - # Account for dilation when padding > kernel/2. Not the case in torch by default. - new_kernel = tuple( - [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] - ) - if isinstance(padding, list) and len(padding) == len(new_kernel): - ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) - - return kernel, strides, padding, dilation - - -# --- Main --- # -# ------------ # - - +# Linear # @handle_exceptions +@handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes +@inputs_to_ivy_arrays @handle_array_function -def conv( +def linear( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], + weight: Union[ivy.Array, ivy.NativeArray], /, *, - transpose: bool = False, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D - input x respectively and filters arrays. + """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. + The operation also supports batching of the weight matrices. This is useful if a + batch of different network parameters are to be represented. Parameters ---------- x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - transpose - True for computing transpose convolution, and False for dilated convolution. - When True, `x_dilations` must be 1 (the default). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output (Default value = None) - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) + The input x to compute linear transformation on. + *[outer_batch_shape,inner_batch_shape,in_features]* + weight + The weight matrix. *[outer_batch_shape,out_features,in_features]* bias - Bias array of shape *[d_out]*. + The bias vector, default is ``None``. *[outer_batch_shape,out_features]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -312,92 +107,179 @@ def conv( Returns ------- ret - The result of the transpose or dilated convolution operation. - """ - if transpose: - return conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - else: - return conv_general_dilated( - x, - filters, - strides, - padding, - dims=dims, - data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, + Result array of the linear transformation. + *[outer_batch_shape,inner_batch_shape,out_features]* + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1., 2., 3.]) + >>> w = ivy.array([[1., 0., 0.]]) + >>> y = ivy.linear(x, w) + >>> print(y) + ivy.array([1.]) + + >>> x = ivy.array([[0.666, -0.4269, 1.911]]) + >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) + >>> y = ivy.zeros((1, 2)) + >>> ivy.linear(x, w, out=y) + >>> print(y) + ivy.array([[0.666, 1.91 ]]) + + >>> x = ivy.array([[1.546, 5.234, 6.487], + ... [0.157, 5.753, 4.52], + ... [5.165, 3.159, 7.101]]) + >>> w = ivy.array([[1.545, 2.547, 3.124], + ... [5.852, 8.753, 6.963]]) + >>> b = ivy.array([-1., 1.]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.linear(x, w, bias=b, out=y) + >>> print(y) + ivy.array([[ 34.98495483, 101.0293808 ], + [ 28.0159359 , 83.74752808], + [ 37.20942307, 108.3205719 ]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [4., 5., 6.]]), + ... b=ivy.array([1.1, 2.2, 3.3])) + >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [-1., 1., 2.]]), + ... b=ivy.array([[0., -1., 1.], + ... [0., 1., 1.]])) + >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) + >>> y = ivy.linear(x, w, bias=b) + >>> print(y) + { + a: ivy.array([[15., 6.], + [33., 12.]]), + b: ivy.array([2.1, 6.5]) + } + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], + ... [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], + ... [7., 13., 17.]])) + >>> w = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.]]) + >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), + ... b=ivy.array([1., 1., 0.])) + >>> ivy.linear(x, w, bias=b, out=x) + >>> print(x) + { + a: ivy.array([[16.4, 35.2, 54.], + [155., 352., 549.]]), + b: ivy.array([[15.1, 32., 47.9], + [85., 196., 306.]]) + } + + """ + outer_batch_shape = list(weight.shape[:-2]) + num_outer_batch_dims = len(outer_batch_shape) + inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) + num_inner_batch_dims = len(inner_batch_shape) + num_out_feats, num_in_feats = list(weight.shape[-2:]) + + # OBS x IBS x OF + y = ivy.matmul( + x, + ivy.swapaxes( + ivy.reshape( + weight, + outer_batch_shape + + [1] * max(num_inner_batch_dims - 1, 0) + + [num_out_feats, num_in_feats], + ), + -1, + -2, + ), + ) + + if ivy.exists(bias): + # OBS x [1]*len(IBS) x OF + bias_broadcast = ivy.reshape( + bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] ) + # OBS x IBS x OF + y = y + bias_broadcast -# Convolutions # + if ivy.exists(out): + return ivy.inplace_update(out, y) + return y + + +linear.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} + + +# Dropout # @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv1d( +def dropout( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + prob: float, /, *, - data_format: str = "NWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int]] = 1, - dilations: Union[int, Tuple[int]] = 1, - bias: Optional[ivy.Array] = None, + scale: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D convolution given 3-D input x and filters arrays. + Randomly setting a fraction of input tensor to zeroes with probability. + + `prob` at each update during training time to prevent possible overfitting. + The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that + overall sum is unchanged at training time and inference time. Parameters ---------- x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. - filters - Convolution filters *[fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]*. + The input array x to perform dropout on. + prob + The probability of zeroing out each array element, float between 0 and 1. + scale + Whether to scale the output by `1/(1-prob)`. Default is ``True``. + dtype + output array data type. If dtype is None, the output array data type + must be inferred from x. Default is ``None``. + training + Turn on dropout if training, turn off otherwise. Default is ``True``. + seed + Set a default seed for random number generating (for reproducibility). Default + is ``None``. + noise_shape + a sequence representing the shape of the binary dropout mask that will be + multiplied with the input. A shape dimension set to None means that a different + mask value will be applied to each element of the input across that dimension. A + dimension set to 1 means the same mask value will be applied to all elements of + the input across that dimension. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -405,7 +287,7 @@ def conv1d( Returns ------- ret - The result of the convolution operation. + Result array after dropout is performed. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -415,95 +297,175 @@ def conv1d( -------- With :class:`ivy.Array` input: - >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC - >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO - >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) - >>> print(result) - ivy.array([[[0.], [3.], [0.]]]) + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3) + >>> print(y) + ivy.array([[ 1.42857146, 2.85714293, 4.28571415], + [ 0. , 7.14285755, 8.5714283 ], + [10. , 11.4285717 , 0. ], + [14.2857151 , 0. , 17.1428566 ]]) - With :class:`ivy.NativeArray` input: - >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) - >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) - >>> result = ivy.conv1d(x, filters, (2,),'VALID') - >>> print(result) - ivy.array([[[3., 1.], - ... [7., 5.]]]) + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + ivy.array([[ 0. , 5.19999981], + [ 0. , 0. ], + [ 0. , 17.39999962]]) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3,scale=False) + >>> print(y) + ivy.array([[ 1., 2., 3.], + [ 4., 5., 0.], + [ 7., 0., 9.], + [10., 11., 0.]]) - >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], - ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), - ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) - >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) - >>> result = ivy.conv1d(x, filters, 3, 'VALID') - >>> print(result) - { - a: ivy.array([[[6., 7.9, 1.2], - ... [15.6, 11.7, 6.1]]]), - ... b: ivy.array([[[15.4, 14.3, 8.8]]]) - } - """ - return current_backend(x).conv1d( - x, - filters, - strides, - padding, - data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5,scale=False) + >>> print(y) + ivy.array([[0., 2.6], + [0., 0. ], + [0., 8.7]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) + { + a: ivy.array([[0., 0., 4.28571415], + [5.71428585, 7.14285755, 0.]]), + b: ivy.array([0., 11.4285717, 12.8571434]) + } + + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 4.4000001, 6.5999999], + [22., 44., 0.]]), + b: ivy.array([[2.49000001, 0.55599999, 8.21000004], + [14., 0., 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) + { + a: ivy.array([[0., 0., 3.], + [4., 5., 0.]]), + b: ivy.array([0., 8., 9.]) + } + + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 2.2, 3.3], + [11., 22., 0.]]), + b: ivy.array([[1.245, 0.278, 4.105], + [7., 0., 0.]]) + } + """ + if prob == 0 or not training: + if dtype is not None: + x = ivy.astype(x, dtype) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + if noise_shape is None: + noise_shape = x.shape + else: + noise_shape = list(noise_shape) + for i, v in enumerate(noise_shape): + if v is None: + noise_shape[i] = x.shape[i] + mask = ivy.where( + ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) + < prob, + 0.0, + 1.0, ) + x = x * mask + if scale: + x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + + +dropout.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} + + +# Attention # @handle_exceptions -@handle_backend_invalid -@handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def conv1d_transpose( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: str, +def scaled_dot_product_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Union[ivy.Array, ivy.NativeArray], + value: Union[ivy.Array, ivy.NativeArray], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NWC", - dilations: Union[int, Tuple[int]] = 1, - bias: Optional[ivy.Array] = None, + scale: Optional[float] = None, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + dropout_p: Optional[float] = 0.0, + is_causal: Optional[bool] = False, + training: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D transpose convolution given 3-D input x and filters arrays. + Apply scaled dot product attention to inputs x using optional mask. Parameters ---------- - x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. - filters - Convolution filters *[fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) - data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]*. + query + The queries input array. The shape of queries input array should be in + *[batch_shape,num_queries,feat_dim]*. The queries input array should have the + same size as keys and values. + key + The keys input array. The shape of keys input array should be in + *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same + size as queries and values. + value + The values input array. The shape of values input should be in + *[batch_shape,num_keys,feat_dim]*. The values input array should have the same + size as queries and keys. + scale + The scale float value. + The scale float value is used to scale the query-key pairs before softmax. + mask + The mask input array. The mask to apply to the query-key values. Default is + None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. + dropout_p + Specifies the dropout probablity, if greater than 0.0, dropout is applied + is_causal + If true, assumes causal attention masking + and errors if both `mask` and `is_causal` are set. + training + If True, dropout is used, otherwise dropout is not activated. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -511,7 +473,9 @@ def conv1d_transpose( Returns ------- ret - The result of the transpose convolution operation. + The output following application of scaled dot-product attention. + The output array is the weighted sum produced by the attention score and value. + The shape of output array is *[batch_shape,num_queries,feat_dim]* . Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -521,227 +485,403 @@ def conv1d_transpose( -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 6) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 64) + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) - >>> y = ivy.zeros((1, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 32) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - With :class:`ivy.NativeArray` input: + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - >>> x = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) - >>> filters = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 512, 32) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) - With one :class:`ivy.Container` input: + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - >>> x = ivy.full((1, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) - >>> print(y.shape) - { - a: ivy.Shape(1, 10, 1), - b: ivy.Shape(1, 10, 1) - } + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - With multiple :class:`ivy.Container` inputs: - - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - { - a: { - c: ivy.Shape(1, 28, 3), - d: ivy.Shape(1, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 3), - d: ivy.Shape(1, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - } - } - """ - return current_backend(x).conv1d_transpose( - x, - filters, - strides, - padding, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - bias=bias, - out=out, - ) + ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def conv2d( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int]] = 1, - dilations: Union[int, Tuple[int, int]] = 1, - bias: Optional[ivy.Array] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 2-D convolution given 4-D input x and filters arrays. + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - Parameters - ---------- - x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. - filters - Convolution filters *[fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", - input data formats, while "channel_last" corresponds to "HWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]*. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + With :class:`ivy.Container` input: - Returns - ------- - ret - The result of the convolution operation. + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) + { + a: ivy.array([[[5.19999981, 1.], + ... [2.59249449, 2.68226194], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[0.2, 1.], + ... [2.19603825, 2.9960382], + ... [4.4000001, 5.5999999]]]) + } - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> mask = ivy.Container( + ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), + ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) + ... ) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) + { + a: ivy.array([[[4.26894283, 5.40236187], + ... [4.39999437, 5.59999037], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[4.35046196, 5.54282808], + ... [4.39989519, 5.5998764], + ... [4.4000001, 5.5999999]]]) + } - Examples - -------- - With :class:`ivy.Array` input: + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - >>> x = ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]]) - >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], - ... [[[0.]],[[1.]], [[0.]]], - ... [[[0.]],[[1.]], [[0.]]]]) - >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) >>> print(result) - ivy.array([[ - [[2.],[4.],[6.]], - [[3.],[6.],[9.]], - [[2.],[4.],[6.]] - ]]) - With one :class:`ivy.Container` input: + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]])) - >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) - >>> print(result) - { - a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) - } + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) + >>> print(out) + ivy.array([[[4.03946018, 5.0280633 ], + ... [4.29981947, 5.29981089], + ... [4.30000019, 5.30000019]]]) - With multiple :class:`ivy.Container` inputs: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), - ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[1, 1, 1], - ... [0, 1, 1], - ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) >>> print(result) { - a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), - b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), - c:ivy.array([[[[2.], [0.], [0.]], - [[1.], [3.], [0.]], - [[0.], [1.], [2.]] - ]]) + a: ivy.array([[[0.40000001, 1.29999995], + ... [2.06345534, 2.9634552], + ... [4.30000019, 5.30000019]]]), + b: ivy.array([[[0.40000001, 1.29999995], + ... [2.19336844, 3.09336829], + ... [4.30000019, 5.30000019]]]) } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[2, 0, 1], - ... [1, 3, 1], - ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... mask=mask, + ... dropout_p=0.1, + ... training=True) >>> print(result) { - a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), - b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) + a: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]), + b: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) } """ - return current_backend(x).conv2d( - x, - filters, - strides, - padding, - data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, + ivy.assertions.check_all( + (not is_causal) or (is_causal and mask is None), + "is_causal and attn_mask cannot be set at the same time", + ) + embed_dim = query.shape[-1] + scale = 1 / (embed_dim**0.5) if not scale else scale + sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale + sim = ivy.dropout(sim, dropout_p, training=training) + if ivy.exists(mask): + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, + ) + elif is_causal: + L = query.shape[-2] # Source sequence length + S = key.shape[-2] # Target sequence length + mask = ivy.tril(ivy.ones((L, S)), k=0) + mask = ivy.astype(mask, ivy.bool) + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, + ) + attn = ivy.softmax(sim, axis=-1) + result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) + return result if not ivy.exists(out) else ivy.inplace_update(out, result) + + +@handle_exceptions +@handle_nestable +@handle_out_argument +# @handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def multi_head_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + /, + *, + num_heads: Optional[int] = 8, + scale: Optional[float] = None, + attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + is_causal: Optional[bool] = False, + return_attention_weights: Optional[bool] = False, + average_attention_weights: Optional[bool] = True, + dropout: Optional[float] = 0.0, + training: Optional[bool] = False, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Apply multi-head attention to inputs x. This is an implementation of multi-headed + attention as described in the paper "Attention is all you Need" (Vaswani et al., + 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each + timestep in `query` attends to the corresponding sequence in `key`, and returns a + fixed-width vector. This layer first projects `query`, `key` and `value`. These are + (effectively) a list of tensors of length `num_attention_heads`, where the + corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, + , key_dim)`, `(batch_size, , + value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are + softmaxed to obtain attention probabilities. The value tensors are then interpolated + by these probabilities, then concatenated back to a single tensor. Finally, the + result tensor with the last dimension as value_dim can take an linear projection and + return. + + Parameters + ---------- + query + query embeddings *[batch_shape,num_queries,query_dim]*. + key + key embeddings *[batch_shape,num_queries,key_dim]*. + value + value embeddings *[batch_shape,num_queries,value_dim]*. + num_heads + The number of attention heads to use. + scale + The value by which to scale the query-key similarity measure before softmax. + attention_mask + The mask to apply to the query-key values. Default is ``None``. + *[batch_shape,num_queries,num_keys]*. + in_proj_weights + The weights used to project query, key and value *[3*E, E]. + q_proj_weights + The weights used to project query if in_proj_weights is None *[new_E, E]. + k_proj_weights + The weights used to project key if in_proj_weights is None *[new_E, E]. + v_proj_weights + The weights used to project value if in_proj_weights is None *[new_E, E]. + out_proj_weights + The weights used to project the output. + in_proj_bias + The bias used when projecting with query, key and value. + out_proj_bias + The bias used when projecting the output. + is_causal + If True, Uses a causal attention mask and ignores provided attention_mask. + return_attention_weights + If True, returns attention_weights alongside the output + as a tuple (output, attenion_weights). Defaults to `False`. + average_attention_weights + If true, indicates that the returned ``attention_weights`` should be averaged + across heads. Otherwise, ``attention_weights`` are provided separately per head. + Note that this flag only has an effect when ``return_attention_weights=True``. + Default: ``True`` (i.e. average weights across heads) + dropout + Specifies the dropout probablity, dropout is applied to attention_weights. + training + If True, dropout is used, otherwise dropout is not activated. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The output following application of multi-head attention. + *[batch_shape,num_queries,out_feat_dim]* if input is batched + otherwise *[num_queries, out_feat_dim] + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + num_dims = query.ndim + ivy.assertions.check_all( + num_dims > 1 and num_dims < 4, + "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" + f" input), got {num_dims}", + ) + if key is None and value is None: + key = value = query + if num_dims == 2: + query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] + if ivy.exists(in_proj_weights): + q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) + elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): + if ivy.exists(in_proj_bias): + b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) + else: + b_q = b_k = b_v = None + q, k, v = ( + ivy.linear(query, q_proj_weights, bias=b_q), + ivy.linear(key, k_proj_weights, bias=b_k), + ivy.linear(value, v_proj_weights, bias=b_v), + ) + else: + q, k, v = query, key, value + batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] + k_seq_length = k.shape[1] + ivy.assertions.check_true( + emb_dim % num_heads == 0, "features must be divisible by number of heads" + ) + dims_per_head = emb_dim // num_heads + # isolate heads + q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 3, 1) + ) + v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + # perform bmm + attn_scores = ivy.matmul(q, k) + # scale + scale = 1 / (dims_per_head**0.5) if not scale else scale + attn_scores *= scale + # apply attention mask + if ivy.exists(attention_mask) or is_causal: + if is_causal: + # create causal mask + attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) + attention_mask = attention_mask.astype("bool") + attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) + # perform softmax + attn_weights = ivy.softmax(attn_scores, axis=-1) + # perform dropout + attn_weights = ivy.dropout(attn_weights, dropout, training=training) + # bmm with values + attention_out = ivy.matmul(attn_weights, v) + attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( + (batch_size, q_seq_length, -1) ) + # proj out if out_proj_weight exists + if ivy.exists(out_proj_weights): + attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) + # if input was unbatched, unbatchify the output + if num_dims == 2: + attention_out = attention_out.squeeze(axis=0) + if return_attention_weights: + if average_attention_weights: + attn_weights = attn_weights.mean(axis=1) + if num_dims == 2: + attn_weights = attn_weights.squeeze(axis=0) + return attention_out, attn_weights + else: + return attention_out + + +# Convolutions # @handle_exceptions @@ -749,47 +889,46 @@ def conv2d( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv2d_transpose( +def conv1d( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: str, + strides: Union[int, Tuple[int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, + data_format: str = "NWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int]] = 1, + dilations: Union[int, Tuple[int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D transpose convolution given 4-D input x and filters arrays. + Compute a 1-D convolution given 3-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. filters - Convolution filters *[fh,fw,d_in,d_out]*. + Convolution filters *[fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to - "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . - x_dilations + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". + x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations The dilation factor for each dimension of input. (Default value = 1) @@ -802,7 +941,7 @@ def conv2d_transpose( Returns ------- ret - The result of the transpose convolution operation. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -811,72 +950,44 @@ def conv2d_transpose( Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 56, 6) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) - >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 128, 64) + >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC + >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO + >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) + >>> print(result) + ivy.array([[[0.], [3.], [0.]]]) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) - >>> y = ivy.zeros((1, 258, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) - >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 258, 32) + With :class:`ivy.NativeArray` input: - With one :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) - >>> print(y.shape) - { - a: ivy.Shape(1, 10, 10, 1), - b: ivy.Shape(1, 10, 10, 1) - } + >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) + >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) + >>> result = ivy.conv1d(x, filters, (2,),'VALID') + >>> print(result) + ivy.array([[[3., 1.], + ... [7., 5.]]]) - With multiple :class:`ivy.Container` inputs: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], + ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), + ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) + >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) + >>> result = ivy.conv1d(x, filters, 3, 'VALID') + >>> print(result) { - a: { - c: ivy.Shape(1, 28, 28, 3), - d: ivy.Shape(1, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 56, 3), - d: ivy.Shape(1, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - } + a: ivy.array([[[6., 7.9, 1.2], + ... [15.6, 11.7, 6.1]]]), + ... b: ivy.array([[[15.4, 14.3, 8.8]]]) } """ - return current_backend(x).conv2d_transpose( + return current_backend(x).conv1d( x, filters, strides, padding, - output_shape=output_shape, data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -888,52 +999,47 @@ def conv2d_transpose( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv3d( - x: Union[ivy.Array, ivy.NativeArray, ivy.Container], - filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], - strides: Union[int, Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], +def conv1d_transpose( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int]], + padding: str, /, *, - data_format: str = "NDHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int, int, int]] = 1, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NWC", + dilations: Union[int, Tuple[int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 3-D convolution given 5-D input x and filters arrays. + Compute a 1-D transpose convolution given 3-D input x and filters arrays. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. + Convolution filters *[fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds - to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). dilations The dilation factor for each dimension of input. (Default value = 1) bias - Bias array of shape *[d_out]* + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -941,7 +1047,7 @@ def conv3d( Returns ------- ret - The result of the convolution operation. + The result of the transpose convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -951,106 +1057,139 @@ def conv3d( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) - >>> filters = ivy.array([[[0.,1.,0.], - ... [0.,1.,0.], - ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) - >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) - >>> print(result) - ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 6) + + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 64) + + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) + >>> y = ivy.zeros((1, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 32) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) + >>> filters = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 512, 32) With one :class:`ivy.Container` input: - >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) - >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.full((1, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) + >>> print(y.shape) { - a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) + a: ivy.Shape(1, 10, 1), + b: ivy.Shape(1, 10, 1) } - With multiple :class:`ivy.Container` input: + With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 3, 5, 5, 1]), - ... b = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 5, 32 ,32, 1]), - ... c = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 32, 32, 32, 1])) - >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 1, 'SAME') - >>> print(result.cont_shapes) + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a: ivy.Shape(1, 3, 5, 5, 3), - b: ivy.Shape(1, 5, 32, 32, 3), - c: ivy.Shape(1, 32, 32, 32, 3) + a: { + c: ivy.Shape(1, 28, 3), + d: ivy.Shape(1, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 3), + d: ivy.Shape(1, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + } } """ - return current_backend(x).conv3d( + return current_backend(x).conv1d_transpose( x, filters, strides, padding, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, ) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv3d_transpose( +def conv2d( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int, int]], - padding: str, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NDHWC", - dilations: Union[int, Tuple[int, int, int]] = 1, + data_format: str = "NHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int]] = 1, + dilations: Union[int, Tuple[int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 3-D transpose convolution given 5-D input x and filters arrays. + Compute a 2-D convolution given 4-D input x and filters arrays. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. + Convolution filters *[fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]* + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", + input data formats, while "channel_last" corresponds to "HWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1058,85 +1197,83 @@ def conv3d_transpose( Returns ------- ret - The result of the transpose convolution operation. + The result of the convolution operation. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 6, 56, 56, 6) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) - >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') - >>> print(y.shape) - ivy.Shape(1, 9, 258, 258, 32) + >>> x = ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]]) + >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], + ... [[[0.]],[[1.]], [[0.]]], + ... [[[0.]],[[1.]], [[0.]]]]) + >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) + >>> print(result) + ivy.array([[ + [[2.],[4.],[6.]], + [[3.],[6.],[9.]], + [[2.],[4.],[6.]] + ]]) - With :class:`ivy.Container` inputs: + With one :class:`ivy.Container` input: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) + >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]])) + >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) + >>> print(result) { - a: { - c: ivy.Shape(1, 6, 28, 28, 3), - d: ivy.Shape(1, 6, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 6, 56, 56, 3), - d: ivy.Shape(1, 6, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - } + a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With multiple :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), + ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[1, 1, 1], + ... [0, 1, 1], + ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') + >>> print(result) { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) + a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), + b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), + c:ivy.array([[[[2.], [0.], [0.]], + [[1.], [3.], [0.]], + [[0.], [1.], [2.]] + ]]) } + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) - >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[2, 0, 1], + ... [1, 3, 1], + ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') + >>> print(result) { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) + a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), + b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) } """ - return current_backend(x).conv3d_transpose( + return current_backend(x).conv2d( x, filters, strides, padding, - output_shape=output_shape, data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1148,58 +1285,50 @@ def conv3d_transpose( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv_general_dilated( +def conv2d_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int]], + padding: str, /, *, - dims: int = 2, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively - and filters arrays. + Compute a 2-D transpose convolution given 4-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. + Convolution filters *[fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) + Either "channel_first" or "channel_last". "channel_first" corresponds to + "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations - The dilation factor for each dimension of filter. (Default value = 1) + The dilation factor for each dimension of input. (Default value = 1) bias Bias array of shape *[d_out]*. out @@ -1210,17 +1339,80 @@ def conv_general_dilated( ------- ret The result of the transpose convolution operation. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 56, 6) + + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) + >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 128, 64) + + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) + >>> y = ivy.zeros((1, 258, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) + >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 258, 32) + + With one :class:`ivy.Container` inputs: + >>> x = ivy.full((1, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) + >>> print(y.shape) + { + a: ivy.Shape(1, 10, 10, 1), + b: ivy.Shape(1, 10, 10, 1) + } + + With multiple :class:`ivy.Container` inputs: + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) + { + a: { + c: ivy.Shape(1, 28, 28, 3), + d: ivy.Shape(1, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 56, 3), + d: ivy.Shape(1, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + } + } """ - return current_backend(x).conv_general_dilated( + return current_backend(x).conv2d_transpose( x, filters, strides, padding, - dims=dims, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1232,91 +1424,14 @@ def conv_general_dilated( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv_general_transpose( +def depthwise_conv2d( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: str, - /, - *, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - feature_group_count: int = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x - respectively and filters arrays. - - Parameters - ---------- - x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output. - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - dilations - The dilation factor for each dimension of input. (Default value = 1) - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - bias - Bias array of shape *[d_out]*. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The result of the transpose convolution operation. - """ - return current_backend(x).conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def depthwise_conv2d( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], /, *, data_format: str = "NHWC", @@ -1443,56 +1558,57 @@ def depthwise_conv2d( ) -# Dropout # - - @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def dropout( - x: Union[ivy.Array, ivy.NativeArray], - prob: float, +@handle_device_shifting +def conv3d( + x: Union[ivy.Array, ivy.NativeArray, ivy.Container], + filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], + strides: Union[int, Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - scale: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, + data_format: str = "NDHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int, int, int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly setting a fraction of input tensor to zeroes with probability. - - `prob` at each update during training time to prevent possible overfitting. - The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that - overall sum is unchanged at training time and inference time. + Compute a 3-D convolution given 5-D input x and filters arrays. Parameters ---------- x - The input array x to perform dropout on. - prob - The probability of zeroing out each array element, float between 0 and 1. - scale - Whether to scale the output by `1/(1-prob)`. Default is ``True``. - dtype - output array data type. If dtype is None, the output array data type - must be inferred from x. Default is ``None``. - training - Turn on dropout if training, turn off otherwise. Default is ``True``. - seed - Set a default seed for random number generating (for reproducibility). Default - is ``None``. - noise_shape - a sequence representing the shape of the binary dropout mask that will be - multiplied with the input. A shape dimension set to None means that a different - mask value will be applied to each element of the input across that dimension. A - dimension set to 1 means the same mask value will be applied to all elements of - the input across that dimension. + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + data_format + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds + to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1500,7 +1616,7 @@ def dropout( Returns ------- ret - Result array after dropout is performed. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1510,142 +1626,106 @@ def dropout( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - ivy.array([[ 1.42857146, 2.85714293, 4.28571415], - [ 0. , 7.14285755, 8.5714283 ], - [10. , 11.4285717 , 0. ], - [14.2857151 , 0. , 17.1428566 ]]) - - - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - ivy.array([[ 0. , 5.19999981], - [ 0. , 0. ], - [ 0. , 17.39999962]]) - - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3,scale=False) - >>> print(y) - ivy.array([[ 1., 2., 3.], - [ 4., 5., 0.], - [ 7., 0., 9.], - [10., 11., 0.]]) - - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5,scale=False) - >>> print(y) - ivy.array([[0., 2.6], - [0., 0. ], - [0., 8.7]]) - - With :class:`ivy.Container` input: + >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) + >>> filters = ivy.array([[[0.,1.,0.], + ... [0.,1.,0.], + ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) + >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) + >>> print(result) + ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - { - a: ivy.array([[0., 0., 4.28571415], - [5.71428585, 7.14285755, 0.]]), - b: ivy.array([0., 11.4285717, 12.8571434]) - } + With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) + >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) + >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 2, 'SAME') + >>> print(result) { - a: ivy.array([[0., 4.4000001, 6.5999999], - [22., 44., 0.]]), - b: ivy.array([[2.49000001, 0.55599999, 8.21000004], - [14., 0., 0.]]) + a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) } - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - { - a: ivy.array([[0., 0., 3.], - [4., 5., 0.]]), - b: ivy.array([0., 8., 9.]) - } + With multiple :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) + >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 3, 5, 5, 1]), + ... b = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 5, 32 ,32, 1]), + ... c = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 32, 32, 32, 1])) + >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 1, 'SAME') + >>> print(result.cont_shapes) { - a: ivy.array([[0., 2.2, 3.3], - [11., 22., 0.]]), - b: ivy.array([[1.245, 0.278, 4.105], - [7., 0., 0.]]) + a: ivy.Shape(1, 3, 5, 5, 3), + b: ivy.Shape(1, 5, 32, 32, 3), + c: ivy.Shape(1, 32, 32, 32, 3) } """ - if prob == 0 or not training: - if dtype is not None: - x = ivy.astype(x, dtype) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) - if noise_shape is None: - noise_shape = x.shape - else: - noise_shape = list(noise_shape) - for i, v in enumerate(noise_shape): - if v is None: - noise_shape[i] = x.shape[i] - mask = ivy.where( - ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) - < prob, - 0.0, - 1.0, + return current_backend(x).conv3d( + x, + filters, + strides, + padding, + data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - x = x * mask - if scale: - x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) -# Linear # @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back @handle_array_function -def linear( +@handle_device_shifting +def conv3d_transpose( x: Union[ivy.Array, ivy.NativeArray], - weight: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int, int, int]], + padding: str, /, *, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NDHWC", + dilations: Union[int, Tuple[int, int, int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. - The operation also supports batching of the weight matrices. This is useful if a - batch of different network parameters are to be represented. + """ + Compute a 3-D transpose convolution given 5-D input x and filters arrays. Parameters ---------- x - The input x to compute linear transformation on. - *[outer_batch_shape,inner_batch_shape,in_features]* - weight - The weight matrix. *[outer_batch_shape,out_features,in_features]* + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) + data_format + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). + dilations + The dilation factor for each dimension of input. (Default value = 1) bias - The bias vector, default is ``None``. *[outer_batch_shape,out_features]* + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1653,303 +1733,228 @@ def linear( Returns ------- ret - Result array of the linear transformation. - *[outer_batch_shape,inner_batch_shape,out_features]* - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The result of the transpose convolution operation. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 0., 0.]]) - >>> y = ivy.linear(x, w) - >>> print(y) - ivy.array([1.]) - - >>> x = ivy.array([[0.666, -0.4269, 1.911]]) - >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) - >>> y = ivy.zeros((1, 2)) - >>> ivy.linear(x, w, out=y) - >>> print(y) - ivy.array([[0.666, 1.91 ]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 6, 56, 56, 6) - >>> x = ivy.array([[1.546, 5.234, 6.487], - ... [0.157, 5.753, 4.52], - ... [5.165, 3.159, 7.101]]) - >>> w = ivy.array([[1.545, 2.547, 3.124], - ... [5.852, 8.753, 6.963]]) - >>> b = ivy.array([-1., 1.]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.linear(x, w, bias=b, out=y) - >>> print(y) - ivy.array([[ 34.98495483, 101.0293808 ], - [ 28.0159359 , 83.74752808], - [ 37.20942307, 108.3205719 ]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) + >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') + >>> print(y.shape) + ivy.Shape(1, 9, 258, 258, 32) - With :class:`ivy.Container` input: + With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [4., 5., 6.]]), - ... b=ivy.array([1.1, 2.2, 3.3])) - >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [-1., 1., 2.]]), - ... b=ivy.array([[0., -1., 1.], - ... [0., 1., 1.]])) - >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) - >>> y = ivy.linear(x, w, bias=b) - >>> print(y) + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a: ivy.array([[15., 6.], - [33., 12.]]), - b: ivy.array([2.1, 6.5]) + a: { + c: ivy.Shape(1, 6, 28, 28, 3), + d: ivy.Shape(1, 6, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 6, 56, 56, 3), + d: ivy.Shape(1, 6, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + } } With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], - ... [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], - ... [7., 13., 17.]])) - >>> w = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.]]) - >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), - ... b=ivy.array([1., 1., 0.])) - >>> ivy.linear(x, w, bias=b, out=x) - >>> print(x) + >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) { - a: ivy.array([[16.4, 35.2, 54.], - [155., 352., 549.]]), - b: ivy.array([[15.1, 32., 47.9], - [85., 196., 306.]]) + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) } - """ - outer_batch_shape = list(weight.shape[:-2]) - num_outer_batch_dims = len(outer_batch_shape) - inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) - num_inner_batch_dims = len(inner_batch_shape) - num_out_feats, num_in_feats = list(weight.shape[-2:]) - # OBS x IBS x OF - y = ivy.matmul( + >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) + >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) + { + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) + } + """ + return current_backend(x).conv3d_transpose( x, - ivy.swapaxes( - ivy.reshape( - weight, - outer_batch_shape - + [1] * max(num_inner_batch_dims - 1, 0) - + [num_out_feats, num_in_feats], - ), - -1, - -2, - ), + filters, + strides, + padding, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + bias=bias, + out=out, ) - if ivy.exists(bias): - # OBS x [1]*len(IBS) x OF - bias_broadcast = ivy.reshape( - bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] - ) - - # OBS x IBS x OF - y = y + bias_broadcast - - if ivy.exists(out): - return ivy.inplace_update(out, y) - return y - - -# LSTM # - @handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def lstm_update( +@handle_device_shifting +def conv_general_dilated( x: Union[ivy.Array, ivy.NativeArray], - init_h: Union[ivy.Array, ivy.NativeArray], - init_c: Union[ivy.Array, ivy.NativeArray], - kernel: Union[ivy.Array, ivy.NativeArray], - recurrent_kernel: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, + dims: int = 2, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Perform long-short term memory update by unrolling time dimension of input array. + Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively + and filters arrays. Parameters ---------- x - input tensor of LSTM layer *[batch_shape, t, in]*. - init_h - initial state tensor for the cell output *[batch_shape, out]*. - init_c - initial state tensor for the cell hidden state *[batch_shape, out]*. - kernel - weights for cell kernel *[in, 4 x out]*. - recurrent_kernel - weights for cell recurrent kernel *[out, 4 x out]*. + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of filter. (Default value = 1) bias - bias for cell kernel *[4 x out]*. (Default value = None) - recurrent_bias - bias for cell recurrent kernel *[4 x out]*. (Default value = None) + Bias array of shape *[d_out]*. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - hidden state for all timesteps *[batch_shape,t,out]* and cell state for last - timestep *[batch_shape,out]* + The result of the transpose convolution operation. """ - # get shapes - x_shape = list(x.shape) - batch_shape = x_shape[:-2] - timesteps = x_shape[-2] - input_channels = x_shape[-1] - x_flat = ivy.reshape(x, (-1, input_channels)) - - # input kernel - Wi = kernel - Wi_x = ivy.reshape( - ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), - batch_shape + [timesteps, -1], + return current_backend(x).conv_general_dilated( + x, + filters, + strides, + padding, + dims=dims, + data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) - - # recurrent kernel - Wh = recurrent_kernel - - # lstm states - ht = init_h - ct = init_c - - # lstm outputs - hts_list = list() - - # unrolled time dimension with lstm steps - for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( - ivy.unstack(Wii_x, axis=-2), - ivy.unstack(Wif_x, axis=-2), - ivy.unstack(Wig_x, axis=-2), - ivy.unstack(Wio_x, axis=-2), - ): - htm1 = ht - ctm1 = ct - - Wh_htm1 = ivy.matmul(htm1, Wh) + ( - recurrent_bias if recurrent_bias is not None else 0 - ) - Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( - Wh_htm1, num_or_size_splits=4, axis=-1 - ) - - it = ivy.sigmoid(Wii_xt + Whi_htm1) - ft = ivy.sigmoid(Wif_xt + Whf_htm1) - gt = ivy.tanh(Wig_xt + Whg_htm1) - ot = ivy.sigmoid(Wio_xt + Who_htm1) - ct = ft * ctm1 + it * gt - ht = ot * ivy.tanh(ct) - - hts_list.append(ivy.expand_dims(ht, axis=-2)) - - return ivy.concat(hts_list, axis=-2), ct @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument -# @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@inputs_to_native_shapes +@to_native_arrays_and_back @handle_array_function -def multi_head_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +@handle_device_shifting +def conv_general_transpose( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: str, /, *, - num_heads: Optional[int] = 8, - scale: Optional[float] = None, - attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - is_causal: Optional[bool] = False, - return_attention_weights: Optional[bool] = False, - average_attention_weights: Optional[bool] = True, - dropout: Optional[float] = 0.0, - training: Optional[bool] = False, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + feature_group_count: int = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Apply multi-head attention to inputs x. This is an implementation of multi-headed - attention as described in the paper "Attention is all you Need" (Vaswani et al., - 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each - timestep in `query` attends to the corresponding sequence in `key`, and returns a - fixed-width vector. This layer first projects `query`, `key` and `value`. These are - (effectively) a list of tensors of length `num_attention_heads`, where the - corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, - , key_dim)`, `(batch_size, , - value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are - softmaxed to obtain attention probabilities. The value tensors are then interpolated - by these probabilities, then concatenated back to a single tensor. Finally, the - result tensor with the last dimension as value_dim can take an linear projection and - return. + Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x + respectively and filters arrays. Parameters ---------- - query - query embeddings *[batch_shape,num_queries,query_dim]*. - key - key embeddings *[batch_shape,num_queries,key_dim]*. - value - value embeddings *[batch_shape,num_queries,value_dim]*. - num_heads - The number of attention heads to use. - scale - The value by which to scale the query-key similarity measure before softmax. - attention_mask - The mask to apply to the query-key values. Default is ``None``. - *[batch_shape,num_queries,num_keys]*. - in_proj_weights - The weights used to project query, key and value *[3*E, E]. - q_proj_weights - The weights used to project query if in_proj_weights is None *[new_E, E]. - k_proj_weights - The weights used to project key if in_proj_weights is None *[new_E, E]. - v_proj_weights - The weights used to project value if in_proj_weights is None *[new_E, E]. - out_proj_weights - The weights used to project the output. - in_proj_bias - The bias used when projecting with query, key and value. - out_proj_bias - The bias used when projecting the output. - is_causal - If True, Uses a causal attention mask and ignores provided attention_mask. - return_attention_weights - If True, returns attention_weights alongside the output - as a tuple (output, attenion_weights). Defaults to `False`. - average_attention_weights - If true, indicates that the returned ``attention_weights`` should be averaged - across heads. Otherwise, ``attention_weights`` are provided separately per head. - Note that this flag only has an effect when ``return_attention_weights=True``. - Default: ``True`` (i.e. average weights across heads) - dropout - Specifies the dropout probablity, dropout is applied to attention_weights. - training - If True, dropout is used, otherwise dropout is not activated. + x + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output. + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + dilations + The dilation factor for each dimension of input. (Default value = 1) + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1957,140 +1962,86 @@ def multi_head_attention( Returns ------- ret - The output following application of multi-head attention. - *[batch_shape,num_queries,out_feat_dim]* if input is batched - otherwise *[num_queries, out_feat_dim] - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The result of the transpose convolution operation. """ - num_dims = query.ndim - ivy.assertions.check_all( - num_dims > 1 and num_dims < 4, - "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" - f" input), got {num_dims}", - ) - if key is None and value is None: - key = value = query - if num_dims == 2: - query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] - if ivy.exists(in_proj_weights): - q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) - elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): - if ivy.exists(in_proj_bias): - b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) - else: - b_q = b_k = b_v = None - q, k, v = ( - ivy.linear(query, q_proj_weights, bias=b_q), - ivy.linear(key, k_proj_weights, bias=b_k), - ivy.linear(value, v_proj_weights, bias=b_v), - ) - else: - q, k, v = query, key, value - batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] - k_seq_length = k.shape[1] - ivy.assertions.check_true( - emb_dim % num_heads == 0, "features must be divisible by number of heads" - ) - dims_per_head = emb_dim // num_heads - # isolate heads - q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 3, 1) - ) - v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - # perform bmm - attn_scores = ivy.matmul(q, k) - # scale - scale = 1 / (dims_per_head**0.5) if not scale else scale - attn_scores *= scale - # apply attention mask - if ivy.exists(attention_mask) or is_causal: - if is_causal: - # create causal mask - attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) - attention_mask = attention_mask.astype("bool") - attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) - # perform softmax - attn_weights = ivy.softmax(attn_scores, axis=-1) - # perform dropout - attn_weights = ivy.dropout(attn_weights, dropout, training=training) - # bmm with values - attention_out = ivy.matmul(attn_weights, v) - attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( - (batch_size, q_seq_length, -1) + return current_backend(x).conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, ) - # proj out if out_proj_weight exists - if ivy.exists(out_proj_weights): - attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) - # if input was unbatched, unbatchify the output - if num_dims == 2: - attention_out = attention_out.squeeze(axis=0) - if return_attention_weights: - if average_attention_weights: - attn_weights = attn_weights.mean(axis=1) - if num_dims == 2: - attn_weights = attn_weights.squeeze(axis=0) - return attention_out, attn_weights - else: - return attention_out - - -# Attention # @handle_exceptions @handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes @handle_array_function -def scaled_dot_product_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Union[ivy.Array, ivy.NativeArray], - value: Union[ivy.Array, ivy.NativeArray], +def conv( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], /, *, - scale: Optional[float] = None, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - dropout_p: Optional[float] = 0.0, - is_causal: Optional[bool] = False, - training: Optional[bool] = False, + transpose: bool = False, + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply scaled dot product attention to inputs x using optional mask. + Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D + input x respectively and filters arrays. Parameters ---------- - query - The queries input array. The shape of queries input array should be in - *[batch_shape,num_queries,feat_dim]*. The queries input array should have the - same size as keys and values. - key - The keys input array. The shape of keys input array should be in - *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same - size as queries and values. - value - The values input array. The shape of values input should be in - *[batch_shape,num_keys,feat_dim]*. The values input array should have the same - size as queries and keys. - scale - The scale float value. - The scale float value is used to scale the query-key pairs before softmax. - mask - The mask input array. The mask to apply to the query-key values. Default is - None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. - dropout_p - Specifies the dropout probablity, if greater than 0.0, dropout is applied - is_causal - If true, assumes causal attention masking - and errors if both `mask` and `is_causal` are set. - training - If True, dropout is used, otherwise dropout is not activated. + x + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + transpose + True for computing transpose convolution, and False for dilated convolution. + When True, `x_dilations` must be 1 (the default). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output (Default value = None) + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -2098,252 +2049,294 @@ def scaled_dot_product_attention( Returns ------- ret - The output following application of scaled dot-product attention. - The output array is the weighted sum produced by the attention score and value. - The shape of output array is *[batch_shape,num_queries,feat_dim]* . - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: + The result of the transpose or dilated convolution operation. + """ + if transpose: + return conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, + ) + else: + return conv_general_dilated( + x, + filters, + strides, + padding, + dims=dims, + data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, + ) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) +# LSTM # - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def lstm_update( + x: Union[ivy.Array, ivy.NativeArray], + init_h: Union[ivy.Array, ivy.NativeArray], + init_c: Union[ivy.Array, ivy.NativeArray], + kernel: Union[ivy.Array, ivy.NativeArray], + recurrent_kernel: Union[ivy.Array, ivy.NativeArray], + /, + *, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[ivy.Array, ivy.Array]: + """ + Perform long-short term memory update by unrolling time dimension of input array. - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + Parameters + ---------- + x + input tensor of LSTM layer *[batch_shape, t, in]*. + init_h + initial state tensor for the cell output *[batch_shape, out]*. + init_c + initial state tensor for the cell hidden state *[batch_shape, out]*. + kernel + weights for cell kernel *[in, 4 x out]*. + recurrent_kernel + weights for cell recurrent kernel *[out, 4 x out]*. + bias + bias for cell kernel *[4 x out]*. (Default value = None) + recurrent_bias + bias for cell recurrent kernel *[4 x out]*. (Default value = None) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + Returns + ------- + ret + hidden state for all timesteps *[batch_shape,t,out]* and cell state for last + timestep *[batch_shape,out]* + """ + # get shapes + x_shape = list(x.shape) + batch_shape = x_shape[:-2] + timesteps = x_shape[-2] + input_channels = x_shape[-1] + x_flat = ivy.reshape(x, (-1, input_channels)) - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) + # input kernel + Wi = kernel + Wi_x = ivy.reshape( + ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), + batch_shape + [timesteps, -1], + ) + Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) - ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) + # recurrent kernel + Wh = recurrent_kernel - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + # lstm states + ht = init_h + ct = init_c - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + # lstm outputs + hts_list = list() - With :class:`ivy.Container` input: + # unrolled time dimension with lstm steps + for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( + ivy.unstack(Wii_x, axis=-2), + ivy.unstack(Wif_x, axis=-2), + ivy.unstack(Wig_x, axis=-2), + ivy.unstack(Wio_x, axis=-2), + ): + htm1 = ht + ctm1 = ct - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) - { - a: ivy.array([[[5.19999981, 1.], - ... [2.59249449, 2.68226194], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[0.2, 1.], - ... [2.19603825, 2.9960382], - ... [4.4000001, 5.5999999]]]) - } + Wh_htm1 = ivy.matmul(htm1, Wh) + ( + recurrent_bias if recurrent_bias is not None else 0 + ) + Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( + Wh_htm1, num_or_size_splits=4, axis=-1 + ) - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> mask = ivy.Container( - ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), - ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) - ... ) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) - { - a: ivy.array([[[4.26894283, 5.40236187], - ... [4.39999437, 5.59999037], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[4.35046196, 5.54282808], - ... [4.39989519, 5.5998764], - ... [4.4000001, 5.5999999]]]) - } + it = ivy.sigmoid(Wii_xt + Whi_htm1) + ft = ivy.sigmoid(Wif_xt + Whf_htm1) + gt = ivy.tanh(Wig_xt + Whg_htm1) + ot = ivy.sigmoid(Wio_xt + Who_htm1) + ct = ft * ctm1 + it * gt + ht = ot * ivy.tanh(ct) - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + hts_list.append(ivy.expand_dims(ht, axis=-2)) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) + return ivy.concat(hts_list, axis=-2), ct - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) - >>> print(out) - ivy.array([[[4.03946018, 5.0280633 ], - ... [4.29981947, 5.29981089], - ... [4.30000019, 5.30000019]]]) +# Helpers # - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) - >>> print(result) - { - a: ivy.array([[[0.40000001, 1.29999995], - ... [2.06345534, 2.9634552], - ... [4.30000019, 5.30000019]]]), - b: ivy.array([[[0.40000001, 1.29999995], - ... [2.19336844, 3.09336829], - ... [4.30000019, 5.30000019]]]) - } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: +def _handle_padding(x, strides, filters, padding): + if isinstance(padding, str) and padding.upper() == "SAME": + if x % strides == 0: + pad = max(filters - strides, 0) + else: + pad = max(filters - (x % strides), 0) + else: + pad = 0 + return pad - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... mask=mask, - ... dropout_p=0.1, - ... training=True) - >>> print(result) - { - a: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]), - b: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) - } - """ - ivy.assertions.check_all( - (not is_causal) or (is_causal and mask is None), - "is_causal and attn_mask cannot be set at the same time", - ) - embed_dim = query.shape[-1] - scale = 1 / (embed_dim**0.5) if not scale else scale - sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale - sim = ivy.dropout(sim, dropout_p, training=training) - if ivy.exists(mask): - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, + +def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): + if isinstance(kernel, int): + kernel = (kernel,) * dims + elif len(kernel) == 1: + kernel = (kernel[0],) * dims + elif (len(kernel) != dims) and (len(kernel) != dims + 2): + raise ValueError( + "The kernel should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" ) - elif is_causal: - L = query.shape[-2] # Source sequence length - S = key.shape[-2] # Target sequence length - mask = ivy.tril(ivy.ones((L, S)), k=0) - mask = ivy.astype(mask, ivy.bool) - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, + + if isinstance(strides, int): + strides = (strides,) * dims + elif len(strides) == 1: + strides = (strides[0],) * dims + elif (len(strides) != dims) and (len(strides) != dims + 2): + raise ValueError( + "The stride should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" ) - attn = ivy.softmax(sim, axis=-1) - result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) - return result if not ivy.exists(out) else ivy.inplace_update(out, result) + if isinstance(padding, int): + padding = [(padding,) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == 1: + padding = [(padding[0],) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == dims: + padding = [(padding[i],) * 2 for i in range(dims)] + elif isinstance(padding, list) and len(padding) == dims: + if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): + raise ValueError("Explicit padding must be a list of tuple of two integers") + if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: + raise ValueError( + f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" + ) -linear.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -dropout.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} + if isinstance(dilation, int): + dilation = (dilation,) * dims + elif len(dilation) == 1: + dilation = (dilation[0],) * dims + elif len(dilation) != dims: + raise ValueError( + f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" + ) + if min(dilation) < 1: + raise ValueError("All values of `dilation` must be positive") + + # Other errors + if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: + raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") + assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" + + # Account for dilation when padding > kernel/2. Not the case in torch by default. + new_kernel = tuple( + [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] + ) + if isinstance(padding, list) and len(padding) == len(new_kernel): + ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) + + return kernel, strides, padding, dilation + + +def _depth_max_pooling_helper( + x_shape, kernel, strides, dims, data_format="channel_last" +): + # Determine depth pooling. + # We assume that the kernel and the data have the same data_format. + depth_pooling = False + CHANNEL_LAST = "channel_last" + channel_idx = -1 if data_format == CHANNEL_LAST else 1 + if len(kernel) == dims + 2: + spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] + if kernel[channel_idx] != 1: + depth_pooling = True + if any(i != 1 for i in spatial_kernel): + raise NotImplementedError( + "MaxPooling supports exactly one of pooling across" + " depth or pooling across width/height." + ) + if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to equal the depth" + " stride" + ) + if x_shape[channel_idx] % kernel[channel_idx] != 0: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to evenly divide" + " the input depth" + ) + kernel = [kernel[channel_idx], *[1] * (dims - 1)] + strides = [strides[channel_idx], *[1] * (dims - 1)] + else: + kernel = spatial_kernel + if len(strides) == dims + 2: + strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] + return kernel, strides, depth_pooling + + +def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): + kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) + if padding == "SAME": + dim_size = dim_size * stride_size + else: + dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) + return dim_size + + +def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): + if dims == 1: + if data_format == "channel_first": + return "NCW" + else: + return "NWC" + if dims == 2: + if data_format == "channel_first": + return "NCHW" + else: + return "NHWC" + elif dims == 3: + if data_format == "channel_first": + return "NCDHW" + else: + return "NDHWC" + + +def _get_num_padded_values(i, p, n, k, s): + """ + Get number of padded values in a specific window. + + Parameters + ---------- + i window index + p total amount of padding + n input size + k kernel size + s stride + + Returns + ------- + number of padded values in a particular window represented by i + """ + current_index = s * i + left_padding = p // 2 + return max(0, left_padding - current_index) + max( + 0, current_index + k - n - left_padding + ) diff --git a/ivy/functional/ivy/linear_algebra.py b/ivy/functional/ivy/linear_algebra.py index 8190c0d4b57d5..612315f0b987a 100644 --- a/ivy/functional/ivy/linear_algebra.py +++ b/ivy/functional/ivy/linear_algebra.py @@ -345,95 +345,6 @@ def det( return current_backend(x).det(x, out=out) -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def diag( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the specified diagonals of the input array, or an array with the input - array's elements as diagonals. - - Parameters - ---------- - x - An array with rank >= 1. - k - An integer that controls which diagonal to consider. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - If x is a 1-D array, the function returns a 2-D square array with the elements - of input as diagonals. - If x is a 2-D array, the function returns a 1-D array with the diagonal elements - of x. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x) - ivy.array([0, 4, 8]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=1) - ivy.array([1, 5]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=-1) - ivy.array([3, 7]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(ivy.diag(x)) - ivy.array([[0, 0, 0], - [0, 4, 0], - [0, 0, 8]]) - """ - return current_backend(x).diag(x, k=k, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -1018,41 +929,6 @@ def inv( return current_backend(x).inv(x, adjoint=adjoint, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lu_factor( - A: Union[ivy.Array, ivy.NativeArray], - /, - *, - pivot: bool = True, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Parameters - ---------- - A - tensor of shape (*, m, n) where * is zero or more batch dimensions. - - pivot - Whether to compute the LU decomposition with partial pivoting, or the regular LU - decomposition. pivot = False not supported on CPU. Default: True. - - out - tuple of two tensors to write the output to. Ignored if None. Default: None. - - Returns - ------- - ret - A named tuple (LU, pivots). - """ - return current_backend(A).lu_factor(A, pivot=pivot, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2439,51 +2315,6 @@ def tensordot( return current_backend(x1, x2).tensordot(x1, x2, axes=axes, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def tensorsolve( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - ndim1 = ivy.get_num_dims(x1) - ndim2 = ivy.get_num_dims(x2) - - if axes is not None: - allaxes = list(range(0, ndim1)) - for k in axes: - allaxes.remove(k) - allaxes.insert(ndim1, k) - - x1 = ivy.matrix_transpose(x1, allaxes) - - old_shape = x1.shape[-(ndim1 - ndim2) :] - - prod = 1 - for k in old_shape: - prod *= k - - if ivy.shape(ivy.flatten(x1))[0] != prod**2: - raise ivy.utils.exceptions.IvyException( - "Input arrays must satisfy the requirement " - "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" - ) - - x1 = ivy.reshape(x1, (prod, prod)) - x2 = ivy.flatten(x2) - res = ivy.solve(x1, x2) - res = ivy.reshape(res, old_shape) - return res - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2612,80 +2443,6 @@ def trace( return current_backend(x).trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def vander( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Generate a Vandermonde matrix. The columns of the output matrix are elementwise - powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, - the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a - geometric progression in each row is named for Alexandre-Theophile Vandermonde. - - Parameters - ---------- - x - 1-D input array. - N - Number of columns in the output. If N is not specified, - a square array is returned (N = len(x)) - increasing - Order of the powers of the columns. If True, the powers increase - from left to right, if False (the default) they are reversed. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Vandermonde matrix. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x) - ivy.array( - [[ 1, 1, 1, 1], - [ 8, 4, 2, 1], - [ 27, 9, 3, 1], - [125, 25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3) - ivy.array( - [[ 1, 1, 1], - [ 4, 2, 1], - [ 9, 3, 1], - [25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3, increasing=True) - ivy.array( - [[ 1, 1, 1], - [ 1, 2, 4], - [ 1, 3, 9], - [ 1, 5, 25]] - ) - """ - return current_backend(x).vander(x, N=N, increasing=increasing, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2900,6 +2657,169 @@ def vector_norm( ) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def diag( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the specified diagonals of the input array, or an array with the input + array's elements as diagonals. + + Parameters + ---------- + x + An array with rank >= 1. + k + An integer that controls which diagonal to consider. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + If x is a 1-D array, the function returns a 2-D square array with the elements + of input as diagonals. + If x is a 2-D array, the function returns a 1-D array with the diagonal elements + of x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x) + ivy.array([0, 4, 8]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=1) + ivy.array([1, 5]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=-1) + ivy.array([3, 7]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(ivy.diag(x)) + ivy.array([[0, 0, 0], + [0, 4, 0], + [0, 0, 8]]) + """ + return current_backend(x).diag(x, k=k, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def vander( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a Vandermonde matrix. The columns of the output matrix are elementwise + powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, + the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a + geometric progression in each row is named for Alexandre-Theophile Vandermonde. + + Parameters + ---------- + x + 1-D input array. + N + Number of columns in the output. If N is not specified, + a square array is returned (N = len(x)) + increasing + Order of the powers of the columns. If True, the powers increase + from left to right, if False (the default) they are reversed. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Vandermonde matrix. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x) + ivy.array( + [[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3) + ivy.array( + [[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3, increasing=True) + ivy.array( + [[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]] + ) + """ + return current_backend(x).vander(x, N=N, increasing=increasing, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2934,3 +2854,84 @@ def vector_to_skew_symmetric_matrix( instances in place of any of the arguments. """ return current_backend(vector).vector_to_skew_symmetric_matrix(vector, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lu_factor( + A: Union[ivy.Array, ivy.NativeArray], + /, + *, + pivot: bool = True, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Parameters + ---------- + A + tensor of shape (*, m, n) where * is zero or more batch dimensions. + + pivot + Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. pivot = False not supported on CPU. Default: True. + + out + tuple of two tensors to write the output to. Ignored if None. Default: None. + + Returns + ------- + ret + A named tuple (LU, pivots). + """ + return current_backend(A).lu_factor(A, pivot=pivot, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def tensorsolve( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = 2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + ndim1 = ivy.get_num_dims(x1) + ndim2 = ivy.get_num_dims(x2) + + if axes is not None: + allaxes = list(range(0, ndim1)) + for k in axes: + allaxes.remove(k) + allaxes.insert(ndim1, k) + + x1 = ivy.matrix_transpose(x1, allaxes) + + old_shape = x1.shape[-(ndim1 - ndim2) :] + + prod = 1 + for k in old_shape: + prod *= k + + if ivy.shape(ivy.flatten(x1))[0] != prod**2: + raise ivy.utils.exceptions.IvyException( + "Input arrays must satisfy the requirement " + "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" + ) + + x1 = ivy.reshape(x1, (prod, prod)) + x2 = ivy.flatten(x2) + res = ivy.solve(x1, x2) + res = ivy.reshape(res, old_shape) + return res + # return current_backend(x1, x2).tensorsolve(x1, x2, axes=axes, out=out) diff --git a/ivy/functional/ivy/losses.py b/ivy/functional/ivy/losses.py index d803eb6fca839..d866cd8bc2487 100644 --- a/ivy/functional/ivy/losses.py +++ b/ivy/functional/ivy/losses.py @@ -12,8 +12,8 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # +# Helpers # +# ------- # def _reduce_loss(red, loss, axis, out): @@ -25,8 +25,64 @@ def _reduce_loss(red, loss, axis, out): return ivy.negative(loss, out=out) -# --- Main --- # -# ------------ # +# Extra # +# ------# + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def cross_entropy( + true: Union[ivy.Array, ivy.NativeArray], + pred: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + epsilon: float = 1e-7, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute cross-entropy between predicted and true discrete distributions. + + Parameters + ---------- + true + input array containing true labels. + pred + input array containing the predicted labels. + axis + the axis along which to compute the cross-entropy. If axis is ``-1``, + the cross-entropy will be computed along the last dimension. Default: ``-1``. + epsilon + a float in [0.0, 1.0] specifying the amount of smoothing when calculating + the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + The cross-entropy loss between the given distributions + + Examples + -------- + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.cross_entropy(x, y)) + ivy.array(1.3862944) + + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.cross_entropy(x, z)) + ivy.array(0.35667497) + """ + ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) + pred = ivy.clip(pred, epsilon, 1 - epsilon) + log_pred = ivy.log(pred) + return _reduce_loss(reduction, log_pred * true, axis, out) @handle_exceptions @@ -212,66 +268,6 @@ def binary_cross_entropy( return _reduce_loss(reduction, loss, axis, out) -# Extra # -# ------# - - -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def cross_entropy( - true: Union[ivy.Array, ivy.NativeArray], - pred: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - epsilon: float = 1e-7, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute cross-entropy between predicted and true discrete distributions. - - Parameters - ---------- - true - input array containing true labels. - pred - input array containing the predicted labels. - axis - the axis along which to compute the cross-entropy. If axis is ``-1``, - the cross-entropy will be computed along the last dimension. Default: ``-1``. - epsilon - a float in [0.0, 1.0] specifying the amount of smoothing when calculating - the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - The cross-entropy loss between the given distributions - - Examples - -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.cross_entropy(x, y)) - ivy.array(1.3862944) - - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.cross_entropy(x, z)) - ivy.array(0.35667497) - """ - ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) - pred = ivy.clip(pred, epsilon, 1 - epsilon) - log_pred = ivy.log(pred) - return _reduce_loss(reduction, log_pred * true, axis, out) - - @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/manipulation.py b/ivy/functional/ivy/manipulation.py index af74e3523eb5f..71d5e65b61c89 100644 --- a/ivy/functional/ivy/manipulation.py +++ b/ivy/functional/ivy/manipulation.py @@ -20,10 +20,6 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # - - def _calculate_out_shape(axis, array_shape): if type(axis) not in (tuple, list): axis = (axis,) @@ -37,141 +33,6 @@ def _calculate_out_shape(axis, array_shape): return out_shape -# --- Main --- # -# ------------ # - - -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def clip( - x: Union[ivy.Array, ivy.NativeArray], - x_min: Union[Number, ivy.Array, ivy.NativeArray], - x_max: Union[Number, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Clips (limits) the values in an array. - - Given an interval, values outside the interval are clipped to the interval edges - (element-wise). For example, if an interval of [0, 1] is specified, values smaller - than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller - or equal to maximum value to return correct results. - - Parameters - ---------- - x - Input array containing elements to clip. - x_min - Minimum value. - x_max - Maximum value. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - An array with the elements of x, but where values < x_min are replaced with - x_min, and those > x_max with x_max. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.zeros_like(x) - >>> ivy.clip(x, 2., 7., out=y) - >>> print(y) - ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) - >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - { - a: ivy.array([1., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> x_min = ivy.Container(a=0, b=-3) - >>> x_max = ivy.Container(a=1, b=-1) - >>> y = ivy.clip(x, x_min,x_max) - >>> print(y) - { - a: ivy.array([0., 1., 1.]), - b: ivy.array([-1., -1., -1.]) - } - - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 0., 1]) - >>> x_max = ivy.array([5., 4., 3.]) - >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> z = ivy.clip(y, x_min, x_max) - >>> print(z) - { - a: ivy.array([3., 1., 2.]), - b: ivy.array([3., 4., 3.]) - } - """ - return current_backend(x).clip(x, x_min, x_max, out=out) - - # Array API Standard # # -------------------# @@ -236,103 +97,6 @@ def concat( return current_backend(xs[0]).concat(xs, axis=axis, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def constant_pad( - x: Union[ivy.Array, ivy.NativeArray], - /, - pad_width: Iterable[Tuple[int]], - *, - value: Number = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Pad an array with a constant value. - - Parameters - ---------- - x - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of - axes of x. - value - The constant value to pad the array with. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Padded array of rank equal to x with shape increased according to pad_width. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3, 4, 5]) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) - >>> print(y) - ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[2.], [3.]]) - >>> y = ivy.zeros((4, 3)) - >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) - >>> print(y) - ivy.array([[5., 5., 5.], - [5., 2., 5.], - [5., 3., 5.], - [5., 5., 5.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), - ... b = ivy.array([3., 4., 5.])) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) - >>> print(y) - { - a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), - b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) - } - """ - return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -671,96 +435,23 @@ def permute_dims( @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion +@handle_view @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def repeat( +def reshape( x: Union[ivy.Array, ivy.NativeArray], /, - repeats: Union[int, Iterable[int]], + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], *, - axis: int = None, + copy: Optional[bool] = None, + order: str = "C", + allowzero: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Repeat values along a given dimension. - - Parameters - ---------- - x - Input array. - repeats - The number of repetitions for each element. repeats is broadcast to fit the - shape of the given axis. - axis - The axis along which to repeat values. By default, use the flattened input - array, and return a flat output array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The repeated output array. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.repeat(x, 2) - >>> print(y) - ivy.array([3, 3, 4, 4, 5, 5]) - - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> y = ivy.repeat(x, [1, 2], axis=0) - >>> print(y) - ivy.array([[1, 2, 3], - [4, 5, 6], - [4, 5, 6]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([0., 1., 2.])) - >>> y = ivy.repeat(x, 2, axis=0) - >>> print(y) - { - a: ivy.array([0., 0., 1., 1., 2., 2.]), - b: ivy.array([0., 0., 1., 1., 2., 2.]) - } - """ - return current_backend(x).repeat(x, repeats, axis=axis, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def reshape( - x: Union[ivy.Array, ivy.NativeArray], - /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - *, - copy: Optional[bool] = None, - order: str = "C", - allowzero: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Give a new shape to an array without changing its data. + Give a new shape to an array without changing its data. Parameters ---------- @@ -980,93 +671,6 @@ def roll( return current_backend(x).roll(x, shift, axis=axis, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def split( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], ivy.Array, ivy.NativeArray] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> List[ivy.Array]: - """ - Split an array into multiple sub-arrays. - - Parameters - ---------- - x - array to be divided into sub-arrays. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - num_or_size_splits - Number of equal arrays to divide the array into along the given axis if an - integer. The size of each split element if a sequence of integers or 1-D array. - Default is to divide into as many 1-dimensional arrays as the axis dimension. - axis - The axis along which to split, default is ``0``. - with_remainder - If the tensor does not split evenly, then store the last remainder entry. - Default is ``False``. - - Returns - ------- - ret - A list of sub-arrays. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.split(x) - >>> print(y) - [ivy.array([1]),ivy.array([2]),ivy.array([3])] - - >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) - >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) - >>> print(y) - [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] - - >>> x = ivy.array([4, 6, 5, 3]) - >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) - >>> print(y) - ivy.array([[4], [6, 5, 3]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([10, 45, 2])) - >>> y = ivy.split(x) - >>> print(y) - { - a:(list[3],shape=[1]) - } - """ - return current_backend(x).split( - x, - copy=copy, - num_or_size_splits=num_or_size_splits, - axis=axis, - with_remainder=with_remainder, - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -1247,6 +851,394 @@ def stack( return res +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def clip( + x: Union[ivy.Array, ivy.NativeArray], + x_min: Union[Number, ivy.Array, ivy.NativeArray], + x_max: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Clips (limits) the values in an array. + + Given an interval, values outside the interval are clipped to the interval edges + (element-wise). For example, if an interval of [0, 1] is specified, values smaller + than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller + or equal to maximum value to return correct results. + + Parameters + ---------- + x + Input array containing elements to clip. + x_min + Minimum value. + x_max + Maximum value. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + An array with the elements of x, but where values < x_min are replaced with + x_min, and those > x_max with x_max. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.zeros_like(x) + >>> ivy.clip(x, 2., 7., out=y) + >>> print(y) + ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) + >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + { + a: ivy.array([1., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> x_min = ivy.Container(a=0, b=-3) + >>> x_max = ivy.Container(a=1, b=-1) + >>> y = ivy.clip(x, x_min,x_max) + >>> print(y) + { + a: ivy.array([0., 1., 1.]), + b: ivy.array([-1., -1., -1.]) + } + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 0., 1]) + >>> x_max = ivy.array([5., 4., 3.]) + >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> z = ivy.clip(y, x_min, x_max) + >>> print(z) + { + a: ivy.array([3., 1., 2.]), + b: ivy.array([3., 4., 3.]) + } + """ + return current_backend(x).clip(x, x_min, x_max, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def constant_pad( + x: Union[ivy.Array, ivy.NativeArray], + /, + pad_width: Iterable[Tuple[int]], + *, + value: Number = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Pad an array with a constant value. + + Parameters + ---------- + x + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of + axes of x. + value + The constant value to pad the array with. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Padded array of rank equal to x with shape increased according to pad_width. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5]) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) + >>> print(y) + ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[2.], [3.]]) + >>> y = ivy.zeros((4, 3)) + >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) + >>> print(y) + ivy.array([[5., 5., 5.], + [5., 2., 5.], + [5., 3., 5.], + [5., 5., 5.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), + ... b = ivy.array([3., 4., 5.])) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) + >>> print(y) + { + a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), + b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) + } + """ + return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def repeat( + x: Union[ivy.Array, ivy.NativeArray], + /, + repeats: Union[int, Iterable[int]], + *, + axis: int = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Repeat values along a given dimension. + + Parameters + ---------- + x + Input array. + repeats + The number of repetitions for each element. repeats is broadcast to fit the + shape of the given axis. + axis + The axis along which to repeat values. By default, use the flattened input + array, and return a flat output array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The repeated output array. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.repeat(x, 2) + >>> print(y) + ivy.array([3, 3, 4, 4, 5, 5]) + + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> y = ivy.repeat(x, [1, 2], axis=0) + >>> print(y) + ivy.array([[1, 2, 3], + [4, 5, 6], + [4, 5, 6]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([0., 1., 2.])) + >>> y = ivy.repeat(x, 2, axis=0) + >>> print(y) + { + a: ivy.array([0., 0., 1., 1., 2., 2.]), + b: ivy.array([0., 0., 1., 1., 2., 2.]) + } + """ + return current_backend(x).repeat(x, repeats, axis=axis, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def split( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], ivy.Array, ivy.NativeArray] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> List[ivy.Array]: + """ + Split an array into multiple sub-arrays. + + Parameters + ---------- + x + array to be divided into sub-arrays. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + num_or_size_splits + Number of equal arrays to divide the array into along the given axis if an + integer. The size of each split element if a sequence of integers or 1-D array. + Default is to divide into as many 1-dimensional arrays as the axis dimension. + axis + The axis along which to split, default is ``0``. + with_remainder + If the tensor does not split evenly, then store the last remainder entry. + Default is ``False``. + + Returns + ------- + ret + A list of sub-arrays. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.split(x) + >>> print(y) + [ivy.array([1]),ivy.array([2]),ivy.array([3])] + + >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) + >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) + >>> print(y) + [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] + + >>> x = ivy.array([4, 6, 5, 3]) + >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) + >>> print(y) + ivy.array([[4], [6, 5, 3]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([10, 45, 2])) + >>> y = ivy.split(x) + >>> print(y) + { + a:(list[3],shape=[1]) + } + """ + return current_backend(x).split( + x, + copy=copy, + num_or_size_splits=num_or_size_splits, + axis=axis, + with_remainder=with_remainder, + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable diff --git a/ivy/functional/ivy/meta.py b/ivy/functional/ivy/meta.py index 66d5452e48832..ab885bb3b207c 100644 --- a/ivy/functional/ivy/meta.py +++ b/ivy/functional/ivy/meta.py @@ -7,11 +7,6 @@ # local from typing import Optional, Union, Callable, Tuple, Any - -# --- Helpers --- # -# --------------- # - - # Extra # # ------# @@ -199,70 +194,6 @@ def _train_task( return final_cost, variables, all_grads -def _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - batched, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, -): - if batched: - return _train_tasks_batched( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) - return _train_tasks_with_for_loop( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) - - def _train_tasks_batched( batch, inner_batch_fn, @@ -405,8 +336,68 @@ def _train_tasks_with_for_loop( return total_cost / num_tasks -# --- Main --- # -# ------------ # +def _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + batched, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, +): + if batched: + return _train_tasks_batched( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) + return _train_tasks_with_for_loop( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) # Public # @@ -535,136 +526,7 @@ def fomaml_step( return cost, grads -# Second Order - - -@handle_exceptions -@handle_array_function -def maml_step( - batch: ivy.Container, - inner_cost_fn: Callable, - outer_cost_fn: Callable, - variables: ivy.Container, - inner_grad_steps: int, - inner_learning_rate: float, - /, - *, - inner_optimization_step: Callable = gradient_descent_update, - inner_batch_fn: Optional[Callable] = None, - outer_batch_fn: Optional[Callable] = None, - average_across_steps: bool = False, - batched: bool = True, - inner_v: Optional[ivy.Container] = None, - keep_inner_v: bool = True, - outer_v: Optional[ivy.Container] = None, - keep_outer_v: bool = True, - return_inner_v: Union[str, bool] = False, - num_tasks: Optional[int] = None, - stop_gradients: bool = True, -) -> Tuple[ivy.Array, ivy.Container, Any]: - """ - Perform step of vanilla second order MAML. - - Parameters - ---------- - batch - The input batch - inner_cost_fn - callable for the inner loop cost function, receiving sub-batch, inner vars and - outer vars - outer_cost_fn - callable for the outer loop cost function, receving task-specific sub-batch, - inner vars and outer vars. If None, the cost from the inner loop will also be - optimized in the outer loop. - variables - Variables to be optimized during the meta step - inner_grad_steps - Number of gradient steps to perform during the inner loop. - inner_learning_rate - The learning rate of the inner loop. - inner_optimization_step - The function used for the inner loop optimization. - Default is ivy.gradient_descent_update. - inner_batch_fn - Function to apply to the task sub-batch, before passing to the inner_cost_fn. - Default is ``None``. - outer_batch_fn - Function to apply to the task sub-batch, before passing to the outer_cost_fn. - Default is ``None``. - average_across_steps - Whether to average the inner loop steps for the outer loop update. - Default is ``False``. - batched - Whether to batch along the time dimension, and run the meta steps in batch. - Default is ``True``. - inner_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_inner_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - outer_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_outer_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - return_inner_v - Either 'first', 'all', or False. 'first' means the variables for the first task - inner loop will also be returned. variables for all tasks will be returned with - 'all'. Default is ``False``. - num_tasks - Number of unique tasks to inner-loop optimize for the meta step. Determined from - batch by default. - stop_gradients - Whether to stop the gradients of the cost. Default is ``True``. - - Returns - ------- - ret - The cost and the gradients with respect to the outer loop variables. - """ - if num_tasks is None: - num_tasks = batch.cont_shape[0] - unique_outer = outer_v is not None - func_ret, grads = ivy.execute_with_gradients( - lambda v: _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables.cont_set_at_key_chains(v) if unique_outer else v, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - 2, - average_across_steps, - batched, - inner_v, - keep_inner_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - False, - ), - ( - variables.cont_at_key_chains(outer_v, ignore_none=True) - if keep_outer_v - else variables.cont_prune_key_chains(outer_v, ignore_none=True) - ), - ) - if isinstance(func_ret, tuple): - grads = grads["0"] if "0" in grads else grads - cost = func_ret[0] - rest = func_ret[1] - else: - cost = func_ret - rest = () - if stop_gradients: - cost = ivy.stop_gradient(cost, preserve_type=False) - return cost, grads.sum(axis=0), rest +fomaml_step.computes_gradients = True @handle_exceptions @@ -803,6 +665,139 @@ def reptile_step( return cost, grads -fomaml_step.computes_gradients = True reptile_step.computes_gradients = True + + +# Second Order + + +@handle_exceptions +@handle_array_function +def maml_step( + batch: ivy.Container, + inner_cost_fn: Callable, + outer_cost_fn: Callable, + variables: ivy.Container, + inner_grad_steps: int, + inner_learning_rate: float, + /, + *, + inner_optimization_step: Callable = gradient_descent_update, + inner_batch_fn: Optional[Callable] = None, + outer_batch_fn: Optional[Callable] = None, + average_across_steps: bool = False, + batched: bool = True, + inner_v: Optional[ivy.Container] = None, + keep_inner_v: bool = True, + outer_v: Optional[ivy.Container] = None, + keep_outer_v: bool = True, + return_inner_v: Union[str, bool] = False, + num_tasks: Optional[int] = None, + stop_gradients: bool = True, +) -> Tuple[ivy.Array, ivy.Container, Any]: + """ + Perform step of vanilla second order MAML. + + Parameters + ---------- + batch + The input batch + inner_cost_fn + callable for the inner loop cost function, receiving sub-batch, inner vars and + outer vars + outer_cost_fn + callable for the outer loop cost function, receving task-specific sub-batch, + inner vars and outer vars. If None, the cost from the inner loop will also be + optimized in the outer loop. + variables + Variables to be optimized during the meta step + inner_grad_steps + Number of gradient steps to perform during the inner loop. + inner_learning_rate + The learning rate of the inner loop. + inner_optimization_step + The function used for the inner loop optimization. + Default is ivy.gradient_descent_update. + inner_batch_fn + Function to apply to the task sub-batch, before passing to the inner_cost_fn. + Default is ``None``. + outer_batch_fn + Function to apply to the task sub-batch, before passing to the outer_cost_fn. + Default is ``None``. + average_across_steps + Whether to average the inner loop steps for the outer loop update. + Default is ``False``. + batched + Whether to batch along the time dimension, and run the meta steps in batch. + Default is ``True``. + inner_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_inner_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + outer_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_outer_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + return_inner_v + Either 'first', 'all', or False. 'first' means the variables for the first task + inner loop will also be returned. variables for all tasks will be returned with + 'all'. Default is ``False``. + num_tasks + Number of unique tasks to inner-loop optimize for the meta step. Determined from + batch by default. + stop_gradients + Whether to stop the gradients of the cost. Default is ``True``. + + Returns + ------- + ret + The cost and the gradients with respect to the outer loop variables. + """ + if num_tasks is None: + num_tasks = batch.cont_shape[0] + unique_outer = outer_v is not None + func_ret, grads = ivy.execute_with_gradients( + lambda v: _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables.cont_set_at_key_chains(v) if unique_outer else v, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + 2, + average_across_steps, + batched, + inner_v, + keep_inner_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + False, + ), + ( + variables.cont_at_key_chains(outer_v, ignore_none=True) + if keep_outer_v + else variables.cont_prune_key_chains(outer_v, ignore_none=True) + ), + ) + if isinstance(func_ret, tuple): + grads = grads["0"] if "0" in grads else grads + cost = func_ret[0] + rest = func_ret[1] + else: + cost = func_ret + rest = () + if stop_gradients: + cost = ivy.stop_gradient(cost, preserve_type=False) + return cost, grads.sum(axis=0), rest + + maml_step.computes_gradients = True diff --git a/ivy/functional/ivy/nest.py b/ivy/functional/ivy/nest.py index c7a052fd7cdb6..05e47e2358582 100644 --- a/ivy/functional/ivy/nest.py +++ b/ivy/functional/ivy/nest.py @@ -11,246 +11,6 @@ from ivy.utils.exceptions import handle_exceptions -@handle_exceptions -def all_nested_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, - /, - include_nests: bool = False, - _index: Optional[Union[int, Sequence[int]]] = None, - _base: bool = True, - extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return indices of all the elements in nest. - - Parameters - ---------- - nest - The nest to check the leaves of. - include_nests - Whether to also include indices of the nests themselves, not only - leaves. Default is ``False``. - _index - The indices detected so far. None at the beginning. Used internally, - do not set manually. - _base - Whether the current function call is the first function call in the - recursive stack. Used internally, do not set manually. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - out - Optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - A set of indices of all elements in nest - - Both the description and the type hints above assumes an array input - for simplicity, but this function is nestable, and therefore also - accepts :class:ivy.Container instances in place of the arguments. - - Examples - -------- - With :class:`Dict` input: - - >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} - >>> y = ivy.all_nested_indices(x) - >>> print(y) - [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] - - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2., 3., 4.]) - >>> y = ivy.all_nested_indices(x, False, out=x) - >>> print(y) - [[]] - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.all_nested_indices(x, True) - >>> print(y) - [['a'], ['b']] - """ - _index = list() if _index is None else _index - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - ind = ivy.argwhere(ivy.ones_like(nest)) - indices = list() - for i in range(len(ind)): - indices.append(_index + ind.to_list()[i]) - return indices - _indices = [ - all_nested_indices( - item, include_nests, _index + [i], False, extra_nest_types - ) - for i, item in enumerate(nest) - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - elif isinstance(nest, dict): - _indices = [ - all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) - for k, v in nest.items() - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - else: - return [_index] - return [index for index in _indices if index] - - -@handle_exceptions -def copy_nest( - nest: Union[ivy.Array, ivy.NativeArray, Iterable], - /, - include_derived: bool = False, - to_mutable: bool = False, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[ivy.Array, ivy.NativeArray, Iterable]: - """ - Copy a nest deeply, but without copying leaves of the nest, only the nest lists, - tuples and dicts are copied. - - Parameters - ---------- - nest - The nest to copy. - include_derived - Whether to also recursive for classes derived from tuple, list and dict. - Default is ``False``. - to_mutable - Whether to convert the nest to a mutable form, changing all tuples to lists. - Default is ``False``. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - - Returns - ------- - ret - The copied nest. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - ivy.array([[1., 2., 3.], - [7., 8., 9.]]) - - With :code:`Iterable` input: - - >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - >>> copied_nest = ivy.copy_nest(nest, include_derived = True) - >>> print(copied_nest) - [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - - >>> nest = ([23, 25, 1337], [63, 98, 6]) - >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) - >>> print(copied_nest) - [[23, 25, 1337], [63, 98, 6]] - - >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} - """ - extra_nest_types = ivy.default(extra_nest_types, ()) - class_instance = type(nest) - check_fn = ( - (lambda x_, t: isinstance(nest, t)) - if include_derived - else (lambda x_, t: type(nest) is t) - ) - if check_fn(nest, tuple): - ret_list = [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - if to_mutable: - return ret_list - if hasattr(nest, "_fields"): - return class_instance(**dict(zip(nest._fields, ret_list))) - return class_instance(tuple(ret_list)) - elif check_fn(nest, list) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - return copy.deepcopy(nest) - return class_instance( - [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - ) - elif check_fn(nest, dict): - class_instance = type(nest) - dict_ = { - k: copy_nest( - v, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for k, v in nest.items() - } - if isinstance(nest, OrderedDict): - return class_instance(**dict_) - return class_instance(dict_) - return nest - - -@handle_exceptions -def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): - """ - Group all unique index chains in a nest. This function is useful for finding all - unique index chains in a nest, and then duplicating the values at those index chains - for functional frameworks. - - Parameters - ---------- - nest - nest to get duplicate index chains for. - - Returns - ------- - list of index chains to duplicate. - """ - all_index_chains = ivy.nested_argwhere(nest, lambda _: True) - duplicates = [] - duplicate_index_chains = {} - for index_chain in all_index_chains: - val = ivy.index_nest(nest, index_chain) - if ivy.is_array(val): - for i in range(len(duplicates)): - if val is duplicates[i]: - duplicate_index_chains[i].append(index_chain) - break - else: - duplicates.append(val) - duplicate_index_chains[len(duplicates) - 1] = [index_chain] - return list(duplicate_index_chains.values()) - - # Extra # # ------# @@ -332,146 +92,140 @@ def index_nest( @handle_exceptions -def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: - if len(index) == 1: - idx = index[0] - if isinstance(nest, list): - nest.insert(idx, value) - else: - nest[index[0]] = value - else: - insert_into_nest_at_index(nest[index[0]], index[1:], value) - - -@handle_exceptions -def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: +def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: """ - Insert a value into the nested item at specified indices with specified values. + Prune a nested object at a specified index. Parameters ---------- nest - The nested object to insert into. - indices - A tuple of tuples of indices for the indices at which to insert - values. - values - The new values for inserting. + The nested object to prune. + index + A tuple of indices for the index at which to prune. """ - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - [ - insert_into_nest_at_index(nest, index, value) - for index, value in zip(indices, values) - ] - - -# noinspection PyShadowingBuiltins + if len(index) == 1: + del nest[index[0]] + else: + prune_nest_at_index(nest[index[0]], index[1:]) @handle_exceptions -def map( - fn: Callable, - constant: Optional[Dict[str, Any]] = None, - unique: Optional[Dict[str, Iterable[Any]]] = None, - mean: bool = False, -) -> List: +def set_nest_at_index( + nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], + index: Sequence[Union[str, int]], + value: Any, + /, + shallow: bool = True, + _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: """ - Apply a function on each item of an iterable x. + Set the value of a nested item at a specified index. Parameters ---------- - fn - The function to map onto x. - constant - keyword arguments which remain constant between each function call. - Default is ``None``. - unique - keyword arguments which are unique for each function call. Default is ``None``. - mean - Whether to compute the mean across the return values, and return this mean. - Default is ``False``. + nest + The nested object to update. + index + A tuple of indices for the index at which to update. + value + The new value for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + _result + Placeholder for the result of the update. do not set this paramter. Returns ------- ret - x following the application of fn to each of its iterated items. + nest with changed value at the given index. Examples -------- - With :code:`int` inputs: - - >>> def special_square(x : float) -> float : return np.square(x) - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x' : [1,2,3]}, - ... mean = False) - >>> print(results) - [1, 4, 9] + With :class:`ivy.Array` inputs: - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x':[0,1,2]}, - ... mean = True) - >>> print(results) - 1.6666666666666667 + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = (1, 1) + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([[1., 2.], [3., 5.]]) - >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = False) - >>> print(results) - [array([1,1]), - array([1,2]), - array([1,3])] + >>> x = ivy.array([1., 2., 3., 4.]) + >>> y = [1] + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([1., 5., 3., 4.]) - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = True) - >>> print(results) - [1. 2.] + With :code:`Dict` input: - With float inputs: + >>> x = {1 : [1, [2, 3]], 2: (4, 5)} + >>> y = (1, 1) + >>> z = 2 + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + {1: [1, 2], 2: (4, 5)} - >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':10., 'b':1.}, - ... unique = {'x':[0.,1.,2.]}, - ... mean = False) - >>> print(results) - [1.0, 11.0, 21.0] + With :code:`List` inputs: - With :class:`ivy.Array` inputs: + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> y = (2, 1, 0) + >>> z = 'H' + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + [['a','b','c'],['d','e','f'],['g',['H','i']]] - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = False) - >>> print(results) - [ivy.array([0., 10., 100.]), - ivy.array([1., 10., 101.])] + With :class:`ivy.Container` input: - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = True) - >>> print(results) - ivy.array([ 0.5, 10. , 100. ]) + >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) + >>> y = ('b',) + >>> z = ivy.array([3., 4.]) + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + { + a: ivy.array([1., 2.]), + b: ivy.array([3., 4.]) + } """ - c = ivy.default(constant, {}) - u = ivy.default(unique, {}) - rets = [ - r - for r in _map( - lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if _result is None: + if shallow: + _result = nest_type(nest) + else: + _result = copy_nest(nest, include_derived=True) + _result = list(_result) if is_tuple else _result + if len(index) == 1: + if shallow: + try: + nest[index[0]] = value + except TypeError: + pass + _result[index[0]] = value + else: + _result[index[0]] = set_nest_at_index( + nest[index[0]], index[1:], value, shallow, _result[index[0]] ) - ] - if mean: - rets = sum(rets) / len(rets) + try: + _result = nest_type(_result) + except TypeError: + _result = nest_type(*_result) + return _result - return rets + +@handle_exceptions +def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: + if len(index) == 1: + idx = index[0] + if isinstance(nest, list): + nest.insert(idx, value) + else: + nest[index[0]] = value + else: + insert_into_nest_at_index(nest[index[0]], index[1:], value) @handle_exceptions @@ -583,130 +337,47 @@ def map_nest_at_index( @handle_exceptions -def map_nest_at_indices( - nest: Iterable, - indices: Tuple, - fn: Callable, +def multi_index_nest( + nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], + indices: Iterable[Iterable[int]], /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: +) -> Iterable[Any]: """ - Map a function to the values of a nested item at the specified indices. + Repeatedly index a nested object, using a tuple of tuples of indices or keys in the + case of dicts. Parameters ---------- nest - The nested object to update. + The nested object to slice. indices - A tuple of tuples of indices for the indices at which to update. - fn - The function to perform on the nest at the given index. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. + A tuple of tuples of indices to apply. Returns ------- ret - nest with applicable of fn on given indices. + The result elements through indexing the nested object. Examples -------- - With :code:`List` inputs: - - >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] - >>> indices = [[0, 4], [1, 5]] - >>> function = lambda x : x + 'b' - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] - With :code:`Tuple` inputs: - >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) - >>> indices = ((0, 2),(1, 0),(1, 2)) - >>> function = abs - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ([-9, 8, 27], [9, -4, 5, 7]) + >>> x = (1, 2) + >>> y = [[0]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [1] - With :code:`Dict` input: + With :class:`ivy.Array` inputs: - >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} - >>> indices = (('a', 2), ('b', 0), ('c', 1)) - >>> function = lambda x : x + 1 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> y = [[0],[1]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [ivy.array([1., 2.], ivy.array([3., 4.])] - With :class:`ivy.Array` inputs: - - >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) - >>> indices = ((0, 1),(1, 1),(1, 2)) - >>> function = lambda x : x ** 2 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ivy.array([[ -9., 64., -17.], - [ 11., 9., 25.]]) - """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if shallow: - result = nest_type(nest) - else: - result = copy_nest(nest, include_derived=True) - result = list(result) if is_tuple else result - for i, index in enumerate(indices): - result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) - try: - result = nest_type(result) - except TypeError: - result = nest_type(*result) - return result - - -@handle_exceptions -def multi_index_nest( - nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], - indices: Iterable[Iterable[int]], - /, -) -> Iterable[Any]: - """ - Repeatedly index a nested object, using a tuple of tuples of indices or keys in the - case of dicts. - - Parameters - ---------- - nest - The nested object to slice. - indices - A tuple of tuples of indices to apply. - - Returns - ------- - ret - The result elements through indexing the nested object. - - Examples - -------- - With :code:`Tuple` inputs: - - >>> x = (1, 2) - >>> y = [[0]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [1] - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> y = [[0],[1]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [ivy.array([1., 2.], ivy.array([3., 4.])] - - With :class:`ivy.Container` input: + With :class:`ivy.Container` input: >>> x = ivy.Container(a=ivy.array([1,2]), ... b=[30,40]) @@ -737,165 +408,322 @@ def multi_index_nest( @handle_exceptions -def nested_any( - nest: Iterable, - fn: Callable, - check_nests: bool = False, - _base: bool = True, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> bool: +def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: """ - Check the leaf nodes of nest x via function fn, and returns True if any evaluate to - True, else False. + Prune a nested object at specified indices. Parameters ---------- nest - The nest to check the leaves of. - fn - The conditon function, returning True or False. - check_nests - Whether to also check the nests for the condition, not only nest leaves. - Default is ``False``. - _base - Whether the current function call is the first function call in the recursive - stack. Used internally, do not set manually. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - - Returns - ------- - ret - A boolean, whether the function evaluates to true for any leaf node. + The nested object to prune. + indices + A tuple of tuples of indices for the indices at which to prune. """ - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - if ivy.any(fn(nest)): - return True - for i, item in enumerate(nest): - if nested_any(item, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif isinstance(nest, dict): - for k, v in nest.items(): - if nested_any(v, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif fn(nest): - return True - return False + # Delete first deeper elements and elements with larger index + indices_sorted = sorted( + indices, + key=str, + reverse=True, + ) + [prune_nest_at_index(nest, index) for index in indices_sorted] @handle_exceptions -def nested_argwhere( - nest: Iterable, - fn: Callable, - check_nests: bool = False, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - _index: Optional[List] = None, - _base: bool = True, - stop_after_n_found: Optional[int] = None, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[Iterable, bool]: +def set_nest_at_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], + indices: Union[List[int], Tuple[int], Iterable[int]], + values: Union[List[int], Tuple[int], Iterable[int]], + /, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: """ - Check the leaf nodes of nested x via function fn, and returns all nest indices where - the method evaluates as True. + Set the value of a nested item at specified indices with specified values. Parameters ---------- nest - The nest to check the leaves of. - fn - The conditon function, returning True or False. - check_nests - Whether to also check the nests for the condition, not only nest leaves. - Default is ``False``. - to_ignore - Types to ignore when deciding whether to go deeper into the nest or not - _index - The indices detected so far. None at the beginning. Used internally, do not set - manually. - _base - Whether the current function call is the first function call in the recursive - stack. Used internally, do not set manually. - stop_after_n_found - to stop after some needed indices are found. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not + The nested object to update. + indices + A tuple of tuples of indices for the indices at which to update. + values + The new values for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. Returns ------- ret - A set of indices for the nest where the function evaluated as True. + nest with updated values at the given indices. Examples -------- - With :code:`List` input: - - >>> nest = [[[1, -2, 3], 19], [[9, -36, 80], -10.19]] - >>> fun = ivy.abs - >>> nested_indices = ivy.nested_argwhere(nest, fn=fun) - >>> print(nested_indices) - [ - [0, 0, 0], [0, 0, 1], - [0, 0, 2], [0, 1], - [1, 0, 0], [1, 0, 1], - [1, 0, 2], [1, 1] - ] + With :code:`List` inputs: + >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] + >>> indices = [[0, 4], [1, 3]] + >>> values = [111, 'x'] + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] - With :code:`Tuple` input: + With :code:`Tuple` inputs: - >>> nest = ([-5, 9, 2], [0.3, 4.]) - >>> fun = ivy.abs - >>> nested_indices = ivy.nested_argwhere(nest, fn=fun, stop_after_n_found=4) - >>> print(nested_indices) - [[0, 0], [0, 1], [0, 2], [1, 0]] + >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] + >>> indices = ((0, 1),(1, 2)) + >>> values = ('ivy', 'x') + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) With :code:`Dict` input: - >>> nest={'a': [2., 0.6, -2.], 'b': [1., 4., 1.9], 'c': [9.4]} - >>> fun = ivy.abs - >>> nested_indices = ivy.nested_argwhere(nest, fn=fun) - >>> print(nested_indices) - [ - ['a', 0], ['a', 1], - ['a', 2], ['b', 0], - ['b', 1], ['b', 2], - ['c', 0] - ] + >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} + >>> indices = (('a', 1), ('b', 2), ('c', 0)) + >>> values = (11., 22., 33.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} + + With :class:`ivy.Array` inputs: + + >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) + >>> indices = ((0, 1),(1, 2)) + >>> values = (11., 22.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + ivy.array([[1., 11., 3.], [4., 5., 22.]]) """ - to_ignore = ivy.default(to_ignore, ()) - extra_nest_types = ivy.default(extra_nest_types, ()) - _index = list() if _index is None else _index - if ( - isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types) - ) and not isinstance(nest, to_ignore): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - cond_met = fn(nest) - ind = ivy.argwhere(cond_met) - _indices = list() - for i in range(len(ind)): - _indices.append(_index + ind.to_list()[i]) - return _indices - n = 0 - _indices = [] - for i, item in enumerate(nest): - ind = ( - nested_argwhere( - item, - fn, - check_nests, - to_ignore, - _index + [i], - False, - stop_after_n_found - n, - extra_nest_types, + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if shallow: + result = nest_type(nest) + else: + result = copy_nest(nest, include_derived=True) + result = list(result) if is_tuple else result + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + for index, value in zip(indices, values): + result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) + try: + result = nest_type(result) + except TypeError: + result = nest_type(*result) + return result + + +@handle_exceptions +def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: + """ + Insert a value into the nested item at specified indices with specified values. + + Parameters + ---------- + nest + The nested object to insert into. + indices + A tuple of tuples of indices for the indices at which to insert + values. + values + The new values for inserting. + """ + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + [ + insert_into_nest_at_index(nest, index, value) + for index, value in zip(indices, values) + ] + + +@handle_exceptions +def map_nest_at_indices( + nest: Iterable, + indices: Tuple, + fn: Callable, + /, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: + """ + Map a function to the values of a nested item at the specified indices. + + Parameters + ---------- + nest + The nested object to update. + indices + A tuple of tuples of indices for the indices at which to update. + fn + The function to perform on the nest at the given index. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + + Returns + ------- + ret + nest with applicable of fn on given indices. + + Examples + -------- + With :code:`List` inputs: + + >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] + >>> indices = [[0, 4], [1, 5]] + >>> function = lambda x : x + 'b' + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] + + With :code:`Tuple` inputs: + + >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) + >>> indices = ((0, 2),(1, 0),(1, 2)) + >>> function = abs + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + ([-9, 8, 27], [9, -4, 5, 7]) + + With :code:`Dict` input: + + >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} + >>> indices = (('a', 2), ('b', 0), ('c', 1)) + >>> function = lambda x : x + 1 + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} + + With :class:`ivy.Array` inputs: + + >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) + >>> indices = ((0, 1),(1, 1),(1, 2)) + >>> function = lambda x : x ** 2 + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + ivy.array([[ -9., 64., -17.], + [ 11., 9., 25.]]) + """ + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if shallow: + result = nest_type(nest) + else: + result = copy_nest(nest, include_derived=True) + result = list(result) if is_tuple else result + for i, index in enumerate(indices): + result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) + try: + result = nest_type(result) + except TypeError: + result = nest_type(*result) + return result + + +@handle_exceptions +def nested_argwhere( + nest: Iterable, + fn: Callable, + check_nests: bool = False, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + _index: Optional[List] = None, + _base: bool = True, + stop_after_n_found: Optional[int] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[Iterable, bool]: + """ + Check the leaf nodes of nested x via function fn, and returns all nest indices where + the method evaluates as True. + + Parameters + ---------- + nest + The nest to check the leaves of. + fn + The conditon function, returning True or False. + check_nests + Whether to also check the nests for the condition, not only nest leaves. + Default is ``False``. + to_ignore + Types to ignore when deciding whether to go deeper into the nest or not + _index + The indices detected so far. None at the beginning. Used internally, do not set + manually. + _base + Whether the current function call is the first function call in the recursive + stack. Used internally, do not set manually. + stop_after_n_found + to stop after some needed indices are found. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + A set of indices for the nest where the function evaluated as True. + + Examples + -------- + With :code:`List` input: + + >>> nest = [[[1, -2, 3], 19], [[9, -36, 80], -10.19]] + >>> fun = ivy.abs + >>> nested_indices = ivy.nested_argwhere(nest, fn=fun) + >>> print(nested_indices) + [ + [0, 0, 0], [0, 0, 1], + [0, 0, 2], [0, 1], + [1, 0, 0], [1, 0, 1], + [1, 0, 2], [1, 1] + ] + + + With :code:`Tuple` input: + + >>> nest = ([-5, 9, 2], [0.3, 4.]) + >>> fun = ivy.abs + >>> nested_indices = ivy.nested_argwhere(nest, fn=fun, stop_after_n_found=4) + >>> print(nested_indices) + [[0, 0], [0, 1], [0, 2], [1, 0]] + + With :code:`Dict` input: + + >>> nest={'a': [2., 0.6, -2.], 'b': [1., 4., 1.9], 'c': [9.4]} + >>> fun = ivy.abs + >>> nested_indices = ivy.nested_argwhere(nest, fn=fun) + >>> print(nested_indices) + [ + ['a', 0], ['a', 1], + ['a', 2], ['b', 0], + ['b', 1], ['b', 2], + ['c', 0] + ] + """ + to_ignore = ivy.default(to_ignore, ()) + extra_nest_types = ivy.default(extra_nest_types, ()) + _index = list() if _index is None else _index + if ( + isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types) + ) and not isinstance(nest, to_ignore): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + cond_met = fn(nest) + ind = ivy.argwhere(cond_met) + _indices = list() + for i in range(len(ind)): + _indices.append(_index + ind.to_list()[i]) + return _indices + n = 0 + _indices = [] + for i, item in enumerate(nest): + ind = ( + nested_argwhere( + item, + fn, + check_nests, + to_ignore, + _index + [i], + False, + stop_after_n_found - n, + extra_nest_types, ) if stop_after_n_found is not None else nested_argwhere( @@ -971,25 +799,230 @@ def nested_argwhere( @handle_exceptions -def nested_map( - x: Union[ivy.Array, ivy.NativeArray, Iterable], +def all_nested_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, /, - fn: Callable, - include_derived: Optional[Union[Dict[type, bool], bool]] = None, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - to_mutable: bool = False, - max_depth: Optional[int] = None, - _depth: int = 0, - _tuple_check_fn: Optional[Callable] = None, - _list_check_fn: Optional[Callable] = None, - _dict_check_fn: Optional[Callable] = None, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: + include_nests: bool = False, + _index: Optional[Union[int, Sequence[int]]] = None, + _base: bool = True, + extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Apply a function on x in a nested manner, whereby all dicts, lists and tuples are - traversed to their lowest leaves before applying the method and returning x. If x is - not nested, the method is applied to x directly. + Return indices of all the elements in nest. + + Parameters + ---------- + nest + The nest to check the leaves of. + include_nests + Whether to also include indices of the nests themselves, not only + leaves. Default is ``False``. + _index + The indices detected so far. None at the beginning. Used internally, + do not set manually. + _base + Whether the current function call is the first function call in the + recursive stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + out + Optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + A set of indices of all elements in nest + + Both the description and the type hints above assumes an array input + for simplicity, but this function is nestable, and therefore also + accepts :class:ivy.Container instances in place of the arguments. + + Examples + -------- + With :class:`Dict` input: + + >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} + >>> y = ivy.all_nested_indices(x) + >>> print(y) + [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] + + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4.]) + >>> y = ivy.all_nested_indices(x, False, out=x) + >>> print(y) + [[]] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.all_nested_indices(x, True) + >>> print(y) + [['a'], ['b']] + """ + _index = list() if _index is None else _index + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + ind = ivy.argwhere(ivy.ones_like(nest)) + indices = list() + for i in range(len(ind)): + indices.append(_index + ind.to_list()[i]) + return indices + _indices = [ + all_nested_indices( + item, include_nests, _index + [i], False, extra_nest_types + ) + for i, item in enumerate(nest) + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + elif isinstance(nest, dict): + _indices = [ + all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) + for k, v in nest.items() + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + else: + return [_index] + return [index for index in _indices if index] + + +# noinspection PyShadowingBuiltins + + +@handle_exceptions +def map( + fn: Callable, + constant: Optional[Dict[str, Any]] = None, + unique: Optional[Dict[str, Iterable[Any]]] = None, + mean: bool = False, +) -> List: + """ + Apply a function on each item of an iterable x. + + Parameters + ---------- + fn + The function to map onto x. + constant + keyword arguments which remain constant between each function call. + Default is ``None``. + unique + keyword arguments which are unique for each function call. Default is ``None``. + mean + Whether to compute the mean across the return values, and return this mean. + Default is ``False``. + + Returns + ------- + ret + x following the application of fn to each of its iterated items. + + Examples + -------- + With :code:`int` inputs: + + >>> def special_square(x : float) -> float : return np.square(x) + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x' : [1,2,3]}, + ... mean = False) + >>> print(results) + [1, 4, 9] + + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x':[0,1,2]}, + ... mean = True) + >>> print(results) + 1.6666666666666667 + + >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = False) + >>> print(results) + [array([1,1]), + array([1,2]), + array([1,3])] + + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = True) + >>> print(results) + [1. 2.] + + With float inputs: + + >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':10., 'b':1.}, + ... unique = {'x':[0.,1.,2.]}, + ... mean = False) + >>> print(results) + [1.0, 11.0, 21.0] + + With :class:`ivy.Array` inputs: + + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = False) + >>> print(results) + [ivy.array([0., 10., 100.]), + ivy.array([1., 10., 101.])] + + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = True) + >>> print(results) + ivy.array([ 0.5, 10. , 100. ]) + """ + c = ivy.default(constant, {}) + u = ivy.default(unique, {}) + rets = [ + r + for r in _map( + lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() + ) + ] + if mean: + rets = sum(rets) / len(rets) + + return rets + + +@handle_exceptions +def nested_map( + x: Union[ivy.Array, ivy.NativeArray, Iterable], + /, + fn: Callable, + include_derived: Optional[Union[Dict[type, bool], bool]] = None, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + to_mutable: bool = False, + max_depth: Optional[int] = None, + _depth: int = 0, + _tuple_check_fn: Optional[Callable] = None, + _list_check_fn: Optional[Callable] = None, + _dict_check_fn: Optional[Callable] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: + """ + Apply a function on x in a nested manner, whereby all dicts, lists and tuples are + traversed to their lowest leaves before applying the method and returning x. If x is + not nested, the method is applied to x directly. Parameters ---------- @@ -1228,45 +1261,210 @@ def nested_map( @handle_exceptions -def nested_multi_map( - func: Callable, - nests: List[Iterable], - index_chains=None, - to_apply=True, - prune_unapplied=False, - index_chain="", - config=None, - to_ivy=True, -): +def nested_any( + nest: Iterable, + fn: Callable, + check_nests: bool = False, + _base: bool = True, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> bool: """ - Apply function to all array values from a collection of identically structured ivy - arrays. + Check the leaf nodes of nest x via function fn, and returns True if any evaluate to + True, else False. Parameters ---------- - func - Function to apply to each nest entry. nest - nests to map. - index_chains - The key-chains to apply or not apply the method to. Default is ``None``. - to_apply - If True, the method will be applied to index_chains, otherwise index_chains will - be skipped. Default is ``True``. - prune_unapplied - Whether to prune index_chains for which the function was not applied, - otherwise the leftmost nest value is used. Default is ``False``. - index_chain - Chain of keys for this dict entry (Default value = '') - config - The configuration for the nests. Default is the same as nest0. - to_ivy - convert the output to ivy_arrays. Default is ``True`` - Returns - ------- - nest containing the result of the function. The structure of the output is the - same as the input with the result of the function applied to each applicable - leaf and the value at that leaf in the first nest for a non-applicable leaf if + The nest to check the leaves of. + fn + The conditon function, returning True or False. + check_nests + Whether to also check the nests for the condition, not only nest leaves. + Default is ``False``. + _base + Whether the current function call is the first function call in the recursive + stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + A boolean, whether the function evaluates to true for any leaf node. + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + if ivy.any(fn(nest)): + return True + for i, item in enumerate(nest): + if nested_any(item, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif isinstance(nest, dict): + for k, v in nest.items(): + if nested_any(v, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif fn(nest): + return True + return False + + +@handle_exceptions +def copy_nest( + nest: Union[ivy.Array, ivy.NativeArray, Iterable], + /, + include_derived: bool = False, + to_mutable: bool = False, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[ivy.Array, ivy.NativeArray, Iterable]: + """ + Copy a nest deeply, but without copying leaves of the nest, only the nest lists, + tuples and dicts are copied. + + Parameters + ---------- + nest + The nest to copy. + include_derived + Whether to also recursive for classes derived from tuple, list and dict. + Default is ``False``. + to_mutable + Whether to convert the nest to a mutable form, changing all tuples to lists. + Default is ``False``. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + The copied nest. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + ivy.array([[1., 2., 3.], + [7., 8., 9.]]) + + With :code:`Iterable` input: + + >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + >>> copied_nest = ivy.copy_nest(nest, include_derived = True) + >>> print(copied_nest) + [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + + >>> nest = ([23, 25, 1337], [63, 98, 6]) + >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) + >>> print(copied_nest) + [[23, 25, 1337], [63, 98, 6]] + + >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + class_instance = type(nest) + check_fn = ( + (lambda x_, t: isinstance(nest, t)) + if include_derived + else (lambda x_, t: type(nest) is t) + ) + if check_fn(nest, tuple): + ret_list = [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + if to_mutable: + return ret_list + if hasattr(nest, "_fields"): + return class_instance(**dict(zip(nest._fields, ret_list))) + return class_instance(tuple(ret_list)) + elif check_fn(nest, list) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + return copy.deepcopy(nest) + return class_instance( + [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + ) + elif check_fn(nest, dict): + class_instance = type(nest) + dict_ = { + k: copy_nest( + v, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for k, v in nest.items() + } + if isinstance(nest, OrderedDict): + return class_instance(**dict_) + return class_instance(dict_) + return nest + + +@handle_exceptions +def nested_multi_map( + func: Callable, + nests: List[Iterable], + index_chains=None, + to_apply=True, + prune_unapplied=False, + index_chain="", + config=None, + to_ivy=True, +): + """ + Apply function to all array values from a collection of identically structured ivy + arrays. + + Parameters + ---------- + func + Function to apply to each nest entry. + nest + nests to map. + index_chains + The key-chains to apply or not apply the method to. Default is ``None``. + to_apply + If True, the method will be applied to index_chains, otherwise index_chains will + be skipped. Default is ``True``. + prune_unapplied + Whether to prune index_chains for which the function was not applied, + otherwise the leftmost nest value is used. Default is ``False``. + index_chain + Chain of keys for this dict entry (Default value = '') + config + The configuration for the nests. Default is the same as nest0. + to_ivy + convert the output to ivy_arrays. Default is ``True`` + Returns + ------- + nest containing the result of the function. The structure of the output is the + same as the input with the result of the function applied to each applicable + leaf and the value at that leaf in the first nest for a non-applicable leaf if prune_unapplied is False else unapplied leaves are pruned. """ nest0 = None @@ -1385,6 +1583,38 @@ def _found_in_index_chains(this_index_chain, index_chains): ) +@handle_exceptions +def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): + """ + Group all unique index chains in a nest. This function is useful for finding all + unique index chains in a nest, and then duplicating the values at those index chains + for functional frameworks. + + Parameters + ---------- + nest + nest to get duplicate index chains for. + + Returns + ------- + list of index chains to duplicate. + """ + all_index_chains = ivy.nested_argwhere(nest, lambda _: True) + duplicates = [] + duplicate_index_chains = {} + for index_chain in all_index_chains: + val = ivy.index_nest(nest, index_chain) + if ivy.is_array(val): + for i in range(len(duplicates)): + if val is duplicates[i]: + duplicate_index_chains[i].append(index_chain) + break + else: + duplicates.append(val) + duplicate_index_chains[len(duplicates) - 1] = [index_chain] + return list(duplicate_index_chains.values()) + + def prune_empty(nest): """ Prune empty nests from a nest. @@ -1420,233 +1650,3 @@ def prune_empty(nest): if not valid and not (ivy.is_array(nest) or isinstance(nest, (int, float, str))): return None return nest - - -@handle_exceptions -def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: - """ - Prune a nested object at a specified index. - - Parameters - ---------- - nest - The nested object to prune. - index - A tuple of indices for the index at which to prune. - """ - if len(index) == 1: - del nest[index[0]] - else: - prune_nest_at_index(nest[index[0]], index[1:]) - - -@handle_exceptions -def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: - """ - Prune a nested object at specified indices. - - Parameters - ---------- - nest - The nested object to prune. - indices - A tuple of tuples of indices for the indices at which to prune. - """ - # Delete first deeper elements and elements with larger index - indices_sorted = sorted( - indices, - key=str, - reverse=True, - ) - [prune_nest_at_index(nest, index) for index in indices_sorted] - - -@handle_exceptions -def set_nest_at_index( - nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], - index: Sequence[Union[str, int]], - value: Any, - /, - shallow: bool = True, - _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: - """ - Set the value of a nested item at a specified index. - - Parameters - ---------- - nest - The nested object to update. - index - A tuple of indices for the index at which to update. - value - The new value for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - _result - Placeholder for the result of the update. do not set this paramter. - - Returns - ------- - ret - nest with changed value at the given index. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = (1, 1) - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([[1., 2.], [3., 5.]]) - - >>> x = ivy.array([1., 2., 3., 4.]) - >>> y = [1] - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([1., 5., 3., 4.]) - - With :code:`Dict` input: - - >>> x = {1 : [1, [2, 3]], 2: (4, 5)} - >>> y = (1, 1) - >>> z = 2 - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - {1: [1, 2], 2: (4, 5)} - - With :code:`List` inputs: - - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> y = (2, 1, 0) - >>> z = 'H' - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - [['a','b','c'],['d','e','f'],['g',['H','i']]] - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) - >>> y = ('b',) - >>> z = ivy.array([3., 4.]) - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - { - a: ivy.array([1., 2.]), - b: ivy.array([3., 4.]) - } - """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if _result is None: - if shallow: - _result = nest_type(nest) - else: - _result = copy_nest(nest, include_derived=True) - _result = list(_result) if is_tuple else _result - if len(index) == 1: - if shallow: - try: - nest[index[0]] = value - except TypeError: - pass - _result[index[0]] = value - else: - _result[index[0]] = set_nest_at_index( - nest[index[0]], index[1:], value, shallow, _result[index[0]] - ) - try: - _result = nest_type(_result) - except TypeError: - _result = nest_type(*_result) - return _result - - -@handle_exceptions -def set_nest_at_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], - indices: Union[List[int], Tuple[int], Iterable[int]], - values: Union[List[int], Tuple[int], Iterable[int]], - /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: - """ - Set the value of a nested item at specified indices with specified values. - - Parameters - ---------- - nest - The nested object to update. - indices - A tuple of tuples of indices for the indices at which to update. - values - The new values for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - - Returns - ------- - ret - nest with updated values at the given indices. - - Examples - -------- - With :code:`List` inputs: - - >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] - >>> indices = [[0, 4], [1, 3]] - >>> values = [111, 'x'] - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] - - With :code:`Tuple` inputs: - - >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] - >>> indices = ((0, 1),(1, 2)) - >>> values = ('ivy', 'x') - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) - - With :code:`Dict` input: - - >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} - >>> indices = (('a', 1), ('b', 2), ('c', 0)) - >>> values = (11., 22., 33.) - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} - - With :class:`ivy.Array` inputs: - - >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) - >>> indices = ((0, 1),(1, 2)) - >>> values = (11., 22.) - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - ivy.array([[1., 11., 3.], [4., 5., 22.]]) - """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if shallow: - result = nest_type(nest) - else: - result = copy_nest(nest, include_derived=True) - result = list(result) if is_tuple else result - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - for index, value in zip(indices, values): - result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) - try: - result = nest_type(result) - except TypeError: - result = nest_type(*result) - return result diff --git a/ivy/functional/ivy/norms.py b/ivy/functional/ivy/norms.py index 48d0cce6435e0..e0804ca4dc668 100644 --- a/ivy/functional/ivy/norms.py +++ b/ivy/functional/ivy/norms.py @@ -13,18 +13,6 @@ from ivy.utils.exceptions import handle_exceptions -layer_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - # Extra # # ------# @@ -150,3 +138,15 @@ def layer_norm( return ivy.multiply(ivy.multiply(x, scale), new_std, out=out) return ivy.multiply(x, new_std, out=out) + + +layer_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} diff --git a/ivy/functional/ivy/random.py b/ivy/functional/ivy/random.py index 70e8c21a5db45..7ed96ae869d1c 100644 --- a/ivy/functional/ivy/random.py +++ b/ivy/functional/ivy/random.py @@ -20,8 +20,8 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # +# Helpers # +# ------- # def _check_bounds_and_get_shape(low, high, shape): @@ -51,17 +51,6 @@ def _check_bounds_and_get_shape(low, high, shape): return ivy.Shape(()) -def _check_shapes_broadcastable(out, inp): - if out is not None: - ivy.utils.assertions.check_shapes_broadcastable(out, inp) - - -def _check_valid_scale(std): - ivy.utils.assertions.check_greater( - std, 0, allow_equal=True, message="std must be non-negative" - ) - - def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_all_or_any_fn( low, @@ -84,51 +73,66 @@ def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_less(low, high) -# --- Main --- # -# ------------ # +def _check_valid_scale(std): + ivy.utils.assertions.check_greater( + std, 0, allow_equal=True, message="std must be non-negative" + ) + + +def _check_shapes_broadcastable(out, inp): + if out is not None: + ivy.utils.assertions.check_shapes_broadcastable(out, inp) + + +# Extra # +# ------# @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function +@infer_dtype @handle_device_shifting @infer_device -def multinomial( - population_size: int, - num_samples: int, - /, +def random_uniform( *, - batch_size: int = 1, - probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - replace: bool = True, + low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, + high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, + shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a multinomial distribution. Specifically, returns a tensor where - each row contains num_samples indices sampled from the multinomial probability - distribution located in the corresponding row of tensor input. + Draws samples from a uniform distribution. Samples are uniformly distributed over + the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In + other words, any value within the given interval is equally likely to be drawn by + uniform. Parameters ---------- - population_size - The size of the population from which to draw samples. - num_samples - Number of independent samples to draw from the population. - batch_size - Number of tensors to generate. Default is 1. - probs - The unnormalized probabilities for all elements in population, - default is uniform *[batch_shape, population_size]* - replace - Whether to replace samples once they've been drawn. Default is ``True``. + low + Lower boundary of the output interval. All values generated will be greater than + or equal to ``low``. If array, must have same shape as ``high``. + high + Upper boundary of the output interval. All the values generated will be less + than ``high``. If array, must have same shape as ``low``. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. + Can only be specified when ``low`` and ``high`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None) + (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out @@ -138,145 +142,67 @@ def multinomial( Returns ------- ret - Drawn samples indices from the multinomial distribution. - - Examples - -------- - >>> y = ivy.multinomial(10, 5) - >>> print(y) - ivy.array([[1, 8, 7, 8, 3]]) - - >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) - >>> print(y) - ivy.array([[3, 9, 7, 5, 1], - [1, 0, 8, 6, 7]]) - - >>> y = ivy.multinomial(10, 5, replace=False) - >>> print(y) - ivy.array([[2, 6, 4, 7, 0]]) - - With :class:`ivy.Array` input: + Drawn samples from the parameterized uniform distribution. - >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) - >>> print(y) - ivy.array([5, 2, 7, 6, 9]) + Functional Examples + ------------------- - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) - >>> print(y) - ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) + >>> ivy.random_uniform() + ivy.array(0.26431865) - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), - ... replace=False) - >>> print(y) - ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) + >>> ivy.random_uniform(shape=3) + ivy.array([0.475, 0.878, 0.861]) - With :class:`ivy.NativeArray` input: + >>> ivy.random_uniform(shape=(2,3)) + ivy.array([[0.929 , 0.545 , 0.789 ], + [0.519 , 0.0435, 0.381 ]]) - >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) - >>> print(y) - ivy.array([5, 7, 4, 2, 1]) + >>> ivy.random_uniform(low=3.0, high=6.0) + ivy.array(3.4608004) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) - >>> print(y) - ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) + ivy.array([[1.85], + [1.81]]) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), - ... replace=False) - >>> print(y) - ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) - """ - return ivy.current_backend().multinomial( - population_size, - num_samples, - batch_size=batch_size, - probs=probs, - replace=replace, - device=device, - seed=seed, - out=out, - ) + >>> z = ivy.zeros(()) + >>> ivy.random_uniform(low=1.0, high=2.0, out=z) + ivy.array(1.8458502) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') + ivy.array([[1.81, 1.8 ], + [1.32, 1.43]]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -@infer_device -def randint( - low: Union[int, ivy.NativeArray, ivy.Array], - high: Union[int, ivy.NativeArray, ivy.Array], - /, - *, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array filled with random integers generated uniformly between low - (inclusive) and high (exclusive). + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', + ... dtype='int32') + ivy.array([[1, 1], + [1, 1]]) - Parameters - ---------- - low - Lowest integer that can be drawn from the distribution. - high - One above the highest integer that can be drawn from the distribution. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn - Can only be specified when ``mean`` and ``std`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. - device - device on which to create the array. 'cuda:0', - 'cuda:1', 'cpu' etc. (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default integer data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. + >>> z = ivy.zeros((1,2)) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', + ... dtype='float64', out=z) + ivy.array([[1.34, 1.02]]) - Returns - ------- - ret - Returns an array with the given shape filled with integers from - the uniform distribution in the “half-open” interval [low, high) + >>> x = ivy.array([4.8, 5.6]) + >>> y = ivy.array([9.8, 7.4]) + >>> ivy.random_uniform(low=x, high=y) + ivy.array([0.475, 0.878]) - Examples - -------- - >>> y = ivy.randint(0, 9, shape=(1,1)) - >>> print(y) - ivy.array([[5]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) + ivy.array([6.67270088, 7.31128597]) - >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) - >>> print(y) - ivy.array([[ 8, 16], - [12, 9]]) + >>> ivy.random_uniform(low=x, high=y, device='cpu') + ivy.array([6.88, 6.75]) - >>> x = ivy.array([1, 2, 3]) - >>> ivy.randint(0, 10, shape=(3,), out=x) - >>> print(x) - ivy.array([2, 6, 7]) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') + ivy.array([8.62, 6.47]) - >>> y = ivy.zeros((3, 3)) - >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) - >>> print(y) - ivy.array([[ 7, 7, 5], - [12, 8, 8], - [ 8, 11, 3]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) + ivy.array([5. , 7.3]) """ - return ivy.current_backend().randint( - low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out + return ivy.current_backend().random_uniform( + low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed ) @@ -394,55 +320,47 @@ def random_normal( ) -# Extra # -# ------# - - @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function -@infer_dtype @handle_device_shifting @infer_device -def random_uniform( +def multinomial( + population_size: int, + num_samples: int, + /, *, - low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, - high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, - shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, + batch_size: int = 1, + probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + replace: bool = True, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a uniform distribution. Samples are uniformly distributed over - the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In - other words, any value within the given interval is equally likely to be drawn by - uniform. + Draws samples from a multinomial distribution. Specifically, returns a tensor where + each row contains num_samples indices sampled from the multinomial probability + distribution located in the corresponding row of tensor input. Parameters ---------- - low - Lower boundary of the output interval. All values generated will be greater than - or equal to ``low``. If array, must have same shape as ``high``. - high - Upper boundary of the output interval. All the values generated will be less - than ``high``. If array, must have same shape as ``low``. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. - Can only be specified when ``low`` and ``high`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. + population_size + The size of the population from which to draw samples. + num_samples + Number of independent samples to draw from the population. + batch_size + Number of tensors to generate. Default is 1. + probs + The unnormalized probabilities for all elements in population, + default is uniform *[batch_shape, population_size]* + replace + Whether to replace samples once they've been drawn. Default is ``True``. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` + (Default value = None) seed A python integer. Used to create a random seed distribution out @@ -452,67 +370,145 @@ def random_uniform( Returns ------- ret - Drawn samples from the parameterized uniform distribution. + Drawn samples indices from the multinomial distribution. - Functional Examples - ------------------- + Examples + -------- + >>> y = ivy.multinomial(10, 5) + >>> print(y) + ivy.array([[1, 8, 7, 8, 3]]) - >>> ivy.random_uniform() - ivy.array(0.26431865) + >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) + >>> print(y) + ivy.array([[3, 9, 7, 5, 1], + [1, 0, 8, 6, 7]]) - >>> ivy.random_uniform(shape=3) - ivy.array([0.475, 0.878, 0.861]) + >>> y = ivy.multinomial(10, 5, replace=False) + >>> print(y) + ivy.array([[2, 6, 4, 7, 0]]) - >>> ivy.random_uniform(shape=(2,3)) - ivy.array([[0.929 , 0.545 , 0.789 ], - [0.519 , 0.0435, 0.381 ]]) + With :class:`ivy.Array` input: - >>> ivy.random_uniform(low=3.0, high=6.0) - ivy.array(3.4608004) + >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) + >>> print(y) + ivy.array([5, 2, 7, 6, 9]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) - ivy.array([[1.85], - [1.81]]) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) + >>> print(y) + ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) - >>> z = ivy.zeros(()) - >>> ivy.random_uniform(low=1.0, high=2.0, out=z) - ivy.array(1.8458502) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), + ... replace=False) + >>> print(y) + ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') - ivy.array([[1.81, 1.8 ], - [1.32, 1.43]]) + With :class:`ivy.NativeArray` input: - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', - ... dtype='int32') - ivy.array([[1, 1], - [1, 1]]) + >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) + >>> print(y) + ivy.array([5, 7, 4, 2, 1]) - >>> z = ivy.zeros((1,2)) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', - ... dtype='float64', out=z) - ivy.array([[1.34, 1.02]]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) + >>> print(y) + ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) - >>> x = ivy.array([4.8, 5.6]) - >>> y = ivy.array([9.8, 7.4]) - >>> ivy.random_uniform(low=x, high=y) - ivy.array([0.475, 0.878]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), + ... replace=False) + >>> print(y) + ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) + """ + return ivy.current_backend().multinomial( + population_size, + num_samples, + batch_size=batch_size, + probs=probs, + replace=replace, + device=device, + seed=seed, + out=out, + ) - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) - ivy.array([6.67270088, 7.31128597]) - >>> ivy.random_uniform(low=x, high=y, device='cpu') - ivy.array([6.88, 6.75]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +@infer_device +def randint( + low: Union[int, ivy.NativeArray, ivy.Array], + high: Union[int, ivy.NativeArray, ivy.Array], + /, + *, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array filled with random integers generated uniformly between low + (inclusive) and high (exclusive). - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') - ivy.array([8.62, 6.47]) + Parameters + ---------- + low + Lowest integer that can be drawn from the distribution. + high + One above the highest integer that can be drawn from the distribution. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn + Can only be specified when ``mean`` and ``std`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. + device + device on which to create the array. 'cuda:0', + 'cuda:1', 'cpu' etc. (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default integer data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) - ivy.array([5. , 7.3]) + Returns + ------- + ret + Returns an array with the given shape filled with integers from + the uniform distribution in the “half-open” interval [low, high) + + Examples + -------- + >>> y = ivy.randint(0, 9, shape=(1,1)) + >>> print(y) + ivy.array([[5]]) + + >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) + >>> print(y) + ivy.array([[ 8, 16], + [12, 9]]) + + >>> x = ivy.array([1, 2, 3]) + >>> ivy.randint(0, 10, shape=(3,), out=x) + >>> print(x) + ivy.array([2, 6, 7]) + + >>> y = ivy.zeros((3, 3)) + >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) + >>> print(y) + ivy.array([[ 7, 7, 5], + [12, 8, 8], + [ 8, 11, 3]]) """ - return ivy.current_backend().random_uniform( - low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed + return ivy.current_backend().randint( + low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out ) diff --git a/ivy/functional/ivy/searching.py b/ivy/functional/ivy/searching.py index c5391a424a431..573c7fc1068ed 100644 --- a/ivy/functional/ivy/searching.py +++ b/ivy/functional/ivy/searching.py @@ -236,81 +236,6 @@ def argmin( ) -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def argwhere( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the indices of all non-zero elements of the input array. - - Parameters - ---------- - x - input array, for which indices are desired. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Indices of non-zero elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> y = ivy.zeros((3, 2), dtype=ivy.int64) - >>> res = ivy.argwhere(x, out=y) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - With a :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0], [1]]), - b: ivy.array([[0], [1]]) - } - - >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0]]), - b: ivy.array([[0], [1]]) - } - """ - return current_backend(x).argwhere(x, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -529,3 +454,78 @@ def where( } """ return current_backend(x1).where(condition, x1, x2, out=out) + + +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def argwhere( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the indices of all non-zero elements of the input array. + + Parameters + ---------- + x + input array, for which indices are desired. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Indices of non-zero elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> y = ivy.zeros((3, 2), dtype=ivy.int64) + >>> res = ivy.argwhere(x, out=y) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + With a :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0], [1]]), + b: ivy.array([[0], [1]]) + } + + >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0]]), + b: ivy.array([[0], [1]]) + } + """ + return current_backend(x).argwhere(x, out=out) diff --git a/ivy/functional/ivy/set.py b/ivy/functional/ivy/set.py index 2c27798531508..1a3e69831ede7 100644 --- a/ivy/functional/ivy/set.py +++ b/ivy/functional/ivy/set.py @@ -148,111 +148,6 @@ def unique_all( return ivy.current_backend(x).unique_all(x, axis=axis, by_value=by_value) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def unique_counts( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Return the unique elements of an input array ``x`` and the corresponding counts for - each unique element in ``x``. - - .. admonition:: Data-dependent output shape - :class: important - - The shapes of two of the output arrays for this function depend on the data - values in the input array; hence, array libraries which build computation graphs - (e.g., JAX, Dask, etc.) may find this function difficult to implement without - knowing array values. Accordingly, such libraries may choose to omit this - function. See :ref:`data-dependent-output-shapes` section for more details. - - .. note:: - Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). - For input arrays having floating-point data types, value-based equality implies - the following behavior. - - - As ``nan`` values compare as ``False``, ``nan`` values should be considered - distinct. - - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be - considered distinct, and the corresponding unique element will be - implementation-dependent (e.g., an implementation could choose to return - ``-0`` if ``-0`` occurs before ``+0``). - - Parameters - ---------- - x - input array. If ``x`` has more than one dimension, the function must flatten - ``x`` and return the unique elements of the flattened array. - - Returns - ------- - ret - a namedtuple ``(values, counts)`` whose - - - first element must have the field name ``values`` and must be an - array containing the unique elements of ``x``. - The array must have the same data type as ``x``. - - second element must have the field name ``counts`` and must be an array - containing the number of times each unique element occurs in ``x``. - The returned array must have same shape as ``values`` and must - have the default array index data type. - - .. note:: - The order of unique elements is not specified and may vary between - implementations. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2,1,3,4,1,3]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) - - >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) - - >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, - 2.29999995]), - counts=ivy.array([3, 1, 1, 1, 1])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), - ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) - >>> y = ivy.unique_counts(x) - >>> print(y) - { - a: (list[2],shape=[4]), - b: (list[2],shape=[4]) - } - """ - return ivy.current_backend(x).unique_counts(x) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -453,3 +348,108 @@ def unique_values( array([0., 1., 2., 3., 4., 5., nan, -0.]) """ return ivy.current_backend(x).unique_values(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def unique_counts( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Return the unique elements of an input array ``x`` and the corresponding counts for + each unique element in ``x``. + + .. admonition:: Data-dependent output shape + :class: important + + The shapes of two of the output arrays for this function depend on the data + values in the input array; hence, array libraries which build computation graphs + (e.g., JAX, Dask, etc.) may find this function difficult to implement without + knowing array values. Accordingly, such libraries may choose to omit this + function. See :ref:`data-dependent-output-shapes` section for more details. + + .. note:: + Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). + For input arrays having floating-point data types, value-based equality implies + the following behavior. + + - As ``nan`` values compare as ``False``, ``nan`` values should be considered + distinct. + - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be + considered distinct, and the corresponding unique element will be + implementation-dependent (e.g., an implementation could choose to return + ``-0`` if ``-0`` occurs before ``+0``). + + Parameters + ---------- + x + input array. If ``x`` has more than one dimension, the function must flatten + ``x`` and return the unique elements of the flattened array. + + Returns + ------- + ret + a namedtuple ``(values, counts)`` whose + + - first element must have the field name ``values`` and must be an + array containing the unique elements of ``x``. + The array must have the same data type as ``x``. + - second element must have the field name ``counts`` and must be an array + containing the number of times each unique element occurs in ``x``. + The returned array must have same shape as ``values`` and must + have the default array index data type. + + .. note:: + The order of unique elements is not specified and may vary between + implementations. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2,1,3,4,1,3]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) + + >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) + + >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, + 2.29999995]), + counts=ivy.array([3, 1, 1, 1, 1])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), + ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) + >>> y = ivy.unique_counts(x) + >>> print(y) + { + a: (list[2],shape=[4]), + b: (list[2],shape=[4]) + } + """ + return ivy.current_backend(x).unique_counts(x) diff --git a/ivy/functional/ivy/sorting.py b/ivy/functional/ivy/sorting.py index 085db2f176276..a3423a308e41a 100644 --- a/ivy/functional/ivy/sorting.py +++ b/ivy/functional/ivy/sorting.py @@ -140,6 +140,113 @@ def argsort( ) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def sort( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a sorted copy of an array. + + Parameters + ---------- + x + input array + axis + axis along which to sort. If set to ``-1``, the function must sort along the + last axis. Default: ``-1``. + descending + direction The direction in which to sort the values + stable + sort stability. If ``True``, + the returned indices must maintain the relative order of ``x`` values which + compare as equal. If ``False``, the returned indices may or may not maintain the + relative order of ``x`` values which compare as equal (i.e., the relative order + of ``x`` values which compare as equal is implementation-dependent). + Default: ``True``. + out + optional output array, for writing the result to. It must have the same shape + as ``x``. + + Returns + ------- + ret + An array with the same dtype and shape as ``x``, with the elements sorted + along the given `axis`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([7, 8, 6]) + >>> y = ivy.sort(x) + >>> print(y) + ivy.array([6, 7, 8]) + + >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) + >>> y = ivy.sort(x, axis=1, descending=True, stable=False) + >>> print(y) + ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) + + >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) + >>> y = ivy.zeros(5) + >>> ivy.sort(x, descending=True, stable=False, out=y) + >>> print(y) + ivy.array([3.2, 2.5, 1.5, 0.7]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.sort(x, out=x) + >>> print(x) + ivy.array([[ 1.1, 2.2, 3.3], + [-6.6, -5.5, -4.4]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) + >>> y = ivy.sort(x, descending=True) + >>> print(y) + { + a: ivy.array([8, 6, 6]), + b: ivy.array([[9., 0.7], [0.4, 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) + >>> y = ivy.sort(x, descending=False, stable=False) + >>> print(y) + { + a: ivy.array([0.7, 1., 3.]), + b: ivy.array([[0.9, 4.], [0.2, 0.6]]) + } + """ + return ivy.current_backend(x).sort( + x, axis=axis, descending=descending, stable=stable, out=out + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -262,110 +369,3 @@ def searchsorted( out=out, ret_dtype=ret_dtype, ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def sort( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a sorted copy of an array. - - Parameters - ---------- - x - input array - axis - axis along which to sort. If set to ``-1``, the function must sort along the - last axis. Default: ``-1``. - descending - direction The direction in which to sort the values - stable - sort stability. If ``True``, - the returned indices must maintain the relative order of ``x`` values which - compare as equal. If ``False``, the returned indices may or may not maintain the - relative order of ``x`` values which compare as equal (i.e., the relative order - of ``x`` values which compare as equal is implementation-dependent). - Default: ``True``. - out - optional output array, for writing the result to. It must have the same shape - as ``x``. - - Returns - ------- - ret - An array with the same dtype and shape as ``x``, with the elements sorted - along the given `axis`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([7, 8, 6]) - >>> y = ivy.sort(x) - >>> print(y) - ivy.array([6, 7, 8]) - - >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) - >>> y = ivy.sort(x, axis=1, descending=True, stable=False) - >>> print(y) - ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) - - >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) - >>> y = ivy.zeros(5) - >>> ivy.sort(x, descending=True, stable=False, out=y) - >>> print(y) - ivy.array([3.2, 2.5, 1.5, 0.7]) - - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.sort(x, out=x) - >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-6.6, -5.5, -4.4]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) - >>> y = ivy.sort(x, descending=True) - >>> print(y) - { - a: ivy.array([8, 6, 6]), - b: ivy.array([[9., 0.7], [0.4, 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) - >>> y = ivy.sort(x, descending=False, stable=False) - >>> print(y) - { - a: ivy.array([0.7, 1., 3.]), - b: ivy.array([[0.9, 4.], [0.2, 0.6]]) - } - """ - return ivy.current_backend(x).sort( - x, axis=axis, descending=descending, stable=stable, out=out - ) diff --git a/ivy/functional/ivy/statistical.py b/ivy/functional/ivy/statistical.py index 81d7c3b8f17d9..aaea746d500d0 100644 --- a/ivy/functional/ivy/statistical.py +++ b/ivy/functional/ivy/statistical.py @@ -16,8 +16,8 @@ from ivy.utils.exceptions import handle_exceptions -# --- Helpers --- # -# --------------- # +# Helpers # +# --------# def _get_promoted_type_of_operands(operands): @@ -31,11 +31,10 @@ def _get_promoted_type_of_operands(operands): return ivy.as_native_dtype(dtype) -# --- Main --- # -# ------------ # +# Array API Standard # +# -------------------# -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -43,138 +42,105 @@ def _get_promoted_type_of_operands(operands): @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumprod( +def min( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative product of the elements along a given axis. + Calculate the minimum value of the input array ``x``. + + .. note:: + When the number of elements over which to compute the minimum value is zero, the + minimum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the maximum possible value + for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, + return ``+infinity``). + + **Special Cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - Input array. + Input array. Should have a real-valued data type. axis - int , axis along which the cumulative product is computed. By default 0. - exclusive - optional bool, Whether to perform the cumprod exclusively. Defaults is False. - reverse - Whether to perform the cumprod from last to first element in the selected - axis. Default is ``False`` (from first to last element) + axis or axes along which minimum values must be computed. By default, the + minimum value must be computed over the entire array. If a tuple of integers, + minimum values must be computed over multiple axes. Default: ``None``. + + keepdims + optional boolean, if ``True``, the reduced axes (dimensions) must be included + in the result as singleton dimensions, and, accordingly, the result must be + compatible with the input array (see :ref:`broadcasting`). Otherwise, + if ``False``, the reduced axes (dimensions) must not be included in the result. + Default: ``False``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Input array with cumulatively multiplied elements along axis. - - Examples - -------- - With :class:`ivy.Array` input: + if the minimum value was computed over the entire array, a zero-dimensional + array containing the minimum value; otherwise, a non-zero-dimensional array + containing the minimum values. The returned array must have the same data type + as ``x``. - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x) - >>> print(y) - ivy.array([2, 6, 24]) - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x, exclusive=True) - >>> print(y) - ivy.array([1, 2, 6]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = ivy.array([[2, 3], - [5, 7], - [11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) - >>> print(y) - ivy.array([[ 1., 2.], - [ 1., 5.], - [ 1., 11.]]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) - >>> print(x) - ivy.array([[1, 1], - [2, 3], - [10, 21]]) + Examples + -------- + With :class:`ivy.Array` input: - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> x.cumprod(axis=0, exclusive=True, out=y) - >>> print(x) - ivy.array([[1., 1.], - [2., 3.], - [10., 21.]]) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.min(x) + >>> print(z) + ivy.array(1) - With :class:`ivy.Container` input: + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array([0, 0, 0]) + >>> y = ivy.min(x, out=z) + >>> print(z) + ivy.array(0) - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.min(x, axis=0, keepdims=True) >>> print(y) - { - a: ivy.array([2, 6, 24]), - b: ivy.array([3, 12, 60]) - } + ivy.array([[0, 1, 2]]) - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x, exclusive=True) + >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.min(x) >>> print(y) - { - a: ivy.array([1, 2, 6]), - b: ivy.array([1, 3, 12]) - } + ivy.array(0) - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) - >>> print(y) - { - a: ivy.array([[1, 2], - [1, 5], - [1, 11]]), - b: ivy.array([[1, 3], - [1, 4], - [1, 5]]) - } + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> x.cumprod(axis=0, exclusive=True, out=x) - >>> print(x) + >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) + >>> z = ivy.min(x) + >>> print(z) { - a: ivy.array([[1, 1], - [2, 3], - [10, 21]]), - b: ivy.array([[1, 1], - [3, 4], - [15, 42]]) + a: ivy.array(1), + b: ivy.array(2) } """ - return current_backend(x).cumprod( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out - ) - - -# Extra # -# ------# + return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -185,141 +151,107 @@ def cumprod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumsum( +def max( x: Union[ivy.Array, ivy.NativeArray], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, + /, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative sum of the elements along a given axis. + Calculate the maximum value of the input array ``x``. + + .. note:: + When the number of elements over which to compute the maximum value is zero, the + maximum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the minimum possible + value for the input array ``x`` data type (e.g., if ``x`` is a floating-point + array, return ``-infinity``). + + **Special Cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- x - Input array. + input array. Should have a numeric data type. axis - Axis along which the cumulative sum is computed. Default is ``0``. - exclusive - Whether to perform cumsum exclusively. Default is ``False``. - reverse - Whether to perform the cumsum from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. + axis or axes along which maximum values must be computed. By default, the + maximum value must be computed over the entire array. If a tuple of integers, + maximum values must be computed over multiple axes. Default: ``None``. + keepdims + if ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cumsum at each - original array elements along the specified axis. + if the maximum value was computed over the entire array, a zero-dimensional + array containing the maximum value; otherwise, a non-zero-dimensional array + containing the maximum values. The returned array must have the same data type + as ``x``. + + + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cumsum(x, exclusive= True, reverse=False) - >>> print(y) - ivy.array([0, 1, 6, 8]) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.max(x) + >>> print(z) + ivy.array(3) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) - >>> print(y) - ivy.array([[7, 7, 2], - [1, 3, 0]]) + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array(0) + >>> y = ivy.max(x, out=z) + >>> print(z) + ivy.array(2) - >>> x = ivy.array([[1, 5, 2], - ... [4, 3, 0]]) - >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.max(x, axis=0, keepdims=True) >>> print(y) - ivy.array([[4, 3, 0], - [0, 0, 0]]) - - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[11, 9, 5], - [14, 11, 5], - [14, 13, 10]]) + ivy.array([[4, 6, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cumsum(x, axis= 0) - >>> print(y) - { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) - } - - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cumsum(x,axis=1,reverse=True, out=y) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.max(x) >>> print(y) { - a: ivy.array([[8, 7, 4]]), - b: ivy.array([[16, 13, 8], - [16, 11, 5]]), - c: ivy.array([[7, 5, 1], - [18, 15, 9], - [5, 5, 3]]) + a: ivy.array(2.), + b: ivy.array(5.) } - >>> x = ivy.Container(a=ivy.array([[0], - ... [5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cumsum(x,axis=0,out=x) - >>> print(x) + >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), + ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) + >>> z = ivy.max(x, axis=1) + >>> print(z) { - a: ivy.array([[0], - [5]]), - b: ivy.array([[6, 8, 7], - [10, 10, 10]]), - c: ivy.array([[1, 2], - [4, 6], - [10, 10]]) + a: ivy.array([3, 2]), + b: ivy.array([4, 2]) } """ - return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) + return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -330,126 +262,108 @@ def cumsum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def einsum( - equation: str, - *operands: Union[ivy.Array, ivy.NativeArray], +def mean( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Sum the product of the elements of the input operands along dimensions specified - using a notation based on the Einstein summation convention. + Calculate the arithmetic mean of the input array ``x``. + + **Special Cases** + + Let ``N`` equal the number of elements over which to compute the arithmetic mean. + - If ``N`` is ``0``, the arithmetic mean is ``NaN``. + - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- - equation - A str describing the contraction, in the same format as numpy.einsum. - operands - seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes - should be consistent with equation. + x + input array. Should have a floating-point data type. + axis + axis or axes along which arithmetic means must be computed. By default, the mean + must be computed over the entire array. If a Sequence of integers, arithmetic + means must be computed over multiple axes. Default: ``None``. + keepdims + bool, if ``True``, the reduced axes (dimensions) must be included in the result + as singleton dimensions, and, accordingly, the result must be compatible with + the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced + axes (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - The array with sums computed. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: + array, if the arithmetic mean was computed over the entire array, a + zero-dimensional array containing the arithmetic mean; otherwise, a + non-zero-dimensional array containing the arithmetic means. The returned + array must have the same data type as ``x``. + .. note:: + While this specification recommends that this function only accept input + arrays having a floating-point data type, specification-compliant array + libraries may choose to accept input arrays having an integer data type. + While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have the + default floating-point data type. - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> y = ivy.einsum('ii', x) - >>> print(y) - ivy.array(12) - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> z = ivy.einsum('ij -> j', x) - >>> print(z) - ivy.array([ 9, 12, 15]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ in the standard. - >>> A = ivy.array([0, 1, 2]) - >>> B = ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]) - >>> C = ivy.einsum('i,ij->i', A, B) - >>> print(C) - ivy.array([ 0, 22, 76]) + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. - >>> A = ivy.array([[1, 1, 1], - ... [2, 2, 2], - ... [5, 5, 5]]) - >>> B = ivy.array([[0, 1, 0], - ... [1, 1, 0], - ... [1, 1, 1]]) - >>> C = ivy.einsum('ij,jk->ik', A, B) - >>> print(C) - ivy.array([[ 2, 3, 1], - [ 4, 6, 2], - [10, 15, 5]]) + Examples + -------- + With :class:`ivy.Array` input: - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i->', A) - >>> print(C) - ivy.array(45) + >>> x = ivy.array([3., 4., 5.]) + >>> y = ivy.mean(x) + >>> print(y) + ivy.array(4.) - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->i', A, B) - >>> print(C) - ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.array(0.) + >>> ivy.mean(x, out=y) + >>> print(y) + ivy.array(1.) - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' - >>> print(C) - ivy.array(510) + >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) + >>> y = ivy.array([0., 0.]) + >>> ivy.mean(x, axis=1, out=y) + >>> print(y) + ivy.array([-1.4, 1.4]) - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,j->ij', A, B) - >>> print(C) - ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], - [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], - [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], - [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], - [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], - [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], - [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], - [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], - [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With :class:`ivy.Container` input: - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]), - ... b=ivy.array([[ 0, 1, 2], - ... [ 4, 5, 6], - ... [ 8, 9, 10]])) - >>> z = ivy.einsum('i,ij->i', x, y) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = ivy.mean(x) + >>> print(y) { - a: ivy.array([0, 22, 76]), - b: ivy.array([0, 15, 54]) + a: ivy.array(0.), + b: ivy.array(0.90000004) } - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), - ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) - >>> y = ivy.einsum('ii', x) + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), + ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) + >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) + >>> ivy.mean(x, axis=0, out=y) >>> print(y) { - a: ivy.array(2), - b: ivy.array(15) + a: ivy.array([1.5, 2.5, 3.5]), + b: ivy.array([4.5, 5.5, 6.5]) } """ - return current_backend(operands[0]).einsum(equation, *operands, out=out) + return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -460,61 +374,72 @@ def einsum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def max( +def prod( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the maximum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the maximum value is zero, the - maximum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the minimum possible - value for the input array ``x`` data type (e.g., if ``x`` is a floating-point - array, return ``-infinity``). + Calculate the product of input array x elements. **Special Cases** - For floating-point operands, + Let ``N`` equal the number of elements over which to compute the product. - - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values - propagate). + - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). + + For both both real-valued and complex floating-point operands, special + cases must be handled as the operation is implemented by successive application + of :func:`ivy.multiply`: Parameters ---------- x input array. Should have a numeric data type. axis - axis or axes along which maximum values must be computed. By default, the - maximum value must be computed over the entire array. If a tuple of integers, - maximum values must be computed over multiple axes. Default: ``None``. + axis or axes along which products must be computed. By default, the product must + be computed over the entire array. If a tuple of integers, products must be + computed over multiple axes. Default: ``None``. keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as + bool, if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + input array (see Broadcasting). Otherwise, if False, the reduced axes (dimensions) must not be included in the result. Default: ``False``. + dtype + data type of the returned array. If None, + if the default data type corresponding to the data type “kind” (integer or + floating-point) of x has a smaller range of values than the data type of x + (e.g., x has data type int64 and the default data type is int32, or x has data + type uint64 and the default data type is int64), the returned array must have + the same data type as x. if x has a floating-point data type, the returned array + must have the default floating-point data type. if x has a signed integer data + type (e.g., int16), the returned array must have the default integer data type. + if x has an unsigned integer data type (e.g., uint16), the returned array must + have an unsigned integer data type having the same number of bits as the default + integer data type (e.g., if the default integer data type is int32, the returned + array must have a uint32 data type). If the data type (either specified or + resolved) differs from the data type of x, the input array should be cast to the + specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - if the maximum value was computed over the entire array, a zero-dimensional - array containing the maximum value; otherwise, a non-zero-dimensional array - containing the maximum values. The returned array must have the same data type - as ``x``. + array, if the product was computed over the entire array, a zero-dimensional + array containing the product; otherwise, a non-zero-dimensional array containing + the products. The returned array must have a data type as described by the dtype + parameter above. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.prod.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -526,41 +451,55 @@ def max( With :class:`ivy.Array` input: >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.max(x) + >>> z = ivy.prod(x) >>> print(z) - ivy.array(3) + ivy.array(6) - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array(0) - >>> y = ivy.max(x, out=z) + >>> x = ivy.array([1, 0, 3]) + >>> z = ivy.prod(x) >>> print(z) - ivy.array(2) + ivy.array(0) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.max(x, axis=0, keepdims=True) + >>> x = ivy.array([[3., 4., 5.]]) + >>> y = ivy.prod(x, keepdims=True) >>> print(y) - ivy.array([[4, 6, 10]]) + ivy.array([60.]) + + >>> x = ivy.array([2., 1.]) + >>> y = ivy.array(0.) + >>> ivy.prod(x, out=y) + >>> print(y) + ivy.array(2.) + + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.prod(x, axis=1) + >>> print(y) + ivy.array([2., 9.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.max(x) + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = ivy.prod(x) >>> print(y) { - a: ivy.array(2.), - b: ivy.array(5.) + a: ivy.array(-0.), + b: ivy.array(0.30800003) } - >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), - ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) - >>> z = ivy.max(x, axis=1) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), + ... b=ivy.array([[ 4., 5.], [5., 6.]])) + >>> y = ivy.prod(x, axis=1, keepdims=True) + >>> print(y) { - a: ivy.array([3, 2]), - b: ivy.array([4, 2]) + a: ivy.array([[2.], + [12.]]), + b: ivy.array([[20.], + [30.]]) } """ - return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).prod( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out + ) @handle_exceptions @@ -571,114 +510,138 @@ def max( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def mean( +def std( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, + correction: Union[int, float] = 0.0, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the arithmetic mean of the input array ``x``. + Calculate the standard deviation of the input array ``x``. **Special Cases** - Let ``N`` equal the number of elements over which to compute the arithmetic mean. - - If ``N`` is ``0``, the arithmetic mean is ``NaN``. - - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values - propagate). + Let ``N`` equal the number of elements over which to compute the standard deviation. + + - If ``N - correction`` is less than or equal to ``0``, + the standard deviation is ``NaN``. + - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - input array. Should have a floating-point data type. + input array. axis - axis or axes along which arithmetic means must be computed. By default, the mean - must be computed over the entire array. If a Sequence of integers, arithmetic - means must be computed over multiple axes. Default: ``None``. + axis or axes along which standard deviations must be computed. By default, the + standard deviation must be computed over the entire array. If a tuple of + integers, standard deviations must be computed over multiple axes. + Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other + than ``0`` has the effect of adjusting the divisor during the calculation of the + standard deviation according to ``N-c`` where ``N`` corresponds to the total + number of elements over which the standard deviation is computed and ``c`` + corresponds to the provided degrees of freedom adjustment. When computing the + standard deviation of a population, setting this parameter to ``0`` is the + standard choice (i.e., the provided array contains data constituting an + entire population). When computing the corrected sample standard deviation, + setting this parameter to ``1`` is the standard choice (i.e., the provided array + contains data sampled from a larger population; this is commonly referred to as + Bessel's correction). + Default: ``0``. keepdims - bool, if ``True``, the reduced axes (dimensions) must be included in the result - as singleton dimensions, and, accordingly, the result must be compatible with - the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced - axes (dimensions) must not be included in the result. Default: ``False``. + if ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - array, if the arithmetic mean was computed over the entire array, a - zero-dimensional array containing the arithmetic mean; otherwise, a - non-zero-dimensional array containing the arithmetic means. The returned - array must have the same data type as ``x``. + if the standard deviation was computed over the entire array, a zero-dimensional + array containing the standard deviation; otherwise, a non-zero-dimensional array + containing the standard deviations. The returned array must have the same data + type as ``x``. + .. note:: While this specification recommends that this function only accept input - arrays having a floating-point data type, specification-compliant array - libraries may choose to accept input arrays having an integer data type. - While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have the - default floating-point data type. - + arrays having a real-valued floating-point data type, specification-compliant + array libraries may choose to accept input arrays having an integer data + type. While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have + the default real-valued floating-point data type. This function conforms to the `Array API Standard `_. This docstring is an extension of the - `docstring `_ in the standard. + `docstring `_ + in the standard. - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.std(x) + >>> print(y) + ivy.array(0.81649661) - >>> x = ivy.array([3., 4., 5.]) - >>> y = ivy.mean(x) + >>> x = ivy.array([-1., 0., 1.]) + >>> z = ivy.std(x, correction=1) + >>> print(z) + ivy.array(1.) + + >>> x = ivy.array([[0., 4.]]) + >>> y = ivy.std(x, keepdims=True) >>> print(y) - ivy.array(4.) + ivy.array([[2.]]) - >>> x = ivy.array([0., 1., 2.]) + >>> x = ivy.array([2., 1.]) >>> y = ivy.array(0.) - >>> ivy.mean(x, out=y) + >>> ivy.std(x, out=y) >>> print(y) - ivy.array(1.) + ivy.array(0.5) - >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) - >>> y = ivy.array([0., 0.]) - >>> ivy.mean(x, axis=1, out=y) + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.std(x, axis=1) >>> print(y) - ivy.array([-1.4, 1.4]) - + ivy.array([0.5, 0. ]) With :class:`ivy.Container` input: >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.mean(x) + >>> y = x.std() >>> print(y) { - a: ivy.array(0.), - b: ivy.array(0.90000004) + a: ivy.array(0.81649661), + b: ivy.array(0.509902) } - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), - ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) - >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) - >>> ivy.mean(x, axis=0, out=y) + >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), + ... b=ivy.array([[ 4., 2.], [2., 1.]])) + >>> y = x.std(axis=1, keepdims=True) >>> print(y) { - a: ivy.array([1.5, 2.5, 3.5]), - b: ivy.array([4.5, 5.5, 6.5]) + a: ivy.array([[1.], + [1.5]]), + b: ivy.array([[1.], + [0.5]]) } """ - return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) - - -# Array API Standard # -# -------------------# + return current_backend(x).std( + x, axis=axis, correction=correction, keepdims=keepdims, out=out + ) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -686,63 +649,81 @@ def mean( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def min( +def sum( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the minimum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the minimum value is zero, the - minimum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the maximum possible value - for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, - return ``+infinity``). + Calculate the sum of the input array x. **Special Cases** - For floating-point operands, - - - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` - (i.e., ``NaN`` values propagate). + Let ``N`` equal the number of elements over which to compute the sum. + - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). + + For floating-point operands, + - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). + + For both real-valued and complex floating-point operands, special cases must + be handled as if the operation is implemented by successive application of + :func:`ivy.add`: Parameters ---------- x - Input array. Should have a real-valued data type. + Input array. Should have a numeric data type. axis - axis or axes along which minimum values must be computed. By default, the - minimum value must be computed over the entire array. If a tuple of integers, - minimum values must be computed over multiple axes. Default: ``None``. + Axis or axes along which sums must be computed. By default, the sum must be + computed over the entire array. If a tuple of integers, sums must be computed + over multiple axes. Default: ``None``. + dtype + Data type of the returned array. If ``None``, + If the default data type corresponding to the data type "kind" (integer or + floating-point) of ``x`` has a smaller range of values than the data type of + ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is + ``int32``, or ``x`` has data type ``uint64`` and the default data type is + ``int64``), the returned array must have the same data type as ``x``. + If ``x`` has a floating-point data type, the returned array must have the + default floating-point data type. + If ``x`` has a signed integer data type (e.g., ``int16``), the returned + array must have the default integer data type. + If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned + array must have an unsigned integer data type having the same number of bits + as the default integer data type (e.g., if the default integer data type is + ``int32``, the returned array must have a ``uint32`` data type). + + If the data type (either specified or resolved) differs from the data type of + ``x``, the input array should be cast to the specified data type before + computing the sum. Default: ``None``. + + .. note:: + keyword argument is intended to help prevent data type overflows. keepdims - optional boolean, if ``True``, the reduced axes (dimensions) must be included - in the result as singleton dimensions, and, accordingly, the result must be - compatible with the input array (see :ref:`broadcasting`). Otherwise, - if ``False``, the reduced axes (dimensions) must not be included in the result. - Default: ``False``. + If ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - if the minimum value was computed over the entire array, a zero-dimensional - array containing the minimum value; otherwise, a non-zero-dimensional array - containing the minimum values. The returned array must have the same data type - as ``x``. + If the sum was computed over the entire array, a zero-dimensional array + containing the sum; otherwise, an array containing the sums. The returned array + must have a data type as described by the ``dtype`` parameter above. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sum.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -753,38 +734,52 @@ def min( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.min(x) - >>> print(z) - ivy.array(1) + >>> x = ivy.array([0.41, 0.89]) + >>> y = ivy.sum(x) + >>> print(y) + ivy.array(1.3) - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array([0, 0, 0]) - >>> y = ivy.min(x, out=z) - >>> print(z) - ivy.array(0) + >>> x = ivy.array([0.5, 0.7, 2.4]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) + >>> print(y) + ivy.array(3.6) >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x, axis=0, keepdims=True) + >>> y = ivy.sum(x, axis = 1, keepdims = False) >>> print(y) - ivy.array([[0, 1, 2]]) + ivy.array([3, 20]) - >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.array([0,0,0]) + >>> ivy.sum(x, axis = 0, keepdims = False, out = y) >>> print(y) - ivy.array(0) + ivy.array([4, 7, 12]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.sum(x) + >>> print(y) + ivy.array(1.9) + + >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) + >>> print(y) + ivy.array(8.) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) - >>> z = ivy.min(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.sum(x) + >>> print(y) { - a: ivy.array(1), - b: ivy.array(2) + a: ivy.array(3.), + b: ivy.array(12.) } """ - return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) @handle_exceptions @@ -795,72 +790,65 @@ def min( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def prod( +def var( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + correction: Union[int, float] = 0.0, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the product of input array x elements. + Calculate the variance of the input array x. **Special Cases** - Let ``N`` equal the number of elements over which to compute the product. + Let N equal the number of elements over which to compute the variance. - - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). + If N - correction is less than or equal to 0, the variance is NaN. - For both both real-valued and complex floating-point operands, special - cases must be handled as the operation is implemented by successive application - of :func:`ivy.multiply`: + If x_i is NaN, the variance is NaN (i.e., NaN values propagate). Parameters ---------- x - input array. Should have a numeric data type. + input array. Should have a floating-point data type. axis - axis or axes along which products must be computed. By default, the product must - be computed over the entire array. If a tuple of integers, products must be - computed over multiple axes. Default: ``None``. + axis or axes along which variances must be computed. By default, the variance + must be computed over the entire array. If a tuple of integers, variances must + be computed over multiple axes. Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other than 0 + has the effect of adjusting the divisor during the calculation of the variance + according to N-c where N corresponds to the total number of elements over which + the variance is computed and c corresponds to the provided degrees of freedom + adjustment. When computing the variance of a population, setting this parameter + to 0 is the standard choice (i.e., the provided array contains data constituting + an entire population). When computing the unbiased sample variance, setting this + parameter to 1 is the standard choice (i.e., the provided array contains data + sampled from a larger population; this is commonly referred to as Bessel's + correction). Default: ``0``. keepdims - bool, if True, the reduced axes (dimensions) must be included in the result as + if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see Broadcasting). Otherwise, if False, the reduced axes (dimensions) must not be included in the result. Default: ``False``. - dtype - data type of the returned array. If None, - if the default data type corresponding to the data type “kind” (integer or - floating-point) of x has a smaller range of values than the data type of x - (e.g., x has data type int64 and the default data type is int32, or x has data - type uint64 and the default data type is int64), the returned array must have - the same data type as x. if x has a floating-point data type, the returned array - must have the default floating-point data type. if x has a signed integer data - type (e.g., int16), the returned array must have the default integer data type. - if x has an unsigned integer data type (e.g., uint16), the returned array must - have an unsigned integer data type having the same number of bits as the default - integer data type (e.g., if the default integer data type is int32, the returned - array must have a uint32 data type). If the data type (either specified or - resolved) differs from the data type of x, the input array should be cast to the - specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - array, if the product was computed over the entire array, a zero-dimensional - array containing the product; otherwise, a non-zero-dimensional array containing - the products. The returned array must have a data type as described by the dtype - parameter above. + if the variance was computed over the entire array, a zero-dimensional array + containing the variance; otherwise, a non-zero-dimensional array containing the + variances. The returned array must have the same data type as x. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.var.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -871,58 +859,46 @@ def prod( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.prod(x) - >>> print(z) - ivy.array(6) - - >>> x = ivy.array([1, 0, 3]) - >>> z = ivy.prod(x) - >>> print(z) - ivy.array(0) - - >>> x = ivy.array([[3., 4., 5.]]) - >>> y = ivy.prod(x, keepdims=True) - >>> print(y) - ivy.array([60.]) - - >>> x = ivy.array([2., 1.]) - >>> y = ivy.array(0.) - >>> ivy.prod(x, out=y) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.var(x) >>> print(y) - ivy.array(2.) + ivy.array(0.07472222) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.prod(x, axis=1) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.array(0.0) + >>> ivy.var(x, out=y) >>> print(y) - ivy.array([2., 9.]) + ivy.array(0.07472222) - With :class:`ivy.Container` input: + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> print(ivy.var(x, axis=1, keepdims=True)) + ivy.array([[0.00666667], + [0.11555555]]) - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.prod(x) + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> y = ivy.var(x, correction=1) >>> print(y) - { - a: ivy.array(-0.), - b: ivy.array(0.30800003) - } + ivy.array(0.08966666) - >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), - ... b=ivy.array([[ 4., 5.], [5., 6.]])) - >>> y = ivy.prod(x, axis=1, keepdims=True) + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), + ... b=ivy.array([0.7, 0.1, 0.9])) + >>> y = ivy.var(x) >>> print(y) { - a: ivy.array([[2.], - [12.]]), - b: ivy.array([[20.], - [30.]]) + a: ivy.array(0.12666667), + b: ivy.array(0.11555555) } """ - return current_backend(x).prod( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out + return current_backend(x).var( + x, axis=axis, correction=correction, keepdims=keepdims, out=out ) +# Extra # +# ------# + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -931,135 +907,141 @@ def prod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def std( +def cumsum( x: Union[ivy.Array, ivy.NativeArray], - /, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, *, - axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, - keepdims: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the standard deviation of the input array ``x``. - - **Special Cases** - - Let ``N`` equal the number of elements over which to compute the standard deviation. - - - If ``N - correction`` is less than or equal to ``0``, - the standard deviation is ``NaN``. - - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` - (i.e., ``NaN`` values propagate). + Return the cumulative sum of the elements along a given axis. Parameters ---------- x - input array. + Input array. axis - axis or axes along which standard deviations must be computed. By default, the - standard deviation must be computed over the entire array. If a tuple of - integers, standard deviations must be computed over multiple axes. - Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other - than ``0`` has the effect of adjusting the divisor during the calculation of the - standard deviation according to ``N-c`` where ``N`` corresponds to the total - number of elements over which the standard deviation is computed and ``c`` - corresponds to the provided degrees of freedom adjustment. When computing the - standard deviation of a population, setting this parameter to ``0`` is the - standard choice (i.e., the provided array contains data constituting an - entire population). When computing the corrected sample standard deviation, - setting this parameter to ``1`` is the standard choice (i.e., the provided array - contains data sampled from a larger population; this is commonly referred to as - Bessel's correction). - Default: ``0``. - keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + Axis along which the cumulative sum is computed. Default is ``0``. + exclusive + Whether to perform cumsum exclusively. Default is ``False``. + reverse + Whether to perform the cumsum from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - if the standard deviation was computed over the entire array, a zero-dimensional - array containing the standard deviation; otherwise, a non-zero-dimensional array - containing the standard deviations. The returned array must have the same data - type as ``x``. - - .. note:: - While this specification recommends that this function only accept input - arrays having a real-valued floating-point data type, specification-compliant - array libraries may choose to accept input arrays having an integer data - type. While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have - the default real-valued floating-point data type. - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Array which holds the result of applying cumsum at each + original array elements along the specified axis. Examples -------- - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.std(x) - >>> print(y) - ivy.array(0.81649661) - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = ivy.std(x, correction=1) - >>> print(z) - ivy.array(1.) + With :class:`ivy.Array` input: - >>> x = ivy.array([[0., 4.]]) - >>> y = ivy.std(x, keepdims=True) + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cumsum(x, exclusive= True, reverse=False) >>> print(y) - ivy.array([[2.]]) + ivy.array([0, 1, 6, 8]) - >>> x = ivy.array([2., 1.]) - >>> y = ivy.array(0.) - >>> ivy.std(x, out=y) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) >>> print(y) - ivy.array(0.5) + ivy.array([[7, 7, 2], + [1, 3, 0]]) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.std(x, axis=1) + >>> x = ivy.array([[1, 5, 2], + ... [4, 3, 0]]) + >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) >>> print(y) - ivy.array([0.5, 0. ]) + ivy.array([[4, 3, 0], + [0, 0, 0]]) + + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[11, 9, 5], + [14, 11, 5], + [14, 13, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = x.std() + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cumsum(x, axis= 0) >>> print(y) { - a: ivy.array(0.81649661), - b: ivy.array(0.509902) + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) } - >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), - ... b=ivy.array([[ 4., 2.], [2., 1.]])) - >>> y = x.std(axis=1, keepdims=True) + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cumsum(x,axis=1,reverse=True, out=y) >>> print(y) { - a: ivy.array([[1.], - [1.5]]), - b: ivy.array([[1.], - [0.5]]) + a: ivy.array([[8, 7, 4]]), + b: ivy.array([[16, 13, 8], + [16, 11, 5]]), + c: ivy.array([[7, 5, 1], + [18, 15, 9], + [5, 5, 3]]) + } + + >>> x = ivy.Container(a=ivy.array([[0], + ... [5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cumsum(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [5]]), + b: ivy.array([[6, 8, 7], + [10, 10, 10]]), + c: ivy.array([[1, 2], + [4, 6], + [10, 10]]) } """ - return current_backend(x).std( - x, axis=axis, correction=correction, keepdims=keepdims, out=out - ) + return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) @handle_exceptions @@ -1070,137 +1052,134 @@ def std( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sum( +def cumprod( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[int, Sequence[int]]] = None, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the sum of the input array x. - - **Special Cases** - - Let ``N`` equal the number of elements over which to compute the sum. - - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). - - For floating-point operands, - - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). - - For both real-valued and complex floating-point operands, special cases must - be handled as if the operation is implemented by successive application of - :func:`ivy.add`: + Return the cumulative product of the elements along a given axis. Parameters ---------- x - Input array. Should have a numeric data type. + Input array. axis - Axis or axes along which sums must be computed. By default, the sum must be - computed over the entire array. If a tuple of integers, sums must be computed - over multiple axes. Default: ``None``. - dtype - Data type of the returned array. If ``None``, - If the default data type corresponding to the data type "kind" (integer or - floating-point) of ``x`` has a smaller range of values than the data type of - ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is - ``int32``, or ``x`` has data type ``uint64`` and the default data type is - ``int64``), the returned array must have the same data type as ``x``. - If ``x`` has a floating-point data type, the returned array must have the - default floating-point data type. - If ``x`` has a signed integer data type (e.g., ``int16``), the returned - array must have the default integer data type. - If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned - array must have an unsigned integer data type having the same number of bits - as the default integer data type (e.g., if the default integer data type is - ``int32``, the returned array must have a ``uint32`` data type). - - If the data type (either specified or resolved) differs from the data type of - ``x``, the input array should be cast to the specified data type before - computing the sum. Default: ``None``. - - .. note:: - keyword argument is intended to help prevent data type overflows. - - keepdims - If ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + int , axis along which the cumulative product is computed. By default 0. + exclusive + optional bool, Whether to perform the cumprod exclusively. Defaults is False. + reverse + Whether to perform the cumprod from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - If the sum was computed over the entire array, a zero-dimensional array - containing the sum; otherwise, an array containing the sums. The returned array - must have a data type as described by the ``dtype`` parameter above. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Input array with cumulatively multiplied elements along axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.41, 0.89]) - >>> y = ivy.sum(x) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x) >>> print(y) - ivy.array(1.3) + ivy.array([2, 6, 24]) - >>> x = ivy.array([0.5, 0.7, 2.4]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x, exclusive=True) >>> print(y) - ivy.array(3.6) + ivy.array([1, 2, 6]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.sum(x, axis = 1, keepdims = False) + >>> x = ivy.array([[2, 3], + [5, 7], + [11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) >>> print(y) - ivy.array([3, 20]) + ivy.array([[ 1., 2.], + [ 1., 5.], + [ 1., 11.]]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.array([0,0,0]) - >>> ivy.sum(x, axis = 0, keepdims = False, out = y) - >>> print(y) - ivy.array([4, 7, 12]) + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) + >>> print(x) + ivy.array([[1, 1], + [2, 3], + [10, 21]]) - With :class:`ivy.NativeArray` input: + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> x.cumprod(axis=0, exclusive=True, out=y) + >>> print(x) + ivy.array([[1., 1.], + [2., 3.], + [10., 21.]]) - >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.sum(x) - >>> print(y) - ivy.array(1.9) + With :class:`ivy.Container` input: - >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x) >>> print(y) - ivy.array(8.) + { + a: ivy.array([2, 6, 24]), + b: ivy.array([3, 12, 60]) + } - With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x, exclusive=True) + >>> print(y) + { + a: ivy.array([1, 2, 6]), + b: ivy.array([1, 3, 12]) + } - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.sum(x) + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) >>> print(y) { - a: ivy.array(3.), - b: ivy.array(12.) + a: ivy.array([[1, 2], + [1, 5], + [1, 11]]), + b: ivy.array([[1, 3], + [1, 4], + [1, 5]]) + } + + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> x.cumprod(axis=0, exclusive=True, out=x) + >>> print(x) + { + a: ivy.array([[1, 1], + [2, 3], + [10, 21]]), + b: ivy.array([[1, 1], + [3, 4], + [15, 42]]) } """ - return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + return current_backend(x).cumprod( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + ) @handle_exceptions @@ -1211,106 +1190,123 @@ def sum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def var( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, - keepdims: bool = False, +def einsum( + equation: str, + *operands: Union[ivy.Array, ivy.NativeArray], out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the variance of the input array x. - - **Special Cases** - - Let N equal the number of elements over which to compute the variance. - - If N - correction is less than or equal to 0, the variance is NaN. - - If x_i is NaN, the variance is NaN (i.e., NaN values propagate). + Sum the product of the elements of the input operands along dimensions specified + using a notation based on the Einstein summation convention. Parameters ---------- - x - input array. Should have a floating-point data type. - axis - axis or axes along which variances must be computed. By default, the variance - must be computed over the entire array. If a tuple of integers, variances must - be computed over multiple axes. Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other than 0 - has the effect of adjusting the divisor during the calculation of the variance - according to N-c where N corresponds to the total number of elements over which - the variance is computed and c corresponds to the provided degrees of freedom - adjustment. When computing the variance of a population, setting this parameter - to 0 is the standard choice (i.e., the provided array contains data constituting - an entire population). When computing the unbiased sample variance, setting this - parameter to 1 is the standard choice (i.e., the provided array contains data - sampled from a larger population; this is commonly referred to as Bessel's - correction). Default: ``0``. - keepdims - if True, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see Broadcasting). Otherwise, if False, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + equation + A str describing the contraction, in the same format as numpy.einsum. + operands + seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes + should be consistent with equation. out optional output array, for writing the result to. Returns ------- ret - if the variance was computed over the entire array, a zero-dimensional array - containing the variance; otherwise, a non-zero-dimensional array containing the - variances. The returned array must have the same data type as x. - - - This method conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + The array with sums computed. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Functional Examples + ------------------- - Examples - -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.var(x) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> y = ivy.einsum('ii', x) >>> print(y) - ivy.array(0.07472222) + ivy.array(12) - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.array(0.0) - >>> ivy.var(x, out=y) - >>> print(y) - ivy.array(0.07472222) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> z = ivy.einsum('ij -> j', x) + >>> print(z) + ivy.array([ 9, 12, 15]) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> print(ivy.var(x, axis=1, keepdims=True)) - ivy.array([[0.00666667], - [0.11555555]]) + >>> A = ivy.array([0, 1, 2]) + >>> B = ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]) + >>> C = ivy.einsum('i,ij->i', A, B) + >>> print(C) + ivy.array([ 0, 22, 76]) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> y = ivy.var(x, correction=1) - >>> print(y) - ivy.array(0.08966666) + >>> A = ivy.array([[1, 1, 1], + ... [2, 2, 2], + ... [5, 5, 5]]) + >>> B = ivy.array([[0, 1, 0], + ... [1, 1, 0], + ... [1, 1, 1]]) + >>> C = ivy.einsum('ij,jk->ik', A, B) + >>> print(C) + ivy.array([[ 2, 3, 1], + [ 4, 6, 2], + [10, 15, 5]]) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i->', A) + >>> print(C) + ivy.array(45) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->i', A, B) + >>> print(C) + ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' + >>> print(C) + ivy.array(510) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,j->ij', A, B) + >>> print(C) + ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], + [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], + [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], + [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], + [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], + [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], + [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], + [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], + [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]), + ... b=ivy.array([[ 0, 1, 2], + ... [ 4, 5, 6], + ... [ 8, 9, 10]])) + >>> z = ivy.einsum('i,ij->i', x, y) + >>> print(z) + { + a: ivy.array([0, 22, 76]), + b: ivy.array([0, 15, 54]) + } With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), - ... b=ivy.array([0.7, 0.1, 0.9])) - >>> y = ivy.var(x) + + >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), + ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) + >>> y = ivy.einsum('ii', x) >>> print(y) { - a: ivy.array(0.12666667), - b: ivy.array(0.11555555) + a: ivy.array(2), + b: ivy.array(15) } """ - return current_backend(x).var( - x, axis=axis, correction=correction, keepdims=keepdims, out=out - ) + return current_backend(operands[0]).einsum(equation, *operands, out=out) diff --git a/ivy/functional/ivy/utility.py b/ivy/functional/ivy/utility.py index cefd9d446c711..09bd432b60291 100644 --- a/ivy/functional/ivy/utility.py +++ b/ivy/functional/ivy/utility.py @@ -237,19 +237,6 @@ def any( return ivy.current_backend(x).any(x, axis=axis, keepdims=keepdims, out=out) -@staticmethod -def load(filepath, format=None, type="module"): - if type == "module": - return ivy.Module.load(filepath) - elif type == "container": - if format is not None: - return ivy.Container.cont_load(filepath, format=format) - else: - return ivy.Container.cont_load(filepath) - else: - raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.") - - # Extra # # ----- # @@ -264,3 +251,16 @@ def save(item, filepath, format=None): item.save(filepath) else: raise ivy.utils.exceptions.IvyException("Unsupported item type for saving.") + + +@staticmethod +def load(filepath, format=None, type="module"): + if type == "module": + return ivy.Module.load(filepath) + elif type == "container": + if format is not None: + return ivy.Container.cont_load(filepath, format=format) + else: + return ivy.Container.cont_load(filepath) + else: + raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.") From 0e1314a97a0d3c6db7fd22bd2889b69ae73f5b2c Mon Sep 17 00:00:00 2001 From: NripeshN Date: Thu, 31 Aug 2023 17:23:16 +0530 Subject: [PATCH 3/3] final --- .pre-commit-config.yaml | 2 +- ivy/functional/backends/jax/activations.py | 32 +- ivy/functional/backends/jax/creation.py | 138 +- ivy/functional/backends/jax/data_type.py | 168 +- ivy/functional/backends/jax/device.py | 146 +- ivy/functional/backends/jax/elementwise.py | 256 +- .../backends/jax/experimental/activations.py | 40 +- .../backends/jax/experimental/creation.py | 77 +- .../backends/jax/experimental/elementwise.py | 284 +- .../backends/jax/experimental/layers.py | 748 +-- .../jax/experimental/linear_algebra.py | 76 +- .../backends/jax/experimental/manipulation.py | 354 +- .../backends/jax/experimental/random.py | 65 +- .../backends/jax/experimental/statistical.py | 416 +- ivy/functional/backends/jax/general.py | 252 +- ivy/functional/backends/jax/gradients.py | 68 +- ivy/functional/backends/jax/layers.py | 156 +- ivy/functional/backends/jax/linear_algebra.py | 92 +- ivy/functional/backends/jax/manipulation.py | 200 +- ivy/functional/backends/jax/random.py | 98 +- ivy/functional/backends/jax/searching.py | 26 +- ivy/functional/backends/jax/sorting.py | 30 +- ivy/functional/backends/jax/statistical.py | 197 +- ivy/functional/backends/mxnet/activations.py | 16 +- ivy/functional/backends/mxnet/creation.py | 80 +- ivy/functional/backends/mxnet/data_type.py | 127 +- ivy/functional/backends/mxnet/device.py | 64 +- ivy/functional/backends/mxnet/elementwise.py | 240 +- .../mxnet/experimental/activations.py | 12 +- .../backends/mxnet/experimental/creation.py | 40 +- .../mxnet/experimental/elementwise.py | 172 +- .../backends/mxnet/experimental/layers.py | 152 +- .../mxnet/experimental/linear_algebra.py | 94 +- .../mxnet/experimental/manipulation.py | 154 +- .../backends/mxnet/experimental/norms.py | 20 +- .../backends/mxnet/experimental/random.py | 38 +- .../mxnet/experimental/statistical.py | 86 +- ivy/functional/backends/mxnet/general.py | 40 +- ivy/functional/backends/mxnet/gradients.py | 30 +- ivy/functional/backends/mxnet/layers.py | 28 +- .../backends/mxnet/linear_algebra.py | 42 +- ivy/functional/backends/mxnet/manipulation.py | 88 +- ivy/functional/backends/mxnet/random.py | 48 +- ivy/functional/backends/mxnet/searching.py | 18 +- ivy/functional/backends/mxnet/sorting.py | 24 +- ivy/functional/backends/mxnet/statistical.py | 72 +- ivy/functional/backends/numpy/activations.py | 104 +- ivy/functional/backends/numpy/creation.py | 140 +- ivy/functional/backends/numpy/data_type.py | 166 +- ivy/functional/backends/numpy/device.py | 100 +- ivy/functional/backends/numpy/elementwise.py | 574 +- .../numpy/experimental/activations.py | 58 +- .../backends/numpy/experimental/creation.py | 127 +- .../numpy/experimental/elementwise.py | 562 +- .../backends/numpy/experimental/layers.py | 1180 ++-- .../numpy/experimental/linear_algebra.py | 114 +- .../numpy/experimental/manipulation.py | 434 +- .../backends/numpy/experimental/random.py | 56 +- .../numpy/experimental/statistical.py | 618 +- ivy/functional/backends/numpy/general.py | 182 +- ivy/functional/backends/numpy/gradients.py | 58 +- ivy/functional/backends/numpy/helpers.py | 50 +- ivy/functional/backends/numpy/layers.py | 120 +- .../backends/numpy/linear_algebra.py | 92 +- ivy/functional/backends/numpy/manipulation.py | 158 +- ivy/functional/backends/numpy/random.py | 73 +- ivy/functional/backends/numpy/searching.py | 26 +- ivy/functional/backends/numpy/sorting.py | 38 +- ivy/functional/backends/numpy/statistical.py | 216 +- ivy/functional/backends/numpy/utility.py | 4 +- ivy/functional/backends/paddle/activations.py | 128 +- ivy/functional/backends/paddle/creation.py | 453 +- ivy/functional/backends/paddle/data_type.py | 139 +- ivy/functional/backends/paddle/device.py | 88 +- ivy/functional/backends/paddle/elementwise.py | 1860 +++--- .../paddle/experimental/activations.py | 72 +- .../backends/paddle/experimental/creation.py | 143 +- .../paddle/experimental/elementwise.py | 590 +- .../backends/paddle/experimental/layers.py | 499 +- .../paddle/experimental/linear_algebra.py | 74 +- .../backends/paddle/experimental/losses.py | 24 +- .../paddle/experimental/manipulation.py | 615 +- .../backends/paddle/experimental/norms.py | 56 +- .../backends/paddle/experimental/random.py | 127 +- .../paddle/experimental/statistical.py | 706 +- ivy/functional/backends/paddle/general.py | 202 +- ivy/functional/backends/paddle/gradients.py | 222 +- ivy/functional/backends/paddle/layers.py | 72 +- .../backends/paddle/linear_algebra.py | 142 +- .../backends/paddle/manipulation.py | 336 +- ivy/functional/backends/paddle/random.py | 111 +- ivy/functional/backends/paddle/searching.py | 50 +- ivy/functional/backends/paddle/sorting.py | 48 +- ivy/functional/backends/paddle/statistical.py | 359 +- .../backends/tensorflow/activations.py | 44 +- .../backends/tensorflow/control_flow_ops.py | 77 +- .../backends/tensorflow/creation.py | 160 +- .../backends/tensorflow/data_type.py | 176 +- ivy/functional/backends/tensorflow/device.py | 126 +- .../backends/tensorflow/elementwise.py | 476 +- .../tensorflow/experimental/activations.py | 44 +- .../tensorflow/experimental/creation.py | 120 +- .../tensorflow/experimental/elementwise.py | 362 +- .../tensorflow/experimental/layers.py | 1494 ++--- .../tensorflow/experimental/linear_algebra.py | 210 +- .../tensorflow/experimental/manipulation.py | 360 +- .../backends/tensorflow/experimental/norms.py | 48 +- .../tensorflow/experimental/random.py | 90 +- .../tensorflow/experimental/statistical.py | 658 +- ivy/functional/backends/tensorflow/general.py | 270 +- .../backends/tensorflow/gradients.py | 180 +- ivy/functional/backends/tensorflow/layers.py | 92 +- .../backends/tensorflow/linear_algebra.py | 60 +- .../backends/tensorflow/manipulation.py | 226 +- ivy/functional/backends/tensorflow/random.py | 82 +- .../backends/tensorflow/searching.py | 46 +- ivy/functional/backends/tensorflow/sorting.py | 46 +- .../backends/tensorflow/statistical.py | 159 +- ivy/functional/backends/torch/activations.py | 118 +- ivy/functional/backends/torch/creation.py | 258 +- ivy/functional/backends/torch/data_type.py | 135 +- ivy/functional/backends/torch/device.py | 107 +- ivy/functional/backends/torch/elementwise.py | 1130 ++-- .../torch/experimental/activations.py | 44 +- .../backends/torch/experimental/creation.py | 147 +- .../torch/experimental/elementwise.py | 340 +- .../backends/torch/experimental/layers.py | 988 +-- .../torch/experimental/linear_algebra.py | 106 +- .../backends/torch/experimental/losses.py | 35 +- .../torch/experimental/manipulation.py | 396 +- .../backends/torch/experimental/norms.py | 118 +- .../backends/torch/experimental/random.py | 104 +- .../torch/experimental/statistical.py | 756 ++- ivy/functional/backends/torch/general.py | 334 +- ivy/functional/backends/torch/gradients.py | 146 +- ivy/functional/backends/torch/layers.py | 124 +- .../backends/torch/linear_algebra.py | 192 +- ivy/functional/backends/torch/manipulation.py | 264 +- ivy/functional/backends/torch/random.py | 97 +- ivy/functional/backends/torch/searching.py | 26 +- ivy/functional/backends/torch/sorting.py | 48 +- ivy/functional/backends/torch/statistical.py | 270 +- ivy/functional/backends/torch/utility.py | 4 +- .../frontends/torch/comparison_ops.py | 2 +- ivy/functional/ivy/activations.py | 368 +- ivy/functional/ivy/constants.py | 81 +- ivy/functional/ivy/control_flow_ops.py | 202 +- ivy/functional/ivy/creation.py | 2441 +++---- ivy/functional/ivy/data_type.py | 2036 +++--- ivy/functional/ivy/device.py | 1391 ++-- ivy/functional/ivy/elementwise.py | 3974 ++++++------ .../ivy/experimental/activations.py | 340 +- ivy/functional/ivy/experimental/creation.py | 990 +-- .../ivy/experimental/elementwise.py | 1244 ++-- ivy/functional/ivy/experimental/general.py | 8 + ivy/functional/ivy/experimental/layers.py | 4041 ++++++------ .../ivy/experimental/linear_algebra.py | 1755 ++--- ivy/functional/ivy/experimental/losses.py | 178 +- .../ivy/experimental/manipulation.py | 3478 +++++----- ivy/functional/ivy/experimental/norms.py | 336 +- ivy/functional/ivy/experimental/random.py | 204 +- .../ivy/experimental/sparse_array.py | 1118 ++-- .../ivy/experimental/statistical.py | 1232 ++-- ivy/functional/ivy/general.py | 5674 ++++++++--------- ivy/functional/ivy/gradients.py | 1686 +++-- ivy/functional/ivy/layers.py | 3223 +++++----- ivy/functional/ivy/linear_algebra.py | 487 +- ivy/functional/ivy/losses.py | 124 +- ivy/functional/ivy/manipulation.py | 824 +-- ivy/functional/ivy/meta.py | 399 +- ivy/functional/ivy/nest.py | 1716 ++--- ivy/functional/ivy/random.py | 466 +- ivy/functional/ivy/searching.py | 150 +- ivy/functional/ivy/set.py | 210 +- ivy/functional/ivy/sorting.py | 214 +- ivy/functional/ivy/statistical.py | 1554 ++--- ivy/functional/ivy/utility.py | 26 +- 177 files changed, 35961 insertions(+), 36076 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f4c4c298876fb..780dba2209b8a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,6 +26,6 @@ repos: # Exclude everything in frontends except __init__.py, and func_wrapper.py exclude: 'ivy/functional/(frontends|backends)/(?!.*/func_wrapper\.py$).*(?!__init__\.py$)' - repo: https://github.com/unifyai/lint-hook - rev: 27646397c5390f644a645f439535b1061b9c0105 + rev: 2e2dc6c06475b5ec47e4b97c30b23dbb4bd01891 hooks: - id: ivy-lint diff --git a/ivy/functional/backends/jax/activations.py b/ivy/functional/backends/jax/activations.py index 6d980aaf6022e..7732873e41b6f 100644 --- a/ivy/functional/backends/jax/activations.py +++ b/ivy/functional/backends/jax/activations.py @@ -23,6 +23,10 @@ def gelu( return jax.nn.gelu(x, approximate) +def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.hard_swish(x) + + def leaky_relu( x: JaxArray, /, @@ -34,6 +38,18 @@ def leaky_relu( return jnp.asarray(jnp.where(x > 0, x, jnp.multiply(x, alpha)), x.dtype) +def log_softmax( + x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None +): + if axis is None: + axis = -1 + return jax.nn.log_softmax(x, axis) + + +def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): + return x * jnp.tanh(jax.nn.softplus(x)) + + def relu( x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None ) -> JaxArray: @@ -78,19 +94,3 @@ def softplus( if threshold is not None: return jnp.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -def log_softmax( - x: JaxArray, /, *, axis: Optional[int] = None, out: Optional[JaxArray] = None -): - if axis is None: - axis = -1 - return jax.nn.log_softmax(x, axis) - - -def mish(x: JaxArray, /, *, out: Optional[JaxArray] = None): - return x * jnp.tanh(jax.nn.softplus(x)) - - -def hardswish(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.hard_swish(x) diff --git a/ivy/functional/backends/jax/creation.py b/ivy/functional/backends/jax/creation.py index 4fbc39f493756..45b8ad89b441f 100644 --- a/ivy/functional/backends/jax/creation.py +++ b/ivy/functional/backends/jax/creation.py @@ -79,6 +79,19 @@ def asarray( return jnp.asarray(obj, dtype=dtype) +def copy_array( + x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + x = ( + jax.core.ShapedArray(x.shape, x.dtype) + if isinstance(x, jax.core.ShapedArray) + else jnp.array(x) + ) + if to_ivy_array: + return ivy.to_ivy(x) + return x + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -127,6 +140,15 @@ def from_dlpack(x, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jax.dlpack.from_dlpack(capsule) +def frombuffer( + buffer: bytes, + dtype: Optional[jnp.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> JaxArray: + return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -234,6 +256,42 @@ def meshgrid( return jnp.meshgrid(*arrays, sparse=sparse, indexing=indexing) +def one_hot( + indices: JaxArray, + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[jnp.dtype] = None, + device: jaxlib.xla_extension.Device, + out: Optional[JaxArray] = None, +) -> JaxArray: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = jnp.float32 + else: + if not on_none: + dtype = jnp.array(on_value).dtype + elif not off_none: + dtype = jnp.array(off_value).dtype + + res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] + res = res.reshape(list(indices.shape) + [depth]) + + if not on_none and not off_none: + res = jnp.where(res == 1, on_value, off_value) + + if axis is not None: + res = jnp.moveaxis(res, -1, axis) + + return res + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -263,6 +321,17 @@ def triu(x: JaxArray, /, *, k: int = 0, out: Optional[JaxArray] = None) -> JaxAr return jnp.triu(x, k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: jaxlib.xla_extension.Device, +) -> Tuple[JaxArray]: + return jnp.triu_indices(n=n_rows, k=k, m=n_cols) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -289,72 +358,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: JaxArray, *, to_ivy_array: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - x = ( - jax.core.ShapedArray(x.shape, x.dtype) - if isinstance(x, jax.core.ShapedArray) - else jnp.array(x) - ) - if to_ivy_array: - return ivy.to_ivy(x) - return x - - -def one_hot( - indices: JaxArray, - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - device: jaxlib.xla_extension.Device, - out: Optional[JaxArray] = None, -) -> JaxArray: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = jnp.float32 - else: - if not on_none: - dtype = jnp.array(on_value).dtype - elif not off_none: - dtype = jnp.array(off_value).dtype - - res = jnp.eye(depth, dtype=dtype)[jnp.array(indices, dtype="int64").reshape(-1)] - res = res.reshape(list(indices.shape) + [depth]) - - if not on_none and not off_none: - res = jnp.where(res == 1, on_value, off_value) - - if axis is not None: - res = jnp.moveaxis(res, -1, axis) - - return res - - -def frombuffer( - buffer: bytes, - dtype: Optional[jnp.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> JaxArray: - return jnp.frombuffer(buffer, dtype=dtype, count=count, offset=offset) - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: jaxlib.xla_extension.Device, -) -> Tuple[JaxArray]: - return jnp.triu_indices(n=n_rows, k=k, m=n_cols) diff --git a/ivy/functional/backends/jax/data_type.py b/ivy/functional/backends/jax/data_type.py index e25e54e745bbe..9a53044c244f3 100644 --- a/ivy/functional/backends/jax/data_type.py +++ b/ivy/functional/backends/jax/data_type.py @@ -9,6 +9,26 @@ from ivy.functional.backends.jax import JaxArray from ivy.functional.ivy.data_type import _handle_nestable_dtype_info +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} ivy_dtype_dict = { jnp.dtype("int8"): "int8", jnp.dtype("int16"): "int16", @@ -41,7 +61,6 @@ jnp.complex128: "complex128", jnp.bool_: "bool", } - native_dtype_dict = { "int8": jnp.dtype("int8"), "int16": jnp.dtype("int16"), @@ -60,27 +79,6 @@ "bool": jnp.dtype("bool"), } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} - class Finfo: def __init__(self, jnp_finfo: jnp.finfo): @@ -110,69 +108,6 @@ def smallest_normal(self): return float(self._jnp_finfo.tiny) -# Array API Standard # -# -------------------# - - -def astype( - x: JaxArray, - dtype: jnp.dtype, - /, - *, - copy: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - ivy.utils.assertions._check_jax_x64_flag(dtype) - if x.dtype == dtype: - return jnp.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: - try: - return jnp.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: JaxArray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return jnp.broadcast_to(x.reshape(-1), shape) - return jnp.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return Finfo(jnp.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype.name - return jnp.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return jnp.result_type(arrays_and_dtypes) - - result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = jnp.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -245,6 +180,45 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: JaxArray, + dtype: jnp.dtype, + /, + *, + copy: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + ivy.utils.assertions._check_jax_x64_flag(dtype) + if x.dtype == dtype: + return jnp.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: JaxArray) -> List[JaxArray]: + try: + return jnp.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: JaxArray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return jnp.broadcast_to(x.reshape(-1), shape) + return jnp.broadcast_to(x, shape) + + def dtype(x: Union[JaxArray, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -264,6 +238,20 @@ def dtype_bits(dtype_in: Union[jnp.dtype, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return Finfo(jnp.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[jnp.dtype, str, JaxArray, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype.name + return jnp.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -271,3 +259,13 @@ def is_native_dtype(dtype_in: Union[jnp.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[JaxArray, jnp.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return jnp.result_type(arrays_and_dtypes) + + result = jnp.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = jnp.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/jax/device.py b/ivy/functional/backends/jax/device.py index 5158c9a0809e5..9f41e96c4be8c 100644 --- a/ivy/functional/backends/jax/device.py +++ b/ivy/functional/backends/jax/device.py @@ -15,54 +15,44 @@ ) -# Helpers # -# --------# +# noinspection PyMethodMayBeStatic +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._save_dir = os.path.join(self._save_dir, "profile") + def start(self): + jax.profiler.start_trace(self._save_dir) -def _to_array(x): - if isinstance(x, jax.interpreters.ad.JVPTracer): - return _to_array(x.primal) - elif isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): - return _to_array(x.aval) - elif isinstance(x, jax.interpreters.batching.BatchTracer): - return _to_array(x.val) - return x + def stop(self): + jax.profiler.stop_trace() + def __enter__(self): + self.start() -# API # -# ----# + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() -def dev( - x: JaxArray, - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, jaxlib.xla_extension.Device]: - if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): - return "" +# --- Helpers --- # +# --------------- # + + +def _dev_is_available(base_dev): try: - dv = _to_array(x).device_buffer.device - dv = dv() - except Exception: - dv = jax.devices()[0] - if as_native: - return dv - return as_ivy_dev(dv) + jax.devices(base_dev) + return True + except RuntimeError: + return False -def to_device( - x: JaxArray, - device: jaxlib.xla_extension.Device, - /, - *, - stream: Optional[int] = None, - out: Optional[JaxArray] = None, -): - if device is not None: - cur_dev = as_native_dev(dev(x)) - if cur_dev != device: - x = jax.device_put(x, as_native_dev(device)) +def _to_array(x): + if isinstance(x, jax.interpreters.ad.JVPTracer): + return _to_array(x.primal) + elif isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): + return _to_array(x.aval) + elif isinstance(x, jax.interpreters.batching.BatchTracer): + return _to_array(x.val) return x @@ -76,6 +66,10 @@ def _to_device(x, device=None): return x +# --- Main --- # +# ------------ # + + def as_ivy_dev(device, /): if isinstance(device, str): return ivy.Device(device) @@ -99,30 +93,44 @@ def as_native_dev(device, /): return jax.devices(device)[idx] -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( - *args, device_shifting_dev=device_shifting_dev, **kwargs - ) - with jax.default_device(device_shifting_dev): - return fn(*args, **kwargs) - - def clear_cached_mem_on_dev(device: str, /): return None -def _dev_is_available(base_dev): +# API # +# ----# + + +def dev( + x: JaxArray, + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, jaxlib.xla_extension.Device]: + if isinstance(x, jax.interpreters.partial_eval.DynamicJaxprTracer): + return "" try: - jax.devices(base_dev) - return True - except RuntimeError: - return False + dv = _to_array(x).device_buffer.device + dv = dv() + except Exception: + dv = jax.devices()[0] + if as_native: + return dv + return as_ivy_dev(dv) def gpu_is_available() -> bool: return _dev_is_available("gpu") +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( + *args, device_shifting_dev=device_shifting_dev, **kwargs + ) + with jax.default_device(device_shifting_dev): + return fn(*args, **kwargs) + + def num_gpus() -> int: try: return len(jax.devices("gpu")) @@ -130,24 +138,20 @@ def num_gpus() -> int: return 0 -def tpu_is_available() -> bool: - return _dev_is_available("tpu") - - -# noinspection PyMethodMayBeStatic -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._save_dir = os.path.join(self._save_dir, "profile") - - def start(self): - jax.profiler.start_trace(self._save_dir) - - def stop(self): - jax.profiler.stop_trace() +def to_device( + x: JaxArray, + device: jaxlib.xla_extension.Device, + /, + *, + stream: Optional[int] = None, + out: Optional[JaxArray] = None, +): + if device is not None: + cur_dev = as_native_dev(dev(x)) + if cur_dev != device: + x = jax.device_put(x, as_native_dev(device)) + return x - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + return _dev_is_available("tpu") diff --git a/ivy/functional/backends/jax/elementwise.py b/ivy/functional/backends/jax/elementwise.py index 6de81f65a16eb..77f304a382779 100644 --- a/ivy/functional/backends/jax/elementwise.py +++ b/ivy/functional/backends/jax/elementwise.py @@ -16,6 +16,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _abs_variant_sign(x): + return jnp.where(x != 0, x / jnp.abs(x), 0) + + +# --- Main --- # +# ------------ # + + def abs( x: Union[float, JaxArray], /, @@ -51,6 +63,16 @@ def add( return jnp.add(x1, x2) +def angle( + z: JaxArray, + /, + *, + deg: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.angle(z, deg=deg) + + def asin(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.arcsin(x) @@ -156,6 +178,10 @@ def cosh(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.cosh(x) +def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.deg2rad(x) + + def divide( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -183,10 +209,28 @@ def equal( return jnp.equal(x1, x2) +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.scipy.special.erf(x) + + def exp(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.exp(x) +def exp2( + x: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.power(2, x) + + def expm1(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.expm1(x) @@ -221,6 +265,29 @@ def fmin( return jnp.fmin(x1, x2) +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def fmod( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmod(x1, x2) + + +def gcd( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.gcd(x1, x2) + + def greater( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -243,6 +310,15 @@ def greater_equal( return jnp.greater_equal(x1, x2) +def imag( + val: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.imag(val) + + def isfinite(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isfinite(x) @@ -268,6 +344,10 @@ def isnan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.isnan(x) +def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.isreal(x) + + def lcm(x1: JaxArray, x2: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: x1, x2 = promote_types_of_inputs(x1, x2) return jnp.lcm(x1, x2) @@ -353,6 +433,34 @@ def logical_xor( return jnp.logical_xor(x1, x2) +def maximum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 >= x2, x1, x2) + return jnp.maximum(x1, x2) + + +def minimum( + x1: Union[float, JaxArray], + x2: Union[float, JaxArray], + /, + *, + use_where: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return jnp.where(x1 <= x2, x1, x2) + return jnp.minimum(x1, x2) + + def multiply( x1: Union[float, JaxArray], x2: Union[float, JaxArray], @@ -411,6 +519,20 @@ def pow( return jnp.power(x1, x2) +def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.rad2deg(x) + + +def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.real(x) + + +def reciprocal( + x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.reciprocal(x) + + @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def remainder( x1: Union[float, JaxArray], @@ -442,10 +564,6 @@ def round( return ret -def _abs_variant_sign(x): - return jnp.where(x != 0, x / jnp.abs(x), 0) - - def sign( x: JaxArray, /, *, np_variant: Optional[bool] = True, out: Optional[JaxArray] = None ) -> JaxArray: @@ -486,6 +604,16 @@ def subtract( return jnp.subtract(x1, x2) +def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.tan(x) + + +def tanh( + x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.tanh(x) + + def trapz( y: JaxArray, /, @@ -498,129 +626,9 @@ def trapz( return jnp.trapz(y, x=x, dx=dx, axis=axis) -def tan(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.tan(x) - - -def tanh( - x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.tanh(x) - - @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def trunc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: if "int" in str(x.dtype): return x else: return jnp.trunc(x) - - -def exp2( - x: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.power(2, x) - - -def imag( - val: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.imag(val) - - -def angle( - z: JaxArray, - /, - *, - deg: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.angle(z, deg=deg) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def erf(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.scipy.special.erf(x) - - -def maximum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 >= x2, x1, x2) - return jnp.maximum(x1, x2) - - -def minimum( - x1: Union[float, JaxArray], - x2: Union[float, JaxArray], - /, - *, - use_where: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return jnp.where(x1 <= x2, x1, x2) - return jnp.minimum(x1, x2) - - -def reciprocal( - x: Union[float, JaxArray], /, *, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.reciprocal(x) - - -def deg2rad(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.deg2rad(x) - - -def rad2deg(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.rad2deg(x) - - -def isreal(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.isreal(x) - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def fmod( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmod(x1, x2) - - -def gcd( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.gcd(x1, x2) - - -def real(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.real(x) diff --git a/ivy/functional/backends/jax/experimental/activations.py b/ivy/functional/backends/jax/experimental/activations.py index 788126f4b44cf..8d270c6c8458a 100644 --- a/ivy/functional/backends/jax/experimental/activations.py +++ b/ivy/functional/backends/jax/experimental/activations.py @@ -8,6 +8,15 @@ import ivy +def elu( + x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None +) -> JaxArray: + ret = jax.nn.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret + + def logit( x: JaxArray, /, @@ -22,6 +31,10 @@ def logit( return jnp.log(x / (1 - x)) +def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jax.nn.log_sigmoid(input) + + def relu6(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: relu6_func = jax.nn.relu6 @@ -38,20 +51,6 @@ def custom_grad_func(x_and_grad, one): return new_func(x).astype(x.dtype) -def thresholded_relu( - x: JaxArray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.where(x > threshold, x, 0).astype(x.dtype) - - -def logsigmoid(input: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jax.nn.log_sigmoid(input) - - def selu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: ret = jax.nn.selu(x).astype(x.dtype) if ivy.exists(out): @@ -66,10 +65,11 @@ def silu(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return ret -def elu( - x: JaxArray, /, *, alpha: float = 1.0, out: Optional[JaxArray] = None +def thresholded_relu( + x: JaxArray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[JaxArray] = None, ) -> JaxArray: - ret = jax.nn.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret + return jnp.where(x > threshold, x, 0).astype(x.dtype) diff --git a/ivy/functional/backends/jax/experimental/creation.py b/ivy/functional/backends/jax/experimental/creation.py index 2bbc0b1b895e6..343a8dfaf9b49 100644 --- a/ivy/functional/backends/jax/experimental/creation.py +++ b/ivy/functional/backends/jax/experimental/creation.py @@ -9,27 +9,23 @@ from ivy.functional.backends.jax import JaxArray import ivy -# Array API Standard # -# ------------------ # - -def vorbis_window( - window_length: JaxArray, +def blackman_window( + size: int, + /, *, - dtype: jnp.dtype = jnp.float32, + periodic: bool = True, + dtype: Optional[jnp.dtype] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.array( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, + if size < 2: + return jnp.ones([size], dtype=dtype) + if periodic: + count = jnp.arange(size) / size + else: + count = jnp.linspace(start=0, stop=size, num=size) + return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( + 0.08 * jnp.cos(2 * jnp.pi * 2 * count) ) @@ -77,6 +73,14 @@ def tril_indices( return jnp.tril_indices(n=n_rows, k=k, m=n_cols) +def trilu( + x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + if upper: + return jnp.triu(x, k) + return jnp.tril(x, k) + + def unsorted_segment_min( data: JaxArray, segment_ids: JaxArray, @@ -104,28 +108,25 @@ def unsorted_segment_sum( return jax.ops.segment_sum(data, segment_ids, num_segments) -def blackman_window( - size: int, - /, +# Array API Standard # +# ------------------ # + + +def vorbis_window( + window_length: JaxArray, *, - periodic: bool = True, - dtype: Optional[jnp.dtype] = None, + dtype: jnp.dtype = jnp.float32, out: Optional[JaxArray] = None, ) -> JaxArray: - if size < 2: - return jnp.ones([size], dtype=dtype) - if periodic: - count = jnp.arange(size) / size - else: - count = jnp.linspace(start=0, stop=size, num=size) - return (0.42 - 0.5 * jnp.cos(2 * jnp.pi * count)) + ( - 0.08 * jnp.cos(2 * jnp.pi * 2 * count) + return jnp.array( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, ) - - -def trilu( - x: JaxArray, /, *, k: int = 0, upper: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - if upper: - return jnp.triu(x, k) - return jnp.tril(x, k) diff --git a/ivy/functional/backends/jax/experimental/elementwise.py b/ivy/functional/backends/jax/experimental/elementwise.py index f6517ca26c95d..c05fecdf8c937 100644 --- a/ivy/functional/backends/jax/experimental/elementwise.py +++ b/ivy/functional/backends/jax/experimental/elementwise.py @@ -19,41 +19,65 @@ jax_ArrayLike = Union[JaxArray, Number] -def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jnp.sinc(x) +# --- Helpers --- # +# --------------- # -@with_supported_dtypes( - {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version -) -def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - return jlax.lgamma(x) +# def gradient( +# x: JaxArray, +# /, +# *, +# spacing: Optional[Union[int, list, tuple]] = 1, +# axis: Optional[Union[int, list, tuple]] = None, +# edge_order: Optional[int] = 1, +# ) -> Union[JaxArray, List[JaxArray]]: +# if type(spacing) == int: +# return jnp.gradient(x, spacing, axis=axis) +# return jnp.gradient(x, *spacing, axis=axis) -def fmax( +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim + + +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis + + +# --- Main --- # +# ------------ # + + +def allclose( x1: JaxArray, x2: JaxArray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[JaxArray] = None, -) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - return jnp.fmax(x1, x2) +) -> bool: + return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) -def float_power( - x1: Union[JaxArray, float, list, tuple], - x2: Union[JaxArray, float, list, tuple], +def conj( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - x1, x2 = promote_types_of_inputs(x1, x2) - if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): - out_dtype = jnp.complex128 - else: - out_dtype = jnp.float64 - return jnp.float_power(x1, x2).astype(out_dtype) + return jnp.conj(x) def copysign( @@ -86,65 +110,6 @@ def count_nonzero( return jnp.array(jnp.count_nonzero(a, axis=axis, keepdims=keepdims), dtype=dtype) -def nansum( - x: JaxArray, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[jnp.dtype] = None, - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - if isinstance(axis, list): - axis = tuple(axis) - return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -def isclose( - a: JaxArray, - b: JaxArray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -def signbit( - x: Union[JaxArray, float, int, list, tuple], - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.signbit(x) - - -def hypot( - x1: JaxArray, - x2: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.hypot(x1, x2) - - -def allclose( - x1: JaxArray, - x2: JaxArray, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[JaxArray] = None, -) -> bool: - return jnp.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - - def diff( x: JaxArray, /, @@ -163,74 +128,54 @@ def diff( return jnp.diff(x, n=n, axis=axis, prepend=prepend, append=append) -def fix( +def digamma( x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.fix(x, out=out) + return js.special.digamma(x) -def nextafter( - x1: JaxArray, - x2: JaxArray, +def fix( + x: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.nextafter(x1, x2) + return jnp.fix(x, out=out) -def zeta( - x: JaxArray, - q: JaxArray, +def float_power( + x1: Union[JaxArray, float, list, tuple], + x2: Union[JaxArray, float, list, tuple], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) - inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) - temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) - temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) - nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) - ret = js.special.zeta(x, q) - ret = ret.at[nan_indices].set(jnp.nan) - ret = ret.at[inf_indices].set(jnp.inf) - return ret - - -# def gradient( -# x: JaxArray, -# /, -# *, -# spacing: Optional[Union[int, list, tuple]] = 1, -# axis: Optional[Union[int, list, tuple]] = None, -# edge_order: Optional[int] = 1, -# ) -> Union[JaxArray, List[JaxArray]]: -# if type(spacing) == int: -# return jnp.gradient(x, spacing, axis=axis) -# return jnp.gradient(x, *spacing, axis=axis) + x1, x2 = promote_types_of_inputs(x1, x2) + if jnp.any(jnp.iscomplex(x1)) or jnp.any(jnp.iscomplex(x2)): + out_dtype = jnp.complex128 + else: + out_dtype = jnp.float64 + return jnp.float_power(x1, x2).astype(out_dtype) -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim +def fmax( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + x1, x2 = promote_types_of_inputs(x1, x2) + return jnp.fmax(x1, x2) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +def frexp( + x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None +) -> Tuple[JaxArray, JaxArray]: + return jnp.frexp(x) def gradient( @@ -422,18 +367,27 @@ def gradient( return outvals -def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: - x, y = promote_types_of_inputs(x, y) - return js.special.xlogy(x, y) +def hypot( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.hypot(x1, x2) -def conj( - x: JaxArray, +def isclose( + a: JaxArray, + b: JaxArray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.conj(x) + return jnp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) def ldexp( @@ -442,10 +396,11 @@ def ldexp( return jnp.ldexp(x1, x2) -def frexp( - x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None -) -> Tuple[JaxArray, JaxArray]: - return jnp.frexp(x) +@with_supported_dtypes( + {"0.4.14 and below": ("float16", "float32", "float64")}, backend_version +) +def lgamma(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jlax.lgamma(x) def modf( @@ -457,10 +412,63 @@ def modf( return jnp.modf(x) -def digamma( +def nansum( x: JaxArray, /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[jnp.dtype] = None, + keepdims: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return js.special.digamma(x) + if isinstance(axis, list): + axis = tuple(axis) + return jnp.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + + +def nextafter( + x1: JaxArray, + x2: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.nextafter(x1, x2) + + +def signbit( + x: Union[JaxArray, float, int, list, tuple], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.signbit(x) + + +def sinc(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + return jnp.sinc(x) + + +def xlogy(x: JaxArray, y: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: + x, y = promote_types_of_inputs(x, y) + return js.special.xlogy(x, y) + + +def zeta( + x: JaxArray, + q: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + temp = jnp.logical_and(jnp.greater(x, 0), jnp.equal(jnp.remainder(x, 2), 0)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + temp = jnp.logical_and(temp, jnp.equal(jnp.remainder(q, 1), 0)) + inf_indices = jnp.logical_or(temp, jnp.equal(x, 1)) + temp = jnp.logical_and(jnp.not_equal(jnp.remainder(x, 2), 0), jnp.greater(x, 1)) + temp = jnp.logical_and(temp, jnp.less_equal(q, 0)) + nan_indices = jnp.logical_or(temp, jnp.less(x, 1)) + ret = js.special.zeta(x, q) + ret = ret.at[nan_indices].set(jnp.nan) + ret = ret.at[inf_indices].set(jnp.inf) + return ret diff --git a/ivy/functional/backends/jax/experimental/layers.py b/ivy/functional/backends/jax/experimental/layers.py index e65256730b950..6ea8306098c5d 100644 --- a/ivy/functional/backends/jax/experimental/layers.py +++ b/ivy/functional/backends/jax/experimental/layers.py @@ -23,6 +23,10 @@ from ivy.functional.backends.jax.experimental.manipulation import _to_nested_tuple +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # determine depth pooling _, _, depth_pooling = _depth_max_pooling_helper( @@ -48,239 +52,8 @@ def _pad_str_to_list(inputs, dims, padding, strides, new_window_shape): return pad_list -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - # This function assumes that param validation is already done - window_shape = tuple(window_shape) - strides = (1,) + strides + (1,) if len(strides) == dim else strides - dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape - if isinstance(dilation, int): - dilation = (1,) + (dilation,) * dim + (1,) - else: - dilation = (1,) + tuple(dilation) + (1,) - - is_single_input = False - if inputs.ndim == len(dims) - 1: - # add singleton batch dimension because lax.reduce_window always - # needs a batch dimension. - inputs = inputs[None] - is_single_input = True - - assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" - - # shape of window after dilation - new_window_shape = tuple( - [ - window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) - for i in range(1, len(dims) - 1) - ] - ) - inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( - inputs, window_shape, strides, dim, data_format="channel_last" - ) - if not depth_pooling: - # manually creating padding list - if isinstance(padding, str): - pad_list = _pad_str_to_list( - inputs, dims, padding, strides, new_window_shape - ) - else: - if isinstance(padding, int): - padding = [(padding,) * 2] * dim - pad_list = [(0, 0)] + list(padding) + [(0, 0)] - - if ceil_mode: - c = [] - for i in range(len(dims) - 2): - pad_list[i + 1], ceil = _padding_ceil_mode( - inputs.shape[i + 1], - new_window_shape[i], - pad_list[i + 1], - strides[i + 1], - True, - ) - c.append(ceil) - - if count_include_pad: - # manually pad inputs with 0 if ceil_mode is True - # because they're not counted in average calculation - if ceil_mode: - ceil = [(0, c[i]) for i in range(len(dims) - 2)] - for i in range(len(dims) - 2): - pad_list[i + 1] = ( - pad_list[i + 1][0], - pad_list[i + 1][1] - ceil[i][1], - ) - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - inputs = jnp.pad( - inputs, - [(0, 0)] + ceil + [(0, 0)], - mode="constant", - constant_values=0.0, - ) - else: - # manually pad inputs with 1s - # because they are counted in average calculation - inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) - pad_list = [(0, 0)] * len(pad_list) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - pad_list = [(0, 0)] * (dim + 2) - - if not ivy.is_array(inputs): - # if dtype is not set here, jax casts it to float64 - inputs = jnp.array(inputs, dtype=jnp.float32) - if not ivy.is_array(init): - init = jnp.array(init, dtype=jnp.float32) - promoted_type = jnp.promote_types(inputs.dtype, init.dtype) - inputs = inputs.astype(promoted_type) - init = init.astype(promoted_type) - y = jlax.reduce_window( - inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation - ) - if is_single_input: - y = jnp.squeeze(y, axis=0) - return y - - -def max_pool1d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCW": - x = jnp.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCW": - res = jnp.transpose(res, (0, 2, 1)) - return res - - -def max_pool2d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCHW": - x = jnp.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCHW": - return jnp.transpose(res, (0, 3, 1, 2)) - - return res - - -def max_pool3d( - x: JaxArray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - if data_format == "NCDHW": - x = jnp.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - res = general_pool( - x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode - ) - - if data_format == "NCDHW": - res = jnp.transpose(res, (0, 4, 1, 2, 3)) - - return res +# --- Main --- # +# ------------ # def avg_pool1d( @@ -512,60 +285,13 @@ def dct( return dct_out -def idct( +def dropout1d( x: JaxArray, + prob: float, /, *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - -def fft( - x: JaxArray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return jnp.fft.fft(x, n, dim, norm) - - -def dropout1d( - x: JaxArray, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", + training: bool = True, + data_format: str = "NWC", out: Optional[JaxArray] = None, ) -> JaxArray: if training: @@ -639,6 +365,225 @@ def dropout3d( return res +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version +) +def embedding( + weights: JaxArray, + indices: JaxArray, + /, + *, + max_norm: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + + embeddings = jnp.take(weights, indices, axis=0) + if max_norm is not None: + norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = jnp.where( + norms > max_norm, embeddings * max_norm / norms, embeddings + ) + embeddings = jnp.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings + ) + return embeddings + + +def fft( + x: JaxArray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return jnp.fft.fft(x, n, dim, norm) + + +def fft2( + x: JaxArray, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_elem_in_list( + norm, + ["backward", "ortho", "forward"], + message=f"Unrecognized normalization mode {norm}", + ) + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " + ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) + + +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + # This function assumes that param validation is already done + window_shape = tuple(window_shape) + strides = (1,) + strides + (1,) if len(strides) == dim else strides + dims = (1,) + window_shape + (1,) if len(window_shape) == dim else window_shape + if isinstance(dilation, int): + dilation = (1,) + (dilation,) * dim + (1,) + else: + dilation = (1,) + tuple(dilation) + (1,) + + is_single_input = False + if inputs.ndim == len(dims) - 1: + # add singleton batch dimension because lax.reduce_window always + # needs a batch dimension. + inputs = inputs[None] + is_single_input = True + + assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})" + + # shape of window after dilation + new_window_shape = tuple( + [ + window_shape[i - 1] + (dilation[i] - 1) * (window_shape[i - 1] - 1) + for i in range(1, len(dims) - 1) + ] + ) + inputs, window_shape, strides, depth_pooling = _determine_depth_max_pooling( + inputs, window_shape, strides, dim, data_format="channel_last" + ) + if not depth_pooling: + # manually creating padding list + if isinstance(padding, str): + pad_list = _pad_str_to_list( + inputs, dims, padding, strides, new_window_shape + ) + else: + if isinstance(padding, int): + padding = [(padding,) * 2] * dim + pad_list = [(0, 0)] + list(padding) + [(0, 0)] + + if ceil_mode: + c = [] + for i in range(len(dims) - 2): + pad_list[i + 1], ceil = _padding_ceil_mode( + inputs.shape[i + 1], + new_window_shape[i], + pad_list[i + 1], + strides[i + 1], + True, + ) + c.append(ceil) + + if count_include_pad: + # manually pad inputs with 0 if ceil_mode is True + # because they're not counted in average calculation + if ceil_mode: + ceil = [(0, c[i]) for i in range(len(dims) - 2)] + for i in range(len(dims) - 2): + pad_list[i + 1] = ( + pad_list[i + 1][0], + pad_list[i + 1][1] - ceil[i][1], + ) + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + inputs = jnp.pad( + inputs, + [(0, 0)] + ceil + [(0, 0)], + mode="constant", + constant_values=0.0, + ) + else: + # manually pad inputs with 1s + # because they are counted in average calculation + inputs = jnp.pad(inputs, pad_list, mode="constant", constant_values=1.0) + pad_list = [(0, 0)] * len(pad_list) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + pad_list = [(0, 0)] * (dim + 2) + + if not ivy.is_array(inputs): + # if dtype is not set here, jax casts it to float64 + inputs = jnp.array(inputs, dtype=jnp.float32) + if not ivy.is_array(init): + init = jnp.array(init, dtype=jnp.float32) + promoted_type = jnp.promote_types(inputs.dtype, init.dtype) + inputs = inputs.astype(promoted_type) + init = init.astype(promoted_type) + y = jlax.reduce_window( + inputs, init, reduce_fn, dims, strides, pad_list, window_dilation=dilation + ) + if is_single_input: + y = jnp.squeeze(y, axis=0) + return y + + +def idct( + x: JaxArray, + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + def ifft( x: JaxArray, dim: int, @@ -671,6 +616,17 @@ def ifft( return jnp.fft.ifft(x, n, dim, norm) +def ifftn( + x: JaxArray, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: str = "backward", + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.fft.ifftn(x, s, axes, norm) + + def interpolate( x: JaxArray, size: Union[Sequence[int], int], @@ -713,19 +669,129 @@ def interpolate( ) -interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (align_corners is None or not align_corners) - and mode - not in [ - "area", - "nearest", - "nd", - "tf_area", - "mitchellcubic", - "gaussian", - "bicubic", - ] -) +def max_pool1d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCW": + x = jnp.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCW": + res = jnp.transpose(res, (0, 2, 1)) + return res + + +def max_pool2d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCHW": + x = jnp.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCHW": + return jnp.transpose(res, (0, 3, 1, 2)) + + return res + + +def max_pool3d( + x: JaxArray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NCDHW": + x = jnp.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + res = general_pool( + x, -jnp.inf, jlax.max, kernel, strides, padding, dims, dilation, ceil_mode + ) + + if data_format == "NCDHW": + res = jnp.transpose(res, (0, 4, 1, 2, 3)) + + return res def reduce_window( @@ -763,79 +829,6 @@ def reduce_window( ) -def fft2( - x: JaxArray, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_elem_in_list( - norm, - ["backward", "ortho", "forward"], - message=f"Unrecognized normalization mode {norm}", - ) - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" - ) - return jnp.fft.fft2(x, s, dim, norm).astype(jnp.complex128) - - -def ifftn( - x: JaxArray, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: str = "backward", - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.fft.ifftn(x, s, axes, norm) - - -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version -) -def embedding( - weights: JaxArray, - indices: JaxArray, - /, - *, - max_norm: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - - embeddings = jnp.take(weights, indices, axis=0) - if max_norm is not None: - norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = jnp.where( - norms > max_norm, embeddings * max_norm / norms, embeddings - ) - embeddings = jnp.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings - ) - return embeddings - - @with_unsupported_dtypes({"0.4.14 and below": ("float16", "complex")}, backend_version) def rfftn( x: JaxArray, @@ -867,3 +860,18 @@ def rfftn( if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") return jnp.fft.rfftn(x, s, axes, norm).astype(jnp.complex128) + + +interpolate.partial_mixed_handler = lambda *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (align_corners is None or not align_corners) + and mode + not in [ + "area", + "nearest", + "nd", + "tf_area", + "mitchellcubic", + "gaussian", + "bicubic", + ] +) diff --git a/ivy/functional/backends/jax/experimental/linear_algebra.py b/ivy/functional/backends/jax/experimental/linear_algebra.py index 897254757e7ba..f1657aceab0af 100644 --- a/ivy/functional/backends/jax/experimental/linear_algebra.py +++ b/ivy/functional/backends/jax/experimental/linear_algebra.py @@ -10,6 +10,28 @@ from ivy.utils.exceptions import IvyNotImplementedException +def adjoint( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return jnp.conjugate(jnp.transpose(x, axes=axes)) + + +def cond( + x: JaxArray, + /, + *, + p: Optional[Union[int, str, None]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.cond(x, p=p) + + def diagflat( x: JaxArray, /, @@ -86,23 +108,14 @@ def diagflat( return ret -def kron( +def dot( a: JaxArray, b: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.kron(a, b) - - -def matrix_exp( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jla.expm(x) + return jnp.dot(a, b) def eig( @@ -120,55 +133,42 @@ def eigvals(x: JaxArray, /) -> JaxArray: return jnp.linalg.eigvals(x) -def adjoint( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return jnp.conjugate(jnp.transpose(x, axes=axes)) - - -def multi_dot( - x: Sequence[JaxArray], +def kron( + a: JaxArray, + b: JaxArray, /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.linalg.multi_dot(x) + return jnp.kron(a, b) -def cond( +def lu_factor( x: JaxArray, /, *, - p: Optional[Union[int, str, None]] = None, + pivot: Optional[bool] = True, out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.cond(x, p=p) +) -> Tuple[JaxArray]: + raise IvyNotImplementedException() -def lu_factor( +def matrix_exp( x: JaxArray, /, *, - pivot: Optional[bool] = True, out: Optional[JaxArray] = None, -) -> Tuple[JaxArray]: - raise IvyNotImplementedException() +) -> JaxArray: + return jla.expm(x) -def dot( - a: JaxArray, - b: JaxArray, +def multi_dot( + x: Sequence[JaxArray], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dot(a, b) + return jnp.linalg.multi_dot(x) dot.support_native_out = True diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index 7258a7ae7ebe7..e38a8343ba23d 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -20,102 +20,180 @@ from ivy.functional.backends.jax import JaxArray -def moveaxis( - a: JaxArray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], +# --- Helpers --- # +# --------------- # + + +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x + + +def _to_nested_tuple(nested_list): + ret = () + if hasattr(nested_list, "__iter__"): + for inner_list in nested_list: + if hasattr(inner_list, "__iter__"): + ret += (tuple(inner_list),) + else: + ret += (inner_list,) + return ret + if ret == (): + return nested_list + + +# --- Main --- # +# ------------ # + + +def atleast_1d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_1d(*arys) + + +def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: + return jnp.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None +) -> List[JaxArray]: + return jnp.atleast_3d(*arys) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return jnp.broadcast_shapes(*shapes) + + +def concat_from_sequence( + input_sequence: Union[Tuple[JaxArray], List[JaxArray]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.moveaxis(a, source, destination) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = jnp.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = jnp.stack(input_sequence, axis=axis) + return ret -def heaviside( - x1: JaxArray, - x2: JaxArray, +def dsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], + /, + *, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def dstack( + arrays: Sequence[JaxArray], /, *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.heaviside(x1, x2) + return jnp.dstack(arrays) -def flipud( - m: JaxArray, +def expand( + x: JaxArray, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.flipud(m) + shape = list(shape) + if len(shape) > len(x.shape): + x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape[i] + return jnp.broadcast_to(x, tuple(shape)) -def vstack( - arrays: Sequence[JaxArray], +def fill_diagonal( + a: JaxArray, + v: Union[int, float], /, *, - out: Optional[JaxArray] = None, + wrap: bool = False, ) -> JaxArray: - return jnp.vstack(arrays) + shape = jnp.array(a.shape) + end = None + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (jnp.cumprod(shape[:-1])).sum() + a = jnp.reshape(a, (-1,)) + a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) + a = jnp.reshape(a, shape) + return a -def hstack( - arrays: Sequence[JaxArray], +def fliplr( + m: JaxArray, /, *, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.hstack(arrays) + return jnp.fliplr(m) -def rot90( +def flipud( m: JaxArray, /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[JaxArray] = None, ) -> JaxArray: - if isinstance(axes, list): - axes = tuple(axes) - return jnp.rot90(m, k, axes) + return jnp.flipud(m) -def top_k( - x: JaxArray, - k: int, +def heaviside( + x1: JaxArray, + x2: JaxArray, /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[JaxArray, JaxArray]] = None, -) -> Tuple[JaxArray, JaxArray]: - k = min(k, x.shape[axis]) - if not largest: - indices = jnp.argsort(x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - else: - indices = jnp.argsort(-x, axis=axis) - indices = jnp.take(indices, jnp.arange(k), axis=axis) - if not sorted: - indices = jnp.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) - val = jnp.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.heaviside(x1, x2) -def fliplr( - m: JaxArray, +def hsplit( + ary: JaxArray, + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) + + +def hstack( + arrays: Sequence[JaxArray], + /, + *, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.fliplr(m) + return jnp.hstack(arrays) def i0( @@ -127,21 +205,16 @@ def i0( return jnp.i0(x) -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x - - -def _to_nested_tuple(nested_list): - ret = () - if hasattr(nested_list, "__iter__"): - for inner_list in nested_list: - if hasattr(inner_list, "__iter__"): - ret += (tuple(inner_list),) - else: - ret += (inner_list,) - return ret - if ret == (): - return nested_list +def moveaxis( + a: JaxArray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.moveaxis(a, source, destination) def pad( @@ -228,57 +301,18 @@ def pad( return ret -def vsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def dsplit( - ary: JaxArray, - indices_or_sections: Union[int, Sequence[int], JaxArray], +def rot90( + m: JaxArray, /, *, copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_1d(*arys) - - -def dstack( - arrays: Sequence[JaxArray], - /, - *, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.dstack(arrays) - - -def atleast_2d(*arys: JaxArray, copy: Optional[bool] = None) -> List[JaxArray]: - return jnp.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[JaxArray, bool, Number], copy: Optional[bool] = None -) -> List[JaxArray]: - return jnp.atleast_3d(*arys) + if isinstance(axes, list): + axes = tuple(axes) + return jnp.rot90(m, k, axes) def take_along_axis( @@ -298,56 +332,28 @@ def take_along_axis( return jnp.take_along_axis(arr, indices, axis, mode=mode) -def hsplit( - ary: JaxArray, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[JaxArray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return jnp.broadcast_shapes(*shapes) - - -def expand( +def top_k( x: JaxArray, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - shape = list(shape) - if len(shape) > len(x.shape): - x = jnp.expand_dims(x, range(len(shape) - len(x.shape))) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape[i] - return jnp.broadcast_to(x, tuple(shape)) - - -def concat_from_sequence( - input_sequence: Union[Tuple[JaxArray], List[JaxArray]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = jnp.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = jnp.stack(input_sequence, axis=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[JaxArray, JaxArray]] = None, +) -> Tuple[JaxArray, JaxArray]: + k = min(k, x.shape[axis]) + if not largest: + indices = jnp.argsort(x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + else: + indices = jnp.argsort(-x, axis=axis) + indices = jnp.take(indices, jnp.arange(k), axis=axis) + if not sorted: + indices = jnp.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", JaxArray), ("indices", JaxArray)]) + val = jnp.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) def unique_consecutive( @@ -393,22 +399,24 @@ def unique_consecutive( ) -def fill_diagonal( - a: JaxArray, - v: Union[int, float], +def vsplit( + ary: JaxArray, + indices_or_sections: Union[int, Sequence[int], JaxArray], /, *, - wrap: bool = False, + copy: Optional[bool] = None, +) -> List[JaxArray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Sequence[JaxArray], + /, + *, + out: Optional[JaxArray] = None, ) -> JaxArray: - shape = jnp.array(a.shape) - end = None - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (jnp.cumprod(shape[:-1])).sum() - a = jnp.reshape(a, (-1,)) - a = a.at[:end:step].set(jnp.array(v).astype(a.dtype)) - a = jnp.reshape(a, shape) - return a + return jnp.vstack(arrays) diff --git a/ivy/functional/backends/jax/experimental/random.py b/ivy/functional/backends/jax/experimental/random.py index cfec5e0f496ce..0b75f787a241c 100644 --- a/ivy/functional/backends/jax/experimental/random.py +++ b/ivy/functional/backends/jax/experimental/random.py @@ -15,26 +15,27 @@ from ivy.func_wrapper import with_unsupported_dtypes from .. import backend_version -# Extra # -# ----- # - -# dirichlet -def dirichlet( - alpha: Union[JaxArray, float, Sequence[float]], - /, +def bernoulli( + probs: Union[float, JaxArray], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + logits: Optional[Union[float, JaxArray]] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: Optional[jaxlib.xla_extension.Device] = None, dtype: Optional[jnp.dtype] = None, seed: Optional[int] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - if seed is not None: + if seed: rng_input = jax.random.PRNGKey(seed) else: RNG_, rng_input = jax.random.split(_getRNG()) _setRNG(RNG_) - return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) + if logits is not None: + probs = jax.nn.softmax(logits, axis=-1) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return jax.random.bernoulli(rng_input, probs, shape=shape) def beta( @@ -56,6 +57,28 @@ def beta( return jax.random.beta(rng_input, a, b, shape, dtype) +# Extra # +# ----- # + + +# dirichlet +def dirichlet( + alpha: Union[JaxArray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[jnp.dtype] = None, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if seed is not None: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.dirichlet(rng_input, alpha, shape=size, dtype=dtype) + + @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, JaxArray], @@ -105,25 +128,3 @@ def poisson( else: ret = jax.random.poisson(rng_input, lam, shape=list_shape).astype(dtype) return ret - - -def bernoulli( - probs: Union[float, JaxArray], - *, - logits: Optional[Union[float, JaxArray]] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: Optional[jaxlib.xla_extension.Device] = None, - dtype: Optional[jnp.dtype] = None, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - if logits is not None: - probs = jax.nn.softmax(logits, axis=-1) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return jax.random.bernoulli(rng_input, probs, shape=shape) diff --git a/ivy/functional/backends/jax/experimental/statistical.py b/ivy/functional/backends/jax/experimental/statistical.py index 64f1bfc031973..b391273e1b628 100644 --- a/ivy/functional/backends/jax/experimental/statistical.py +++ b/ivy/functional/backends/jax/experimental/statistical.py @@ -9,6 +9,198 @@ from ..statistical import _infer_dtype +# --- Helpers --- # +# --------------- # + + +def __find_cummax_indices( + x: JaxArray, + axis: int = 0, +) -> JaxArray: + n, indice, indices = 0, [], [] + + if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + else: + z_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices = indices.at[y_index].set(multi_index[0]) + else: + indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) + else: + n, indices = 0, [] + for idx, y in enumerate(x): + if idx == 0 or x[n] <= y: + n = idx + indices.append(n) + + return jnp.asarray(indices, dtype="int64") + + +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] + + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices + + +# --- Main --- # +# ------------ # + + +def bincount( + x: JaxArray, + /, + *, + weights: Optional[JaxArray] = None, + minlength: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + if weights is not None: + ret = jnp.bincount(x, weights=weights, minlength=minlength) + ret = ret.astype(weights.dtype) + else: + ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) + return ret + + +def corrcoef( + x: JaxArray, + /, + *, + y: Optional[JaxArray] = None, + rowvar: bool = True, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.corrcoef(x, y=y, rowvar=rowvar) + + +def cov( + x1: JaxArray, + x2: JaxArray = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[JaxArray] = None, + aweights: Optional[JaxArray] = None, + dtype: Optional[jnp.dtype] = None, +) -> JaxArray: + if not dtype: + x1 = jnp.asarray(x1, dtype=jnp.float64) + + if jnp.ndim(x1) > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if jnp.ndim(x2) > 2: + raise ValueError("x2 has more than 2 dimensions") + + if fweights is not None: + fweights = jnp.asarray(fweights, dtype=jnp.int64) + + return jnp.cov( + m=x1, + y=x2, + rowvar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, + ) + + +def cummax( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> Tuple[JaxArray, JaxArray]: + if x.dtype in (jnp.bool_, jnp.float16): + x = x.astype(jnp.float64) + elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): + x = x.astype(jnp.int64) + elif x.dtype in (jnp.complex128, jnp.complex64): + x = jnp.real(x).astype(jnp.float64) + + if exclusive or (reverse and exclusive): + if exclusive and reverse: + indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) + x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + x, indices = jnp.concatenate( + (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 + ), jnp.concatenate( + (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 + ) + x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) + res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + indices = __find_cummax_indices(x, axis=axis) + res = jlax.cummax(x, axis=axis) + return res, indices + + if reverse: + y = jnp.flip(x, axis=axis) + indices = __find_cummax_indices(y, axis=axis) + indices = jnp.flip(indices, axis=axis) + else: + indices = __find_cummax_indices(x, axis=axis) + return jlax.cummax(x, axis, reverse=reverse), indices + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cummin( + x: JaxArray, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if axis < 0: + axis = axis + len(x.shape) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + return jlax.cummin(x, axis, reverse=reverse).astype(dtype) + + @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16",)}, backend_version, @@ -120,6 +312,16 @@ def histogram( return ret +def igamma( + a: JaxArray, + /, + *, + x: JaxArray, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jlax.igamma(a=a, x=x) + + @with_unsupported_dtypes( {"0.4.14 and below": ("complex64", "complex128")}, backend_version ) @@ -162,34 +364,6 @@ def nanmean( return jnp.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) -def quantile( - a: JaxArray, - q: Union[float, JaxArray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - interpolation = "nearest" if interpolation == "nearest_jax" else interpolation - return jnp.quantile( - a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out - ) - - -def corrcoef( - x: JaxArray, - /, - *, - y: Optional[JaxArray] = None, - rowvar: bool = True, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.corrcoef(x, y=y, rowvar=rowvar) - - def nanmedian( input: JaxArray, /, @@ -219,184 +393,18 @@ def nanmedian( ) -def bincount( - x: JaxArray, - /, - *, - weights: Optional[JaxArray] = None, - minlength: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - if weights is not None: - ret = jnp.bincount(x, weights=weights, minlength=minlength) - ret = ret.astype(weights.dtype) - else: - ret = jnp.bincount(x, minlength=minlength).astype(x.dtype) - return ret - - -def cov( - x1: JaxArray, - x2: JaxArray = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[JaxArray] = None, - aweights: Optional[JaxArray] = None, - dtype: Optional[jnp.dtype] = None, -) -> JaxArray: - if not dtype: - x1 = jnp.asarray(x1, dtype=jnp.float64) - - if jnp.ndim(x1) > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if jnp.ndim(x2) > 2: - raise ValueError("x2 has more than 2 dimensions") - - if fweights is not None: - fweights = jnp.asarray(fweights, dtype=jnp.int64) - - return jnp.cov( - m=x1, - y=x2, - rowvar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - ) - - -def cummax( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> Tuple[JaxArray, JaxArray]: - if x.dtype in (jnp.bool_, jnp.float16): - x = x.astype(jnp.float64) - elif x.dtype in (jnp.int16, jnp.int8, jnp.uint8): - x = x.astype(jnp.int64) - elif x.dtype in (jnp.complex128, jnp.complex64): - x = jnp.real(x).astype(jnp.float64) - - if exclusive or (reverse and exclusive): - if exclusive and reverse: - indices = __find_cummax_indices(jnp.flip(x, axis=axis), axis=axis) - x = jlax.cummax(jnp.flip(x, axis=axis), axis=axis) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - x, indices = jnp.concatenate( - (jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1 - ), jnp.concatenate( - (jnp.zeros_like(indices[..., -1:]), indices[..., :-1]), -1 - ) - x, indices = jnp.swapaxes(x, axis, -1), jnp.swapaxes(indices, axis, -1) - res, indices = jnp.flip(x, axis=axis), jnp.flip(indices, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - indices = __find_cummax_indices(x, axis=axis) - res = jlax.cummax(x, axis=axis) - return res, indices - - if reverse: - y = jnp.flip(x, axis=axis) - indices = __find_cummax_indices(y, axis=axis) - indices = jnp.flip(indices, axis=axis) - else: - indices = __find_cummax_indices(x, axis=axis) - return jlax.cummax(x, axis, reverse=reverse), indices - - -def __find_cummax_indices( - x: JaxArray, - axis: int = 0, -) -> JaxArray: - n, indice, indices = 0, [], [] - - if isinstance(x[0], JaxArray) and len(x[0].shape) >= 1: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - else: - z_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices = jnp.zeros(jnp.asarray(indices.shape), dtype=x.dtype) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices = indices.at[y_index].set(multi_index[0]) - else: - indices = indices.at[y_index].set(n1[tuple(multi_index[1:])]) - else: - n, indices = 0, [] - for idx, y in enumerate(x): - if idx == 0 or x[n] <= y: - n = idx - indices.append(n) - - return jnp.asarray(indices, dtype="int64") - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cummin( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if axis < 0: - axis = axis + len(x.shape) - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - return jlax.cummin(x, axis, reverse=reverse).astype(dtype) - - -def igamma( +def quantile( a: JaxArray, + q: Union[float, JaxArray], /, *, - x: JaxArray, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, out: Optional[JaxArray] = None, ) -> JaxArray: - return jlax.igamma(a=a, x=x) + axis = tuple(axis) if isinstance(axis, list) else axis + interpolation = "nearest" if interpolation == "nearest_jax" else interpolation + return jnp.quantile( + a, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ) diff --git a/ivy/functional/backends/jax/general.py b/ivy/functional/backends/jax/general.py index 9221319cfeaed..1315cef38980a 100644 --- a/ivy/functional/backends/jax/general.py +++ b/ivy/functional/backends/jax/general.py @@ -22,33 +22,8 @@ from . import backend_version -def container_types(): - flat_mapping_spec = importlib.util.find_spec( - "FlatMapping", "haiku._src.data_structures" - ) - if not flat_mapping_spec: - from haiku._src.data_structures import FlatMapping - else: - FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) - return [FlatMapping] - - -def current_backend_str() -> str: - return "jax" - - -def is_native_array(x, /, *, exclusive=False): - if exclusive: - return isinstance(x, NativeArray) - return isinstance( - x, - ( - NativeArray, - jax.interpreters.ad.JVPTracer, - jax.core.ShapedArray, - jax.interpreters.partial_eval.DynamicJaxprTracer, - ), - ) +# --- Helpers --- # +# --------------- # def _mask_to_index(query, x): @@ -61,63 +36,35 @@ def _mask_to_index(query, x): return jnp.where(query), expected_shape -def get_item( - x: JaxArray, - /, - query: Union[JaxArray, Tuple], - *, - copy: bool = None, -) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return jnp.array([], dtype=x.dtype) - else: - return jnp.expand_dims(x, 0) - query, _ = _mask_to_index(query, x) - elif isinstance(query, list): - query = (query,) - return x.__getitem__(query) +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view -def set_item( - x: JaxArray, - query: Union[JaxArray, Tuple], - val: JaxArray, - /, - *, - copy: Optional[bool] = False, -) -> JaxArray: - if ivy.is_array(query) and ivy.is_bool_dtype(query): - query, expected_shape = _mask_to_index(query, x) - val = _broadcast_to(val, expected_shape)._data - ret = x.at[query].set(val) - if copy: - return ret - return ivy.inplace_update(x, _to_device(ret)) +# --- Main --- # +# ------------ # def array_equal(x0: JaxArray, x1: JaxArray, /) -> bool: return bool(jnp.array_equal(x0, x1)) -@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) -def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return np.array(_to_array(x)) - else: - return np.asarray(_to_array(x)) - - -def to_scalar(x: JaxArray, /) -> Number: - if isinstance(x, Number): - return x +def container_types(): + flat_mapping_spec = importlib.util.find_spec( + "FlatMapping", "haiku._src.data_structures" + ) + if not flat_mapping_spec: + from haiku._src.data_structures import FlatMapping else: - return _to_array(x).item() + FlatMapping = importlib.util.module_from_spec(flat_mapping_spec) + return [FlatMapping] -def to_list(x: JaxArray, /) -> list: - return _to_array(x).tolist() +def current_backend_str() -> str: + return "jax" def gather( @@ -152,6 +99,36 @@ def gather( return result +def gather_nd( + params: JaxArray, + indices: JaxArray, + /, + *, + batch_dims: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = jnp.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -183,34 +160,23 @@ def gather_nd_helper(params, indices): return ret -def gather_nd( - params: JaxArray, - indices: JaxArray, +def get_item( + x: JaxArray, /, + query: Union[JaxArray, Tuple], *, - batch_dims: int = 0, - out: Optional[JaxArray] = None, + copy: bool = None, ) -> JaxArray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return jnp.array([], dtype=x.dtype) else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = jnp.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result + return jnp.expand_dims(x, 0) + query, _ = _mask_to_index(query, x) + elif isinstance(query, list): + query = (query,) + return x.__getitem__(query) def get_num_dims(x: JaxArray, /, *, as_array: bool = False) -> Union[JaxArray, int]: @@ -288,18 +254,40 @@ def inplace_update( return val -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view - - def inplace_variables_supported(): return False +def is_native_array(x, /, *, exclusive=False): + if exclusive: + return isinstance(x, NativeArray) + return isinstance( + x, + ( + NativeArray, + jax.interpreters.ad.JVPTracer, + jax.core.ShapedArray, + jax.interpreters.partial_eval.DynamicJaxprTracer, + ), + ) + + +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def isin( + elements: JaxArray, + test_elements: JaxArray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> JaxArray: + return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) + + +def itemsize(x: JaxArray) -> int: + return x.itemsize + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -342,9 +330,6 @@ def scatter_flat( return target -scatter_flat.support_native_out = True - - def scatter_nd( indices: JaxArray, updates: JaxArray, @@ -392,7 +377,21 @@ def scatter_nd( return target -scatter_nd.support_native_out = True +def set_item( + x: JaxArray, + query: Union[JaxArray, Tuple], + val: JaxArray, + /, + *, + copy: Optional[bool] = False, +) -> JaxArray: + if ivy.is_array(query) and ivy.is_bool_dtype(query): + query, expected_shape = _mask_to_index(query, x) + val = _broadcast_to(val, expected_shape)._data + ret = x.at[query].set(val) + if copy: + return ret + return ivy.inplace_update(x, _to_device(ret)) def shape( @@ -407,6 +406,25 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: JaxArray, /) -> list: + return _to_array(x).tolist() + + +@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) +def to_numpy(x: JaxArray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return np.array(_to_array(x)) + else: + return np.asarray(_to_array(x)) + + +def to_scalar(x: JaxArray, /) -> Number: + if isinstance(x, Number): + return x + else: + return _to_array(x).item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -418,17 +436,5 @@ def vmap( ) -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def isin( - elements: JaxArray, - test_elements: JaxArray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> JaxArray: - return jnp.isin(elements, test_elements, assume_unique=assume_unique, invert=invert) - - -def itemsize(x: JaxArray) -> int: - return x.itemsize +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True diff --git a/ivy/functional/backends/jax/gradients.py b/ivy/functional/backends/jax/gradients.py index a5e6d789a070b..dcd85bc2b395e 100644 --- a/ivy/functional/backends/jax/gradients.py +++ b/ivy/functional/backends/jax/gradients.py @@ -18,19 +18,8 @@ ) -# ToDo: modify these functions to track whether variable() has been called -def variable(x, /): - return x - - -def is_variable(x, /, *, exclusive=False): - if exclusive: - return False - return isinstance(x, NativeArray) - - -def variable_data(x: JaxArray, /) -> JaxArray: - return x +# --- Helpers --- # +# --------------- # def _forward_fn( @@ -68,6 +57,10 @@ def _forward_fn( return ret_values +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: JaxArray, @@ -127,21 +120,18 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - value, grad = jax.value_and_grad(grad_fn)(xs) - return ivy.to_ivy(value), ivy.to_ivy(grad) - +def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): + grad_fn = lambda x_in: ivy.to_native(func(x_in)) + callback_fn = lambda x_in: ivy.to_ivy( + jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) + ) return callback_fn -def stop_gradient( - x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None -) -> JaxArray: - return jlax.stop_gradient(x) +def is_variable(x, /, *, exclusive=False): + if exclusive: + return False + return isinstance(x, NativeArray) def jac(func: Callable): @@ -158,9 +148,27 @@ def jac(func: Callable): return callback_fn -def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0): - grad_fn = lambda x_in: ivy.to_native(func(x_in)) - callback_fn = lambda x_in: ivy.to_ivy( - jax.grad(grad_fn, argnums)(ivy.to_native(x_in)) - ) +def stop_gradient( + x: JaxArray, /, *, preserve_type: bool = True, out: Optional[JaxArray] = None +) -> JaxArray: + return jlax.stop_gradient(x) + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + value, grad = jax.value_and_grad(grad_fn)(xs) + return ivy.to_ivy(value), ivy.to_ivy(grad) + return callback_fn + + +# ToDo: modify these functions to track whether variable() has been called +def variable(x, /): + return x + + +def variable_data(x: JaxArray, /) -> JaxArray: + return x diff --git a/ivy/functional/backends/jax/layers.py b/ivy/functional/backends/jax/layers.py index ff4c264439da7..c27e92a2680eb 100644 --- a/ivy/functional/backends/jax/layers.py +++ b/ivy/functional/backends/jax/layers.py @@ -16,41 +16,18 @@ ) -def _transpose_padding_helper(k, s, padding, dilation, diff=0): - k = (k - 1) * dilation + 1 - if padding == "SAME": - pad_len = k + s - 2 - pad_len -= diff - if s > k - 1: - pad_a = k - 1 - else: - pad_a = int(jnp.ceil(pad_len / 2)) - else: - pad_len = k + s - 2 + max(k - s, 0) - pad_a = k - 1 - pad_b = pad_len - pad_a - return pad_a, pad_b +# --- Helpers --- # +# --------------- # -def _get_tranpose_padding( - x_shape, filter_shape, strides, padding, dims, dilations, output_shape -): - new_shape = [ - _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) - for i in range(dims) - ] - if output_shape is None: - output_shape = [x_shape[0], *new_shape, filter_shape[-1]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] - shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] - pad_list = [ - _transpose_padding_helper( - filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] - ) - for i in range(dims) - ] - return pad_list +def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): + first = True if filter_format == "channel_first" else False + if dims == 1: + return "WIO" if not first else "OIW" + if dims == 2: + return "HWIO" if not first else "OIHW" + elif dims == 3: + return "DHWIO" if not first else "OIDHW" def _get_new_padding_before_conv( @@ -94,6 +71,47 @@ def _get_new_padding_before_conv( return padding +def _get_tranpose_padding( + x_shape, filter_shape, strides, padding, dims, dilations, output_shape +): + new_shape = [ + _deconv_length(x_shape[i], strides[i], filter_shape[i], padding, dilations[i]) + for i in range(dims) + ] + if output_shape is None: + output_shape = [x_shape[0], *new_shape, filter_shape[-1]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + list(output_shape) + [filter_shape[-1]] + shape_diff = [-(output_shape[1 + i] - new_shape[i]) for i in range(dims)] + pad_list = [ + _transpose_padding_helper( + filter_shape[i], strides[i], padding, dilations[i], shape_diff[i] + ) + for i in range(dims) + ] + return pad_list + + +def _transpose_padding_helper(k, s, padding, dilation, diff=0): + k = (k - 1) * dilation + 1 + if padding == "SAME": + pad_len = k + s - 2 + pad_len -= diff + if s > k - 1: + pad_a = k - 1 + else: + pad_a = int(jnp.ceil(pad_len / 2)) + else: + pad_len = k + s - 2 + max(k - s, 0) + pad_a = k - 1 + pad_b = pad_len - pad_a + return pad_a, pad_b + + +# --- Main --- # +# ------------ # + + def conv1d( x: JaxArray, filters: JaxArray, @@ -231,37 +249,6 @@ def conv2d_transpose( return res -def depthwise_conv2d( - x: JaxArray, - filters: JaxArray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[JaxArray] = None, -) -> JaxArray: - strides = [strides] * 2 if isinstance(strides, int) else strides - strides = [strides[1], strides[2]] if len(strides) == 4 else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters - cn = filters.shape[-1] - filters = jnp.expand_dims(filters, -2) - return jlax.conv_general_dilated( - x, - filters, - strides, - padding, - None, - dilations, - (data_format, "HWIO", data_format), - feature_group_count=cn, - ) - - def conv3d( x: JaxArray, filters: JaxArray, @@ -330,16 +317,6 @@ def conv3d_transpose( return res -def _get_filter_dataformat(dims: int = 2, filter_format: str = "channel_last"): - first = True if filter_format == "channel_first" else False - if dims == 1: - return "WIO" if not first else "OIW" - if dims == 2: - return "HWIO" if not first else "OIHW" - elif dims == 3: - return "DHWIO" if not first else "OIDHW" - - def conv_general_dilated( x: JaxArray, filters: JaxArray, @@ -455,3 +432,34 @@ def conv_general_transpose( if data_format == "channel_first": return jnp.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +def depthwise_conv2d( + x: JaxArray, + filters: JaxArray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[JaxArray] = None, +) -> JaxArray: + strides = [strides] * 2 if isinstance(strides, int) else strides + strides = [strides[1], strides[2]] if len(strides) == 4 else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + filters = jnp.squeeze(filters, 3) if filters.ndim == 4 else filters + cn = filters.shape[-1] + filters = jnp.expand_dims(filters, -2) + return jlax.conv_general_dilated( + x, + filters, + strides, + padding, + None, + dilations, + (data_format, "HWIO", data_format), + feature_group_count=cn, + ) diff --git a/ivy/functional/backends/jax/linear_algebra.py b/ivy/functional/backends/jax/linear_algebra.py index 383a9d032abc6..64c80f2339090 100644 --- a/ivy/functional/backends/jax/linear_algebra.py +++ b/ivy/functional/backends/jax/linear_algebra.py @@ -58,13 +58,19 @@ def det(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray: return jnp.linalg.det(x) -@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) -def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] - ) - eigenvalues, eigenvectors = jnp.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def diag( + x: JaxArray, + /, + *, + k: int = 0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.diag(x, k=k) @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) @@ -92,15 +98,13 @@ def diagonal( return ret -def tensorsolve( - x1: JaxArray, - x2: JaxArray, - /, - *, - axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.linalg.tensorsolve(x1, x2, axes) +@with_unsupported_dtypes({"0.4.14 and below": ("float16", "bfloat16")}, backend_version) +def eig(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> Tuple[JaxArray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", JaxArray), ("eigenvectors", JaxArray)] + ) + eigenvalues, eigenvectors = jnp.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) @with_unsupported_dtypes( @@ -388,6 +392,17 @@ def tensordot( return jnp.tensordot(x1, x2, axes) +def tensorsolve( + x1: JaxArray, + x2: JaxArray, + /, + *, + axes: Optional[Union[int, Tuple[Sequence[int], Sequence[int]]]] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.linalg.tensorsolve(x1, x2, axes) + + @with_unsupported_dtypes( {"0.4.14 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -404,6 +419,21 @@ def trace( return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) +@with_unsupported_dtypes( + {"0.4.14 and below": ("bfloat16", "float16", "complex")}, + backend_version, +) +def vander( + x: JaxArray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) def vecdot( x1: JaxArray, x2: JaxArray, /, *, axis: int = -1, out: Optional[JaxArray] = None @@ -438,36 +468,6 @@ def vector_norm( return jnp.sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def diag( - x: JaxArray, - /, - *, - k: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.diag(x, k=k) - - -@with_unsupported_dtypes( - {"0.4.14 and below": ("bfloat16", "float16", "complex")}, - backend_version, -) -def vander( - x: JaxArray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes( { "0.4.14 and below": ( diff --git a/ivy/functional/backends/jax/manipulation.py b/ivy/functional/backends/jax/manipulation.py index 61e638b9509cb..4ebc42a93b9c7 100644 --- a/ivy/functional/backends/jax/manipulation.py +++ b/ivy/functional/backends/jax/manipulation.py @@ -12,10 +12,61 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x +# --- Main --- # +# ------------ # + + +def clip( + x: JaxArray, + x_min: Union[Number, JaxArray], + x_max: Union[Number, JaxArray], + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + if ( + hasattr(x_min, "dtype") + and hasattr(x_max, "dtype") + and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) + ): + if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( + jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float32) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + elif ( + jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) + ) and ( + jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) + or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) + ): + promoted_type = jnp.promote_types(x.dtype, jnp.float64) + promoted_type = jnp.promote_types(promoted_type, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x = x.astype(promoted_type) + else: + promoted_type = jnp.promote_types(x.dtype, x_min.dtype) + promoted_type = jnp.promote_types(promoted_type, x_max.dtype) + x.astype(promoted_type) + # jnp.clip isn't used because of inconsistent gradients + x = jnp.where(x > x_max, x_max, x) + return jnp.where(x < x_min, x_min, x) + + # Array API Standard # # -------------------# @@ -42,6 +93,18 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) +@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) +def constant_pad( + x: JaxArray, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) + + def expand_dims( x: JaxArray, /, @@ -79,6 +142,17 @@ def permute_dims( return jnp.transpose(x, axes) +def repeat( + x: JaxArray, + /, + repeats: Union[int, Iterable[int]], + *, + axis: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.repeat(x, repeats, axis) + + def reshape( x: JaxArray, /, @@ -114,38 +188,6 @@ def roll( return jnp.roll(x, shift, axis) -def squeeze( - x: JaxArray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - else: - ret = jnp.squeeze(x, axis=axis) - return ret - - -def stack( - arrays: Union[Tuple[JaxArray], List[JaxArray]], - /, - *, - axis: int = 0, - out: Optional[JaxArray] = None, -) -> JaxArray: - try: - return jnp.stack(arrays, axis=axis) - except ValueError as error: - raise ivy.utils.exceptions.IvyIndexError(error) - - # Extra # # ------# @@ -184,76 +226,54 @@ def split( return jnp.split(x, num_or_size_splits, axis) -def repeat( +def squeeze( x: JaxArray, /, - repeats: Union[int, Iterable[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.repeat(x, repeats, axis) - - -def tile( - x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.tile(x, repeats) + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + else: + ret = jnp.squeeze(x, axis=axis) + return ret -def clip( - x: JaxArray, - x_min: Union[Number, JaxArray], - x_max: Union[Number, JaxArray], +def stack( + arrays: Union[Tuple[JaxArray], List[JaxArray]], /, *, + axis: int = 0, out: Optional[JaxArray] = None, ) -> JaxArray: - if ( - hasattr(x_min, "dtype") - and hasattr(x_max, "dtype") - and (x.dtype != x_min.dtype or x.dtype != x_max.dtype) - ): - if (jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype)) and ( - jnp.int16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint16 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float32) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - elif ( - jnp.float16 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.float32 in (x.dtype, x_min.dtype, x_max.dtype) - ) and ( - jnp.int32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint32 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.uint64 in (x.dtype, x_min.dtype, x_max.dtype) - or jnp.int64 in (x.dtype, x_min.dtype, x_max.dtype) - ): - promoted_type = jnp.promote_types(x.dtype, jnp.float64) - promoted_type = jnp.promote_types(promoted_type, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x = x.astype(promoted_type) - else: - promoted_type = jnp.promote_types(x.dtype, x_min.dtype) - promoted_type = jnp.promote_types(promoted_type, x_max.dtype) - x.astype(promoted_type) - # jnp.clip isn't used because of inconsistent gradients - x = jnp.where(x > x_max, x_max, x) - return jnp.where(x < x_min, x_min, x) + try: + return jnp.stack(arrays, axis=axis) + except ValueError as error: + raise ivy.utils.exceptions.IvyIndexError(error) -@with_unsupported_dtypes({"0.4.14 and below": ("uint64",)}, backend_version) -def constant_pad( +def swapaxes( x: JaxArray, + axis0: int, + axis1: int, /, - pad_width: List[List[int]], *, - value: Number = 0.0, + copy: Optional[bool] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) + return jnp.swapaxes(x, axis0, axis1) + + +def tile( + x: JaxArray, /, repeats: Iterable[int], *, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.tile(x, repeats) def unstack( @@ -278,15 +298,3 @@ def zero_pad( x: JaxArray, /, pad_width: List[List[int]], *, out: Optional[JaxArray] = None ): return jnp.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=0) - - -def swapaxes( - x: JaxArray, - axis0: int, - axis1: int, - /, - *, - copy: Optional[bool] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.swapaxes(x, axis0, axis1) diff --git a/ivy/functional/backends/jax/random.py b/ivy/functional/backends/jax/random.py index 53aedb4e811e4..38b599f8ef405 100644 --- a/ivy/functional/backends/jax/random.py +++ b/ivy/functional/backends/jax/random.py @@ -26,12 +26,8 @@ def __init__(self): self.key = jax.random.PRNGKey(0) -RNG = RNGWrapper() - - -def _setRNG(key): - global RNG - RNG.key = key +# --- Helpers --- # +# --------------- # def _getRNG(): @@ -39,47 +35,13 @@ def _getRNG(): return RNG.key -def random_uniform( - *, - low: Union[float, JaxArray] = 0.0, - high: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - shape = _check_bounds_and_get_shape(low, high, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.uniform( - rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 - ).astype(dtype) +def _setRNG(key): + global RNG + RNG.key = key -def random_normal( - *, - mean: Union[float, JaxArray] = 0.0, - std: Union[float, JaxArray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: jaxlib.xla_extension.Device, - dtype: jnp.dtype, - seed: Optional[int] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - - if seed: - rng_input = jax.random.PRNGKey(seed) - else: - RNG_, rng_input = jax.random.split(_getRNG()) - _setRNG(RNG_) - return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, backend_version) @@ -152,6 +114,49 @@ def randint( return jax.random.randint(rng_input, shape, low, high, dtype) +def random_normal( + *, + mean: Union[float, JaxArray] = 0.0, + std: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.normal(rng_input, shape, dtype=dtype) * std + mean + + +def random_uniform( + *, + low: Union[float, JaxArray] = 0.0, + high: Union[float, JaxArray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: jaxlib.xla_extension.Device, + dtype: jnp.dtype, + seed: Optional[int] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + shape = _check_bounds_and_get_shape(low, high, shape).shape + + if seed: + rng_input = jax.random.PRNGKey(seed) + else: + RNG_, rng_input = jax.random.split(_getRNG()) + _setRNG(RNG_) + return jax.random.uniform( + rng_input, shape, minval=low, maxval=high, dtype=jnp.float32 + ).astype(dtype) + + def seed(*, seed_value: int = 0) -> None: _setRNG(jax.random.PRNGKey(seed_value)) return @@ -176,3 +181,6 @@ def shuffle( # jax.random.shuffle is deprecated; identical behaviour reproduced with # jax.random.permutation return jax.random.permutation(key=rng_input, x=x, axis=axis, independent=True) + + +RNG = RNGWrapper() diff --git a/ivy/functional/backends/jax/searching.py b/ivy/functional/backends/jax/searching.py index 794436773c74a..e5f6a320ef428 100644 --- a/ivy/functional/backends/jax/searching.py +++ b/ivy/functional/backends/jax/searching.py @@ -64,6 +64,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: JaxArray, + /, + *, + out: Optional[JaxArray] = None, +) -> JaxArray: + return jnp.argwhere(x) + + def nonzero( x: JaxArray, /, @@ -90,16 +103,3 @@ def where( ) -> JaxArray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(jnp.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: JaxArray, - /, - *, - out: Optional[JaxArray] = None, -) -> JaxArray: - return jnp.argwhere(x) diff --git a/ivy/functional/backends/jax/sorting.py b/ivy/functional/backends/jax/sorting.py index 5111eca178530..5ccd39bef98f6 100644 --- a/ivy/functional/backends/jax/sorting.py +++ b/ivy/functional/backends/jax/sorting.py @@ -26,20 +26,15 @@ def argsort( ) -def sort( - x: JaxArray, +# msort +@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) +def msort( + a: Union[JaxArray, list, tuple], /, *, - axis: int = -1, - descending: bool = False, - stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - kind = "stable" if stable else "quicksort" - ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) - if descending: - ret = jnp.asarray(jnp.flip(ret, axis=axis)) - return ret + return jnp.sort(a, axis=0, kind="mergesort") def searchsorted( @@ -79,12 +74,17 @@ def searchsorted( return ret.astype(ret_dtype) -# msort -@with_unsupported_dtypes({"0.4.14 and below": ("complex",)}, backend_version) -def msort( - a: Union[JaxArray, list, tuple], +def sort( + x: JaxArray, /, *, + axis: int = -1, + descending: bool = False, + stable: bool = True, out: Optional[JaxArray] = None, ) -> JaxArray: - return jnp.sort(a, axis=0, kind="mergesort") + kind = "stable" if stable else "quicksort" + ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind)) + if descending: + ret = jnp.asarray(jnp.flip(ret, axis=axis)) + return ret diff --git a/ivy/functional/backends/jax/statistical.py b/ivy/functional/backends/jax/statistical.py index b03be2a367141..b6af324bafde4 100644 --- a/ivy/functional/backends/jax/statistical.py +++ b/ivy/functional/backends/jax/statistical.py @@ -9,20 +9,102 @@ from ivy.functional.backends.jax import JaxArray from . import backend_version -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -def min( + +def _infer_dtype(dtype: jnp.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) +def cumprod( x: JaxArray, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[jnp.dtype] = None, out: Optional[JaxArray] = None, ) -> JaxArray: - axis = tuple(axis) if isinstance(axis, list) else axis - return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return jnp.cumprod(x, axis, dtype=dtype) + elif exclusive and reverse: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + return jnp.flip(x, axis=(axis,)) + + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumprod(x, -1, dtype=dtype) + return jnp.swapaxes(x, axis, -1) + else: + x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) + return jnp.flip(x, axis=axis) + + +def cumsum( + x: JaxArray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is jnp.bool_: + dtype = ivy.default_int_dtype(as_native=True) + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.swapaxes(x, axis, -1) + res = jnp.flip(x, axis=axis) + elif exclusive: + x = jnp.swapaxes(x, axis, -1) + x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = jnp.cumsum(x, -1, dtype=dtype) + res = jnp.swapaxes(x, axis, -1) + elif reverse: + x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) + res = jnp.flip(x, axis=axis) + return res + return jnp.cumsum(x, axis, dtype=dtype) + + +def einsum( + equation: str, *operands: JaxArray, out: Optional[JaxArray] = None +) -> JaxArray: + return jnp.einsum(equation, *operands) def max( @@ -49,11 +131,20 @@ def mean( return jnp.mean(x, axis=axis, keepdims=keepdims) -def _infer_dtype(dtype: jnp.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +# Array API Standard # +# -------------------# + + +def min( + x: JaxArray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[JaxArray] = None, +) -> JaxArray: + axis = tuple(axis) if isinstance(axis, list) else axis + return jnp.min(a=jnp.asarray(x), axis=axis, keepdims=keepdims) def prod( @@ -133,85 +224,3 @@ def var( x.dtype, copy=False, ) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version) -def cumprod( - x: JaxArray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return jnp.cumprod(x, axis, dtype=dtype) - elif exclusive and reverse: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - return jnp.flip(x, axis=(axis,)) - - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumprod(x, -1, dtype=dtype) - return jnp.swapaxes(x, axis, -1) - else: - x = jnp.cumprod(jnp.flip(x, axis=(axis,)), axis=axis, dtype=dtype) - return jnp.flip(x, axis=axis) - - -def cumsum( - x: JaxArray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[jnp.dtype] = None, - out: Optional[JaxArray] = None, -) -> JaxArray: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is jnp.bool_: - dtype = ivy.default_int_dtype(as_native=True) - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.swapaxes(x, axis, -1) - res = jnp.flip(x, axis=axis) - elif exclusive: - x = jnp.swapaxes(x, axis, -1) - x = jnp.concatenate((jnp.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = jnp.cumsum(x, -1, dtype=dtype) - res = jnp.swapaxes(x, axis, -1) - elif reverse: - x = jnp.cumsum(jnp.flip(x, axis=axis), axis=axis, dtype=dtype) - res = jnp.flip(x, axis=axis) - return res - return jnp.cumsum(x, axis, dtype=dtype) - - -def einsum( - equation: str, *operands: JaxArray, out: Optional[JaxArray] = None -) -> JaxArray: - return jnp.einsum(equation, *operands) diff --git a/ivy/functional/backends/mxnet/activations.py b/ivy/functional/backends/mxnet/activations.py index 8efd486936334..a6f3fc0e9b713 100644 --- a/ivy/functional/backends/mxnet/activations.py +++ b/ivy/functional/backends/mxnet/activations.py @@ -32,6 +32,14 @@ def leaky_relu( return mx.nd.LeakyReLU(x, slope=alpha) +def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): + raise IvyNotImplementedException() + + +def mish(x: None, /, *, out: Optional[None] = None) -> None: + raise IvyNotImplementedException() + + def relu(x: None, /, *, complex_mode="jax", out: Optional[None] = None) -> None: return mx.nd.relu(x) @@ -71,11 +79,3 @@ def softplus( if threshold is not None: return mx.nd.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -def log_softmax(x: None, /, *, axis: Optional[int] = None, out: Optional[None] = None): - raise IvyNotImplementedException() - - -def mish(x: None, /, *, out: Optional[None] = None) -> None: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/creation.py b/ivy/functional/backends/mxnet/creation.py index 632b1daca8f9e..417bea7a45a17 100644 --- a/ivy/functional/backends/mxnet/creation.py +++ b/ivy/functional/backends/mxnet/creation.py @@ -60,7 +60,15 @@ def asarray( return ret -array = asarray +def copy_array( + x: Union[(None, mx.ndarray.NDArray)], + *, + to_ivy_array: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() def empty( @@ -107,6 +115,15 @@ def from_dlpack( raise IvyNotImplementedException() +def frombuffer( + buffer: bytes, + dtype: Optional[None] = float, + count: Optional[int] = (-1), + offset: Optional[int] = 0, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def full( shape: Union[(ivy.NativeShape, Sequence[int])], fill_value: Union[(int, float, bool)], @@ -154,6 +171,21 @@ def meshgrid( raise IvyNotImplementedException() +def one_hot( + indices: Union[(None, mx.ndarray.NDArray)], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[None] = None, + device: str, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def ones( shape: Optional[ivy.NativeShape] = None, *, @@ -195,6 +227,12 @@ def triu( raise IvyNotImplementedException() +def triu_indices( + n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str +) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + def zeros( *size: Union[(int, Sequence[int])], shape: Optional[ivy.NativeShape] = None, @@ -220,42 +258,4 @@ def zeros_like( return ivy.to_device(ret, device) -def copy_array( - x: Union[(None, mx.ndarray.NDArray)], - *, - to_ivy_array: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() - - -def one_hot( - indices: Union[(None, mx.ndarray.NDArray)], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[None] = None, - device: str, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def frombuffer( - buffer: bytes, - dtype: Optional[None] = float, - count: Optional[int] = (-1), - offset: Optional[int] = 0, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def triu_indices( - n_rows: int, n_cols: Optional[int] = None, k: int = 0, /, *, device: str -) -> Tuple[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() +array = asarray diff --git a/ivy/functional/backends/mxnet/data_type.py b/ivy/functional/backends/mxnet/data_type.py index 1ec0e1174ceec..6d15845c7c19f 100644 --- a/ivy/functional/backends/mxnet/data_type.py +++ b/ivy/functional/backends/mxnet/data_type.py @@ -5,6 +5,18 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from ivy.utils.exceptions import IvyNotImplementedException +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "u1": "uint8", +} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int32"): "int32", @@ -34,19 +46,6 @@ "bool": np.bool_, } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "u1": "uint8", -} - class Finfo: def __init__(self, mx_finfo: mx.np.finfo): @@ -84,57 +83,6 @@ def __repr__(self): return repr(self._mx_finfo) -def astype( - x: Union[(None, mx.ndarray.NDArray)], - dtype: Union[(None, str)], - /, - *, - copy: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return mx.nd.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays( - *arrays: Union[(None, mx.ndarray.NDArray)] -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_to( - x: Union[(None, mx.ndarray.NDArray)], - /, - shape: Union[(ivy.NativeShape, Sequence[int])], - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -@_handle_nestable_dtype_info -def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: - if isinstance(type, mx.ndarray.NDArray): - type = type.dtype - return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: - # using np.iinfo as mx use np dtypes and mxnet iinfo not provided - if isinstance(type, mx.ndarray.NDArray): - type = type.asnumpy().dtype - return np.iinfo(ivy.as_native_dtype(type)) - - -def result_type( - *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] -) -> ivy.Dtype: - raise IvyNotImplementedException() - - def as_ivy_dtype( dtype_in: Union[(str, int, float, complex, bool, np.dtype)], / ) -> ivy.Dtype: @@ -189,6 +137,36 @@ def as_native_dtype(dtype_in: Union[(None, str, bool, int, float, np.dtype)]) -> ) +def astype( + x: Union[(None, mx.ndarray.NDArray)], + dtype: Union[(None, str)], + /, + *, + copy: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return mx.nd.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays( + *arrays: Union[(None, mx.ndarray.NDArray)] +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def broadcast_to( + x: Union[(None, mx.ndarray.NDArray)], + /, + shape: Union[(ivy.NativeShape, Sequence[int])], + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def dtype( x: Union[(None, mx.ndarray.NDArray, np.ndarray)], *, as_native: bool = False ) -> ivy.Dtype: @@ -201,5 +179,26 @@ def dtype_bits(dtype_in: Union[(None, str, np.dtype)], /) -> int: raise IvyNotImplementedException() +@_handle_nestable_dtype_info +def finfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> Finfo: + if isinstance(type, mx.ndarray.NDArray): + type = type.dtype + return Finfo(mx.np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[str, mx.ndarray.NDArray, np.dtype], /) -> np.iinfo: + # using np.iinfo as mx use np dtypes and mxnet iinfo not provided + if isinstance(type, mx.ndarray.NDArray): + type = type.asnumpy().dtype + return np.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[(None, str)], /) -> bool: raise IvyNotImplementedException() + + +def result_type( + *arrays_and_dtypes: Union[(None, mx.ndarray.NDArray, None)] +) -> ivy.Dtype: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/device.py b/ivy/functional/backends/mxnet/device.py index 309621811a71e..2a47dc7848916 100644 --- a/ivy/functional/backends/mxnet/device.py +++ b/ivy/functional/backends/mxnet/device.py @@ -11,23 +11,21 @@ from ivy.utils.exceptions import IvyNotImplementedException -def dev( - x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False -) -> Union[(ivy.Device, str)]: - if as_native: - return x.context - return as_ivy_dev(x.context) +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + raise IvyNotImplementedException() + def start(self): + raise IvyNotImplementedException() -def to_device( - x: Union[(None, mx.ndarray.NDArray)], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return x.as_in_context(as_native_dev(device)) + def stop(self): + raise IvyNotImplementedException() + + def __enter__(self): + raise IvyNotImplementedException() + + def __exit__(self, exc_type, exc_val, exc_tb): + raise IvyNotImplementedException() def as_ivy_dev(device): @@ -62,8 +60,12 @@ def clear_cached_mem_on_dev(device: str, /): raise IvyNotImplementedException() -def num_gpus() -> int: - return mx.context.num_gpus() +def dev( + x: Union[(None, mx.ndarray.NDArray)], /, *, as_native: bool = False +) -> Union[(ivy.Device, str)]: + if as_native: + return x.context + return as_ivy_dev(x.context) def gpu_is_available() -> bool: @@ -72,22 +74,20 @@ def gpu_is_available() -> bool: return False -def tpu_is_available() -> bool: - return False - - -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - raise IvyNotImplementedException() +def num_gpus() -> int: + return mx.context.num_gpus() - def start(self): - raise IvyNotImplementedException() - def stop(self): - raise IvyNotImplementedException() +def to_device( + x: Union[(None, mx.ndarray.NDArray)], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return x.as_in_context(as_native_dev(device)) - def __enter__(self): - raise IvyNotImplementedException() - def __exit__(self, exc_type, exc_val, exc_tb): - raise IvyNotImplementedException() +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/mxnet/elementwise.py b/ivy/functional/backends/mxnet/elementwise.py index 42a8a2788b605..f3a98ade3f1cf 100644 --- a/ivy/functional/backends/mxnet/elementwise.py +++ b/ivy/functional/backends/mxnet/elementwise.py @@ -46,6 +46,16 @@ def add( return mx.nd.add(mx.nd.multiply(x1, alpha), x2) +def angle( + input: Union[(None, mx.ndarray.NDArray)], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def asin( x: Union[(None, mx.ndarray.NDArray)], /, @@ -178,6 +188,15 @@ def cosh( return mx.nd.cosh(x) +def deg2rad( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def divide( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -203,6 +222,15 @@ def equal( raise IvyNotImplementedException() +def erf( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def exp( x: Union[(None, mx.ndarray.NDArray)], /, @@ -212,6 +240,15 @@ def exp( return mx.nd.exp(x) +def exp2( + x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def expm1( x: Union[(None, mx.ndarray.NDArray)], /, @@ -250,6 +287,16 @@ def fmin( raise IvyNotImplementedException() +def fmod( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def greater( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -270,6 +317,15 @@ def greater_equal( return mx.nd.greater_equal(x1, x2) +def imag( + val: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def isfinite( x: Union[(None, mx.ndarray.NDArray)], /, @@ -299,6 +355,15 @@ def isnan( raise IvyNotImplementedException() +def isreal( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def lcm( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -424,6 +489,28 @@ def logical_xor( return mx.nd.logical_xor(x1, x2) +def maximum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def minimum( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], + /, + *, + use_where: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def multiply( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -472,6 +559,33 @@ def pow( return mx.nd.power(x1, x2) +def rad2deg( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def real( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def reciprocal( + x: Union[(float, None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.reciprocal(x) + + def remainder( x1: Union[(float, None, mx.ndarray.NDArray)], x2: Union[(float, None, mx.ndarray.NDArray)], @@ -573,18 +687,6 @@ def subtract( return mx.nd.subtract(x1, x2) -def trapz( - y: Union[(None, mx.ndarray.NDArray)], - /, - *, - x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - dx: float = 1.0, - axis: int = (-1), - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def tan( x: Union[(None, mx.ndarray.NDArray)], /, @@ -604,121 +706,19 @@ def tanh( return mx.nd.tanh(x) -def trunc( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def imag( - val: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def angle( - input: Union[(None, mx.ndarray.NDArray)], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def exp2( - x: Union[(None, mx.ndarray.NDArray, float, list, tuple)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def erf( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def maximum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def minimum( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], - /, - *, - use_where: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def reciprocal( - x: Union[(float, None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.reciprocal(x) - - -def deg2rad( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def rad2deg( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def isreal( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def fmod( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def trapz( + y: Union[(None, mx.ndarray.NDArray)], /, *, + x: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + dx: float = 1.0, + axis: int = (-1), out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def real( +def trunc( x: Union[(None, mx.ndarray.NDArray)], /, *, diff --git a/ivy/functional/backends/mxnet/experimental/activations.py b/ivy/functional/backends/mxnet/experimental/activations.py index 2ab8f7443e0af..74eac3b3ada9c 100644 --- a/ivy/functional/backends/mxnet/experimental/activations.py +++ b/ivy/functional/backends/mxnet/experimental/activations.py @@ -14,9 +14,7 @@ def logit( raise IvyNotImplementedException() -def thresholded_relu( - x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None -) -> None: +def logsigmoid(input: None) -> None: raise IvyNotImplementedException() @@ -24,13 +22,15 @@ def relu6(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def logsigmoid(input: None) -> None: +def selu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def selu(x: None, /, *, out: Optional[None] = None) -> None: +def silu(x: None, /, *, out: Optional[None] = None) -> None: raise IvyNotImplementedException() -def silu(x: None, /, *, out: Optional[None] = None) -> None: +def thresholded_relu( + x: None, /, *, threshold: Union[(int, float)] = 0, out: Optional[None] = None +) -> None: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/creation.py b/ivy/functional/backends/mxnet/experimental/creation.py index 633a0125504d0..3942f38dae26c 100644 --- a/ivy/functional/backends/mxnet/experimental/creation.py +++ b/ivy/functional/backends/mxnet/experimental/creation.py @@ -5,42 +5,44 @@ from ivy.utils.exceptions import IvyNotImplementedException -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def blackman_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def kaiser_bessel_derived_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def vorbis_window( - window_length: Union[(None, mx.ndarray.NDArray)], +def kaiser_bessel_derived_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, *, - dtype: None = np.float32, + dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def hann_window( - size: int, - /, - *, +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[None] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: @@ -53,12 +55,10 @@ def tril_indices( raise IvyNotImplementedException() -def blackman_window( - size: int, - /, +def vorbis_window( + window_length: Union[(None, mx.ndarray.NDArray)], *, - periodic: bool = True, - dtype: Optional[None] = None, + dtype: None = np.float32, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/elementwise.py b/ivy/functional/backends/mxnet/experimental/elementwise.py index ce69fe9e91ed3..c6b93dab5e8a7 100644 --- a/ivy/functional/backends/mxnet/experimental/elementwise.py +++ b/ivy/functional/backends/mxnet/experimental/elementwise.py @@ -7,41 +7,21 @@ from .. import backend_version -@with_supported_dtypes( - {"1.9.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.np.log(mx.npx.gamma(x)) - - -def sinc( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def fmax( +def allclose( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> bool: raise IvyNotImplementedException() -def float_power( - x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], - x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def conj( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -71,21 +51,23 @@ def count_nonzero( raise IvyNotImplementedException() -def nansum( - x: Union[(None, mx.ndarray.NDArray)], +def diff( + x: Union[(None, mx.ndarray.NDArray, list, tuple)], /, *, - axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, - dtype: Optional[None] = None, - keepdims: bool = False, + n: int = 1, + axis: int = (-1), + prepend: Optional[ + Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] + ] = None, + append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def gcd( - x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], - x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], +def fix( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -93,21 +75,19 @@ def gcd( raise IvyNotImplementedException() -def isclose( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def float_power( + x1: Union[(None, mx.ndarray.NDArray, float, list, tuple)], + x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def signbit( - x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], +def fmax( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -115,9 +95,20 @@ def signbit( raise IvyNotImplementedException() -def hypot( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def frexp( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[ + Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] + ] = None, +) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: + raise IvyNotImplementedException() + + +def gcd( + x1: Union[(None, mx.ndarray.NDArray, int, list, tuple)], + x2: Union[(None, mx.ndarray.NDArray, float, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -125,31 +116,43 @@ def hypot( raise IvyNotImplementedException() -def allclose( +def gradient( + x: None, + /, + *, + spacing: Union[(int, list, tuple)] = 1, + axis: Optional[Union[(int, list, tuple)]] = None, + edge_order: int = 1, +) -> Union[(None, List[None])]: + raise IvyNotImplementedException() + + +def hypot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> bool: +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def fix( - x: Union[(None, mx.ndarray.NDArray)], +def isclose( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def nextafter( +def ldexp( x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray, int)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -157,24 +160,21 @@ def nextafter( raise IvyNotImplementedException() -def diff( - x: Union[(None, mx.ndarray.NDArray, list, tuple)], +@with_supported_dtypes( + {"1.9.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( + x: Union[(None, mx.ndarray.NDArray)], /, *, - n: int = 1, - axis: int = (-1), - prepend: Optional[ - Union[(None, mx.ndarray.NDArray, int, float, list, tuple)] - ] = None, - append: Optional[Union[(None, mx.ndarray.NDArray, int, float, list, tuple)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + return mx.np.log(mx.npx.gamma(x)) -def zeta( +def modf( x: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -182,20 +182,21 @@ def zeta( raise IvyNotImplementedException() -def gradient( - x: None, +def nansum( + x: Union[(None, mx.ndarray.NDArray)], /, *, - spacing: Union[(int, list, tuple)] = 1, - axis: Optional[Union[(int, list, tuple)]] = None, - edge_order: int = 1, -) -> Union[(None, List[None])]: + axis: Optional[Union[(Tuple[(int, ...)], int)]] = None, + dtype: Optional[None] = None, + keepdims: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def xlogy( - x: Union[(None, mx.ndarray.NDArray)], - y: Union[(None, mx.ndarray.NDArray)], +def nextafter( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -203,8 +204,8 @@ def xlogy( raise IvyNotImplementedException() -def conj( - x: Union[(None, mx.ndarray.NDArray)], +def signbit( + x: Union[(None, mx.ndarray.NDArray, float, int, list, tuple)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -212,9 +213,8 @@ def conj( raise IvyNotImplementedException() -def ldexp( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray, int)], +def sinc( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -222,19 +222,19 @@ def ldexp( raise IvyNotImplementedException() -def frexp( +def xlogy( x: Union[(None, mx.ndarray.NDArray)], + y: Union[(None, mx.ndarray.NDArray)], /, *, - out: Optional[ - Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])] - ] = None, -) -> Union[(Tuple[(None, None)], Tuple[(mx.ndarray.NDArray, mx.ndarray.NDArray)])]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def modf( +def zeta( x: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, diff --git a/ivy/functional/backends/mxnet/experimental/layers.py b/ivy/functional/backends/mxnet/experimental/layers.py index 84fab7a4ba311..4d9b6a781898e 100644 --- a/ivy/functional/backends/mxnet/experimental/layers.py +++ b/ivy/functional/backends/mxnet/experimental/layers.py @@ -7,70 +7,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def general_pool( - inputs, - init, - reduce_fn, - window_shape, - strides, - padding, - dim, - dilation=1, - ceil_mode=False, - count_include_pad=False, -): - raise IvyNotImplementedException() - - -def max_pool1d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: Union[str, int, Tuple[int]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool2d( - x: mx.nd.NDArray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: Union[str, int, Tuple[int], Tuple[int, int]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - -def max_pool3d( - x: mx.nd.NDArray, - kernel: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - strides: Union[ - int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] - ], - padding: Union[str, int, Tuple[int], Tuple[int, int, int]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, - ceil_mode: bool = False, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - def avg_pool1d( x: mx.nd.NDArray, kernel: Union[int, Tuple[int]], @@ -131,18 +67,6 @@ def dct( raise IvyNotImplementedException() -def fft( - x: mx.nd.NDArray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[mx.nd.NDArray] = None, -) -> mx.nd.NDArray: - raise IvyNotImplementedException() - - def dropout1d( x: mx.nd.NDArray, prob: float, @@ -179,6 +103,33 @@ def dropout3d( raise IvyNotImplementedException() +def fft( + x: mx.nd.NDArray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def general_pool( + inputs, + init, + reduce_fn, + window_shape, + strides, + padding, + dim, + dilation=1, + ceil_mode=False, + count_include_pad=False, +): + raise IvyNotImplementedException() + + def ifft( x: mx.nd.NDArray, dim: int, @@ -231,3 +182,52 @@ def interpolate( out: Optional[mx.nd.NDArray] = None, ): raise IvyNotImplementedException() + + +def max_pool1d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: Union[str, int, Tuple[int]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool2d( + x: mx.nd.NDArray, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: Union[str, int, Tuple[int], Tuple[int, int]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int], Tuple[int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() + + +def max_pool3d( + x: mx.nd.NDArray, + kernel: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + strides: Union[ + int, Tuple[int], Tuple[int, int, int], Tuple[int, int, int, int, int] + ], + padding: Union[str, int, Tuple[int], Tuple[int, int, int]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int], Tuple[int, int, int]] = 1, + ceil_mode: bool = False, + out: Optional[mx.nd.NDArray] = None, +) -> mx.nd.NDArray: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/linear_algebra.py b/ivy/functional/backends/mxnet/experimental/linear_algebra.py index 8212bc3c54f76..824597f227276 100644 --- a/ivy/functional/backends/mxnet/experimental/linear_algebra.py +++ b/ivy/functional/backends/mxnet/experimental/linear_algebra.py @@ -4,24 +4,22 @@ from ivy.utils.exceptions import IvyNotImplementedException -def eigh_tridiagonal( - alpha: Union[(None, mx.ndarray.NDArray)], - beta: Union[(None, mx.ndarray.NDArray)], +def adjoint( + x: Union[(None, mx.ndarray.NDArray)], /, *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] - ] = None, - tol: Optional[float] = None, -) -> Union[ - ( - None, - mx.ndarray.NDArray, - Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], - ) -]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def cond( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + p: Optional[Union[(None, int, str)]] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -39,31 +37,43 @@ def diagflat( raise IvyNotImplementedException() -def kron( - a: Union[(None, mx.ndarray.NDArray)], - b: Union[(None, mx.ndarray.NDArray)], +def dot( + a: mx.ndarray.NDArray, + b: mx.ndarray.NDArray, /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + out: Optional[mx.ndarray.NDArray] = None, +) -> mx.ndarray.NDArray: + return mx.symbol.dot(a, b, out=out) -def matrix_exp( +def eig( x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> Tuple[None]: raise IvyNotImplementedException() -def eig( - x: Union[(None, mx.ndarray.NDArray)], +def eigh_tridiagonal( + alpha: Union[(None, mx.ndarray.NDArray)], + beta: Union[(None, mx.ndarray.NDArray)], /, *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Tuple[None]: + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[(Tuple[(int, int)], List[int], None, mx.ndarray.NDArray)] + ] = None, + tol: Optional[float] = None, +) -> Union[ + ( + None, + mx.ndarray.NDArray, + Tuple[(Union[(None, mx.ndarray.NDArray)], Union[(None, mx.ndarray.NDArray)])], + ) +]: raise IvyNotImplementedException() @@ -73,8 +83,9 @@ def eigvals( raise IvyNotImplementedException() -def adjoint( - x: Union[(None, mx.ndarray.NDArray)], +def kron( + a: Union[(None, mx.ndarray.NDArray)], + b: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -82,33 +93,22 @@ def adjoint( raise IvyNotImplementedException() -def multi_dot( - x: Sequence[Union[(None, mx.ndarray.NDArray)]], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: - raise IvyNotImplementedException() - - -def cond( +def matrix_exp( x: Union[(None, mx.ndarray.NDArray)], /, *, - p: Optional[Union[(None, int, str)]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dot( - a: mx.ndarray.NDArray, - b: mx.ndarray.NDArray, +def multi_dot( + x: Sequence[Union[(None, mx.ndarray.NDArray)]], /, *, - out: Optional[mx.ndarray.NDArray] = None, -) -> mx.ndarray.NDArray: - return mx.symbol.dot(a, b, out=out) + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> None: + raise IvyNotImplementedException() dot.support_native_out = True diff --git a/ivy/functional/backends/mxnet/experimental/manipulation.py b/ivy/functional/backends/mxnet/experimental/manipulation.py index 66b46e813ae47..aa8ac9077a4a0 100644 --- a/ivy/functional/backends/mxnet/experimental/manipulation.py +++ b/ivy/functional/backends/mxnet/experimental/manipulation.py @@ -5,39 +5,50 @@ from ivy.utils.exceptions import IvyNotImplementedException -def moveaxis( - a: Union[(None, mx.ndarray.NDArray)], - source: Union[(int, Sequence[int])], - destination: Union[(int, Sequence[int])], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +def atleast_1d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def heaviside( - x1: Union[(None, mx.ndarray.NDArray)], - x2: Union[(None, mx.ndarray.NDArray)], +def atleast_2d( + *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def atleast_3d( + *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None +) -> List[Union[(None, mx.ndarray.NDArray)]]: + raise IvyNotImplementedException() + + +def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: + raise IvyNotImplementedException() + + +def concat_from_sequence( + input_sequence: Union[(Tuple[None], List[None])], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def flipud( - m: Union[(None, mx.ndarray.NDArray)], +def dsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def vstack( +def dstack( arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, @@ -46,41 +57,28 @@ def vstack( raise IvyNotImplementedException() -def hstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def expand( + x: Union[(None, mx.ndarray.NDArray)], + shape: Union[(List[int], List[Tuple])], /, *, + copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def rot90( +def fliplr( m: Union[(None, mx.ndarray.NDArray)], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[(int, int)] = (0, 1), - out: Union[(None, mx.ndarray.NDArray)] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def top_k( - x: None, - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[(None, None)]] = None, -) -> Tuple[(None, None)]: - raise IvyNotImplementedException() - - -def fliplr( +def flipud( m: Union[(None, mx.ndarray.NDArray)], /, *, @@ -90,8 +88,9 @@ def fliplr( raise IvyNotImplementedException() -def i0( - x: Union[(None, mx.ndarray.NDArray)], +def heaviside( + x1: Union[(None, mx.ndarray.NDArray)], + x2: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -99,7 +98,7 @@ def i0( raise IvyNotImplementedException() -def vsplit( +def hsplit( ary: Union[(None, mx.ndarray.NDArray)], indices_or_sections: Union[(int, Tuple[(int, ...)])], /, @@ -109,24 +108,17 @@ def vsplit( raise IvyNotImplementedException() -def dsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def hstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def atleast_1d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def dstack( - arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], +def i0( + x: Union[(None, mx.ndarray.NDArray)], /, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -134,15 +126,27 @@ def dstack( raise IvyNotImplementedException() -def atleast_2d( - *arys: Union[(None, mx.ndarray.NDArray)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: +def moveaxis( + a: Union[(None, mx.ndarray.NDArray)], + source: Union[(int, Sequence[int])], + destination: Union[(int, Sequence[int])], + /, + *, + copy: Optional[bool] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def atleast_3d( - *arys: Union[(None, mx.ndarray.NDArray, bool, Number)], copy: Optional[bool] = None -) -> List[Union[(None, mx.ndarray.NDArray)]]: +def rot90( + m: Union[(None, mx.ndarray.NDArray)], + /, + *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[(int, int)] = (0, 1), + out: Union[(None, mx.ndarray.NDArray)] = None, +) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -158,37 +162,33 @@ def take_along_axis( raise IvyNotImplementedException() -def hsplit( - ary: Union[(None, mx.ndarray.NDArray)], - indices_or_sections: Union[(int, Tuple[(int, ...)])], +def top_k( + x: None, + k: int, /, *, - copy: Optional[bool] = None, -) -> List[Union[(None, mx.ndarray.NDArray)]]: - raise IvyNotImplementedException() - - -def broadcast_shapes(*shapes: Union[(List[int], List[Tuple])]) -> Tuple[(int, ...)]: + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[(None, None)]] = None, +) -> Tuple[(None, None)]: raise IvyNotImplementedException() -def expand( - x: Union[(None, mx.ndarray.NDArray)], - shape: Union[(List[int], List[Tuple])], +def vsplit( + ary: Union[(None, mx.ndarray.NDArray)], + indices_or_sections: Union[(int, Tuple[(int, ...)])], /, *, copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: +) -> List[Union[(None, mx.ndarray.NDArray)]]: raise IvyNotImplementedException() -def concat_from_sequence( - input_sequence: Union[(Tuple[None], List[None])], +def vstack( + arrays: Union[(Sequence[None], Sequence[mx.ndarray.NDArray])], /, *, - new_axis: int = 0, - axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/norms.py b/ivy/functional/backends/mxnet/experimental/norms.py index 19fea95289a8d..ee8cb935f70f0 100644 --- a/ivy/functional/backends/mxnet/experimental/norms.py +++ b/ivy/functional/backends/mxnet/experimental/norms.py @@ -4,16 +4,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def l2_normalize( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: Optional[int] = None, - out: Optional[None] = None, -) -> None: - raise IvyNotImplementedException() - - def batch_norm( x: Union[(None, mx.ndarray.NDArray)], mean: Union[(None, mx.ndarray.NDArray)], @@ -58,6 +48,16 @@ def instance_norm( raise IvyNotImplementedException() +def l2_normalize( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[int] = None, + out: Optional[None] = None, +) -> None: + raise IvyNotImplementedException() + + def lp_normalize( x: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/experimental/random.py b/ivy/functional/backends/mxnet/experimental/random.py index 8c3c4dfc0554d..36c313891b167 100644 --- a/ivy/functional/backends/mxnet/experimental/random.py +++ b/ivy/functional/backends/mxnet/experimental/random.py @@ -5,14 +5,15 @@ from ivy.utils.exceptions import IvyNotImplementedException -def dirichlet( - alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], - /, +def bernoulli( + probs: Union[(float, None, mx.ndarray.NDArray)], *, - size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + logits: Union[(float, None, mx.ndarray.NDArray)] = None, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + device: str, + dtype: None, seed: Optional[int] = None, - dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -31,6 +32,18 @@ def beta( raise IvyNotImplementedException() +def dirichlet( + alpha: Union[(None, mx.ndarray.NDArray, float, Sequence[float])], + /, + *, + size: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + seed: Optional[int] = None, + dtype: Optional[None] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def gamma( alpha: Union[(float, None, mx.ndarray.NDArray)], beta: Union[(float, None, mx.ndarray.NDArray)], @@ -55,16 +68,3 @@ def poisson( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def bernoulli( - probs: Union[(float, None, mx.ndarray.NDArray)], - *, - logits: Union[(float, None, mx.ndarray.NDArray)] = None, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - device: str, - dtype: None, - seed: Optional[int] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/experimental/statistical.py b/ivy/functional/backends/mxnet/experimental/statistical.py index 252e27dd08e8c..985ab019ce0e1 100644 --- a/ivy/functional/backends/mxnet/experimental/statistical.py +++ b/ivy/functional/backends/mxnet/experimental/statistical.py @@ -4,6 +4,43 @@ from ivy.utils.exceptions import IvyNotImplementedException +def bincount( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + minlength: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def corrcoef( + x: None, + /, + *, + y: None, + rowvar: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> None: + raise IvyNotImplementedException() + + +def cov( + x1: None, + x2: None = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[None] = None, + aweights: Optional[None] = None, + dtype: Optional[type] = None, +) -> None: + raise IvyNotImplementedException() + + def histogram( a: None, /, @@ -44,30 +81,6 @@ def nanmean( raise IvyNotImplementedException() -def quantile( - a: Union[(None, mx.ndarray.NDArray)], - q: Union[(None, float)], - /, - *, - axis: Optional[Union[(int, Sequence[int])]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def corrcoef( - x: None, - /, - *, - y: None, - rowvar: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> None: - raise IvyNotImplementedException() - - def nanmedian( input: Union[(None, mx.ndarray.NDArray)], /, @@ -79,27 +92,14 @@ def nanmedian( raise IvyNotImplementedException() -def bincount( - x: Union[(None, mx.ndarray.NDArray)], +def quantile( + a: Union[(None, mx.ndarray.NDArray)], + q: Union[(None, float)], /, *, - weights: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - minlength: int = 0, + axis: Optional[Union[(int, Sequence[int])]] = None, + interpolation: str = "linear", + keepdims: bool = False, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def cov( - x1: None, - x2: None = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[None] = None, - aweights: Optional[None] = None, - dtype: Optional[type] = None, -) -> None: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/general.py b/ivy/functional/backends/mxnet/general.py index c469b517c47c1..34eebe8331057 100644 --- a/ivy/functional/backends/mxnet/general.py +++ b/ivy/functional/backends/mxnet/general.py @@ -5,10 +5,26 @@ from ivy.utils.exceptions import IvyNotImplementedException +def container_types(): + return [] + + def current_backend_str() -> str: return "mxnet" +def gather( + x: mx.ndarray.NDArray, + indices: mx.ndarray.NDArray, + /, + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[mx.ndarray.NDArray] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def is_native_array( x: Union[(None, mx.ndarray.NDArray)], /, @@ -21,6 +37,10 @@ def is_native_array( return isinstance(x, mx.ndarray.NDArray) or isinstance(x, np.ndarray) +def itemsize(x: mx.ndarray.NDArray, /) -> int: + return x.asnumpy().itemsize + + def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: if copy: if x.shape == (): @@ -29,23 +49,3 @@ def to_numpy(x: mx.ndarray.NDArray, /, *, copy: bool = True) -> np.ndarray: return x.copy().asnumpy() else: return x.asnumpy() - - -def itemsize(x: mx.ndarray.NDArray, /) -> int: - return x.asnumpy().itemsize - - -def container_types(): - return [] - - -def gather( - x: mx.ndarray.NDArray, - indices: mx.ndarray.NDArray, - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[mx.ndarray.NDArray] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/gradients.py b/ivy/functional/backends/mxnet/gradients.py index dd3e9041601be..28189a4412a33 100644 --- a/ivy/functional/backends/mxnet/gradients.py +++ b/ivy/functional/backends/mxnet/gradients.py @@ -8,18 +8,6 @@ from ivy.utils.exceptions import IvyNotImplementedException -def variable(x, /): - return x - - -def is_variable(x, /, *, exclusive=False): - return isinstance(x, mx.ndarray.NDArray) - - -def variable_data(x, /): - raise IvyNotImplementedException() - - def execute_with_gradients( func, xs, @@ -32,17 +20,29 @@ def execute_with_gradients( raise IvyNotImplementedException() -def value_and_grad(func): +def grad(func, argnums=0): raise IvyNotImplementedException() +def is_variable(x, /, *, exclusive=False): + return isinstance(x, mx.ndarray.NDArray) + + def jac(func): raise IvyNotImplementedException() -def grad(func, argnums=0): +def stop_gradient(x, /, *, preserve_type=True, out=None): raise IvyNotImplementedException() -def stop_gradient(x, /, *, preserve_type=True, out=None): +def value_and_grad(func): + raise IvyNotImplementedException() + + +def variable(x, /): + return x + + +def variable_data(x, /): raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/layers.py b/ivy/functional/backends/mxnet/layers.py index 1ae0560d0c356..3937442087024 100644 --- a/ivy/functional/backends/mxnet/layers.py +++ b/ivy/functional/backends/mxnet/layers.py @@ -66,20 +66,6 @@ def conv2d_transpose( raise IvyNotImplementedException() -def depthwise_conv2d( - x: Union[(None, mx.ndarray.NDArray)], - filters: Union[(None, mx.ndarray.NDArray)], - strides: Union[(int, Tuple[(int, int)])], - padding: Union[(str, int, Sequence[Tuple[(int, int)]])], - /, - *, - data_format: str = "NHWC", - dilations: Union[(int, Tuple[(int, int)])] = 1, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def conv3d( x: Union[(None, mx.ndarray.NDArray)], filters: Union[(None, mx.ndarray.NDArray)], @@ -145,3 +131,17 @@ def conv_general_transpose( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def depthwise_conv2d( + x: Union[(None, mx.ndarray.NDArray)], + filters: Union[(None, mx.ndarray.NDArray)], + strides: Union[(int, Tuple[(int, int)])], + padding: Union[(str, int, Sequence[Tuple[(int, int)]])], + /, + *, + data_format: str = "NHWC", + dilations: Union[(int, Tuple[(int, int)])] = 1, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/linear_algebra.py b/ivy/functional/backends/mxnet/linear_algebra.py index 5077fd3a24bc2..52682499037c2 100644 --- a/ivy/functional/backends/mxnet/linear_algebra.py +++ b/ivy/functional/backends/mxnet/linear_algebra.py @@ -38,6 +38,16 @@ def det( return mx.nd.linalg.det(x) +def diag( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + k: int = 0, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def diagonal( x: Union[(None, mx.ndarray.NDArray)], /, @@ -249,6 +259,17 @@ def trace( raise IvyNotImplementedException() +def vander( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def vecdot( x1: Union[(None, mx.ndarray.NDArray)], x2: Union[(None, mx.ndarray.NDArray)], @@ -273,27 +294,6 @@ def vector_norm( raise IvyNotImplementedException() -def diag( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - k: int = 0, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def vander( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def vector_to_skew_symmetric_matrix( vector: Union[(None, mx.ndarray.NDArray)], /, diff --git a/ivy/functional/backends/mxnet/manipulation.py b/ivy/functional/backends/mxnet/manipulation.py index feb262db2a152..49dd4aa56f818 100644 --- a/ivy/functional/backends/mxnet/manipulation.py +++ b/ivy/functional/backends/mxnet/manipulation.py @@ -6,6 +6,17 @@ from ivy.utils.exceptions import IvyNotImplementedException +def clip( + x: Union[(None, mx.ndarray.NDArray)], + x_min: Union[(Number, None, mx.ndarray.NDArray)], + x_max: Union[(Number, None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + return mx.nd.clip(x, x_min, x_max) + + def concat( xs: Union[(Tuple[(None, ...)], List[None])], /, @@ -16,6 +27,12 @@ def concat( raise IvyNotImplementedException() +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): + raise IvyNotImplementedException() + + def expand_dims( x: Union[(None, mx.ndarray.NDArray)], /, @@ -49,46 +66,36 @@ def permute_dims( raise IvyNotImplementedException() -def reshape( +def repeat( x: Union[(None, mx.ndarray.NDArray)], /, - shape: Union[(ivy.NativeShape, Sequence[int])], + repeats: Union[(int, List[int])], *, - copy: Optional[bool] = None, - order: str = "C", - allowzero: bool = True, + axis: int = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def roll( +def reshape( x: Union[(None, mx.ndarray.NDArray)], /, - shift: Union[(int, Sequence[int])], + shape: Union[(ivy.NativeShape, Sequence[int])], *, - axis: Optional[Union[(int, Sequence[int])]] = None, + copy: Optional[bool] = None, + order: str = "C", + allowzero: bool = True, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def squeeze( +def roll( x: Union[(None, mx.ndarray.NDArray)], /, + shift: Union[(int, Sequence[int])], *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.squeeze(x, axis=axis) - - -def stack( - arrays: Union[(Tuple[None], List[None])], - /, - *, - axis: int = 0, + axis: Optional[Union[(int, Sequence[int])]] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -106,36 +113,24 @@ def split( raise IvyNotImplementedException() -def repeat( +def squeeze( x: Union[(None, mx.ndarray.NDArray)], /, - repeats: Union[(int, List[int])], *, - axis: int = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() + return mx.nd.squeeze(x, axis=axis) -def tile( - x: Union[(None, mx.ndarray.NDArray)], +def stack( + arrays: Union[(Tuple[None], List[None])], /, - repeats: Sequence[int], *, + axis: int = 0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.tile(x, repeats) - - -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): - raise IvyNotImplementedException() - - -def zero_pad( - x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None -): raise IvyNotImplementedException() @@ -151,15 +146,14 @@ def swapaxes( raise IvyNotImplementedException() -def clip( +def tile( x: Union[(None, mx.ndarray.NDArray)], - x_min: Union[(Number, None, mx.ndarray.NDArray)], - x_max: Union[(Number, None, mx.ndarray.NDArray)], /, + repeats: Sequence[int], *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: - return mx.nd.clip(x, x_min, x_max) + return mx.nd.tile(x, repeats) def unstack( @@ -171,3 +165,9 @@ def unstack( keepdims: bool = False, ) -> List[None]: raise IvyNotImplementedException() + + +def zero_pad( + x, /, pad_width, *, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None +): + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/random.py b/ivy/functional/backends/mxnet/random.py index 4f1e25f4763d5..97a4c3959de99 100644 --- a/ivy/functional/backends/mxnet/random.py +++ b/ivy/functional/backends/mxnet/random.py @@ -11,12 +11,14 @@ from ivy.utils.exceptions import IvyNotImplementedException -def random_uniform( +def multinomial( + population_size: int, + num_samples: int, + /, *, - low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, - shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, - dtype: None, + batch_size: int = 1, + probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, + replace: bool = True, device: str, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, @@ -24,42 +26,40 @@ def random_uniform( raise IvyNotImplementedException() -def random_normal( +def randint( + low: Union[(float, None, mx.ndarray.NDArray)], + high: Union[(float, None, mx.ndarray.NDArray)], + /, *, - mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, - std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, - dtype: None, - seed: Optional[int] = None, device: str, + dtype: Optional[Union[(None, ivy.Dtype)]] = None, + seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def multinomial( - population_size: int, - num_samples: int, - /, +def random_normal( *, - batch_size: int = 1, - probs: Optional[Union[(None, mx.ndarray.NDArray)]] = None, - replace: bool = True, - device: str, + mean: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + std: Union[(float, None, mx.ndarray.NDArray)] = 1.0, + shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + dtype: None, seed: Optional[int] = None, + device: str, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() -def randint( - low: Union[(float, None, mx.ndarray.NDArray)], - high: Union[(float, None, mx.ndarray.NDArray)], - /, +def random_uniform( *, - shape: Optional[Union[(ivy.NativeShape, Sequence[int])]] = None, + low: Union[(float, None, mx.ndarray.NDArray)] = 0.0, + high: Union[(float, None, mx.ndarray.NDArray)] = 1.0, + shape: Optional[Union[(ivy.NativeShape, Sequence[int], None)]] = None, + dtype: None, device: str, - dtype: Optional[Union[(None, ivy.Dtype)]] = None, seed: Optional[int] = None, out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: diff --git a/ivy/functional/backends/mxnet/searching.py b/ivy/functional/backends/mxnet/searching.py index ed20dee3fccc9..8e758dae93f55 100644 --- a/ivy/functional/backends/mxnet/searching.py +++ b/ivy/functional/backends/mxnet/searching.py @@ -33,6 +33,15 @@ def argmin( raise IvyNotImplementedException() +def argwhere( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def nonzero( x: Union[(None, mx.ndarray.NDArray)], /, @@ -53,12 +62,3 @@ def where( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def argwhere( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/sorting.py b/ivy/functional/backends/mxnet/sorting.py index 356c4b9414403..85954c0a9e91b 100644 --- a/ivy/functional/backends/mxnet/sorting.py +++ b/ivy/functional/backends/mxnet/sorting.py @@ -17,18 +17,6 @@ def argsort( raise IvyNotImplementedException() -def sort( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: int = (-1), - descending: bool = False, - stable: bool = True, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - def msort( a: Union[(None, mx.ndarray.NDArray, list, tuple)], /, @@ -49,3 +37,15 @@ def searchsorted( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def sort( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: int = (-1), + descending: bool = False, + stable: bool = True, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/mxnet/statistical.py b/ivy/functional/backends/mxnet/statistical.py index 67a24d9922d61..c688052b9fc2e 100644 --- a/ivy/functional/backends/mxnet/statistical.py +++ b/ivy/functional/backends/mxnet/statistical.py @@ -7,12 +7,34 @@ import ivy -def min( +def cumprod( x: Union[(None, mx.ndarray.NDArray)], /, *, - axis: Optional[Union[(int, Sequence[int])]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def cumsum( + x: Union[(None, mx.ndarray.NDArray)], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[None] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + +def einsum( + equation: str, + *operands: Union[(None, mx.ndarray.NDArray)], out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() @@ -40,6 +62,17 @@ def mean( raise IvyNotImplementedException() +def min( + x: Union[(None, mx.ndarray.NDArray)], + /, + *, + axis: Optional[Union[(int, Sequence[int])]] = None, + keepdims: bool = False, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + raise IvyNotImplementedException() + + def prod( x: Union[(None, mx.ndarray.NDArray)], /, @@ -110,36 +143,3 @@ def var( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() - - -def cumprod( - x: Union[(None, mx.ndarray.NDArray)], - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def cumsum( - x: Union[(None, mx.ndarray.NDArray)], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[None] = None, - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() - - -def einsum( - equation: str, - *operands: Union[(None, mx.ndarray.NDArray)], - out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, -) -> Union[(None, mx.ndarray.NDArray)]: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/numpy/activations.py b/ivy/functional/backends/numpy/activations.py index 2e230498b9574..cb2c4d61211f4 100644 --- a/ivy/functional/backends/numpy/activations.py +++ b/ivy/functional/backends/numpy/activations.py @@ -10,13 +10,25 @@ @_scalar_output_to_0d_array -def relu( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +def gelu( + x: np.ndarray, + /, + *, + approximate: bool = False, + complex_mode="jax", + out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.maximum(x, 0, out=out, dtype=x.dtype) + if approximate: + ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) + else: + ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) + return ivy.astype(ret, x.dtype, copy=False) -relu.support_native_out = True +@_scalar_output_to_0d_array +def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) + return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) def leaky_relu( @@ -31,19 +43,36 @@ def leaky_relu( @_scalar_output_to_0d_array -def gelu( - x: np.ndarray, - /, - *, - approximate: bool = False, - complex_mode="jax", - out: Optional[np.ndarray] = None, +def log_softmax( + x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None ) -> np.ndarray: - if approximate: - ret = 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x))) - else: - ret = 0.5 * x * (1 + ivy.erf(x / np.sqrt(2))) - return ivy.astype(ret, x.dtype, copy=False) + if axis is None: + axis = -1 + x_max = np.max(x, axis=axis, keepdims=True) + if x_max.ndim > 0: + x_max[~np.isfinite(x_max)] = 0 + elif not np.isfinite(x_max): + x_max = 0 + exp_tmp = np.exp(x - x_max) + + with np.errstate(divide="ignore"): + s = np.sum(exp_tmp, axis=axis, keepdims=True) + ret = np.log(s) + + ret = x - x_max - ret + return ret + + +@_scalar_output_to_0d_array +def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return x * np.tanh(np.log1p(np.exp(x))) + + +@_scalar_output_to_0d_array +def relu( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.maximum(x, 0, out=out, dtype=x.dtype) def sigmoid(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -60,9 +89,6 @@ def softmax( return np.divide(exp_x, np.sum(exp_x, axis=axis, keepdims=True), out=out) -softmax.support_native_out = True - - @_scalar_output_to_0d_array def softplus( x: np.ndarray, @@ -94,45 +120,9 @@ def softplus( return res.astype(x.dtype) +relu.support_native_out = True +softmax.support_native_out = True softplus.support_native_out = True - - -@_scalar_output_to_0d_array -def log_softmax( - x: np.ndarray, /, *, axis: Optional[int] = None, out: Optional[np.ndarray] = None -) -> np.ndarray: - if axis is None: - axis = -1 - x_max = np.max(x, axis=axis, keepdims=True) - if x_max.ndim > 0: - x_max[~np.isfinite(x_max)] = 0 - elif not np.isfinite(x_max): - x_max = 0 - exp_tmp = np.exp(x - x_max) - - with np.errstate(divide="ignore"): - s = np.sum(exp_tmp, axis=axis, keepdims=True) - ret = np.log(s) - - ret = x - x_max - ret - return ret - - log_softmax.support_native_out = True - - -@_scalar_output_to_0d_array -def mish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return x * np.tanh(np.log1p(np.exp(x))) - - mish.support_native_out = True - - -@_scalar_output_to_0d_array -def hardswish(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - max_x_3 = np.maximum(x + 3, 0, dtype=x.dtype) - return (x * np.minimum(max_x_3, 6, out=out, dtype=x.dtype) / 6).astype(x.dtype) - - hardswish.support_native_out = True diff --git a/ivy/functional/backends/numpy/creation.py b/ivy/functional/backends/numpy/creation.py index 7600f3cf9dba0..0eb26a54646d9 100644 --- a/ivy/functional/backends/numpy/creation.py +++ b/ivy/functional/backends/numpy/creation.py @@ -64,6 +64,17 @@ def asarray( return np.copy(ret) if copy else ret +def copy_array( + x: np.ndarray, + *, + to_ivy_array: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if to_ivy_array: + return ivy.to_ivy(x.copy()) + return x.copy() + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -107,6 +118,17 @@ def from_dlpack(x, /, *, out: Optional[np.ndarray] = None): return np.from_dlpack(x) +def frombuffer( + buffer: bytes, + dtype: Optional[np.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> np.ndarray: + if isinstance(dtype, list): + dtype = np.dtype(dtype[0]) + return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -170,68 +192,6 @@ def meshgrid( return np.meshgrid(*arrays, sparse=sparse, indexing=indexing) -def ones( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.ones(shape, dtype), device=device) - - -def ones_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.ones_like(x, dtype=dtype), device=device) - - -def tril( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tril(x, k) - - -def triu( - x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.triu(x, k) - - -def zeros( - shape: Union[ivy.NativeShape, Sequence[int]], - *, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return _to_device(np.zeros(shape, dtype), device=device) - - -def zeros_like( - x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None -) -> np.ndarray: - return _to_device(np.zeros_like(x, dtype=dtype), device=device) - - -# Extra # -# ------# - - -array = asarray - - -def copy_array( - x: np.ndarray, - *, - to_ivy_array: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if to_ivy_array: - return ivy.to_ivy(x.copy()) - return x.copy() - - def one_hot( indices: np.ndarray, depth: int, @@ -268,15 +228,32 @@ def one_hot( return res -def frombuffer( - buffer: bytes, - dtype: Optional[np.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, +def ones( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(dtype, list): - dtype = np.dtype(dtype[0]) - return np.frombuffer(buffer, dtype=dtype, count=count, offset=offset) + return _to_device(np.ones(shape, dtype), device=device) + + +def ones_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.ones_like(x, dtype=dtype), device=device) + + +def tril( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tril(x, k) + + +def triu( + x: np.ndarray, /, *, k: int = 0, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.triu(x, k) def triu_indices( @@ -290,3 +267,26 @@ def triu_indices( return tuple( _to_device(np.asarray(np.triu_indices(n=n_rows, k=k, m=n_cols)), device=device) ) + + +def zeros( + shape: Union[ivy.NativeShape, Sequence[int]], + *, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return _to_device(np.zeros(shape, dtype), device=device) + + +def zeros_like( + x: np.ndarray, /, *, dtype: np.dtype, device: str, out: Optional[np.ndarray] = None +) -> np.ndarray: + return _to_device(np.zeros_like(x, dtype=dtype), device=device) + + +# Extra # +# ------# + + +array = asarray diff --git a/ivy/functional/backends/numpy/data_type.py b/ivy/functional/backends/numpy/data_type.py index 317cc71c8608d..172989ff82968 100644 --- a/ivy/functional/backends/numpy/data_type.py +++ b/ivy/functional/backends/numpy/data_type.py @@ -9,6 +9,26 @@ from ivy.functional.ivy.data_type import _handle_nestable_dtype_info from . import backend_version +char_rep_dtype_dict = { + "?": "bool", + "i": int, + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "f": float, + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c": complex, + "c8": "complex64", + "c16": "complex128", + "u": "uint32", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", +} ivy_dtype_dict = { np.dtype("int8"): "int8", np.dtype("int16"): "int16", @@ -40,7 +60,6 @@ np.complex128: "complex128", np.bool_: "bool", } - native_dtype_dict = { "int8": np.dtype("int8"), "int16": np.dtype("int16"), @@ -58,27 +77,6 @@ "bool": np.dtype("bool"), } -char_rep_dtype_dict = { - "?": "bool", - "i": int, - "i1": "int8", - "i2": "int16", - "i4": "int32", - "i8": "int64", - "f": float, - "f2": "float16", - "f4": "float32", - "f8": "float64", - "c": complex, - "c8": "complex64", - "c16": "complex128", - "u": "uint32", - "u1": "uint8", - "u2": "uint16", - "u4": "uint32", - "u8": "uint64", -} - class Finfo: def __init__(self, np_finfo: np.finfo): @@ -108,68 +106,6 @@ def smallest_normal(self): return float(self._np_finfo.tiny) -# Array API Standard # -# -------------------# - - -def astype( - x: np.ndarray, - dtype: np.dtype, - /, - *, - copy: bool = True, - out: Optional[ivy.Array] = None, -) -> np.ndarray: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return np.copy(x) if copy else x - return x.astype(dtype) - - -def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: - try: - return np.broadcast_arrays(*arrays) - except ValueError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def broadcast_to( - x: np.ndarray, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return np.broadcast_to(x.reshape([-1]), shape) - return np.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: - if isinstance(type, np.ndarray): - type = type.dtype - return Finfo(np.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: - if isinstance(type, np.ndarray): - type = type.dtype - return np.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return np.result_type(arrays_and_dtypes) - result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) - for i in range(2, len(arrays_and_dtypes)): - result = np.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -238,6 +174,45 @@ def as_native_dtype(dtype_in: Union[np.dtype, str, bool, int, float], /) -> np.d ) +# Array API Standard # +# -------------------# + + +def astype( + x: np.ndarray, + dtype: np.dtype, + /, + *, + copy: bool = True, + out: Optional[ivy.Array] = None, +) -> np.ndarray: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return np.copy(x) if copy else x + return x.astype(dtype) + + +def broadcast_arrays(*arrays: np.ndarray) -> List[np.ndarray]: + try: + return np.broadcast_arrays(*arrays) + except ValueError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def broadcast_to( + x: np.ndarray, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return np.broadcast_to(x.reshape([-1]), shape) + return np.broadcast_to(x, shape) + + def dtype(x: np.ndarray, *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.to_native(x).dtype @@ -257,6 +232,20 @@ def dtype_bits(dtype_in: Union[np.dtype, str], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[np.dtype, str, np.ndarray], /) -> Finfo: + if isinstance(type, np.ndarray): + type = type.dtype + return Finfo(np.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[np.dtype, str, np.ndarray], /) -> np.iinfo: + if isinstance(type, np.ndarray): + type = type.dtype + return np.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -264,3 +253,12 @@ def is_native_dtype(dtype_in: Union[np.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[np.ndarray, np.dtype]) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return np.result_type(arrays_and_dtypes) + result = np.result_type(arrays_and_dtypes[0], arrays_and_dtypes[1]) + for i in range(2, len(arrays_and_dtypes)): + result = np.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/numpy/device.py b/ivy/functional/backends/numpy/device.py index b10c8119f8cc1..4576b42831bc5 100644 --- a/ivy/functional/backends/numpy/device.py +++ b/ivy/functional/backends/numpy/device.py @@ -11,34 +11,30 @@ from ivy.functional.ivy.device import Profiler as BaseProfiler -def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: - if as_native: - return "cpu" - return as_ivy_dev("cpu") - - -def as_ivy_dev(device: str, /): - return ivy.Device("cpu") - - -def as_native_dev(device: str, /): - return "cpu" - - -def clear_cached_mem_on_dev(device: str, /): - return None +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper numpy profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + def start(self): + self._start_time = time.perf_counter() -def tpu_is_available() -> bool: - return False + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + def __enter__(self): + self.start() -def num_gpus() -> int: - return 0 + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() -def gpu_is_available() -> bool: - return False +# --- Helpers --- # +# --------------- # # private version of to_device to be used in backend implementations @@ -60,6 +56,40 @@ def _to_device(x: np.ndarray, device=None) -> np.ndarray: return x +# --- Main --- # +# ------------ # + + +def as_ivy_dev(device: str, /): + return ivy.Device("cpu") + + +def as_native_dev(device: str, /): + return "cpu" + + +def clear_cached_mem_on_dev(device: str, /): + return None + + +def dev(x: np.ndarray, /, *, as_native: bool = False) -> Union[ivy.Device, str]: + if as_native: + return "cpu" + return as_ivy_dev("cpu") + + +def gpu_is_available() -> bool: + return False + + +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return fn(*args, **kwargs) + + +def num_gpus() -> int: + return 0 + + def to_device( x: np.ndarray, device: str, @@ -85,27 +115,5 @@ def to_device( return x -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return fn(*args, **kwargs) - - -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper numpy profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None - - def start(self): - self._start_time = time.perf_counter() - - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) - - def __enter__(self): - self.start() - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/numpy/elementwise.py b/ivy/functional/backends/numpy/elementwise.py index 6c4459bbaaf4f..5d643eb77e7a3 100644 --- a/ivy/functional/backends/numpy/elementwise.py +++ b/ivy/functional/backends/numpy/elementwise.py @@ -10,6 +10,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _abs_variant_sign(x): + return np.divide(x, np.abs(x), where=x != 0) + + +# --- Main --- # +# ------------ # + + @_scalar_output_to_0d_array def abs( x: Union[float, np.ndarray], @@ -20,25 +32,16 @@ def abs( return np.absolute(x, out=out) -abs.support_native_out = True - - @_scalar_output_to_0d_array def acos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccos(x, out=out) -acos.support_native_out = True - - @_scalar_output_to_0d_array def acosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arccosh(x, out=out) -acosh.support_native_out = True - - @_scalar_output_to_0d_array def add( x1: Union[float, np.ndarray], @@ -55,7 +58,14 @@ def add( return np.add(x1, x2, out=out) -add.support_native_out = True +def angle( + z: np.ndarray, + /, + *, + deg: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.angle(z, deg=deg) @_scalar_output_to_0d_array @@ -63,25 +73,16 @@ def asin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsin(x, out=out) -asin.support_native_out = True - - @_scalar_output_to_0d_array def asinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arcsinh(x, out=out) -asinh.support_native_out = True - - @_scalar_output_to_0d_array def atan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctan(x, out=out) -atan.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def atan2( @@ -91,17 +92,11 @@ def atan2( return np.arctan2(x1, x2, out=out) -atan2.support_native_out = True - - @_scalar_output_to_0d_array def atanh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.arctanh(x, out=out) -atanh.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_and( @@ -115,9 +110,6 @@ def bitwise_and( return np.bitwise_and(x1, x2, out=out) -bitwise_and.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_invert( @@ -126,9 +118,6 @@ def bitwise_invert( return np.invert(x, out=out) -bitwise_invert.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_left_shift( @@ -142,9 +131,6 @@ def bitwise_left_shift( return np.left_shift(x1, x2, out=out) -bitwise_left_shift.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_or( @@ -158,9 +144,6 @@ def bitwise_or( return np.bitwise_or(x1, x2, out=out) -bitwise_or.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_right_shift( @@ -174,9 +157,6 @@ def bitwise_right_shift( return np.right_shift(x1, x2, out=out) -bitwise_right_shift.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def bitwise_xor( @@ -190,9 +170,6 @@ def bitwise_xor( return np.bitwise_xor(x1, x2, out=out) -bitwise_xor.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) @_scalar_output_to_0d_array def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -205,24 +182,21 @@ def ceil(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -ceil.support_native_out = True - - @_scalar_output_to_0d_array def cos(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cos(x, out=out) -cos.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) @_scalar_output_to_0d_array def cosh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.cosh(x, out=out) -cosh.support_native_out = True +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.deg2rad(x, out=out) @_scalar_output_to_0d_array @@ -242,9 +216,6 @@ def divide( return ret -divide.support_native_out = True - - @_scalar_output_to_0d_array def equal( x1: Union[float, np.ndarray], @@ -257,15 +228,36 @@ def equal( return np.equal(x1, x2, out=out) -equal.support_native_out = True +# Extra # +# ------# @_scalar_output_to_0d_array -def exp(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.exp(x, out=out) +def erf(x, /, *, out: Optional[np.ndarray] = None): + a1 = 0.254829592 + a2 = -0.284496736 + a3 = 1.421413741 + a4 = -1.453152027 + a5 = 1.061405429 + p = 0.3275911 + sign = np.sign(x) + x = np.abs(x) -exp.support_native_out = True + # A&S formula 7.1.26 + t = 1.0 / (1.0 + p * x) + y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) + ret = sign * y + if hasattr(x, "dtype"): + ret = np.asarray(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + +@_scalar_output_to_0d_array +def exp(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.exp(x, out=out) def exp2( @@ -277,17 +269,11 @@ def exp2( return np.exp2(x, out=out) -exp2.support_native_out = True - - @_scalar_output_to_0d_array def expm1(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.expm1(x, out=out) -expm1.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: @@ -300,9 +286,6 @@ def floor(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -floor.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def floor_divide( @@ -337,7 +320,32 @@ def fmin( ) -fmin.support_native_out = True +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def fmod( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmod( + x1, + x2, + out=None, + ) + + +def gcd( + x1: Union[np.ndarray, int, list, tuple], + x2: Union[np.ndarray, float, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.gcd(x1, x2, out=out) @_scalar_output_to_0d_array @@ -352,9 +360,6 @@ def greater( return np.greater(x1, x2, out=out) -greater.support_native_out = True - - @_scalar_output_to_0d_array def greater_equal( x1: Union[float, np.ndarray], @@ -367,7 +372,13 @@ def greater_equal( return np.greater_equal(x1, x2, out=out) -greater_equal.support_native_out = True +def imag( + val: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.imag(val) @_scalar_output_to_0d_array @@ -375,9 +386,6 @@ def isfinite(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarra return np.isfinite(x, out=out) -isfinite.support_native_out = True - - @_scalar_output_to_0d_array def isinf( x: np.ndarray, @@ -401,7 +409,9 @@ def isnan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.isnan(x, out=out) -isnan.support_native_out = True +@_scalar_output_to_0d_array +def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.isreal(x) @_scalar_output_to_0d_array @@ -420,9 +430,6 @@ def lcm( ) -lcm.support_native_out = True - - @_scalar_output_to_0d_array def less( x1: Union[float, np.ndarray], @@ -435,9 +442,6 @@ def less( return np.less(x1, x2, out=out) -less.support_native_out = True - - @_scalar_output_to_0d_array def less_equal( x1: Union[float, np.ndarray], @@ -450,41 +454,26 @@ def less_equal( return np.less_equal(x1, x2, out=out) -less_equal.support_native_out = True - - @_scalar_output_to_0d_array def log(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log(x, out=out) -log.support_native_out = True - - @_scalar_output_to_0d_array def log10(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log10(x, out=out) -log10.support_native_out = True - - @_scalar_output_to_0d_array def log1p(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log1p(x, out=out) -log1p.support_native_out = True - - @_scalar_output_to_0d_array def log2(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.log2(x, out=out) -log2.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def logaddexp( @@ -494,9 +483,6 @@ def logaddexp( return np.logaddexp(x1, x2, out=out) -logaddexp.support_native_out = True - - def logaddexp2( x1: Union[np.ndarray, int, list, tuple], x2: Union[np.ndarray, int, list, tuple], @@ -511,9 +497,6 @@ def logaddexp2( return np.logaddexp2(x1, x2, out=out) -logaddexp2.support_native_out = True - - @_scalar_output_to_0d_array def logical_and( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -521,17 +504,11 @@ def logical_and( return np.logical_and(x1, x2, out=out) -logical_and.support_native_out = True - - @_scalar_output_to_0d_array def logical_not(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.logical_not(x, out=out) -logical_not.support_native_out = True - - @_scalar_output_to_0d_array def logical_or( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -539,9 +516,6 @@ def logical_or( return np.logical_or(x1, x2, out=out) -logical_or.support_native_out = True - - @_scalar_output_to_0d_array def logical_xor( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None @@ -549,7 +523,40 @@ def logical_xor( return np.logical_xor(x1, x2, out=out) -logical_xor.support_native_out = True +@_scalar_output_to_0d_array +def maximum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +): + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 >= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.maximum(x1, x2, out=out) + + +@_scalar_output_to_0d_array +def minimum( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, + use_where: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + ret = np.where(x1 <= x2, x1, x2) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + return np.minimum(x1, x2, out=out) @_scalar_output_to_0d_array @@ -564,7 +571,17 @@ def multiply( return np.multiply(x1, x2, out=out) -multiply.support_native_out = True +def nan_to_num( + x: np.ndarray, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) @_scalar_output_to_0d_array @@ -574,9 +591,6 @@ def negative( return np.negative(x, out=out) -negative.support_native_out = True - - @_scalar_output_to_0d_array def not_equal( x1: Union[float, np.ndarray], @@ -589,9 +603,6 @@ def not_equal( return np.not_equal(x1, x2, out=out) -not_equal.support_native_out = True - - @_scalar_output_to_0d_array def positive( x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None @@ -599,9 +610,6 @@ def positive( return np.positive(x, out=out) -positive.support_native_out = True - - @_scalar_output_to_0d_array def pow( x1: Union[float, np.ndarray], @@ -614,16 +622,31 @@ def pow( return np.power(x1, x2, out=out) -pow.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def remainder( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, +def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.rad2deg(x, out=out) + + +def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.real(x) + + +@_scalar_output_to_0d_array +def reciprocal( + x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + numerator = np.ones_like(x) + return np.true_divide(numerator, x, out=out) + + +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def remainder( + x1: Union[float, np.ndarray], + x2: Union[float, np.ndarray], + /, + *, modulus: bool = True, out: Optional[np.ndarray] = None, ) -> np.ndarray: @@ -637,9 +660,6 @@ def remainder( return np.remainder(x1, x2, out=out) -remainder.support_native_out = True - - @_scalar_output_to_0d_array def round( x: np.ndarray, /, *, decimals: int = 0, out: Optional[np.ndarray] = None @@ -653,13 +673,6 @@ def round( return ret -round.support_native_out = True - - -def _abs_variant_sign(x): - return np.divide(x, np.abs(x), where=x != 0) - - @_scalar_output_to_0d_array def sign( x: np.ndarray, @@ -673,41 +686,26 @@ def sign( return np.sign(x, out=out) -sign.support_native_out = True - - @_scalar_output_to_0d_array def sin(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sin(x, out=out) -sin.support_native_out = True - - @_scalar_output_to_0d_array def sinh(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sinh(x, out=out) -sinh.support_native_out = True - - @_scalar_output_to_0d_array def sqrt(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.sqrt(x, out=out) -sqrt.support_native_out = True - - @_scalar_output_to_0d_array def square(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.square(x, out=out) -square.support_native_out = True - - @_scalar_output_to_0d_array def subtract( x1: Union[float, np.ndarray], @@ -725,7 +723,16 @@ def subtract( return np.subtract(x1, x2, out=out) -subtract.support_native_out = True +@_scalar_output_to_0d_array +def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.tan(x, out=out) + + +@_scalar_output_to_0d_array +def tanh( + x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tanh(x, out=out) @_scalar_output_to_0d_array @@ -741,27 +748,6 @@ def trapz( return np.trapz(y, x=x, dx=dx, axis=axis) -trapz.support_native_out = False - - -@_scalar_output_to_0d_array -def tan(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.tan(x, out=out) - - -tan.support_native_out = True - - -@_scalar_output_to_0d_array -def tanh( - x: np.ndarray, /, *, complex_mode="jax", out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tanh(x, out=out) - - -tanh.support_native_out = True - - @_scalar_output_to_0d_array def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: if "int" in str(x.dtype): @@ -773,192 +759,74 @@ def trunc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret +abs.support_native_out = True +acos.support_native_out = True +acosh.support_native_out = True +add.support_native_out = True +asin.support_native_out = True +asinh.support_native_out = True +atan.support_native_out = True +atan2.support_native_out = True +atanh.support_native_out = True +bitwise_and.support_native_out = True +bitwise_invert.support_native_out = True +bitwise_left_shift.support_native_out = True +bitwise_or.support_native_out = True +bitwise_right_shift.support_native_out = True +bitwise_xor.support_native_out = True +ceil.support_native_out = True +cos.support_native_out = True +cosh.support_native_out = True +divide.support_native_out = True +equal.support_native_out = True +exp.support_native_out = True +exp2.support_native_out = True +expm1.support_native_out = True +floor.support_native_out = True +fmin.support_native_out = True +greater.support_native_out = True +greater_equal.support_native_out = True +isfinite.support_native_out = True +isnan.support_native_out = True +lcm.support_native_out = True +less.support_native_out = True +less_equal.support_native_out = True +log.support_native_out = True +log10.support_native_out = True +log1p.support_native_out = True +log2.support_native_out = True +logaddexp.support_native_out = True +logaddexp2.support_native_out = True +logical_and.support_native_out = True +logical_not.support_native_out = True +logical_or.support_native_out = True +logical_xor.support_native_out = True +multiply.support_native_out = True +negative.support_native_out = True +not_equal.support_native_out = True +positive.support_native_out = True +pow.support_native_out = True +remainder.support_native_out = True +round.support_native_out = True +sign.support_native_out = True +sin.support_native_out = True +sinh.support_native_out = True +sqrt.support_native_out = True +square.support_native_out = True +subtract.support_native_out = True +trapz.support_native_out = False +tan.support_native_out = True +tanh.support_native_out = True trunc.support_native_out = True - - -# Extra # -# ------# - - -@_scalar_output_to_0d_array -def erf(x, /, *, out: Optional[np.ndarray] = None): - a1 = 0.254829592 - a2 = -0.284496736 - a3 = 1.421413741 - a4 = -1.453152027 - a5 = 1.061405429 - p = 0.3275911 - - sign = np.sign(x) - x = np.abs(x) - - # A&S formula 7.1.26 - t = 1.0 / (1.0 + p * x) - y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) - ret = sign * y - if hasattr(x, "dtype"): - ret = np.asarray(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - erf.support_native_out = True - - -@_scalar_output_to_0d_array -def maximum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -): - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 >= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.maximum(x1, x2, out=out) - - maximum.support_native_out = True - - -@_scalar_output_to_0d_array -def minimum( - x1: Union[float, np.ndarray], - x2: Union[float, np.ndarray], - /, - *, - use_where: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - ret = np.where(x1 <= x2, x1, x2) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - return np.minimum(x1, x2, out=out) - - minimum.support_native_out = True - - -@_scalar_output_to_0d_array -def reciprocal( - x: Union[float, np.ndarray], /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - numerator = np.ones_like(x) - return np.true_divide(numerator, x, out=out) - - reciprocal.support_native_out = True - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def deg2rad(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.deg2rad(x, out=out) - - deg2rad.support_native_out = True - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def rad2deg(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.rad2deg(x, out=out) - - rad2deg.support_native_out = True - - -@_scalar_output_to_0d_array -def isreal(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.isreal(x) - - isreal.support_native_out = False - - -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def fmod( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmod( - x1, - x2, - out=None, - ) - - fmod.support_native_out = True - - -def angle( - z: np.ndarray, - /, - *, - deg: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.angle(z, deg=deg) - - angle.support_native_out = False - - -def gcd( - x1: Union[np.ndarray, int, list, tuple], - x2: Union[np.ndarray, float, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.gcd(x1, x2, out=out) - - gcd.support_native_out = True - - -def imag( - val: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.imag(val) - - imag.support_native_out = False - - -def nan_to_num( - x: np.ndarray, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nan_to_num(x, copy=copy, nan=nan, posinf=posinf, neginf=neginf) - - nan_to_num.support_native_out = False - - -def real(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.real(x) diff --git a/ivy/functional/backends/numpy/experimental/activations.py b/ivy/functional/backends/numpy/experimental/activations.py index d3f8282a22de2..6a0217411f7a9 100644 --- a/ivy/functional/backends/numpy/experimental/activations.py +++ b/ivy/functional/backends/numpy/experimental/activations.py @@ -10,6 +10,17 @@ from . import backend_version +@_scalar_output_to_0d_array +def elu( + x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None +) -> np.ndarray: + # exp = np.expm1(x) + ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ret + + def logit( x: np.ndarray, /, @@ -28,18 +39,10 @@ def logit( return ret +@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) @_scalar_output_to_0d_array -def thresholded_relu( - x: np.ndarray, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.where(x > threshold, x, 0).astype(x.dtype) - - -thresholded_relu.support_native_out = True +def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return -(np.log1p(np.exp(-(input)))) @_scalar_output_to_0d_array @@ -47,15 +50,6 @@ def relu6(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.minimum(np.maximum(x, 0, dtype=x.dtype), 6, out=out, dtype=x.dtype) -relu6.support_native_out = True - - -@with_unsupported_dtypes({"1.25.2 and below": ("bool",)}, backend_version) -@_scalar_output_to_0d_array -def logsigmoid(input: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return -(np.log1p(np.exp(-(input)))) - - @_scalar_output_to_0d_array def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: alpha = 1.6732632423543772848170429916717 @@ -66,9 +60,6 @@ def selu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return ret -selu.support_native_out = True - - @_scalar_output_to_0d_array def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: ret = np.asarray(x * (1 / (1 + np.exp(-x)))) @@ -80,18 +71,19 @@ def silu(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.asarray(x * (1 / (1 + np.exp(-x)))).astype(x.dtype) -silu.support_native_out = True - - @_scalar_output_to_0d_array -def elu( - x: np.ndarray, /, *, alpha: float = 1.0, out: Optional[np.ndarray] = None +def thresholded_relu( + x: np.ndarray, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - # exp = np.expm1(x) - ret = np.where(x > 0, x, np.multiply(alpha, np.expm1(x))).astype(x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ret + return np.where(x > threshold, x, 0).astype(x.dtype) +thresholded_relu.support_native_out = True +relu6.support_native_out = True +selu.support_native_out = True +silu.support_native_out = True elu.support_native_out = True diff --git a/ivy/functional/backends/numpy/experimental/creation.py b/ivy/functional/backends/numpy/experimental/creation.py index c2e01f77c55e8..a003729e658cc 100644 --- a/ivy/functional/backends/numpy/experimental/creation.py +++ b/ivy/functional/backends/numpy/experimental/creation.py @@ -7,37 +7,26 @@ from ivy.functional.backends.numpy.device import _to_device import ivy -# Array API Standard # -# -------------------# - -def vorbis_window( - window_length: np.ndarray, +def blackman_window( + size: int, + /, *, - dtype: np.dtype = np.float32, + periodic: bool = True, + dtype: Optional[np.dtype] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - result = [] - for i in range(1, window_length * 2, 2): - temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) - result.append(round(temp, 8)) - return np.array(result, dtype=dtype) - - -vorbis_window.support_native_out = False - + if size < 2: + return np.ones([size], dtype=dtype) + if periodic: + count = np.arange(size) / size + else: + count = np.linspace(start=0, stop=size, num=size) -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: str, -) -> Tuple[np.ndarray, ...]: - return tuple( - _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) - ) + return ( + (0.42 - 0.5 * np.cos(2 * np.pi * count)) + + (0.08 * np.cos(2 * np.pi * 2 * count)) + ).astype(dtype) def hann_window( @@ -57,7 +46,12 @@ def hann_window( return (0.5 - 0.5 * np.cos(2 * np.pi * count)).astype(dtype) -hann_window.support_native_out = False +def indices( + dimensions: Sequence, + dtype: np.dtype = np.int64, + sparse: bool = False, +) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: + return np.indices(dimensions, dtype=dtype, sparse=sparse) def kaiser_window( @@ -76,15 +70,30 @@ def kaiser_window( return np.kaiser(M=window_length + 1, beta=beta)[:-1].astype(dtype) -kaiser_window.support_native_out = False +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[np.ndarray, ...]: + return tuple( + _to_device(np.asarray(np.tril_indices(n=n_rows, k=k, m=n_cols)), device=device) + ) -def indices( - dimensions: Sequence, - dtype: np.dtype = np.int64, - sparse: bool = False, -) -> Union[np.ndarray, Tuple[np.ndarray, ...]]: - return np.indices(dimensions, dtype=dtype, sparse=sparse) +def trilu( + x: np.ndarray, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if upper: + return np.triu(x, k) + return np.tril(x, k) def unsorted_segment_min( @@ -113,30 +122,6 @@ def unsorted_segment_min( return res -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if size < 2: - return np.ones([size], dtype=dtype) - if periodic: - count = np.arange(size) / size - else: - count = np.linspace(start=0, stop=size, num=size) - - return ( - (0.42 - 0.5 * np.cos(2 * np.pi * count)) - + (0.08 * np.cos(2 * np.pi * 2 * count)) - ).astype(dtype) - - -blackman_window.support_native_out = False - - def unsorted_segment_sum( data: np.ndarray, segment_ids: np.ndarray, @@ -160,14 +145,24 @@ def unsorted_segment_sum( return res -def trilu( - x: np.ndarray, - /, +# Array API Standard # +# -------------------# + + +def vorbis_window( + window_length: np.ndarray, *, - k: int = 0, - upper: bool = True, + dtype: np.dtype = np.float32, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if upper: - return np.triu(x, k) - return np.tril(x, k) + result = [] + for i in range(1, window_length * 2, 2): + temp = np.sin(ivy.pi / 2 * (np.sin(ivy.pi * i / (window_length * 2)) ** 2)) + result.append(round(temp, 8)) + return np.array(result, dtype=dtype) + + +vorbis_window.support_native_out = False +hann_window.support_native_out = False +kaiser_window.support_native_out = False +blackman_window.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/elementwise.py b/ivy/functional/backends/numpy/experimental/elementwise.py index cf91193091296..6e2519d9e9467 100644 --- a/ivy/functional/backends/numpy/experimental/elementwise.py +++ b/ivy/functional/backends/numpy/experimental/elementwise.py @@ -9,49 +9,84 @@ from . import backend_version -@_scalar_output_to_0d_array -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: - return np.sinc(x).astype(x.dtype) +# --- LGAMMA --- # +LANCZOS_N = 13 +kBaseLanczosCoeff = 0.99999999999980993227684700473478 +kLanczosCoefficients = np.array( + [ + 676.520368121885098567009190444019, + -1259.13921672240287047156078755283, + 771.3234287776530788486528258894, + -176.61502916214059906584551354, + 12.507343278686904814458936853, + -0.13857109526572011689554707, + 9.984369578019570859563e-6, + 1.50563273514931155834e-7, + ] +) +# ---digamma---# +kLanczosGamma = 7 # aka g +lanczos_den_coeffs = np.array( + [ + 0.0, + 39916800.0, + 120543840.0, + 150917976.0, + 105258076.0, + 45995730.0, + 13339535.0, + 2637558.0, + 357423.0, + 32670.0, + 1925.0, + 66.0, + 1.0, + ] +) +lanczos_g = 6.024680040776729583740234375 +lanczos_num_coeffs = np.array( + [ + 23531376880.410759688572007674451636754734846804940, + 42919803642.649098768957899047001988850926355848959, + 35711959237.355668049440185451547166705960488635843, + 17921034426.037209699919755754458931112671403265390, + 6039542586.3520280050642916443072979210699388420708, + 1439720407.3117216736632230727949123939715485786772, + 248874557.86205415651146038641322942321632125127801, + 31426415.585400194380614231628318205362874684987640, + 2876370.6289353724412254090516208496135991145378768, + 186056.26539522349504029498971604569928220784236328, + 8071.6720023658162106380029022722506138218516325024, + 210.82427775157934587250973392071336271166969580291, + 2.5066282746310002701649081771338373386264310793408, + ] +) @_scalar_output_to_0d_array -def fmax( +def allclose( x1: np.ndarray, x2: np.ndarray, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[np.ndarray] = None, -) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.fmax( - x1, - x2, - out=None, - where=True, - casting="same_kind", - order="K", - dtype=None, - subok=True, - ) - - -fmax.support_native_out = True +) -> bool: + return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) -@_scalar_output_to_0d_array -def float_power( - x1: Union[np.ndarray, float, list, tuple], - x2: Union[np.ndarray, float, list, tuple], +def conj( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - x1, x2 = promote_types_of_inputs(x1, x2) - return np.float_power(x1, x2, out=out) - - -float_power.support_native_out = True + ret = np.conj(x, out=out) + if x.dtype == bool: + return ret.astype("bool") + return ret @_scalar_output_to_0d_array @@ -69,9 +104,6 @@ def copysign( return np.copysign(x1, x2, out=out) -copysign.support_native_out = True - - @_scalar_output_to_0d_array def count_nonzero( a: np.ndarray, @@ -90,206 +122,221 @@ def count_nonzero( return ret.astype(dtype) -count_nonzero.support_native_out = False - - -def nansum( - x: np.ndarray, +def diff( + x: Union[np.ndarray, list, tuple], /, *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[np.dtype] = None, - keepdims: bool = False, + n: int = 1, + axis: int = -1, + prepend: Optional[Union[np.ndarray, int, float, list, tuple]] = None, + append: Optional[Union[np.ndarray, int, float, list, tuple]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - - -nansum.support_native_out = True + prepend = prepend if prepend is not None else np._NoValue + append = append if append is not None else np._NoValue + return np.diff(x, n=n, axis=axis, prepend=prepend, append=append) -def isclose( - a: np.ndarray, - b: np.ndarray, +def digamma( + x: np.ndarray, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - if np.isscalar(ret): - return np.array(ret, dtype="bool") - return ret - - -isclose.support_native_out = False + # Using `np.errstate` to ignore divide by zero error + # to maintain the same behaviour as other frameworks. + with np.errstate(divide="ignore", invalid="ignore"): + x = np.asarray(x, dtype=x.dtype) + zero = np.zeros_like(x) + one_half = 0.5 * np.ones_like(x) + one = np.ones_like(x) + pi = np.pi * np.ones_like(x) + lanczos_gamma = kLanczosGamma * np.ones_like(x) + lanczos_gamma_plus_one_half = (kLanczosGamma + 0.5) * np.ones_like(x) + log_lanczos_gamma_plus_one_half = np.log(kLanczosGamma + 0.5) * np.ones_like(x) + base_lanczos_coeff = kBaseLanczosCoeff * np.ones_like(x) + need_to_reflect = x < one_half + z = np.where(need_to_reflect, -x, x - one) + num = zero + denom = base_lanczos_coeff + for i in range(len(kLanczosCoefficients)): + lanczos_coefficient = kLanczosCoefficients[i] * np.ones_like(x) + index = i * np.ones_like(x) + num = num - lanczos_coefficient / ((z + index + one) * (z + index + one)) + denom = denom + lanczos_coefficient / (z + index + one) -def signbit( - x: Union[np.ndarray, float, int, list, tuple], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.signbit(x, out=out) + t = lanczos_gamma_plus_one_half + z + log_t = log_lanczos_gamma_plus_one_half + np.log1p( + z / lanczos_gamma_plus_one_half + ) + y = log_t + num / denom - lanczos_gamma / t + reduced_x = x + np.abs(np.floor(x + 0.5)) + reflection = y - pi * np.cos(pi * reduced_x) / np.sin(pi * reduced_x) + real_result = np.where(need_to_reflect, reflection, y) -signbit.support_native_out = True + return np.where( + np.logical_and(x <= zero, x == np.floor(x)), np.nan, real_result + ) -def hypot( - x1: np.ndarray, - x2: np.ndarray, +def fix( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.hypot(x1, x2) + return np.fix(x, out=out) -def diff( - x: Union[np.ndarray, list, tuple], +@_scalar_output_to_0d_array +def float_power( + x1: Union[np.ndarray, float, list, tuple], + x2: Union[np.ndarray, float, list, tuple], /, *, - n: int = 1, - axis: int = -1, - prepend: Optional[Union[np.ndarray, int, float, list, tuple]] = None, - append: Optional[Union[np.ndarray, int, float, list, tuple]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - prepend = prepend if prepend is not None else np._NoValue - append = append if append is not None else np._NoValue - return np.diff(x, n=n, axis=axis, prepend=prepend, append=append) - - -diff.support_native_out = False + x1, x2 = promote_types_of_inputs(x1, x2) + return np.float_power(x1, x2, out=out) @_scalar_output_to_0d_array -def allclose( +def fmax( x1: np.ndarray, x2: np.ndarray, /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[np.ndarray] = None, -) -> bool: - return np.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) +) -> np.ndarray: + x1, x2 = promote_types_of_inputs(x1, x2) + return np.fmax( + x1, + x2, + out=None, + where=True, + casting="same_kind", + order="K", + dtype=None, + subok=True, + ) -allclose.support_native_out = False +def frexp( + x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None +) -> Tuple[np.ndarray, np.ndarray]: + if out is None: + return np.frexp(x, out=(None, None)) + else: + return np.frexp(x, out=out) -def fix( +def gradient( x: np.ndarray, /, *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.fix(x, out=out) - - -fix.support_native_out = True + spacing: Union[int, list, tuple] = 1, + axis: Optional[Union[int, list, tuple]] = None, + edge_order: int = 1, +) -> Union[np.ndarray, List[np.ndarray]]: + if type(spacing) in (int, float): + return np.gradient(x, spacing, axis=axis, edge_order=edge_order) + return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) -def nextafter( +def hypot( x1: np.ndarray, x2: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.nextafter(x1, x2) - - -nextafter.support_natvie_out = True + return np.hypot(x1, x2) -def zeta( - x: np.ndarray, - q: np.ndarray, +def isclose( + a: np.ndarray, + b: np.ndarray, /, *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) - inf_indices = np.logical_or(temp, np.equal(x, 1)) - temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) - temp = np.logical_and(temp, np.less_equal(q, 0)) - nan_indices = np.logical_or(temp, np.less(x, 1)) - n, res = 1, 1 / q**x - while n < 10000: - term = 1 / (q + n) ** x - n, res = n + 1, res + term - ret = np.round(res, decimals=4) - ret[nan_indices] = np.nan - ret[inf_indices] = np.inf + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ret = np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + if np.isscalar(ret): + return np.array(ret, dtype="bool") return ret -zeta.support_native_out = False - - -def gradient( - x: np.ndarray, - /, - *, - spacing: Union[int, list, tuple] = 1, - axis: Optional[Union[int, list, tuple]] = None, - edge_order: int = 1, -) -> Union[np.ndarray, List[np.ndarray]]: - if type(spacing) in (int, float): - return np.gradient(x, spacing, axis=axis, edge_order=edge_order) - return np.gradient(x, *spacing, axis=axis, edge_order=edge_order) - +def lanczos_sum(x): + num = 0.0 + den = 0.0 -def xlogy( - x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None -) -> np.ndarray: - x, y = promote_types_of_inputs(x, y) - if (x == 0).all(): - return 0.0 + if x < 5.0: + for i in range(LANCZOS_N - 1, -1, -1): + num = num * x + lanczos_num_coeffs[i] + den = den * x + lanczos_den_coeffs[i] else: - return x * np.log(y) + for i in range(LANCZOS_N): + num = num / x + lanczos_num_coeffs[i] + den = den / x + lanczos_den_coeffs[i] + return num / den -def conj( - x: np.ndarray, + +def ldexp( + x1: np.ndarray, + x2: Union[np.ndarray, int, list, tuple], /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ret = np.conj(x, out=out) - if x.dtype == bool: - return ret.astype("bool") - return ret + return np.ldexp(x1, x2, out=out) -def ldexp( - x1: np.ndarray, - x2: Union[np.ndarray, int, list, tuple], +# TODO: Replace with native lgamma implementation when available +def lgamma( + x: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.ldexp(x1, x2, out=out) + def func(x): + if not np.isfinite(x): + if np.isnan(x): + return x # lgamma(nan) = nan + else: + return np.inf # lgamma(+-inf) = +inf + if x == np.floor(x) and x <= 2.0: + if x <= 0.0: + return np.inf # lgamma(n) = inf for integers n <= 0 + else: + return 0.0 # lgamma(1) = lgamma(2) = 0.0 -def frexp( - x: np.ndarray, /, *, out: Optional[Tuple[np.ndarray, np.ndarray]] = None -) -> Tuple[np.ndarray, np.ndarray]: - if out is None: - return np.frexp(x, out=(None, None)) - else: - return np.frexp(x, out=out) + absx = np.abs(x) + if absx < 1e-20: + return -np.log(absx) + + # Lanczos' formula + r = np.log(lanczos_sum(absx)) - lanczos_g + r += (absx - 0.5) * (np.log(absx + lanczos_g - 0.5) - 1) + + if x < 0.0: + # Use reflection formula to get value for negative x. + r = np.log(np.pi) - np.log(np.abs(sinpi(absx))) - np.log(absx) - r + + if np.isinf(r): + raise OverflowError("Range error in lgamma") + + return r + + # Vectorize 'func' for element-wise operations on 'x', output matching 'x' dtype. + vfunc = np.vectorize(func, otypes=[x.dtype]) + return vfunc(x) def modf( @@ -303,104 +350,43 @@ def modf( return np.modf(x) -# ---digamma---# -kLanczosGamma = 7 # aka g -kBaseLanczosCoeff = 0.99999999999980993227684700473478 -kLanczosCoefficients = np.array( - [ - 676.520368121885098567009190444019, - -1259.13921672240287047156078755283, - 771.3234287776530788486528258894, - -176.61502916214059906584551354, - 12.507343278686904814458936853, - -0.13857109526572011689554707, - 9.984369578019570859563e-6, - 1.50563273514931155834e-7, - ] -) - - -def digamma( +def nansum( x: np.ndarray, /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[np.dtype] = None, + keepdims: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - # Using `np.errstate` to ignore divide by zero error - # to maintain the same behaviour as other frameworks. - with np.errstate(divide="ignore", invalid="ignore"): - x = np.asarray(x, dtype=x.dtype) - zero = np.zeros_like(x) - one_half = 0.5 * np.ones_like(x) - one = np.ones_like(x) - pi = np.pi * np.ones_like(x) - lanczos_gamma = kLanczosGamma * np.ones_like(x) - lanczos_gamma_plus_one_half = (kLanczosGamma + 0.5) * np.ones_like(x) - log_lanczos_gamma_plus_one_half = np.log(kLanczosGamma + 0.5) * np.ones_like(x) - base_lanczos_coeff = kBaseLanczosCoeff * np.ones_like(x) - need_to_reflect = x < one_half - z = np.where(need_to_reflect, -x, x - one) + if isinstance(axis, list): + axis = tuple(axis) + return np.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) - num = zero - denom = base_lanczos_coeff - for i in range(len(kLanczosCoefficients)): - lanczos_coefficient = kLanczosCoefficients[i] * np.ones_like(x) - index = i * np.ones_like(x) - num = num - lanczos_coefficient / ((z + index + one) * (z + index + one)) - denom = denom + lanczos_coefficient / (z + index + one) - t = lanczos_gamma_plus_one_half + z - log_t = log_lanczos_gamma_plus_one_half + np.log1p( - z / lanczos_gamma_plus_one_half - ) - y = log_t + num / denom - lanczos_gamma / t +def nextafter( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nextafter(x1, x2) - reduced_x = x + np.abs(np.floor(x + 0.5)) - reflection = y - pi * np.cos(pi * reduced_x) / np.sin(pi * reduced_x) - real_result = np.where(need_to_reflect, reflection, y) - return np.where( - np.logical_and(x <= zero, x == np.floor(x)), np.nan, real_result - ) +def signbit( + x: Union[np.ndarray, float, int, list, tuple], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.signbit(x, out=out) -# --- LGAMMA --- # -LANCZOS_N = 13 -lanczos_g = 6.024680040776729583740234375 -lanczos_num_coeffs = np.array( - [ - 23531376880.410759688572007674451636754734846804940, - 42919803642.649098768957899047001988850926355848959, - 35711959237.355668049440185451547166705960488635843, - 17921034426.037209699919755754458931112671403265390, - 6039542586.3520280050642916443072979210699388420708, - 1439720407.3117216736632230727949123939715485786772, - 248874557.86205415651146038641322942321632125127801, - 31426415.585400194380614231628318205362874684987640, - 2876370.6289353724412254090516208496135991145378768, - 186056.26539522349504029498971604569928220784236328, - 8071.6720023658162106380029022722506138218516325024, - 210.82427775157934587250973392071336271166969580291, - 2.5066282746310002701649081771338373386264310793408, - ] -) -lanczos_den_coeffs = np.array( - [ - 0.0, - 39916800.0, - 120543840.0, - 150917976.0, - 105258076.0, - 45995730.0, - 13339535.0, - 2637558.0, - 357423.0, - 32670.0, - 1925.0, - 66.0, - 1.0, - ] -) +@_scalar_output_to_0d_array +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def sinc(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: + return np.sinc(x).astype(x.dtype) def sinpi(x): @@ -424,59 +410,49 @@ def sinpi(x): return np.copysign(1.0, x) * r -def lanczos_sum(x): - num = 0.0 - den = 0.0 - - if x < 5.0: - for i in range(LANCZOS_N - 1, -1, -1): - num = num * x + lanczos_num_coeffs[i] - den = den * x + lanczos_den_coeffs[i] +def xlogy( + x: np.ndarray, y: np.ndarray, /, *, out: Optional[np.ndarray] = None +) -> np.ndarray: + x, y = promote_types_of_inputs(x, y) + if (x == 0).all(): + return 0.0 else: - for i in range(LANCZOS_N): - num = num / x + lanczos_num_coeffs[i] - den = den / x + lanczos_den_coeffs[i] - - return num / den + return x * np.log(y) -# TODO: Replace with native lgamma implementation when available -def lgamma( +def zeta( x: np.ndarray, + q: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - def func(x): - if not np.isfinite(x): - if np.isnan(x): - return x # lgamma(nan) = nan - else: - return np.inf # lgamma(+-inf) = +inf - - if x == np.floor(x) and x <= 2.0: - if x <= 0.0: - return np.inf # lgamma(n) = inf for integers n <= 0 - else: - return 0.0 # lgamma(1) = lgamma(2) = 0.0 - - absx = np.abs(x) - if absx < 1e-20: - return -np.log(absx) - - # Lanczos' formula - r = np.log(lanczos_sum(absx)) - lanczos_g - r += (absx - 0.5) * (np.log(absx + lanczos_g - 0.5) - 1) - - if x < 0.0: - # Use reflection formula to get value for negative x. - r = np.log(np.pi) - np.log(np.abs(sinpi(absx))) - np.log(absx) - r - - if np.isinf(r): - raise OverflowError("Range error in lgamma") + temp = np.logical_and(np.greater(x, 0), np.equal(np.remainder(x, 2), 0)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + temp = np.logical_and(temp, np.equal(np.remainder(q, 1), 0)) + inf_indices = np.logical_or(temp, np.equal(x, 1)) + temp = np.logical_and(np.not_equal(np.remainder(x, 2), 0), np.greater(x, 1)) + temp = np.logical_and(temp, np.less_equal(q, 0)) + nan_indices = np.logical_or(temp, np.less(x, 1)) + n, res = 1, 1 / q**x + while n < 10000: + term = 1 / (q + n) ** x + n, res = n + 1, res + term + ret = np.round(res, decimals=4) + ret[nan_indices] = np.nan + ret[inf_indices] = np.inf + return ret - return r - # Vectorize 'func' for element-wise operations on 'x', output matching 'x' dtype. - vfunc = np.vectorize(func, otypes=[x.dtype]) - return vfunc(x) +fmax.support_native_out = True +float_power.support_native_out = True +copysign.support_native_out = True +count_nonzero.support_native_out = False +nansum.support_native_out = True +isclose.support_native_out = False +signbit.support_native_out = True +diff.support_native_out = False +allclose.support_native_out = False +fix.support_native_out = True +nextafter.support_natvie_out = True +zeta.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/layers.py b/ivy/functional/backends/numpy/experimental/layers.py index 757e434c7b834..4360cf4099908 100644 --- a/ivy/functional/backends/numpy/experimental/layers.py +++ b/ivy/functional/backends/numpy/experimental/layers.py @@ -19,6 +19,10 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): kernel, strides, depth_pooling = _depth_max_pooling_helper( x.shape, kernel, strides, dims=dims, data_format=data_format @@ -28,82 +32,77 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def max_pool1d( +def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): + if isinstance(padding, str): + pad_specific = [ + _handle_padding(x_shape[i], strides[i], kernel[i], padding) + for i in range(dim) + ] + padding = [ + (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) + for i in range(dim) + ] + else: + pad_specific = [sum(padding[i]) for i in range(dim)] + + c = [] + if ceil_mode: + for i in range(dim): + padding[i], c_i = _padding_ceil_mode( + x_shape[i], kernel[i], padding[i], strides[i], True + ) + c.append(c_i) + pad_specific[i] = sum(padding[i]) + return padding, pad_specific, c + + +# --- Main --- # +# ------------ # + + +def avg_pool1d( x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, /, *, data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, + count_include_pad: bool = False, ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + if isinstance(kernel, int): + kernel = [kernel] + elif len(kernel) == 1: + kernel = [kernel[0]] - if data_format == "NCW": + if isinstance(strides, int): + strides = [strides] + elif len(strides) == 1: + strides = [strides[0]] + + if data_format in ("NCW", "NCL"): x = np.swapaxes(x, 1, 2) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" + x_shape = x.shape[1:-1] + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 1 ) - x_shape = x.shape[1:2] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - if dilation[0] > 1: - filters = _add_dilations(filters, dilation[0], axis=0, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_list = [ - (pad_w // 2, pad_w - pad_w // 2), - ] - if ceil_mode: - pad_list[0] = _padding_ceil_mode( - x_shape[0], kernel[0], pad_list[0], strides[0] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) x_shape = x.shape new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] + new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] new_strides = ( x.strides[0], x.strides[1] * strides[0], @@ -111,101 +110,79 @@ def max_pool1d( x.strides[2], ) - # B x OW x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - # B x OW x KW x I - sub_matrices = np.where( - filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf - ) + res = np.mean(sub_matrices, axis=2) - res = sub_matrices.max(axis=(2)) + if (not count_include_pad or ceil_mode) and any(pad_specific): + if not count_include_pad: + num_padded_values = np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[0], + "n": x.shape[1] - pad_specific[0], + "k": kernel[0], + "s": strides[0], + }, + unique={ + "i": np.arange(res.shape[1]), + }, + ), + dtype=res.dtype, + ) + else: + num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) + num_padded_values[-1] = c[0] + res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) + + if data_format in ("NCW", "NCL"): + return res.swapaxes(1, 2) - if depth_pooling: - res = np.swapaxes(res, 1, 2) - if data_format == "NCW": - res = np.swapaxes(res, 1, 2) return res -def max_pool2d( +def avg_pool2d( x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, /, *, data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, + count_include_pad: bool = False, ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + if isinstance(kernel, int): + kernel = [kernel] * 2 + elif len(kernel) == 1: + kernel = [kernel[0]] * 2 + + if isinstance(strides, int): + strides = [strides] * 2 + elif len(strides) == 1: + strides = [strides[0]] * 2 if data_format == "NCHW": x = np.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) x_shape = list(x.shape[1:3]) - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_list = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] - ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + padding, pad_specific, c = _get_padded_values( + x_shape, kernel, strides, padding, ceil_mode, 2 + ) + x = np.pad( + x, + [ + (0, 0), + *padding, + (0, 0), + ], + constant_values=0.0, + ) x_shape = x.shape new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 @@ -219,195 +196,90 @@ def max_pool2d( x.strides[2], x.strides[3], ) - # B x OH x OW x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - # B x OH x OW x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf - ) - # B x OH x OW x O - res = sub_matrices.max(axis=(3, 4)) - - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 1)) - if data_format == "NCHW": - return np.transpose(res, (0, 3, 1, 2)) - return res - - -def max_pool3d( - x: np.ndarray, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCDHW": - x = np.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - x_shape = x.shape[1:4] - filters = np.ones((list(kernel)), dtype=x.dtype) - if not depth_pooling: - for j in range(dims): - if dilation[j] > 1: - filters = _add_dilations(filters, dilation[j], axis=j, values=0) - kernel = list(filters.shape) - pad_list = padding - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) - pad_list = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - pad_list = list(pad_list) - if ceil_mode: - for i in range(dims): - pad_list[i] = _padding_ceil_mode( - x_shape[i], kernel[i], pad_list[i], strides[i] + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override + else: + res = np.mean(sub_matrices, axis=(3, 4)) + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): + if not count_include_pad: + num_padded_values = [ + np.array( + ivy.map( + _get_num_padded_values, + constant={ + "p": pad_specific[i], + "n": x.shape[i + 1] - pad_specific[i], + "k": kernel[i], + "s": strides[i], + }, + unique={ + "i": np.arange(res.shape[i + 1]), + }, + ), + dtype=res.dtype, ) - - x = np.pad( - x, - [ - (0, 0), - *pad_list, - (0, 0), - ], - "constant", - constant_values=-math.inf, + for i in range(2) + ] + else: + num_padded_values = [] + for i in range(2): + num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) + num_pad[-1] = c[i] + num_padded_values.append(num_pad) + num_padded_values1 = num_padded_values[0][:, None] + num_padded_values2 = num_padded_values[1][None, :] + num_padded_values = ( + num_padded_values1 * kernel[1] + + num_padded_values2 * kernel[0] + - num_padded_values1 * num_padded_values2 ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - x_shape = x.shape - new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 - new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[3] * strides[2], - x.strides[1], - x.strides[2], - x.strides[3], - x.strides[4], - ) - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OD x OH x OW x KD x KH x KW x I - sub_matrices = np.where( - filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf - ) - - # B x OD x OH x OW x O - res = sub_matrices.max(axis=(4, 5, 6)) + kernel_mul = np.prod(kernel) + res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) - if depth_pooling: - res = np.transpose(res, (0, 2, 3, 4, 1)) - if data_format == "NCDHW": - return np.transpose(res, (0, 4, 1, 2, 3)) + if data_format == "NCHW": + return np.transpose(res, (0, 3, 1, 2)) return res -def _get_padded_values(x_shape, kernel, strides, padding, ceil_mode, dim): - if isinstance(padding, str): - pad_specific = [ - _handle_padding(x_shape[i], strides[i], kernel[i], padding) - for i in range(dim) - ] - padding = [ - (pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2) - for i in range(dim) - ] - else: - pad_specific = [sum(padding[i]) for i in range(dim)] - - c = [] - if ceil_mode: - for i in range(dim): - padding[i], c_i = _padding_ceil_mode( - x_shape[i], kernel[i], padding[i], strides[i], True - ) - c.append(c_i) - pad_specific[i] = sum(padding[i]) - return padding, pad_specific, c - - -def avg_pool1d( +def avg_pool3d( x: np.ndarray, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], padding: str, /, *, - data_format: str = "NWC", + data_format: str = "NDHWC", count_include_pad: bool = False, ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: if isinstance(kernel, int): - kernel = [kernel] + kernel = [kernel] * 3 elif len(kernel) == 1: - kernel = [kernel[0]] + kernel = [kernel[0]] * 3 if isinstance(strides, int): - strides = [strides] + strides = [strides] * 3 elif len(strides) == 1: - strides = [strides[0]] + strides = [strides[0]] * 3 - if data_format in ("NCW", "NCL"): - x = np.swapaxes(x, 1, 2) + if data_format == "NCDHW": + x = np.transpose(x, (0, 2, 3, 4, 1)) - x_shape = x.shape[1:-1] + x_shape = list(x.shape[1:4]) padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 1 + x_shape, kernel, strides, padding, ceil_mode, 3 ) x = np.pad( @@ -421,228 +293,36 @@ def avg_pool1d( ) x_shape = x.shape - new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_shape = [x_shape[0], new_w, kernel[0]] + [x_shape[-1]] + new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 + new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] new_strides = ( x.strides[0], x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[3] * strides[2], x.strides[1], x.strides[2], + x.strides[3], + x.strides[4], ) - + # B x OH x OW x KH x KW x I sub_matrices = np.lib.stride_tricks.as_strided( x, new_shape, new_strides, writeable=False ) - res = np.mean(sub_matrices, axis=2) + # B x OH x OW x O + if divisor_override is not None: + res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override + else: + res = np.mean(sub_matrices, axis=(4, 5, 6)) - if (not count_include_pad or ceil_mode) and any(pad_specific): - if not count_include_pad: - num_padded_values = np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[0], - "n": x.shape[1] - pad_specific[0], - "k": kernel[0], - "s": strides[0], - }, - unique={ - "i": np.arange(res.shape[1]), - }, - ), - dtype=res.dtype, - ) - else: - num_padded_values = np.zeros(res.shape[1], dtype=res.dtype) - num_padded_values[-1] = c[0] - res = (kernel[0] * res) / (kernel[0] - num_padded_values[:, None]) - - if data_format in ("NCW", "NCL"): - return res.swapaxes(1, 2) - - return res - - -def avg_pool2d( - x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 2 - elif len(kernel) == 1: - kernel = [kernel[0]] * 2 - - if isinstance(strides, int): - strides = [strides] * 2 - elif len(strides) == 1: - strides = [strides[0]] * 2 - - if data_format == "NCHW": - x = np.transpose(x, (0, 2, 3, 1)) - - x_shape = list(x.shape[1:3]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 2 - ) - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) - - x_shape = x.shape - new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[1], - x.strides[2], - x.strides[3], - ) - # B x OH x OW x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(3, 4)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(3, 4)) - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): - if not count_include_pad: - num_padded_values = [ - np.array( - ivy.map( - _get_num_padded_values, - constant={ - "p": pad_specific[i], - "n": x.shape[i + 1] - pad_specific[i], - "k": kernel[i], - "s": strides[i], - }, - unique={ - "i": np.arange(res.shape[i + 1]), - }, - ), - dtype=res.dtype, - ) - for i in range(2) - ] - else: - num_padded_values = [] - for i in range(2): - num_pad = np.zeros(res.shape[i + 1], dtype=res.dtype) - num_pad[-1] = c[i] - num_padded_values.append(num_pad) - num_padded_values1 = num_padded_values[0][:, None] - num_padded_values2 = num_padded_values[1][None, :] - num_padded_values = ( - num_padded_values1 * kernel[1] - + num_padded_values2 * kernel[0] - - num_padded_values1 * num_padded_values2 - ) - kernel_mul = np.prod(kernel) - res = (kernel_mul * res) / (kernel_mul - np.expand_dims(num_padded_values, -1)) - - if data_format == "NCHW": - return np.transpose(res, (0, 3, 1, 2)) - return res - - -def avg_pool3d( - x: np.ndarray, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(kernel, int): - kernel = [kernel] * 3 - elif len(kernel) == 1: - kernel = [kernel[0]] * 3 - - if isinstance(strides, int): - strides = [strides] * 3 - elif len(strides) == 1: - strides = [strides[0]] * 3 - - if data_format == "NCDHW": - x = np.transpose(x, (0, 2, 3, 4, 1)) - - x_shape = list(x.shape[1:4]) - padding, pad_specific, c = _get_padded_values( - x_shape, kernel, strides, padding, ceil_mode, 3 - ) - - x = np.pad( - x, - [ - (0, 0), - *padding, - (0, 0), - ], - constant_values=0.0, - ) - - x_shape = x.shape - new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 - new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 - new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 - new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] - new_strides = ( - x.strides[0], - x.strides[1] * strides[0], - x.strides[2] * strides[1], - x.strides[3] * strides[2], - x.strides[1], - x.strides[2], - x.strides[3], - x.strides[4], - ) - # B x OH x OW x KH x KW x I - sub_matrices = np.lib.stride_tricks.as_strided( - x, new_shape, new_strides, writeable=False - ) - - # B x OH x OW x O - if divisor_override is not None: - res = np.sum(sub_matrices, axis=(4, 5, 6)) / divisor_override - else: - res = np.mean(sub_matrices, axis=(4, 5, 6)) - - if ( - (not count_include_pad or ceil_mode) - and any(pad_specific) - and not divisor_override - ): + if ( + (not count_include_pad or ceil_mode) + and any(pad_specific) + and not divisor_override + ): if not count_include_pad: num_padded_values = [ np.array( @@ -687,43 +367,6 @@ def avg_pool3d( return res -def fft( - x: np.ndarray, - dim: int, - /, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.dtype in [np.uint64, np.int64, np.float64, np.complex128]: - out_dtype = np.complex128 - else: - out_dtype = np.complex64 - return np.fft.fft(x, n, dim, norm).astype(out_dtype) - - @with_supported_dtypes({"1.25.2 and below": ("float32", "float64")}, backend_version) def dct( x: np.ndarray, @@ -820,20 +463,6 @@ def dct( return dct_out.astype(np.float32) if cast_final else dct_out -def idct( - x: np.ndarray, - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - def dropout1d( x: np.ndarray, prob: float, @@ -911,22 +540,48 @@ def dropout3d( return res -def ifft( - x: np.ndarray, - dim: int, +@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) +def embedding( + weights: np.ndarray, + indices: np.ndarray, + /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + max_norm: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + + embeddings = np.take(weights, indices, axis=0) + if max_norm is not None: + norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) + embeddings = np.where( + norms > max_norm, embeddings * max_norm / norms, embeddings ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( + embeddings = np.where( + norms < -max_norm, embeddings * -max_norm / norms, embeddings + ) + return embeddings + + +def fft( + x: np.ndarray, + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( f"Invalid dim {dim}, expecting ranging" " from {-len(x.shape)} to {len(x.shape)-1} " ) @@ -940,7 +595,11 @@ def ifft( ) if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return np.asarray(np.fft.ifft(x, n, dim, norm), dtype=x.dtype) + if x.dtype in [np.uint64, np.int64, np.float64, np.complex128]: + out_dtype = np.complex128 + else: + out_dtype = np.complex64 + return np.fft.fft(x, n, dim, norm).astype(out_dtype) def fft2( @@ -978,6 +637,52 @@ def fft2( return np.fft.fft2(x, s, dim, norm).astype(np.complex128) +def idct( + x: np.ndarray, + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def ifft( + x: np.ndarray, + dim: int, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return np.asarray(np.fft.ifft(x, n, dim, norm), dtype=x.dtype) + + def ifftn( x: np.ndarray, s: Optional[Union[int, Tuple[int]]] = None, @@ -989,29 +694,332 @@ def ifftn( return np.fft.ifftn(x, s, axes, norm).astype(x.dtype) -@with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) -def embedding( - weights: np.ndarray, - indices: np.ndarray, +def max_pool1d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - max_norm: Optional[int] = None, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, out: Optional[np.ndarray] = None, ) -> np.ndarray: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims ) - embeddings = np.take(weights, indices, axis=0) - if max_norm is not None: - norms = np.linalg.norm(embeddings, axis=-1, keepdims=True) - embeddings = np.where( - norms > max_norm, embeddings * max_norm / norms, embeddings + if data_format == "NCW": + x = np.swapaxes(x, 1, 2) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides ) - embeddings = np.where( - norms < -max_norm, embeddings * -max_norm / norms, embeddings + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - return embeddings + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = x.shape[1:2] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + if dilation[0] > 1: + filters = _add_dilations(filters, dilation[0], axis=0, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_w = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_list = [ + (pad_w // 2, pad_w - pad_w // 2), + ] + if ceil_mode: + pad_list[0] = _padding_ceil_mode( + x_shape[0], kernel[0], pad_list[0], strides[0] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_w = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_shape = [x_shape[0], new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[1], + x.strides[2], + ) + + # B x OW x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OW x KW x I + sub_matrices = np.where( + filters.reshape([1] * 2 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + res = sub_matrices.max(axis=(2)) + + if depth_pooling: + res = np.swapaxes(res, 1, 2) + if data_format == "NCW": + res = np.swapaxes(res, 1, 2) + return res + + +def max_pool2d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCHW": + x = np.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = list(x.shape[1:3]) + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_list = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_h = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_w = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_shape = [x_shape[0], new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[1], + x.strides[2], + x.strides[3], + ) + + # B x OH x OW x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OH x OW x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 3 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + # B x OH x OW x O + res = sub_matrices.max(axis=(3, 4)) + + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 1)) + if data_format == "NCHW": + return np.transpose(res, (0, 3, 1, 2)) + return res + + +def max_pool3d( + x: np.ndarray, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + + if data_format == "NCDHW": + x = np.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel + ) + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) + + x_shape = x.shape[1:4] + filters = np.ones((list(kernel)), dtype=x.dtype) + if not depth_pooling: + for j in range(dims): + if dilation[j] > 1: + filters = _add_dilations(filters, dilation[j], axis=j, values=0) + kernel = list(filters.shape) + pad_list = padding + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], kernel[2], padding) + pad_list = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + pad_list = list(pad_list) + if ceil_mode: + for i in range(dims): + pad_list[i] = _padding_ceil_mode( + x_shape[i], kernel[i], pad_list[i], strides[i] + ) + + x = np.pad( + x, + [ + (0, 0), + *pad_list, + (0, 0), + ], + "constant", + constant_values=-math.inf, + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + x_shape = x.shape + new_d = (x_shape[1] - kernel[0]) // strides[0] + 1 + new_h = (x_shape[2] - kernel[1]) // strides[1] + 1 + new_w = (x_shape[3] - kernel[2]) // strides[2] + 1 + new_shape = [x_shape[0], new_d, new_h, new_w] + list(kernel) + [x_shape[-1]] + new_strides = ( + x.strides[0], + x.strides[1] * strides[0], + x.strides[2] * strides[1], + x.strides[3] * strides[2], + x.strides[1], + x.strides[2], + x.strides[3], + x.strides[4], + ) + # B x OD x OH x OW x KD x KH x KW x I + sub_matrices = np.lib.stride_tricks.as_strided( + x, new_shape, new_strides, writeable=False + ) + + # B x OD x OH x OW x KD x KH x KW x I + sub_matrices = np.where( + filters.reshape([1] * 4 + list(kernel) + [1]), sub_matrices, -math.inf + ) + + # B x OD x OH x OW x O + res = sub_matrices.max(axis=(4, 5, 6)) + + if depth_pooling: + res = np.transpose(res, (0, 2, 3, 4, 1)) + if data_format == "NCDHW": + return np.transpose(res, (0, 4, 1, 2, 3)) + return res def rfftn( diff --git a/ivy/functional/backends/numpy/experimental/linear_algebra.py b/ivy/functional/backends/numpy/experimental/linear_algebra.py index 915b0b4ad9111..c3579d21a6485 100644 --- a/ivy/functional/backends/numpy/experimental/linear_algebra.py +++ b/ivy/functional/backends/numpy/experimental/linear_algebra.py @@ -10,6 +10,28 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size +def adjoint( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + _check_valid_dimension_size(x) + axes = list(range(len(x.shape))) + axes[-1], axes[-2] = axes[-2], axes[-1] + return np.conjugate(np.transpose(x, axes=axes)) + + +def cond( + x: np.ndarray, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[np.ndarray] = None, +) -> Any: + return np.linalg.cond(x, p=p) + + def diagflat( x: np.ndarray, /, @@ -84,34 +106,14 @@ def diagflat( return ret -def kron( +def dot( a: np.ndarray, b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.kron(a, b) - - -kron.support_native_out = False - - -@with_supported_dtypes( - {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, - backend_version, -) -def matrix_exp( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - eig_vals, eig_vecs = np.linalg.eig(x) - exp_diag = np.exp(eig_vals) - exp_diag_mat = np.diag(exp_diag) - exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) - return exp_mat.astype(x.dtype) + return np.dot(a, b, out=out) def eig( @@ -126,9 +128,6 @@ def eig( return e.astype(complex), v.astype(complex) -eig.support_native_out = False - - def eigvals(x: np.ndarray, /) -> np.ndarray: if ivy.dtype(x) == ivy.float16: x = x.astype(np.float32) @@ -136,64 +135,55 @@ def eigvals(x: np.ndarray, /) -> np.ndarray: return e.astype(complex) -eigvals.support_native_out = False - - -def adjoint( - x: np.ndarray, +def kron( + a: np.ndarray, + b: np.ndarray, /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - _check_valid_dimension_size(x) - axes = list(range(len(x.shape))) - axes[-1], axes[-2] = axes[-2], axes[-1] - return np.conjugate(np.transpose(x, axes=axes)) - - -def multi_dot( - x: Sequence[np.ndarray], - /, - *, - out: Optional[np.array] = None, -) -> np.ndarray: - return np.linalg.multi_dot(x, out=out) - - -multi_dot.support_native_out = True + return np.kron(a, b) -def cond( +def lu_factor( x: np.ndarray, /, *, - p: Optional[Union[None, int, str]] = None, + pivot: Optional[bool] = True, out: Optional[np.ndarray] = None, -) -> Any: - return np.linalg.cond(x, p=p) - - -cond.support_native_out = False +) -> Tuple[np.ndarray]: + raise IvyNotImplementedException() -def lu_factor( +@with_supported_dtypes( + {"1.25.2 and below": ("float32", "float64", "complex64", "complex128")}, + backend_version, +) +def matrix_exp( x: np.ndarray, /, *, - pivot: Optional[bool] = True, out: Optional[np.ndarray] = None, -) -> Tuple[np.ndarray]: - raise IvyNotImplementedException() +) -> np.ndarray: + eig_vals, eig_vecs = np.linalg.eig(x) + exp_diag = np.exp(eig_vals) + exp_diag_mat = np.diag(exp_diag) + exp_mat = eig_vecs @ exp_diag_mat @ np.linalg.inv(eig_vecs) + return exp_mat.astype(x.dtype) -def dot( - a: np.ndarray, - b: np.ndarray, +def multi_dot( + x: Sequence[np.ndarray], /, *, - out: Optional[np.ndarray] = None, + out: Optional[np.array] = None, ) -> np.ndarray: - return np.dot(a, b, out=out) + return np.linalg.multi_dot(x, out=out) +kron.support_native_out = False +eig.support_native_out = False +eigvals.support_native_out = False +multi_dot.support_native_out = True +cond.support_native_out = False dot.support_native_out = True diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 53436987e2187..0c03925ae8aa2 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -19,103 +19,152 @@ from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array -def moveaxis( - a: np.ndarray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.moveaxis(a, source, destination) +# --- Helpers --- # +# --------------- # -moveaxis.support_native_out = False +def _flat_array_to_1_dim_array(x): + return x.reshape((1,)) if x.shape == () else x -def heaviside( - x1: np.ndarray, - x2: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.heaviside( - x1, - x2, - out=out, - ) +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = np.full(new_shape, padding_value, dtype=operand.dtype) + src_indices = np.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) -heaviside.support_native_out = True + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = np.pad(padded, pad_width, constant_values=padding_value) + return padded -def flipud( - m: np.ndarray, - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.flipud(m) +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] -flipud.support_native_out = False +# --- Main --- # +# ------------ # -def vstack( - arrays: Sequence[np.ndarray], + +def atleast_1d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_1d(*arys) + + +def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: + return np.atleast_2d(*arys) + + +def atleast_3d( + *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None +) -> List[np.ndarray]: + return np.atleast_3d(*arys) + + +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return np.broadcast_shapes(*shapes) + + +def concat_from_sequence( + input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], /, *, + new_axis: int = 0, + axis: int = 0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.vstack(arrays) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = np.concatenate(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = np.stack(input_sequence, axis=axis) + return ret -def hstack( +def dsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + + +def dstack( arrays: Sequence[np.ndarray], /, *, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.hstack(arrays) + return np.dstack(arrays) -def rot90( - m: np.ndarray, +def expand( + x: np.ndarray, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.rot90(m, k, axes) + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) + return np.broadcast_to(x, tuple(shape)) -def top_k( - x: np.ndarray, - k: int, +def fill_diagonal( + a: np.ndarray, + v: Union[int, float], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[np.ndarray, np.ndarray]] = None, -) -> Tuple[np.ndarray, np.ndarray]: - k = min(k, x.shape[axis]) - if not largest: - indices = np.argsort(x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - else: - indices = np.argsort(-x, axis=axis) - indices = np.take(indices, np.arange(k), axis=axis) - if not sorted: - indices = np.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) - val = np.take_along_axis(x, indices, axis=axis) - return topk_res(val, indices) + wrap: bool = False, +) -> np.ndarray: + np.fill_diagonal(a, v, wrap=wrap) + return a def fliplr( @@ -128,71 +177,70 @@ def fliplr( return np.fliplr(m) -fliplr.support_native_out = False - - -def i0( - x: np.ndarray, +def flipud( + m: np.ndarray, /, *, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.i0(x) - + return np.flipud(m) -i0.support_native_out = False +def heaviside( + x1: np.ndarray, + x2: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.heaviside( + x1, + x2, + out=out, + ) -def _flat_array_to_1_dim_array(x): - return x.reshape((1,)) if x.shape == () else x +def hsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] +def hstack( + arrays: Sequence[np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.hstack(arrays) -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = np.full(new_shape, padding_value, dtype=operand.dtype) - src_indices = np.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array +def i0( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.i0(x) - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = np.pad(padded, pad_width, constant_values=padding_value) - return padded +def moveaxis( + a: np.ndarray, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.moveaxis(a, source, destination) def pad( @@ -268,57 +316,16 @@ def pad( ) -def vsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Sequence[int], np.ndarray], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - - -def dsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], +def rot90( + m: np.ndarray, /, *, copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_1d(*arys) - - -def dstack( - arrays: Sequence[np.ndarray], - /, - *, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.dstack(arrays) - - -def atleast_2d(*arys: np.ndarray, copy: Optional[bool] = None) -> List[np.ndarray]: - return np.atleast_2d(*arys) - - -def atleast_3d( - *arys: Union[np.ndarray, bool, Number], copy: Optional[bool] = None -) -> List[np.ndarray]: - return np.atleast_3d(*arys) + return np.rot90(m, k, axes) @_scalar_output_to_0d_array @@ -366,63 +373,28 @@ def take_along_axis( return np.take_along_axis(arr, indices, axis) -def hsplit( - ary: np.ndarray, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[np.ndarray]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -take_along_axis.support_native_out = False - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return np.broadcast_shapes(*shapes) - - -broadcast_shapes.support_native_out = False - - -def expand( +def top_k( x: np.ndarray, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = int(np.prod(x.shape) / np.prod([s for s in shape if s > 0])) - return np.broadcast_to(x, tuple(shape)) - - -expand.support_native_out = False - - -def concat_from_sequence( - input_sequence: Union[Tuple[np.ndarray], List[np.ndarray]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = np.concatenate(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = np.stack(input_sequence, axis=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[np.ndarray, np.ndarray]] = None, +) -> Tuple[np.ndarray, np.ndarray]: + k = min(k, x.shape[axis]) + if not largest: + indices = np.argsort(x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + else: + indices = np.argsort(-x, axis=axis) + indices = np.take(indices, np.arange(k), axis=axis) + if not sorted: + indices = np.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", np.ndarray), ("indices", np.ndarray)]) + val = np.take_along_axis(x, indices, axis=axis) + return topk_res(val, indices) def unique_consecutive( @@ -468,12 +440,34 @@ def unique_consecutive( ) -def fill_diagonal( - a: np.ndarray, - v: Union[int, float], +def vsplit( + ary: np.ndarray, + indices_or_sections: Union[int, Sequence[int], np.ndarray], /, *, - wrap: bool = False, + copy: Optional[bool] = None, +) -> List[np.ndarray]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Sequence[np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, ) -> np.ndarray: - np.fill_diagonal(a, v, wrap=wrap) - return a + return np.vstack(arrays) + + +moveaxis.support_native_out = False +heaviside.support_native_out = True +flipud.support_native_out = False +fliplr.support_native_out = False +i0.support_native_out = False +take_along_axis.support_native_out = False +broadcast_shapes.support_native_out = False +expand.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/random.py b/ivy/functional/backends/numpy/experimental/random.py index 16293d4eac407..43ff14af9c947 100644 --- a/ivy/functional/backends/numpy/experimental/random.py +++ b/ivy/functional/backends/numpy/experimental/random.py @@ -10,24 +10,23 @@ ) -# dirichlet -def dirichlet( - alpha: Union[np.ndarray, float, Sequence[float]], - /, +def bernoulli( + probs: Union[float, np.ndarray], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + logits: Optional[Union[float, np.ndarray]] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, dtype: Optional[np.dtype] = None, seed: Optional[int] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - size = size if size is not None else len(alpha) - dtype = dtype if dtype is not None else np.float64 if seed is not None: np.random.seed(seed) - return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) - - -dirichlet.support_native_out = False + if logits is not None: + probs = np.asarray(ivy.softmax(logits), dtype=dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) def beta( @@ -47,6 +46,23 @@ def beta( return np.asarray(np.random.beta(alpha, beta, shape), dtype=dtype) +# dirichlet +def dirichlet( + alpha: Union[np.ndarray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[np.dtype] = None, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + size = size if size is not None else len(alpha) + dtype = dtype if dtype is not None else np.float64 + if seed is not None: + np.random.seed(seed) + return np.asarray(np.random.dirichlet(alpha, size=size), dtype=dtype) + + def gamma( alpha: Union[float, np.ndarray], beta: Union[float, np.ndarray], @@ -89,20 +105,4 @@ def poisson( return np.asarray(ret, dtype=dtype) -def bernoulli( - probs: Union[float, np.ndarray], - *, - logits: Optional[Union[float, np.ndarray]] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, - dtype: Optional[np.dtype] = None, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if seed is not None: - np.random.seed(seed) - if logits is not None: - probs = np.asarray(ivy.softmax(logits), dtype=dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return np.asarray(np.random.binomial(1, p=probs, size=shape), dtype=dtype) +dirichlet.support_native_out = False diff --git a/ivy/functional/backends/numpy/experimental/statistical.py b/ivy/functional/backends/numpy/experimental/statistical.py index 12c3565148ec7..6efebf5a93ba6 100644 --- a/ivy/functional/backends/numpy/experimental/statistical.py +++ b/ivy/functional/backends/numpy/experimental/statistical.py @@ -8,194 +8,87 @@ from copy import deepcopy -@with_unsupported_dtypes( - {"1.25.2 and below": ("bfloat16",)}, - backend_version, -) -def histogram( - a: np.ndarray, - /, - *, - bins: Optional[Union[int, np.ndarray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[np.dtype] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[np.ndarray] = None, - density: Optional[bool] = False, - out: Optional[np.ndarray] = None, -) -> Tuple[np.ndarray]: - min_a = np.min(a) - max_a = np.max(a) - if isinstance(bins, np.ndarray) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - elif isinstance(bins, int): - range = (min_a, max_a) - bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) - range = None - if bins.size < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - bins_out = bins.copy() - if extend_lower_interval and min_a < bins[0]: - bins[0] = min_a - if extend_upper_interval and max_a > bins[-1]: - bins[-1] = max_a - if a.ndim > 0 and axis is not None: - inverted_shape_dims = list(np.flip(np.arange(a.ndim))) - if isinstance(axis, int): - axis = [axis] - shape_axes = 1 - for dimension in axis: - inverted_shape_dims.remove(dimension) - inverted_shape_dims.append(dimension) - shape_axes *= a.shape[dimension] - a_along_axis_1d = ( - a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) - ) - if weights is None: - ret = [] - for a_1d in a_along_axis_1d: - ret_1d = np.histogram( - a_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - else: - weights_along_axis_1d = ( - weights.transpose(inverted_shape_dims) - .flatten() - .reshape((-1, shape_axes)) - ) - ret = [] - for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): - ret_1d = np.histogram( - a_1d, - weights=weights_1d, - bins=bins, - range=range, - # TODO: waiting tensorflow version support to density - # density=density, - )[0] - ret.append(ret_1d) - out_shape = list(a.shape) - for dimension in sorted(axis, reverse=True): - del out_shape[dimension] - out_shape.insert(0, len(bins) - 1) - ret = np.array(ret) - ret = ret.flatten() - index = np.zeros(len(out_shape), dtype=int) - ret_shaped = np.zeros(out_shape) - dim = 0 - i = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - while list(index) != list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - dim_full_flag = False - while index[dim] == out_shape[dim] - 1: - index[dim] = 0 - dim += 1 - dim_full_flag = True - index[dim] += 1 - i += 1 - if dim_full_flag: - dim = 0 - if list(index) == list(np.array(out_shape) - 1): - ret_shaped[tuple(index)] = ret[i] - ret = ret_shaped - else: - ret = np.histogram( - a=a, bins=bins, range=range, weights=weights, density=density - )[0] - if dtype: - ret = ret.astype(dtype) - bins_out = np.array(bins_out).astype(dtype) - # TODO: weird error when returning bins: return ret, bins_out - return ret +# --- Helpers --- # +# --------------- # -def median( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if out is not None: - out = np.reshape(out, input.shape) - ret = np.median( - input, - axis=axis, - keepdims=keepdims, - out=out, - ) - if input.dtype in [np.uint64, np.int64, np.float64]: - return ret.astype(np.float64) - elif input.dtype in [np.float16]: - return ret.astype(input.dtype) - else: - return ret.astype(np.float32) - - -median.support_native_out = True - - -def nanmean( - a: np.ndarray, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, +def __find_cummax_indices( + x: np.ndarray, + axis: int = 0, ) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) + indices = [] + if x[0] is np.ndarray: + if axis >= 1: + for ret1 in x: + indice = __find_cummax_indices(ret1, axis=axis - 1) + indices.append(indice) + else: + indice_list = __get_index(x.tolist()) + indices, n1 = x.copy(), {} + indices.fill(0) + indice_list = sorted(indice_list, key=lambda i: i[1]) + for y, y_index in indice_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + elif ( + y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + else: + indices[y_index] = n1[tuple(multi_index[1:])] + else: + n = 0 + for index1, ret1 in enumerate(x): + if x[n] <= ret1 or index1 == 0: + n = index1 + indices.append(n) + return np.array(indices, dtype=np.int64) -nanmean.support_native_out = True +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] -def _validate_quantile(q): - if isinstance(q, float): - q = np.asarray(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) else: - if not (np.all(0 <= q) and np.all(q <= 1)): - return False - return True - - -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] + indices.append((lst, tuple(prefix))) + return indices - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + axis = tuple(axis) if isinstance(axis, list) else axis - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis + return np.quantile( + x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out + ).astype(x.dtype) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) def _handle_axis(a, q, fn, keepdims=False, axis=None): @@ -261,89 +154,39 @@ def _quantile(a, q, axis=None): return out.astype(ret_dtype) -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - axis = tuple(axis) if isinstance(axis, list) else axis - - return np.quantile( - x, q, axis=axis, method=interpolation, keepdims=keepdims, out=out - ).astype(x.dtype) - else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - - -def quantile( - a: np.ndarray, - q: Union[float, np.ndarray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, - interpolation: str = "linear", - out: Optional[np.ndarray] = None, -) -> np.ndarray: - # quantile method in numpy backend, always return an array with dtype=float64. - # in other backends, the output is the same dtype as the input. - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - out=out, - ) +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + if len(axis) == 0: + raise ValueError("Axis can't be empty!") -def corrcoef( - x: np.ndarray, - /, - *, - y: Optional[np.ndarray] = None, - rowvar: bool = True, - dtype: np.dtype = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - dtype = dtype if dtype is not None else np.float64 + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") - return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis -@with_unsupported_dtypes( - {"1.25.0 and below": ("bfloat16",)}, - backend_version, -) -def nanmedian( - input: np.ndarray, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out - ) +def _validate_quantile(q): + if isinstance(q, float): + q = np.asarray(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (np.all(0 <= q) and np.all(q <= 1)): + return False + return True -nanmedian.support_native_out = True +# --- Main --- # +# ------------ # def bincount( @@ -363,7 +206,18 @@ def bincount( return ret -bincount.support_native_out = False +def corrcoef( + x: np.ndarray, + /, + *, + y: Optional[np.ndarray] = None, + rowvar: bool = True, + dtype: np.dtype = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + dtype = dtype if dtype is not None else np.float64 + + return np.corrcoef(x, y=y, rowvar=rowvar, dtype=dtype) def cov( @@ -390,9 +244,6 @@ def cov( ) -cov.support_native_out = False - - def cummax( x: np.ndarray, /, @@ -438,58 +289,6 @@ def cummax( return np.maximum.accumulate(x, axis=axis, dtype=x.dtype), indices -def __find_cummax_indices( - x: np.ndarray, - axis: int = 0, -) -> np.ndarray: - indices = [] - if x[0] is np.ndarray: - if axis >= 1: - for ret1 in x: - indice = __find_cummax_indices(ret1, axis=axis - 1) - indices.append(indice) - - else: - indice_list = __get_index(x.tolist()) - indices, n1 = x.copy(), {} - indices.fill(0) - indice_list = sorted(indice_list, key=lambda i: i[1]) - for y, y_index in indice_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - elif ( - y >= x[tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:]))] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - else: - indices[y_index] = n1[tuple(multi_index[1:])] - else: - n = 0 - for index1, ret1 in enumerate(x): - if x[n] <= ret1 or index1 == 0: - n = index1 - indices.append(n) - return np.array(indices, dtype=np.int64) - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) def cummin( x: np.ndarray, @@ -513,6 +312,121 @@ def cummin( return np.flip(x, axis=axis) +@with_unsupported_dtypes( + {"1.25.2 and below": ("bfloat16",)}, + backend_version, +) +def histogram( + a: np.ndarray, + /, + *, + bins: Optional[Union[int, np.ndarray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[np.dtype] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[np.ndarray] = None, + density: Optional[bool] = False, + out: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray]: + min_a = np.min(a) + max_a = np.max(a) + if isinstance(bins, np.ndarray) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + elif isinstance(bins, int): + range = (min_a, max_a) + bins = np.linspace(start=range[0], stop=range[1], num=bins + 1, dtype=a.dtype) + range = None + if bins.size < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + bins_out = bins.copy() + if extend_lower_interval and min_a < bins[0]: + bins[0] = min_a + if extend_upper_interval and max_a > bins[-1]: + bins[-1] = max_a + if a.ndim > 0 and axis is not None: + inverted_shape_dims = list(np.flip(np.arange(a.ndim))) + if isinstance(axis, int): + axis = [axis] + shape_axes = 1 + for dimension in axis: + inverted_shape_dims.remove(dimension) + inverted_shape_dims.append(dimension) + shape_axes *= a.shape[dimension] + a_along_axis_1d = ( + a.transpose(inverted_shape_dims).flatten().reshape((-1, shape_axes)) + ) + if weights is None: + ret = [] + for a_1d in a_along_axis_1d: + ret_1d = np.histogram( + a_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + else: + weights_along_axis_1d = ( + weights.transpose(inverted_shape_dims) + .flatten() + .reshape((-1, shape_axes)) + ) + ret = [] + for a_1d, weights_1d in zip(a_along_axis_1d, weights_along_axis_1d): + ret_1d = np.histogram( + a_1d, + weights=weights_1d, + bins=bins, + range=range, + # TODO: waiting tensorflow version support to density + # density=density, + )[0] + ret.append(ret_1d) + out_shape = list(a.shape) + for dimension in sorted(axis, reverse=True): + del out_shape[dimension] + out_shape.insert(0, len(bins) - 1) + ret = np.array(ret) + ret = ret.flatten() + index = np.zeros(len(out_shape), dtype=int) + ret_shaped = np.zeros(out_shape) + dim = 0 + i = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + while list(index) != list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + dim_full_flag = False + while index[dim] == out_shape[dim] - 1: + index[dim] = 0 + dim += 1 + dim_full_flag = True + index[dim] += 1 + i += 1 + if dim_full_flag: + dim = 0 + if list(index) == list(np.array(out_shape) - 1): + ret_shaped[tuple(index)] = ret[i] + ret = ret_shaped + else: + ret = np.histogram( + a=a, bins=bins, range=range, weights=weights, density=density + )[0] + if dtype: + ret = ret.astype(dtype) + bins_out = np.array(bins_out).astype(dtype) + # TODO: weird error when returning bins: return ret, bins_out + return ret + + def igamma( a: np.ndarray, /, @@ -528,3 +442,89 @@ def igamma_cal(a, x): igamma_vec = np.vectorize(igamma_cal) return igamma_vec(a, x).astype(a.dtype) + + +def median( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if out is not None: + out = np.reshape(out, input.shape) + ret = np.median( + input, + axis=axis, + keepdims=keepdims, + out=out, + ) + if input.dtype in [np.uint64, np.int64, np.float64]: + return ret.astype(np.float64) + elif input.dtype in [np.float16]: + return ret.astype(input.dtype) + else: + return ret.astype(np.float32) + + +def nanmean( + a: np.ndarray, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + return np.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype, out=out) + + +@with_unsupported_dtypes( + {"1.25.0 and below": ("bfloat16",)}, + backend_version, +) +def nanmedian( + input: np.ndarray, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + ) + + +def quantile( + a: np.ndarray, + q: Union[float, np.ndarray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + interpolation: str = "linear", + out: Optional[np.ndarray] = None, +) -> np.ndarray: + # quantile method in numpy backend, always return an array with dtype=float64. + # in other backends, the output is the same dtype as the input. + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + out=out, + ) + + +median.support_native_out = True +nanmean.support_native_out = True +nanmedian.support_native_out = True +bincount.support_native_out = False +cov.support_native_out = False diff --git a/ivy/functional/backends/numpy/general.py b/ivy/functional/backends/numpy/general.py index e4c5a8d75ccac..9d1128a1418fa 100644 --- a/ivy/functional/backends/numpy/general.py +++ b/ivy/functional/backends/numpy/general.py @@ -29,49 +29,6 @@ def current_backend_str() -> str: return "numpy" -@_scalar_output_to_0d_array -def get_item( - x: np.ndarray, - /, - query: Union[np.ndarray, Tuple], - *, - copy: bool = None, -) -> np.ndarray: - return x.__getitem__(query) - - -@_scalar_output_to_0d_array -def set_item( - x: np.ndarray, - query: Union[np.ndarray, Tuple], - val: np.ndarray, - /, - *, - copy: Optional[bool] = False, -) -> np.ndarray: - if copy: - x = np.copy(x) - x.__setitem__(query, val) - return x - - -def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: - if copy: - return x.copy() - else: - return x - - -def to_scalar(x: np.ndarray, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - -def to_list(x: np.ndarray, /) -> list: - return x.tolist() - - def gather( params: np.ndarray, indices: np.ndarray, @@ -104,6 +61,36 @@ def gather( return _to_device(result) +def gather_nd( + params: np.ndarray, + indices: np.ndarray, + /, + *, + batch_dims: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, np.asarray(i, indices.dtype)) + result.append(r) + result = np.array(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return _to_device(result) + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -135,34 +122,15 @@ def gather_nd_helper(params, indices): return res -def gather_nd( - params: np.ndarray, - indices: np.ndarray, +@_scalar_output_to_0d_array +def get_item( + x: np.ndarray, /, + query: Union[np.ndarray, Tuple], *, - batch_dims: int = 0, - out: Optional[np.ndarray] = None, + copy: bool = None, ) -> np.ndarray: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, np.asarray(i, indices.dtype)) - result.append(r) - result = np.array(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return _to_device(result) + return x.__getitem__(query) def get_num_dims(x, /, *, as_array=False): @@ -244,6 +212,27 @@ def is_native_array(x, /, *, exclusive=False): return False +@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) +def isin( + elements: np.ndarray, + test_elements: np.ndarray, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> np.ndarray: + return np.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + +def itemsize(x: np.ndarray) -> int: + return x.itemsize + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -288,9 +277,6 @@ def scatter_flat( return target -scatter_flat.support_native_out = True - - def scatter_nd( indices: np.ndarray, updates: np.ndarray, @@ -332,7 +318,19 @@ def scatter_nd( return _to_device(target) -scatter_nd.support_native_out = True +@_scalar_output_to_0d_array +def set_item( + x: np.ndarray, + query: Union[np.ndarray, Tuple], + val: np.ndarray, + /, + *, + copy: Optional[bool] = False, +) -> np.ndarray: + if copy: + x = np.copy(x) + x.__setitem__(query, val) + return x def shape( @@ -347,6 +345,23 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: np.ndarray, /) -> list: + return x.tolist() + + +def to_numpy(x: np.ndarray, /, *, copy: bool = True) -> np.ndarray: + if copy: + return x.copy() + else: + return x + + +def to_scalar(x: np.ndarray, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -433,25 +448,6 @@ def _vmap(*args): return _vmap -@with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) -def isin( - elements: np.ndarray, - test_elements: np.ndarray, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> np.ndarray: - return np.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True isin.support_native_out = True - - -def itemsize(x: np.ndarray) -> int: - return x.itemsize diff --git a/ivy/functional/backends/numpy/gradients.py b/ivy/functional/backends/numpy/gradients.py index 87be747e053b8..73621665e3d54 100644 --- a/ivy/functional/backends/numpy/gradients.py +++ b/ivy/functional/backends/numpy/gradients.py @@ -6,24 +6,6 @@ import ivy -def variable(x, /): - logging.warning( - "NumPy does not support autograd, declaring a 'variable' " - "is identical to declaring an 'array' when using numpy backend." - ) - return x - - -def is_variable(x, /, *, exclusive=False): - # NumPy does not support autograd, checking if x is a variable does have any meaning - # for NumPy. Return False. - return False - - -def variable_data(x, /): - return x - - def execute_with_gradients( func, xs, @@ -42,23 +24,29 @@ def execute_with_gradients( return func_ret, None -def value_and_grad(func): +def grad(func, argnums=0): logging.warning( - "NumPy does not support autograd, 'value_and_grad' " + "NumPy does not support autograd, 'grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grads = ivy.nested_map( + grad = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return y, grads + return grad return grad_fn +def is_variable(x, /, *, exclusive=False): + # NumPy does not support autograd, checking if x is a variable does have any meaning + # for NumPy. Return False. + return False + + def jac(func): logging.warning( "NumPy does not support autograd, 'jac' " @@ -74,26 +62,38 @@ def grad_fn(xs): return grad_fn -def grad(func, argnums=0): +def stop_gradient(x, /, *, preserve_type=True, out=None): logging.warning( - "NumPy does not support autograd, 'grad' " + "NumPy does not support autograd, 'stop_gradient' " + "has no effect on the array, as gradients are not supported in the first place." + ) + return x + + +def value_and_grad(func): + logging.warning( + "NumPy does not support autograd, 'value_and_grad' " "has no effect on the array, as gradients are not supported in the first place." ) def grad_fn(xs): - grad = ivy.nested_map( + grads = ivy.nested_map( xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False ) y = func(xs) y = ivy.to_ivy(y) - return grad + return y, grads return grad_fn -def stop_gradient(x, /, *, preserve_type=True, out=None): +def variable(x, /): logging.warning( - "NumPy does not support autograd, 'stop_gradient' " - "has no effect on the array, as gradients are not supported in the first place." + "NumPy does not support autograd, declaring a 'variable' " + "is identical to declaring an 'array' when using numpy backend." ) return x + + +def variable_data(x, /): + return x diff --git a/ivy/functional/backends/numpy/helpers.py b/ivy/functional/backends/numpy/helpers.py index b5ae02d3a09f0..110cbebe2a53c 100644 --- a/ivy/functional/backends/numpy/helpers.py +++ b/ivy/functional/backends/numpy/helpers.py @@ -1,23 +1,27 @@ -import functools -from typing import Callable -import numpy as np - - -def _scalar_output_to_0d_array(function: Callable) -> Callable: - """ - Convert scalar outputs to 0d arrays. - - Sometimes NumPy functions return scalars e.g. `np.add` does when the - inputs are both 0 dimensional. - - We use this wrapper to handle such cases, and convert scalar outputs - to 0d arrays, since the array API standard dictates outputs must be - arrays. - """ - - @functools.wraps(function) - def new_function(*args, **kwargs): - ret = function(*args, **kwargs) - return np.asarray(ret) if np.isscalar(ret) else ret - - return new_function +import functools +from typing import Callable +import numpy as np + + +# --- Helpers --- # +# --------------- # + + +def _scalar_output_to_0d_array(function: Callable) -> Callable: + """ + Convert scalar outputs to 0d arrays. + + Sometimes NumPy functions return scalars e.g. `np.add` does when the + inputs are both 0 dimensional. + + We use this wrapper to handle such cases, and convert scalar outputs + to 0d arrays, since the array API standard dictates outputs must be + arrays. + """ + + @functools.wraps(function) + def new_function(*args, **kwargs): + ret = function(*args, **kwargs) + return np.asarray(ret) if np.isscalar(ret) else ret + + return new_function diff --git a/ivy/functional/backends/numpy/layers.py b/ivy/functional/backends/numpy/layers.py index 52ad223304735..ca278076265d0 100644 --- a/ivy/functional/backends/numpy/layers.py +++ b/ivy/functional/backends/numpy/layers.py @@ -14,6 +14,10 @@ ) +# --- Helpers --- # +# --------------- # + + def _add_dilations(x, dilations, axis, values=0): return np.insert( x, @@ -113,6 +117,10 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters +# --- Main --- # +# ------------ # + + def conv1d( x: np.ndarray, filters: np.ndarray, @@ -280,62 +288,6 @@ def conv2d_transpose( return res -def depthwise_conv2d( - x: np.ndarray, - filters: np.ndarray, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[np.ndarray] = None, -): - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if isinstance(padding, int): - padding = [(padding, padding)] * 2 - if data_format == "NHWC": - x = np.transpose(x, (3, 0, 1, 2)) - else: - x = np.transpose(x, (1, 0, 2, 3)) - filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters - filters = np.transpose(filters, (2, 0, 1)) - filters = np.expand_dims(filters, (-1, -2)) - filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) - filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) - if isinstance(padding, str): - if padding == "VALID": - out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) - out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) - else: - out_height = np.ceil(float(x.shape[2]) / float(strides[0])) - out_width = np.ceil(float(x.shape[3]) / float(strides[1])) - else: - out_height = np.ceil( - float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) - / float(strides[0]) - ) - out_width = np.ceil( - float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) - / float(strides[1]) - ) - if data_format == "NHWC": - outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) - else: - outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) - x = np.expand_dims(x, -1) - for i in range(x.shape[0]): - output = conv2d( - x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations - ) - if data_format == "NHWC": - outputs = np.append(outputs, output, axis=-1) - else: - outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) - return outputs - - def conv3d( x: np.ndarray, filters: np.ndarray, @@ -550,3 +502,59 @@ def conv_general_transpose( if data_format == "channel_first": return np.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +def depthwise_conv2d( + x: np.ndarray, + filters: np.ndarray, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[np.ndarray] = None, +): + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if isinstance(padding, int): + padding = [(padding, padding)] * 2 + if data_format == "NHWC": + x = np.transpose(x, (3, 0, 1, 2)) + else: + x = np.transpose(x, (1, 0, 2, 3)) + filters = np.squeeze(filters, 3) if filters.ndim == 4 else filters + filters = np.transpose(filters, (2, 0, 1)) + filters = np.expand_dims(filters, (-1, -2)) + filter_h = filters.shape[1] + (filters.shape[1] - 1) * (dilations[0] - 1) + filter_w = filters.shape[2] + (filters.shape[2] - 1) * (dilations[1] - 1) + if isinstance(padding, str): + if padding == "VALID": + out_height = np.ceil(float(x.shape[2] - filter_h + 1) / float(strides[0])) + out_width = np.ceil(float(x.shape[3] - filter_w + 1) / float(strides[1])) + else: + out_height = np.ceil(float(x.shape[2]) / float(strides[0])) + out_width = np.ceil(float(x.shape[3]) / float(strides[1])) + else: + out_height = np.ceil( + float(x.shape[2] - filter_h + padding[0][0] + padding[0][1] + 1) + / float(strides[0]) + ) + out_width = np.ceil( + float(x.shape[3] - filter_w + padding[1][0] + padding[1][1] + 1) + / float(strides[1]) + ) + if data_format == "NHWC": + outputs = np.empty([x.shape[1], int(out_height), int(out_width), 0], x.dtype) + else: + outputs = np.empty([x.shape[1], 0, int(out_height), int(out_width)], x.dtype) + x = np.expand_dims(x, -1) + for i in range(x.shape[0]): + output = conv2d( + x[i], filters[i], strides, padding, data_format="NHWC", dilations=dilations + ) + if data_format == "NHWC": + outputs = np.append(outputs, output, axis=-1) + else: + outputs = np.append(outputs, np.transpose(output, (0, 3, 1, 2)), axis=1) + return outputs diff --git a/ivy/functional/backends/numpy/linear_algebra.py b/ivy/functional/backends/numpy/linear_algebra.py index 9f9fcdbc17f1b..0928d01919819 100644 --- a/ivy/functional/backends/numpy/linear_algebra.py +++ b/ivy/functional/backends/numpy/linear_algebra.py @@ -52,6 +52,20 @@ def det(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray: return np.linalg.det(x) +# Extra # +# ----- # + + +def diag( + x: np.ndarray, + /, + *, + k: int = 0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.diag(x, k=k) + + def diagonal( x: np.ndarray, /, @@ -64,6 +78,15 @@ def diagonal( return np.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) +@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) +def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] + ) + eigenvalues, eigenvectors = np.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def eigh( x: np.ndarray, /, *, UPLO: str = "L", out: Optional[np.ndarray] = None @@ -137,9 +160,6 @@ def matmul( return ret -matmul.support_native_out = True - - @_scalar_output_to_0d_array @with_unsupported_dtypes({"1.25.2 and below": ("float16", "bfloat16")}, backend_version) def matrix_norm( @@ -214,9 +234,6 @@ def outer( return np.outer(x1, x2, out=out) -outer.support_native_out = True - - @with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) def pinv( x: np.ndarray, @@ -303,27 +320,27 @@ def svdvals(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> np.ndarray return np.linalg.svd(x, compute_uv=False) -def tensorsolve( +def tensordot( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + axes: Union[int, Tuple[List[int], List[int]]] = 2, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.linalg.tensorsolve(x1, x2, axes=axes) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return np.tensordot(x1, x2, axes=axes) -def tensordot( +def tensorsolve( x1: np.ndarray, x2: np.ndarray, /, *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return np.tensordot(x1, x2, axes=axes) + return np.linalg.tensorsolve(x1, x2, axes=axes) @_scalar_output_to_0d_array @@ -340,7 +357,16 @@ def trace( return np.trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) -trace.support_native_out = True +@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) +def vander( + x: np.ndarray, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.vander(x, N=N, increasing=increasing).astype(x.dtype) def vecdot( @@ -355,15 +381,6 @@ def vecdot( return np.tensordot(x1, x2, axes=(axis, axis)) -@with_unsupported_dtypes({"1.25.2 and below": ("float16",)}, backend_version) -def eig(x: np.ndarray, /, *, out: Optional[np.ndarray] = None) -> Tuple[np.ndarray]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", np.ndarray), ("eigenvectors", np.ndarray)] - ) - eigenvalues, eigenvectors = np.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - def vector_norm( x: np.ndarray, /, @@ -404,32 +421,6 @@ def vector_norm( return res -# Extra # -# ----- # - - -def diag( - x: np.ndarray, - /, - *, - k: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.diag(x, k=k) - - -@with_unsupported_dtypes({"1.24.0 and below": ("complex",)}, backend_version) -def vander( - x: np.ndarray, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.vander(x, N=N, increasing=increasing).astype(x.dtype) - - @with_unsupported_dtypes( { "1.25.2 and below": ( @@ -459,4 +450,7 @@ def vector_to_skew_symmetric_matrix( return np.concatenate((row1, row2, row3), -2, out=out) +matmul.support_native_out = True +outer.support_native_out = True +trace.support_native_out = True vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/numpy/manipulation.py b/ivy/functional/backends/numpy/manipulation.py index f258de2b5c59d..8e614f3d7335a 100644 --- a/ivy/functional/backends/numpy/manipulation.py +++ b/ivy/functional/backends/numpy/manipulation.py @@ -10,10 +10,42 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _flat_array_to_1_dim_array(x): return x.reshape((1,)) if x.shape == () else x +# --- Main --- # +# ------------ # + + +def as_strided( + x: np.ndarray, + shape: Union[ivy.NativeShape, Sequence[int]], + strides: Sequence[int], + /, +) -> np.ndarray: + return np.lib.stride_tricks.as_strided( + x, + shape=shape, + strides=strides, + ) + + +def clip( + x: np.ndarray, + x_min: Union[Number, np.ndarray], + x_max: Union[Number, np.ndarray], + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) + + # Array API Standard # # -------------------# @@ -41,7 +73,15 @@ def concat( return ivy.astype(ret, highest_dtype, copy=False) -concat.support_native_out = True +def constant_pad( + x: np.ndarray, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) def expand_dims( @@ -88,6 +128,18 @@ def permute_dims( return np.transpose(x, axes) +@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) +def repeat( + x: np.ndarray, + /, + repeats: Union[int, List[int]], + *, + axis: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.repeat(x, repeats, axis) + + def reshape( x: np.ndarray, /, @@ -118,38 +170,6 @@ def roll( return np.roll(x, shift, axis) -def squeeze( - x: np.ndarray, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if isinstance(axis, list): - axis = tuple(axis) - if x.shape == (): - if axis is None or axis == 0 or axis == -1: - return x - raise ivy.utils.exceptions.IvyException( - "tried to squeeze a zero-dimensional input by axis {}".format(axis) - ) - return np.squeeze(x, axis=axis) - - -def stack( - arrays: Union[Tuple[np.ndarray], List[np.ndarray]], - /, - *, - axis: int = 0, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.stack(arrays, axis, out=out) - - -stack.support_native_out = True - - # Extra # # ------# @@ -191,39 +211,33 @@ def split( return np.split(x, num_or_size_splits, axis) -@with_unsupported_dtypes({"1.25.2 and below": ("uint64",)}, backend_version) -def repeat( +def squeeze( x: np.ndarray, /, - repeats: Union[int, List[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.repeat(x, repeats, axis) - - -def tile( - x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None -) -> np.ndarray: - return np.tile(x, repeats) + if isinstance(axis, list): + axis = tuple(axis) + if x.shape == (): + if axis is None or axis == 0 or axis == -1: + return x + raise ivy.utils.exceptions.IvyException( + "tried to squeeze a zero-dimensional input by axis {}".format(axis) + ) + return np.squeeze(x, axis=axis) -def constant_pad( - x: np.ndarray, +def stack( + arrays: Union[Tuple[np.ndarray], List[np.ndarray]], /, - pad_width: List[List[int]], *, - value: Number = 0.0, + axis: int = 0, out: Optional[np.ndarray] = None, ) -> np.ndarray: - return np.pad(_flat_array_to_1_dim_array(x), pad_width, constant_values=value) - - -def zero_pad( - x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None -): - return np.pad(_flat_array_to_1_dim_array(x), pad_width) + return np.stack(arrays, axis, out=out) def swapaxes( @@ -238,6 +252,12 @@ def swapaxes( return np.swapaxes(x, axis0, axis1) +def tile( + x: np.ndarray, /, repeats: Sequence[int], *, out: Optional[np.ndarray] = None +) -> np.ndarray: + return np.tile(x, repeats) + + def unstack( x: np.ndarray, /, @@ -259,28 +279,12 @@ def unstack( return [np.squeeze(item, axis) for item in x_split] -def clip( - x: np.ndarray, - x_min: Union[Number, np.ndarray], - x_max: Union[Number, np.ndarray], - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.asarray(np.clip(x, x_min, x_max, out=out), dtype=x.dtype) +def zero_pad( + x: np.ndarray, /, pad_width: List[List[int]], *, out: Optional[np.ndarray] = None +): + return np.pad(_flat_array_to_1_dim_array(x), pad_width) +concat.support_native_out = True +stack.support_native_out = True clip.support_native_out = True - - -def as_strided( - x: np.ndarray, - shape: Union[ivy.NativeShape, Sequence[int]], - strides: Sequence[int], - /, -) -> np.ndarray: - return np.lib.stride_tricks.as_strided( - x, - shape=shape, - strides=strides, - ) diff --git a/ivy/functional/backends/numpy/random.py b/ivy/functional/backends/numpy/random.py index bba1591aa039a..2f8db47cedd99 100644 --- a/ivy/functional/backends/numpy/random.py +++ b/ivy/functional/backends/numpy/random.py @@ -14,42 +14,6 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, np.ndarray] = 0.0, - high: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, - dtype: np.dtype, - device: str, - out: Optional[np.ndarray] = None, - seed: Optional[int] = None, -) -> np.ndarray: - if seed: - np.random.seed(seed) - shape = _check_bounds_and_get_shape(low, high, shape).shape - return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) - - -def random_normal( - *, - mean: Union[float, np.ndarray] = 0.0, - std: Union[float, np.ndarray] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: np.dtype, - seed: Optional[int] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - if seed: - np.random.seed(seed) - return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) - @with_unsupported_dtypes({"1.25.2 and below": ("bfloat16",)}, backend_version) def multinomial( @@ -110,6 +74,43 @@ def randint( return np.random.randint(low, high, shape, dtype=dtype) +def random_normal( + *, + mean: Union[float, np.ndarray] = 0.0, + std: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: np.dtype, + seed: Optional[int] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + if seed: + np.random.seed(seed) + return np.asarray(np.random.normal(mean, std, shape), dtype=dtype) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, np.ndarray] = 0.0, + high: Union[float, np.ndarray] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], np.ndarray]] = None, + dtype: np.dtype, + device: str, + out: Optional[np.ndarray] = None, + seed: Optional[int] = None, +) -> np.ndarray: + if seed: + np.random.seed(seed) + shape = _check_bounds_and_get_shape(low, high, shape).shape + return np.asarray(np.random.uniform(low, high, shape), dtype=dtype) + + def seed(*, seed_value: int = 0) -> None: np.random.seed(seed_value) return diff --git a/ivy/functional/backends/numpy/searching.py b/ivy/functional/backends/numpy/searching.py index c5b2b7194025c..604c06c861b77 100644 --- a/ivy/functional/backends/numpy/searching.py +++ b/ivy/functional/backends/numpy/searching.py @@ -60,6 +60,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: np.ndarray, + /, + *, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + return np.argwhere(x) + + def nonzero( x: np.ndarray, /, @@ -95,16 +108,3 @@ def where( ) -> np.ndarray: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return ivy.astype(np.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: np.ndarray, - /, - *, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - return np.argwhere(x) diff --git a/ivy/functional/backends/numpy/sorting.py b/ivy/functional/backends/numpy/sorting.py index aa3d84358220b..0d5aa21213abe 100644 --- a/ivy/functional/backends/numpy/sorting.py +++ b/ivy/functional/backends/numpy/sorting.py @@ -25,22 +25,6 @@ def argsort( ) -def sort( - x: np.ndarray, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - kind = "stable" if stable else "quicksort" - ret = np.asarray(np.sort(x, axis=axis, kind=kind)) - if descending: - ret = np.asarray((np.flip(ret, axis))) - return ret - - # msort @with_unsupported_dtypes({"1.25.2 and below": ("complex",)}, backend_version) def msort( @@ -49,9 +33,6 @@ def msort( return np.msort(a) -msort.support_native_out = False - - def searchsorted( x: np.ndarray, v: np.ndarray, @@ -89,3 +70,22 @@ def searchsorted( else: ret = np.searchsorted(x, v, side=side, sorter=sorter) return ret.astype(ret_dtype) + + +def sort( + x: np.ndarray, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + kind = "stable" if stable else "quicksort" + ret = np.asarray(np.sort(x, axis=axis, kind=kind)) + if descending: + ret = np.asarray((np.flip(ret, axis))) + return ret + + +msort.support_native_out = False diff --git a/ivy/functional/backends/numpy/statistical.py b/ivy/functional/backends/numpy/statistical.py index 419ef9355d336..248318c8104ac 100644 --- a/ivy/functional/backends/numpy/statistical.py +++ b/ivy/functional/backends/numpy/statistical.py @@ -10,23 +10,100 @@ from ivy.utils.einsum_parser import legalise_einsum_expr -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -def min( +def _infer_dtype(dtype: np.dtype): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) +def cumprod( x: np.ndarray, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[np.dtype] = None, out: Optional[np.ndarray] = None, ) -> np.ndarray: - axis = tuple(axis) if isinstance(axis, list) else axis - return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + else: + dtype = _infer_dtype(x.dtype) + if not (exclusive or reverse): + return np.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + return np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumprod(x, -1, dtype=dtype) + return np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) + return np.flip(x, axis=axis) -min.support_native_out = True +def cumsum( + x: np.ndarray, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if dtype is None: + if x.dtype == "bool": + dtype = ivy.default_int_dtype(as_native=True) + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + + if exclusive or reverse: + if exclusive and reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.swapaxes(x, axis, -1) + res = np.flip(x, axis=axis) + elif exclusive: + x = np.swapaxes(x, axis, -1) + x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = np.cumsum(x, -1, dtype=dtype) + res = np.swapaxes(x, axis, -1) + elif reverse: + x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) + res = np.flip(x, axis=axis) + return res + return np.cumsum(x, axis, dtype=dtype, out=out) + + +@_scalar_output_to_0d_array +def einsum( + equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None +) -> np.ndarray: + equation = legalise_einsum_expr(*[equation, *operands]) + return np.einsum(equation, *operands, out=out) def max( @@ -41,9 +118,6 @@ def max( return np.asarray(np.amax(a=x, axis=axis, keepdims=keepdims, out=out)) -max.support_native_out = True - - @_scalar_output_to_0d_array def mean( x: np.ndarray, @@ -59,14 +133,20 @@ def mean( ) -mean.support_native_out = True +# Array API Standard # +# -------------------# -def _infer_dtype(dtype: np.dtype): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +def min( + x: np.ndarray, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + axis = tuple(axis) if isinstance(axis, list) else axis + return np.asarray(np.amin(a=x, axis=axis, keepdims=keepdims, out=out)) def prod( @@ -85,9 +165,6 @@ def prod( return np.asarray(np.prod(a=x, axis=axis, dtype=dtype, keepdims=keepdims, out=out)) -prod.support_native_out = True - - def std( x: np.ndarray, /, @@ -101,9 +178,6 @@ def std( return np.asarray(np.std(x, axis=axis, ddof=correction, keepdims=keepdims, out=out)) -std.support_native_out = True - - def sum( x: np.ndarray, /, @@ -127,9 +201,6 @@ def sum( ) -sum.support_native_out = True - - @_scalar_output_to_0d_array def var( x: np.ndarray, @@ -164,94 +235,13 @@ def var( ) +min.support_native_out = True +max.support_native_out = True +mean.support_native_out = True +prod.support_native_out = True +std.support_native_out = True +sum.support_native_out = True var.support_native_out = True - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"1.25.2 and below": "bfloat16"}, backend_version) -def cumprod( - x: np.ndarray, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - else: - dtype = _infer_dtype(x.dtype) - if not (exclusive or reverse): - return np.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - return np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumprod(x, -1, dtype=dtype) - return np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumprod(np.flip(x, axis=axis), axis=axis, dtype=dtype) - return np.flip(x, axis=axis) - - cumprod.support_native_out = True - - -def cumsum( - x: np.ndarray, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[np.dtype] = None, - out: Optional[np.ndarray] = None, -) -> np.ndarray: - if dtype is None: - if x.dtype == "bool": - dtype = ivy.default_int_dtype(as_native=True) - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - - if exclusive or reverse: - if exclusive and reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.swapaxes(x, axis, -1) - res = np.flip(x, axis=axis) - elif exclusive: - x = np.swapaxes(x, axis, -1) - x = np.concatenate((np.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = np.cumsum(x, -1, dtype=dtype) - res = np.swapaxes(x, axis, -1) - elif reverse: - x = np.cumsum(np.flip(x, axis=axis), axis=axis, dtype=dtype) - res = np.flip(x, axis=axis) - return res - return np.cumsum(x, axis, dtype=dtype, out=out) - - cumsum.support_native_out = True - - -@_scalar_output_to_0d_array -def einsum( - equation: str, *operands: np.ndarray, out: Optional[np.ndarray] = None -) -> np.ndarray: - equation = legalise_einsum_expr(*[equation, *operands]) - return np.einsum(equation, *operands, out=out) - - einsum.support_native_out = True diff --git a/ivy/functional/backends/numpy/utility.py b/ivy/functional/backends/numpy/utility.py index de5597f952940..023c9a9386f17 100644 --- a/ivy/functional/backends/numpy/utility.py +++ b/ivy/functional/backends/numpy/utility.py @@ -20,9 +20,6 @@ def all( raise ivy.utils.exceptions.IvyIndexError(error) -all.support_native_out = True - - def any( x: np.ndarray, /, @@ -34,4 +31,5 @@ def any( return np.asarray(np.any(x, axis=axis, keepdims=keepdims, out=out)) +all.support_native_out = True any.support_native_out = True diff --git a/ivy/functional/backends/paddle/activations.py b/ivy/functional/backends/paddle/activations.py index d090f69a66a4b..6979b59f13f8b 100644 --- a/ivy/functional/backends/paddle/activations.py +++ b/ivy/functional/backends/paddle/activations.py @@ -30,14 +30,38 @@ ] -def relu( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version +) +def gelu( + x: paddle.Tensor, + /, + *, + approximate: bool = False, + complex_mode="jax", + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: + if paddle.is_complex(x): + sqrt_2_over_pi = 0.7978845608 + # the other magic number comes directly from the formula in + # https://doi.org/10.48550/arXiv.1606.08415 + return ( + 0.5 + * x + * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) + ) if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return paddle.complex(F.relu(x.real()), F.relu(x.imag())) - return F.relu(x.cast("float32")).cast(x.dtype) - return F.relu(x) + return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) + return F.gelu(x, approximate=approximate) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def hardswish( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return F.hardswish(x) @with_unsupported_device_and_dtypes( @@ -62,28 +86,47 @@ def leaky_relu( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) -def gelu( +def log_softmax( x: paddle.Tensor, /, *, - approximate: bool = False, - complex_mode="jax", + axis: Optional[int] = None, out: Optional[paddle.Tensor] = None, +): + if axis is None: + axis = -1 + x_max = paddle_backend.max(x, axis=axis, keepdims=True) + x_max = paddle_backend.where( + paddle_backend.isfinite(x_max), + x_max, + paddle.zeros(shape=x_max.shape).astype(x_max.dtype), + ) + exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) + + s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) + ret = paddle_backend.log(s) + ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) + return ret + + +def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in unsupported_dtypes: + if paddle.is_complex(x): + return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) + return F.mish(x.cast("float32")).cast(x.dtype) + return F.mish(x) + + +def relu( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if paddle.is_complex(x): - sqrt_2_over_pi = 0.7978845608 - # the other magic number comes directly from the formula in - # https://doi.org/10.48550/arXiv.1606.08415 - return ( - 0.5 - * x - * (1 + paddle_backend.tanh(sqrt_2_over_pi * (x + 0.044715 * x * x * x))) - ) if x.dtype in unsupported_dtypes: - return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype) - return F.gelu(x, approximate=approximate) + if paddle.is_complex(x): + return paddle.complex(F.relu(x.real()), F.relu(x.imag())) + return F.relu(x.cast("float32")).cast(x.dtype) + return F.relu(x) def sigmoid( @@ -142,46 +185,3 @@ def softplus( if threshold is not None: return ivy.where(x_beta > threshold, x, res).astype(x.dtype) return res.astype(x.dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def log_softmax( - x: paddle.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -): - if axis is None: - axis = -1 - x_max = paddle_backend.max(x, axis=axis, keepdims=True) - x_max = paddle_backend.where( - paddle_backend.isfinite(x_max), - x_max, - paddle.zeros(shape=x_max.shape).astype(x_max.dtype), - ) - exp_tmp = paddle_backend.exp(paddle_backend.subtract(x, x_max)) - - s = paddle_backend.sum(exp_tmp, axis=axis, keepdims=True) - ret = paddle_backend.log(s) - ret = paddle_backend.subtract(paddle_backend.subtract(x, x_max), ret) - return ret - - -def mish(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in unsupported_dtypes: - if paddle.is_complex(x): - return x * paddle_backend.tanh(paddle_backend.log1p(paddle_backend.exp(x))) - return F.mish(x.cast("float32")).cast(x.dtype) - return F.mish(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def hardswish( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return F.hardswish(x) diff --git a/ivy/functional/backends/paddle/creation.py b/ivy/functional/backends/paddle/creation.py index 7b01cae902d53..c216ae2522833 100644 --- a/ivy/functional/backends/paddle/creation.py +++ b/ivy/functional/backends/paddle/creation.py @@ -25,6 +25,114 @@ from . import backend_version from paddle.device import core + +# --- Helpers --- # +# --------------- # + + +def _differentiable_linspace(start, stop, num, *, dtype=None): + start = ivy.to_native(start) + num = paddle.to_tensor(num, stop_gradient=False) + if num == 1: + return paddle_backend.expand_dims(start, axis=0) + n_m_1 = paddle_backend.subtract(num, 1) + increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) + increment_tiled = paddle_backend.repeat(increment, n_m_1) + increments = paddle_backend.multiply( + increment_tiled, + paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), + ) + if isinstance(start, int) or start.ndim == 0: + start = paddle_backend.expand_dims(start, axis=0) + res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) + return res.cast(dtype) + + +def _linspace_helper(start, stop, num, axis=None, *, dtype=None): + num = num.detach().item() if isinstance(num, paddle.Tensor) else num + start_is_array = isinstance(start, paddle.Tensor) + stop_is_array = isinstance(stop, paddle.Tensor) + linspace_method = paddle.linspace + sos_shape = [] + if start_is_array: + start_shape = start.shape + sos_shape = start_shape + if num == 1: + if axis is not None: + return paddle_backend.expand_dims(start, axis=axis) + else: + return paddle_backend.expand_dims(start, axis=-1) + start = start.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not start.stop_gradient else paddle.linspace + ) + if stop_is_array: + stop_shape = stop.shape + sos_shape = stop_shape + if num == 1: + return ( + paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start + ) + stop = stop.reshape((-1,)) + linspace_method = ( + _differentiable_linspace if not stop.stop_gradient else paddle.linspace + ) + if start_is_array and stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=-1) + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = paddle_backend.subtract(stop, start) + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [ + linspace_method(strt, stp, num) + for strt, stp in zip( + paddle_backend.unstack(start, keepdims=True), + paddle_backend.unstack(stop, keepdims=True), + ) + ] + elif start_is_array and not stop_is_array: + if num < start.shape[0]: + start = paddle_backend.expand_dims(start, axis=axis) + diff = stop - start + inc = diff / (num - 1) + res = [start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(paddle.ones(start.shape).astype(start.dtype) * stop) + else: + res = [linspace_method(strt, stop, num) for strt in start] + elif not start_is_array and stop_is_array: + if num < stop.shape[0]: + stop = paddle_backend.expand_dims(stop, axis=-1) + diff = stop - start + inc = diff / (num - 1) + res = [paddle.ones(stop.shape).astype(stop.dtype) * start] + res += [start + inc * i for i in range(1, num - 1)] + res.append(stop) + else: + res = [linspace_method(start, stp, num) for stp in stop] + else: + return linspace_method(start, stop, num, dtype=dtype) + res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) + if axis is not None: + ndim = res.ndim + perm = list(range(ndim - 1)) + perm.insert(axis % (ndim + 1), ndim - 1) + res = paddle_backend.permute_dims(res, perm) + return res + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +# --- Main --- # +# ------------ # + + # Array API Standard # # -------------------# @@ -108,6 +216,21 @@ def asarray( return paddle.to_tensor(obj, dtype=dtype, place=device) +def copy_array( + x: paddle.Tensor, + *, + to_ivy_array: Optional[bool] = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if 0 in x.shape: + new_arr = paddle.empty(x.shape, dtype=x.dtype) + else: + new_arr = x.clone() + if to_ivy_array: + return ivy.to_ivy(new_arr) + return new_arr + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -198,6 +321,46 @@ def from_dlpack(x, /, *, out: Optional[paddle.Tensor] = None): return paddle.utils.dlpack.from_dlpack(x_d) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def frombuffer( + buffer: bytes, + dtype: Optional[paddle.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> paddle.Tensor: + dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) + if str(dtype) == "bool": + dtype_bytes = 1 + dtype_str = str(dtype) + struct_format = { + "bool": "?", + "int8": "b", + "int16": "h", + "int32": "i", + "int64": "q", + "uint8": "B", + "float16": "e", + "float32": "f", + "float64": "d", + } + ret = [] + for i in range(0, len(buffer), dtype_bytes): + x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) + ret = ret + list(x) + if offset > 0: + offset = int(offset / dtype_bytes) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + ret = paddle.to_tensor(ret, dtype=dtype) + + return ret + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -237,105 +400,6 @@ def full_like( ) -def _linspace_helper(start, stop, num, axis=None, *, dtype=None): - num = num.detach().item() if isinstance(num, paddle.Tensor) else num - start_is_array = isinstance(start, paddle.Tensor) - stop_is_array = isinstance(stop, paddle.Tensor) - linspace_method = paddle.linspace - sos_shape = [] - if start_is_array: - start_shape = start.shape - sos_shape = start_shape - if num == 1: - if axis is not None: - return paddle_backend.expand_dims(start, axis=axis) - else: - return paddle_backend.expand_dims(start, axis=-1) - start = start.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not start.stop_gradient else paddle.linspace - ) - if stop_is_array: - stop_shape = stop.shape - sos_shape = stop_shape - if num == 1: - return ( - paddle_backend.ones(stop_shape[:axis] + [1] + stop_shape[axis:]) * start - ) - stop = stop.reshape((-1,)) - linspace_method = ( - _differentiable_linspace if not stop.stop_gradient else paddle.linspace - ) - if start_is_array and stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=-1) - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = paddle_backend.subtract(stop, start) - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [ - linspace_method(strt, stp, num) - for strt, stp in zip( - paddle_backend.unstack(start, keepdims=True), - paddle_backend.unstack(stop, keepdims=True), - ) - ] - elif start_is_array and not stop_is_array: - if num < start.shape[0]: - start = paddle_backend.expand_dims(start, axis=axis) - diff = stop - start - inc = diff / (num - 1) - res = [start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(paddle.ones(start.shape).astype(start.dtype) * stop) - else: - res = [linspace_method(strt, stop, num) for strt in start] - elif not start_is_array and stop_is_array: - if num < stop.shape[0]: - stop = paddle_backend.expand_dims(stop, axis=-1) - diff = stop - start - inc = diff / (num - 1) - res = [paddle.ones(stop.shape).astype(stop.dtype) * start] - res += [start + inc * i for i in range(1, num - 1)] - res.append(stop) - else: - res = [linspace_method(start, stp, num) for stp in stop] - else: - return linspace_method(start, stop, num, dtype=dtype) - res = paddle_backend.concat(res, axis=-1).reshape(sos_shape + [num]) - if axis is not None: - ndim = res.ndim - perm = list(range(ndim - 1)) - perm.insert(axis % (ndim + 1), ndim - 1) - res = paddle_backend.permute_dims(res, perm) - return res - - -def _differentiable_linspace(start, stop, num, *, dtype=None): - start = ivy.to_native(start) - num = paddle.to_tensor(num, stop_gradient=False) - if num == 1: - return paddle_backend.expand_dims(start, axis=0) - n_m_1 = paddle_backend.subtract(num, 1) - increment = paddle_backend.divide(paddle_backend.subtract(stop, start), n_m_1) - increment_tiled = paddle_backend.repeat(increment, n_m_1) - increments = paddle_backend.multiply( - increment_tiled, - paddle.linspace(1, n_m_1, n_m_1.cast(paddle.int32), dtype=dtype), - ) - if isinstance(start, int) or start.ndim == 0: - start = paddle_backend.expand_dims(start, axis=0) - res = paddle_backend.concat((start, paddle_backend.add(start, increments)), axis=0) - return res.cast(dtype) - - -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint16", "bfloat16", "float16")}}, backend_version ) @@ -441,6 +505,58 @@ def meshgrid( return res +def one_hot( + indices: paddle.Tensor, + depth: int, + /, + *, + on_value: Optional[paddle.Tensor] = None, + off_value: Optional[paddle.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[paddle.dtype] = None, + device: core.Place, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + on_none = on_value is None + off_none = off_value is None + expand_ret = False + if indices.ndim == 0: + expand_ret = True + indices = indices.cast("int64").unsqueeze(0) + if dtype is None: + if on_none and off_none: + dtype = paddle.float32 + else: + if not on_none: + dtype = paddle.to_tensor(on_value).dtype + elif not off_none: + dtype = paddle.to_tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = ( + paddle.to_tensor(1.0, dtype="float32") + if on_none + else paddle.to_tensor(on_value, dtype="float32") + ) + off_value = ( + paddle.to_tensor(0.0, dtype="float32") + if off_none + else paddle.to_tensor(off_value, dtype="float32") + ) + + res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) + + if not on_none or not off_none: + res = paddle.where(res == 1, on_value, off_value) + + if axis is not None: + res = paddle.moveaxis(res, -1, axis) + if expand_ret: + res = res.squeeze() + return res.cast(dtype) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -500,6 +616,22 @@ def triu( return paddle.triu(x=x, diagonal=k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: Optional[int] = 0, + /, + *, + device: core.Place, +) -> Tuple[paddle.Tensor]: + # special case due to inconsistent behavior when n_cols=1 and n_rows=0 + if n_cols == 1 and n_rows == 0: + return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( + [], place=device, dtype="int64" + ) + return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -526,126 +658,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: paddle.Tensor, - *, - to_ivy_array: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if 0 in x.shape: - new_arr = paddle.empty(x.shape, dtype=x.dtype) - else: - new_arr = x.clone() - if to_ivy_array: - return ivy.to_ivy(new_arr) - return new_arr - - -def one_hot( - indices: paddle.Tensor, - depth: int, - /, - *, - on_value: Optional[paddle.Tensor] = None, - off_value: Optional[paddle.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[paddle.dtype] = None, - device: core.Place, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - on_none = on_value is None - off_none = off_value is None - expand_ret = False - if indices.ndim == 0: - expand_ret = True - indices = indices.cast("int64").unsqueeze(0) - if dtype is None: - if on_none and off_none: - dtype = paddle.float32 - else: - if not on_none: - dtype = paddle.to_tensor(on_value).dtype - elif not off_none: - dtype = paddle.to_tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = ( - paddle.to_tensor(1.0, dtype="float32") - if on_none - else paddle.to_tensor(on_value, dtype="float32") - ) - off_value = ( - paddle.to_tensor(0.0, dtype="float32") - if off_none - else paddle.to_tensor(off_value, dtype="float32") - ) - - res = paddle.nn.functional.one_hot(indices.cast(paddle.int64), depth) - - if not on_none or not off_none: - res = paddle.where(res == 1, on_value, off_value) - - if axis is not None: - res = paddle.moveaxis(res, -1, axis) - if expand_ret: - res = res.squeeze() - return res.cast(dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def frombuffer( - buffer: bytes, - dtype: Optional[paddle.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> paddle.Tensor: - dtype_bytes = int(ivy.Dtype(dtype).dtype_bits / 8) - if str(dtype) == "bool": - dtype_bytes = 1 - dtype_str = str(dtype) - struct_format = { - "bool": "?", - "int8": "b", - "int16": "h", - "int32": "i", - "int64": "q", - "uint8": "B", - "float16": "e", - "float32": "f", - "float64": "d", - } - ret = [] - for i in range(0, len(buffer), dtype_bytes): - x = struct.unpack(struct_format[dtype_str], buffer[i : i + dtype_bytes]) - ret = ret + list(x) - if offset > 0: - offset = int(offset / dtype_bytes) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - ret = paddle.to_tensor(ret, dtype=dtype) - - return ret - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: Optional[int] = 0, - /, - *, - device: core.Place, -) -> Tuple[paddle.Tensor]: - # special case due to inconsistent behavior when n_cols=1 and n_rows=0 - if n_cols == 1 and n_rows == 0: - return paddle.to_tensor([], place=device, dtype="int64"), paddle.to_tensor( - [], place=device, dtype="int64" - ) - return tuple(paddle.triu_indices(n_rows, col=n_cols, offset=k, dtype="int64")) diff --git a/ivy/functional/backends/paddle/data_type.py b/ivy/functional/backends/paddle/data_type.py index 1a67457c36079..bb4550a92dcc8 100644 --- a/ivy/functional/backends/paddle/data_type.py +++ b/ivy/functional/backends/paddle/data_type.py @@ -22,7 +22,6 @@ paddle.complex128: "complex128", paddle.bool: "bool", } - native_dtype_dict = { "int8": paddle.int8, "int16": paddle.int16, @@ -102,6 +101,51 @@ def __repr__(self): ) +# Extra # +# ------# + + +def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: + if dtype_in is int: + return ivy.default_int_dtype() + if dtype_in is float: + return ivy.default_float_dtype() + if dtype_in is complex: + return ivy.default_complex_dtype() + if dtype_in is bool: + return ivy.Dtype("bool") + if isinstance(dtype_in, str): + if dtype_in in native_dtype_dict: + return ivy.Dtype(dtype_in) + else: + raise ivy.utils.exceptions.IvyException( + "Cannot convert to ivy dtype." + f" {dtype_in} is not supported by Paddle backend." + ) + return ivy.Dtype(ivy_dtype_dict[dtype_in]) + + +def as_native_dtype( + dtype_in: Union[paddle.dtype, str, bool, int, float] +) -> paddle.dtype: + if dtype_in is int: + return ivy.default_int_dtype(as_native=True) + if dtype_in is float: + return ivy.default_float_dtype(as_native=True) + if dtype_in is complex: + return ivy.default_complex_dtype(as_native=True) + if dtype_in is bool: + return paddle.bool + if not isinstance(dtype_in, str): + return dtype_in + if dtype_in in native_dtype_dict.keys(): + return native_dtype_dict[ivy.Dtype(dtype_in)] + else: + raise ivy.utils.exceptions.IvyException( + f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." + ) + + # Array API Standard # # -------------------# @@ -172,6 +216,26 @@ def broadcast_to( return paddle.broadcast_to(x, shape) +def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: + if as_native: + return ivy.to_native(x).dtype + return as_ivy_dtype(x.dtype) + + +def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: + dtype_str = as_ivy_dtype(dtype_in) + if "bool" in dtype_str: + return 1 + return int( + dtype_str.replace("paddle.", "") + .replace("uint", "") + .replace("int", "") + .replace("bfloat", "") + .replace("float", "") + .replace("complex", "") + ) + + @_handle_nestable_dtype_info def finfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Finfo: if isinstance(type, paddle.Tensor): @@ -195,75 +259,6 @@ def iinfo(type: Union[paddle.dtype, str, paddle.Tensor], /) -> Iinfo: return Iinfo(np.iinfo(type)) -def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: - return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) - - -# Extra # -# ------# - - -def as_ivy_dtype(dtype_in: Union[paddle.dtype, str, bool, int, float], /) -> ivy.Dtype: - if dtype_in is int: - return ivy.default_int_dtype() - if dtype_in is float: - return ivy.default_float_dtype() - if dtype_in is complex: - return ivy.default_complex_dtype() - if dtype_in is bool: - return ivy.Dtype("bool") - if isinstance(dtype_in, str): - if dtype_in in native_dtype_dict: - return ivy.Dtype(dtype_in) - else: - raise ivy.utils.exceptions.IvyException( - "Cannot convert to ivy dtype." - f" {dtype_in} is not supported by Paddle backend." - ) - return ivy.Dtype(ivy_dtype_dict[dtype_in]) - - -def as_native_dtype( - dtype_in: Union[paddle.dtype, str, bool, int, float] -) -> paddle.dtype: - if dtype_in is int: - return ivy.default_int_dtype(as_native=True) - if dtype_in is float: - return ivy.default_float_dtype(as_native=True) - if dtype_in is complex: - return ivy.default_complex_dtype(as_native=True) - if dtype_in is bool: - return paddle.bool - if not isinstance(dtype_in, str): - return dtype_in - if dtype_in in native_dtype_dict.keys(): - return native_dtype_dict[ivy.Dtype(dtype_in)] - else: - raise ivy.utils.exceptions.IvyException( - f"Cannot convert to Paddle dtype. {dtype_in} is not supported by Paddle." - ) - - -def dtype(x: paddle.Tensor, *, as_native: bool = False) -> ivy.Dtype: - if as_native: - return ivy.to_native(x).dtype - return as_ivy_dtype(x.dtype) - - -def dtype_bits(dtype_in: Union[paddle.dtype, str], /) -> int: - dtype_str = as_ivy_dtype(dtype_in) - if "bool" in dtype_str: - return 1 - return int( - dtype_str.replace("paddle.", "") - .replace("uint", "") - .replace("int", "") - .replace("bfloat", "") - .replace("float", "") - .replace("complex", "") - ) - - def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -271,3 +266,7 @@ def is_native_dtype(dtype_in: Union[paddle.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[paddle.Tensor, paddle.dtype]) -> ivy.Dtype: + return ivy.promote_types(arrays_and_dtypes[0].dtype, arrays_and_dtypes[1].dtype) diff --git a/ivy/functional/backends/paddle/device.py b/ivy/functional/backends/paddle/device.py index c8e9d9ba385d9..f8295eaaf3d1d 100644 --- a/ivy/functional/backends/paddle/device.py +++ b/ivy/functional/backends/paddle/device.py @@ -13,29 +13,26 @@ from paddle.device import core -# API # -# ----# +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + # ToDO: add proper Paddle profiler + super(Profiler, self).__init__(save_dir) + os.makedirs(save_dir, exist_ok=True) + self._start_time = None + def start(self): + self._start_time = time.perf_counter() -def dev( - x: paddle.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, core.Place]: - return x.place if as_native else as_ivy_dev(x.place) + def stop(self): + time_taken = time.perf_counter() - self._start_time + with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: + f.write("took {} seconds to complete".format(time_taken)) + def __enter__(self): + self.start() -def to_device( - x: paddle.Tensor, - device: core.Place, - /, - *, - stream: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - device = as_native_dev(device) - if device.is_cpu_place(): - return x.cpu() - elif device.is_gpu_place(): - return x.cuda(device.gpu_device_id()) + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() def as_ivy_dev(device: core.Place, /): @@ -73,29 +70,30 @@ def as_native_dev( return native_dev -def clear_mem_on_dev(device: core.Place, /): +def clear_cached_mem_on_dev(device: str, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -def clear_cached_mem_on_dev(device: str, /): +def clear_mem_on_dev(device: core.Place, /): device = as_native_dev(device) if device.is_gpu_place(): paddle.device.cuda.empty_cache() -def num_gpus() -> int: - return paddle.device.cuda.device_count() +# API # +# ----# -def gpu_is_available() -> bool: - return bool(paddle.device.cuda.device_count()) +def dev( + x: paddle.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, core.Place]: + return x.place if as_native else as_ivy_dev(x.place) -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - return False +def gpu_is_available() -> bool: + return bool(paddle.device.cuda.device_count()) def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): @@ -112,23 +110,25 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return ret -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - # ToDO: add proper Paddle profiler - super(Profiler, self).__init__(save_dir) - os.makedirs(save_dir, exist_ok=True) - self._start_time = None +def num_gpus() -> int: + return paddle.device.cuda.device_count() - def start(self): - self._start_time = time.perf_counter() - def stop(self): - time_taken = time.perf_counter() - self._start_time - with open(os.path.join(self._save_dir, "profile.log"), "w+") as f: - f.write("took {} seconds to complete".format(time_taken)) +def to_device( + x: paddle.Tensor, + device: core.Place, + /, + *, + stream: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + device = as_native_dev(device) + if device.is_cpu_place(): + return x.cpu() + elif device.is_gpu_place(): + return x.cuda(device.gpu_device_id()) - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + return False diff --git a/ivy/functional/backends/paddle/elementwise.py b/ivy/functional/backends/paddle/elementwise.py index f4664e17d527a..cbba2fecde68e 100644 --- a/ivy/functional/backends/paddle/elementwise.py +++ b/ivy/functional/backends/paddle/elementwise.py @@ -12,187 +12,155 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. +def _determine_sqrt_dtype_cast( + dtype: Type[paddle.Tensor], +) -> Tuple[Optional[str], Optional[str]]: + """ + Determine the appropriate casting dtype for sqrt operations. + + Returns: + (intermediate_dtype, output_dtype) + """ + + cast_and_return_float32_dtype = { + paddle.int8, + paddle.int16, + paddle.int32, + paddle.uint8, + paddle.bool, + } + + if dtype in cast_and_return_float32_dtype: + return "float32", "float32" + elif dtype == paddle.int64: + return "float64", "float64" + elif dtype == paddle.float16: + return "float32", "float16" + elif dtype == paddle.bfloat16: + return "float32", "bfloat16" + else: + return None, None + + def _elementwise_helper(x1, x2): x1, x2 = ivy.promote_types_of_inputs(x1, x2) x1, x2 = paddle_backend.broadcast_arrays(x1, x2) return x1, x2, x1.dtype -def add( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], +# --- Main --- # +# ------------ # + + +def abs( + x: Union[float, paddle.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() + if x.dtype in [ paddle.int8, + paddle.int16, paddle.uint8, paddle.float16, - paddle.bool, paddle.bfloat16, + paddle.bool, ]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if alpha not in (1, None): - x2 = paddle_backend.multiply(x2, alpha) - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.add(x1, x2).astype(ret_dtype) - - -def bitwise_xor( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_xor(x1, x2) - - -def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: - return paddle.expm1(x) - return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) - - -def bitwise_invert( - x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.bitwise_not(x) + return paddle.abs(x.astype("float32")).astype(x.dtype) + return paddle.abs(x) @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex64", - "complex128", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def isfinite( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.isfinite(x) - - -def isinf( - x: paddle.Tensor, - /, - *, - detect_positive: bool = True, - detect_negative: bool = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if detect_negative and detect_positive: - return paddle.isinf(x) - - if detect_negative: - return paddle_backend.equal(x, float("-inf")) - - if detect_positive: - return paddle_backend.equal(x, float("inf")) - - return paddle.zeros(shape=x.shape, dtype=bool) +def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.acos(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa + s1 = paddle_backend.sqrt(1 - x) + s2 = paddle_backend.sqrt(1 + x) + return paddle.complex( + 2.0 * paddle.atan2(s1.real(), s2.real()), + paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), + ) + return paddle.acos(x) -def equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - diff = paddle_backend.subtract(x1, x2) - ret = paddle_backend.logical_and( - paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) - ) - # ret result is sufficient for all cases except where the value is +/-INF of NaN - return paddle_backend.where( - paddle_backend.isnan(diff), - ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), - ret, - ) +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.acosh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa + s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) + s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) + return paddle.complex( + paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), + 2.0 * paddle.atan2(s1.imag(), s2.real()), + ) + return paddle.acosh(x) -def less_equal( +def add( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - if paddle.is_complex(x1): - real = paddle.less_equal(x1.real(), x2.real()) - imag = paddle.less_equal(x1.imag(), x2.imag()) - return paddle_backend.logical_and(real, imag) - return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_equal(x1, x2) - - -def bitwise_and( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_and(x1, x2) - - -def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ + if x1.dtype in [ paddle.int8, - paddle.int16, - paddle.int32, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, paddle.bool, + paddle.bfloat16, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) - return paddle.ceil(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.ceil(x.astype("float64")).astype(x_dtype) - return paddle.ceil(x) + x1, x2 = x1.astype("float32"), x2.astype("float32") + if alpha not in (1, None): + x2 = paddle_backend.multiply(x2, alpha) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.add(x1, x2).astype(ret_dtype) -def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - x_dtype = x.dtype - if x_dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) - return paddle.floor(x.astype("float32")).astype(x_dtype) - elif x_dtype == paddle.int64: - return paddle.floor(x.astype("float64")).astype(x_dtype) - return paddle.floor(x) +def angle( + input: paddle.Tensor, + /, + *, + deg: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + result = paddle.angle(input) + if deg: + result = paddle.rad2deg(result) + return result @with_unsupported_device_and_dtypes( @@ -243,16 +211,10 @@ def asinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def sign( - x: paddle.Tensor, - /, - *, - np_variant: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: +def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -260,69 +222,33 @@ def sign( paddle.int64, paddle.uint8, paddle.float16, - paddle.bfloat16, - paddle.bool, ]: - return paddle.sgn(x.astype("float32")).astype(x.dtype) - return paddle.sgn(x) - - -# TODO: Remove `float16` from the list once paddle add it's supporting kernel to `CPU`. -def _determine_sqrt_dtype_cast( - dtype: Type[paddle.Tensor], -) -> Tuple[Optional[str], Optional[str]]: - """ - Determine the appropriate casting dtype for sqrt operations. - - Returns: - (intermediate_dtype, output_dtype) - """ - - cast_and_return_float32_dtype = { - paddle.int8, - paddle.int16, - paddle.int32, - paddle.uint8, - paddle.bool, - } - - if dtype in cast_and_return_float32_dtype: - return "float32", "float32" - elif dtype == paddle.int64: - return "float64", "float64" - elif dtype == paddle.float16: - return "float32", "float16" - elif dtype == paddle.bfloat16: - return "float32", "bfloat16" - else: - return None, None - - -def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - """Calculate the square root with type handling.""" - - if paddle.is_complex(x): - angle = paddle.angle(x) - return paddle.complex( - paddle.cos(angle / 2), paddle.sin(angle / 2) - ) * paddle.sqrt(paddle.abs(x)) - - if x.dtype in {paddle.float32, paddle.float64}: - return paddle.sqrt(x) + ret_dtype = x.dtype + return paddle.atan(x.astype("float32")).astype(ret_dtype) + if x.dtype in [paddle.complex64, paddle.complex128]: + atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) + return paddle.atan(x) - intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) - if intermediate_dtype: - result = paddle.sqrt(x.astype(intermediate_dtype)) - return result.astype(output_dtype) - raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + backend_version, +) +def atan2( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.atan2(x1, x2).astype(ret_dtype) @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -332,22 +258,83 @@ def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.float16, ]: ret_dtype = x.dtype - return paddle.cosh(x.astype("float32")).astype(ret_dtype) + return paddle.atanh(x.astype("float32")).astype(ret_dtype) if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) - ) - return paddle.cosh(x) + return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) + return paddle.atanh(x) -def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def bitwise_and( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_and(x1, x2) + + +def bitwise_invert( + x: Union[int, bool, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.bitwise_not(x) + + +def bitwise_left_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( + ret_dtype + ) + + +def bitwise_or( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_or(x1, x2) + + +def bitwise_right_shift( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( + ret_dtype + ) + + +def bitwise_xor( + x1: Union[int, bool, paddle.Tensor], + x2: Union[int, bool, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + return paddle.bitwise_xor(x1, x2) + + +def ceil(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ paddle.int8, paddle.int16, paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, @@ -355,15 +342,18 @@ def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.bool, ]: if paddle.is_complex(x): - base = paddle.to_tensor(10.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log10(x.astype("float32")).astype(x.dtype) - return paddle.log10(x) + return paddle.complex(paddle.ceil(x.real()), paddle.ceil(x.imag())) + return paddle.ceil(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.ceil(x.astype("float64")).astype(x_dtype) + return paddle.ceil(x) -def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def cos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -371,20 +361,24 @@ def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - base = paddle.to_tensor(2.0).squeeze() - return paddle_backend.divide( - paddle_backend.log(x), paddle_backend.log(base) - ).astype(x.dtype) - return paddle.log2(x.astype("float32")).astype(x.dtype) - return paddle.log2(x) + ret_dtype = x.dtype + return paddle.cos(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.cos(re) * paddle.cosh(im), + -paddle.sin(re) * paddle.sinh(im), + ) + return paddle.cos(x) -def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def cosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -392,32 +386,27 @@ def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.int64, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, ]: - if paddle.is_complex(x): - return paddle_backend.log(x + 1) - return paddle.log1p(x.astype("float32")).astype(x.dtype) - return paddle.log1p(x) + ret_dtype = x.dtype + return paddle.cosh(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.cosh(re) * paddle.cos(im), paddle.sinh(re) * paddle.sin(im) + ) + return paddle.cosh(x) -def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) - return paddle.isnan(x.astype("float32")) - return paddle.isnan(x) +def deg2rad( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: + return paddle.deg2rad(x.astype("float32")).astype(x.dtype) + return paddle.deg2rad(x) -def less( +def divide( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -425,17 +414,14 @@ def less( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - real = paddle.less_than(x1.real(), x2.real()) - imag = paddle.less_than(x1.imag(), x2.imag()) - return logical_and(real, imag) - return paddle.less_than(x1.astype("float32"), x2.astype("float32")) - - return paddle.less_than(x1, x2) + if x1.dtype in [paddle.float16, paddle.bfloat16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if not (ivy.is_float_dtype(ret_dtype) or ivy.is_complex_dtype(ret_dtype)): + ret_dtype = ivy.default_float_dtype(as_native=True) + return (x1 / x2).astype(ret_dtype) -def multiply( +def equal( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -443,49 +429,81 @@ def multiply( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.multiply(x1, x2).astype(ret_dtype) - + diff = paddle_backend.subtract(x1, x2) + ret = paddle_backend.logical_and( + paddle_backend.less_equal(diff, 0), paddle_backend.greater_equal(diff, 0) + ) + # ret result is sufficient for all cases except where the value is +/-INF of NaN + return paddle_backend.where( + paddle_backend.isnan(diff), + ~paddle_backend.logical_or(paddle_backend.isnan(x1), paddle_backend.isnan(x2)), + ret, + ) + + +# Extra # +# ------# + @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, backend_version, ) -def cos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + # TODO: add support for complex x, supported in scipy only atm + if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: + return paddle.erf(x.astype("float32")).astype(x.dtype) + return paddle.erf(x) + + +def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + return paddle.exp(x) + if paddle.is_complex(x): + return paddle.multiply( + paddle.exp(x.real()), + paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), + ) + return paddle_backend.pow(math.e, x).astype(x.dtype) + + +def exp2( + x: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.pow(2, x) + + +def expm1(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float16, paddle.float32, paddle.float64]: + return paddle.expm1(x) + return paddle_backend.subtract(paddle_backend.exp(x), 1.0).astype(x.dtype) + + +def floor(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + x_dtype = x.dtype + if x_dtype in [ paddle.int8, paddle.int16, paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - ret_dtype = x.dtype - return paddle.cos(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.cos(re) * paddle.cosh(im), - -paddle.sin(re) * paddle.sinh(im), - ) - return paddle.cos(x) - - -def logical_not( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: if paddle.is_complex(x): - return paddle.logical_and( - paddle.logical_not(x.real()), paddle.logical_not(x.imag()) - ) - return paddle.logical_not(x.astype("float32")) - return paddle.logical_not(x) + return paddle.complex(paddle.floor(x.real()), paddle.floor(x.imag())) + return paddle.floor(x.astype("float32")).astype(x_dtype) + elif x_dtype == paddle.int64: + return paddle.floor(x.astype("float64")).astype(x_dtype) + return paddle.floor(x) -def divide( +def floor_divide( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -493,11 +511,9 @@ def divide( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.float16, paddle.bfloat16]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if not (ivy.is_float_dtype(ret_dtype) or ivy.is_complex_dtype(ret_dtype)): - ret_dtype = ivy.default_float_dtype(as_native=True) - return (x1 / x2).astype(ret_dtype) + if x1.dtype in [paddle.int32, paddle.int64]: + return paddle.floor_divide(x1, x2) + return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) @with_supported_dtypes( @@ -516,6 +532,32 @@ def fmin( return paddle.fmin(x1, x2) +def fmod( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) + return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version +) +def gcd( + x1: Union[paddle.Tensor, int, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.gcd(x1, x2) + + def greater( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], @@ -553,179 +595,121 @@ def greater_equal( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, backend_version, ) -def acos(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.acos(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L178 # noqa - s1 = paddle_backend.sqrt(1 - x) - s2 = paddle_backend.sqrt(1 + x) - return paddle.complex( - 2.0 * paddle.atan2(s1.real(), s2.real()), - paddle.asinh(s2.real() * s1.imag() - s2.imag() * s1.real()), - ) - return paddle.acos(x) +def imag( + val: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.imag(val) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex64", + "complex128", + "bool", + ) + } + }, backend_version, ) -def logical_xor( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def isfinite( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # with the XOR logic - # if paddle.is_complex(x1): - # return paddle.logical_xor( - # paddle.logical_xor(x1.real(), x2.real()), - # paddle.logical_xor(x1.imag(), x2.imag()), - # ) - return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_xor(x1, x2) + return paddle.isfinite(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def logical_and( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def isinf( + x: paddle.Tensor, + /, + *, + detect_positive: bool = True, + detect_negative: bool = True, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - # this logic works well when both inputs are complex but when one of them - # is casted from real to complex, the imaginary part is zero which messes - # if paddle.is_complex(x1): - # return paddle.logical_and( - # paddle.logical_and(x1.real(), x2.real()), - # paddle.logical_and(x1.imag(), x2.imag()), - # ) - return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_and(x1, x2) + if detect_negative and detect_positive: + return paddle.isinf(x) + if detect_negative: + return paddle_backend.equal(x, float("-inf")) -def logical_or( - x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x1): - return paddle.logical_or( - paddle.logical_or(x1.real(), x2.real()), - paddle.logical_or(x1.imag(), x2.imag()), - ) - return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) - return paddle.logical_or(x1, x2) + if detect_positive: + return paddle_backend.equal(x, float("inf")) + return paddle.zeros(shape=x.shape, dtype=bool) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def acosh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + +def isnan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, - paddle.int32, - paddle.int64, paddle.uint8, - paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - return paddle.acosh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - # From https://github.com/python/cpython/blob/39ef93edb9802dccdb6555d4209ac2e60875a011/Modules/cmathmodule.c#L221 # noqa - s1 = paddle_backend.sqrt(paddle.complex(x.real() - 1, x.imag())) - s2 = paddle_backend.sqrt(paddle.complex(x.real() + 1, x.imag())) - return paddle.complex( - paddle.asinh(s1.real() * s2.real() + s1.imag() * s2.imag()), - 2.0 * paddle.atan2(s1.imag(), s2.real()), - ) - return paddle.acosh(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.sin(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) - ) - return paddle.sin(x) - - -def negative( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - if x.dtype == paddle.bool: - return paddle.logical_not(x) - return paddle.neg(x) + if paddle.is_complex(x): + return paddle.logical_or(paddle.isnan(x.real()), paddle.isnan(x.imag())) + return paddle.isnan(x.astype("float32")) + return paddle.isnan(x) -def not_equal( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def isreal( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - return paddle.logical_not(paddle_backend.equal(x1, x2)) + if paddle.is_complex(x): + return paddle.logical_not(x.imag().astype(bool)) + else: + return paddle.ones_like(x, dtype="bool") @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, backend_version, ) -def tanh( - x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +def lcm( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - return paddle.tanh(x.astype("float32")).astype(x.dtype) - if paddle.is_complex(x): - tanh_a = paddle.tanh(paddle.real(x)) - tan_b = paddle.tan(paddle.imag(x)) - return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) - return paddle.tanh(x) + x1_dtype = x1.dtype + x2_dtype = x2.dtype + if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): + return paddle.cast( + paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), + paddle.int16, + ) + elif x1_dtype != x2_dtype: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.lcm(x1, x2) -def floor_divide( +def less( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, @@ -733,27 +717,36 @@ def floor_divide( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int32, paddle.int64]: - return paddle.floor_divide(x1, x2) - return paddle_backend.floor(paddle_backend.divide(x1, x2)).astype(ret_dtype) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + real = paddle.less_than(x1.real(), x2.real()) + imag = paddle.less_than(x1.imag(), x2.imag()) + return logical_and(real, imag) + return paddle.less_than(x1.astype("float32"), x2.astype("float32")) + return paddle.less_than(x1, x2) -def bitwise_or( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], + +def less_equal( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.bitwise_or(x1, x2) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + if paddle.is_complex(x1): + real = paddle.less_equal(x1.real(), x2.real()) + imag = paddle.less_equal(x1.imag(), x2.imag()) + return paddle_backend.logical_and(real, imag) + return paddle.less_equal(x1.astype("float32"), x2.astype("float32")) + return paddle.less_equal(x1, x2) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + +def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, @@ -761,112 +754,43 @@ def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, ]: - ret_dtype = x.dtype - return paddle.sinh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - re = x.real() - im = x.imag() - return paddle.complex( - paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) - ) - return paddle.sinh(x) - - -def positive( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor( - x, dtype=ivy.default_dtype(item=x, as_native=True) - ).squeeze() - return x.clone() - - -def square( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.square(x) - if paddle.is_complex(x): - return paddle.complex( - paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), - 2.0 * paddle.real(x) * paddle.imag(x), - ) - return paddle_backend.pow(x, 2).astype(x.dtype) + if paddle.is_complex(x): + return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) + return paddle.log(x.astype("float32")).astype(x.dtype) + return paddle.log(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version -) -def pow( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ +def log10(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, + paddle.complex64, + paddle.complex128, paddle.bool, ]: - return paddle.pow(x1.astype("float32"), x2.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x1): - # https://math.stackexchange.com/questions/476968/complex-power-of-a-complex-number - r = paddle.abs(x1) - theta = paddle.angle(x1) - res_mag = paddle.pow(r, x2.real()) / paddle.exp(x2.imag() * theta) - res_ang = paddle.log(r) * x2.imag() + theta * x2.real() - result = res_mag * paddle.complex(paddle.cos(res_ang), paddle.sin(res_ang)) - return result.astype(ret_dtype) - return paddle.pow(x1, x2) - - -def round( - x: paddle.Tensor, /, *, decimals: int = 0, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - def _np_round(x, decimals): - # this is a logic to mimic np.round behaviour - # which rounds odd numbers up and even numbers down at limits like 0.5 - eps = 1e-6 * paddle.sign(x) - - # check if the integer is even or odd - candidate_ints = paddle_backend.remainder(paddle_backend.trunc(x), 2.0).astype( - bool - ) - # check if the fraction is exactly half - candidate_fractions = paddle_backend.equal( - paddle_backend.abs(paddle_backend.subtract(x, paddle_backend.trunc(x))), - 0.5, - ) - x = paddle_backend.where( - paddle.logical_and(~candidate_ints, candidate_fractions), - x - eps, - x, - ) - factor = paddle_backend.pow(10.0, decimals).astype(x.dtype) - factor_denom = ivy.where(ivy.isinf(x), 1.0, factor) - return paddle_backend.divide( - paddle.round(paddle_backend.multiply(x, factor)), factor_denom - ) - - if x.dtype not in [paddle.float32, paddle.float64]: if paddle.is_complex(x): - return paddle.complex( - _np_round(x.real(), decimals), _np_round(x.imag(), decimals) - ) - return _np_round(x.astype("float32"), decimals).astype(x.dtype) - return _np_round(x, decimals).astype(x.dtype) + base = paddle.to_tensor(10.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log10(x.astype("float32")).astype(x.dtype) + return paddle.log10(x) -def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: +def log1p(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, @@ -874,79 +798,30 @@ def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle paddle.bool, ]: if paddle.is_complex(x): - return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) - return paddle.trunc(x.astype("float32")).astype(x.dtype) - return paddle.trunc(x) + return paddle_backend.log(x + 1) + return paddle.log1p(x.astype("float32")).astype(x.dtype) + return paddle.log1p(x) -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32")}, - backend_version, -) -def trapz( - y: paddle.Tensor, - /, - *, - x: Optional[paddle.Tensor] = None, - dx: Optional[float] = 1.0, - axis: Optional[int] = -1, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x is None: - d = dx - else: - if x.ndim == 1: - d = paddle.diff(x) - # reshape to correct shape - shape = [1] * y.ndim - shape[axis] = d.shape[0] - d = d.reshape(shape) - else: - d = paddle.diff(x, axis=axis) - - slice1 = [slice(None)] * y.ndim - slice2 = [slice(None)] * y.ndim - - slice1[axis] = slice(1, None) - slice2[axis] = slice(None, -1) - - with ivy.ArrayMode(False): - if y.shape[axis] < 2: - return ivy.zeros_like(ivy.squeeze(y, axis=axis)) - ret = ivy.sum( - ivy.divide( - ivy.multiply( - d, - ivy.add( - ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) - ), - ), - 2.0, - ), - axis=axis, - ) - - return ret - - -def abs( - x: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if not isinstance(x, paddle.Tensor): - x = paddle.to_tensor(x, dtype=ivy.default_dtype(item=x)).squeeze() +def log2(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [ paddle.int8, paddle.int16, + paddle.int32, + paddle.int64, paddle.uint8, paddle.float16, - paddle.bfloat16, + paddle.complex64, + paddle.complex128, paddle.bool, ]: - return paddle.abs(x.astype("float32")).astype(x.dtype) - return paddle.abs(x) + if paddle.is_complex(x): + base = paddle.to_tensor(2.0).squeeze() + return paddle_backend.divide( + paddle_backend.log(x), paddle_backend.log(base) + ).astype(x.dtype) + return paddle.log2(x.astype("float32")).astype(x.dtype) + return paddle.log2(x) @with_unsupported_device_and_dtypes( @@ -977,264 +852,224 @@ def logaddexp2( @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.real(x) +def logical_and( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # if paddle.is_complex(x1): + # return paddle.logical_and( + # paddle.logical_and(x1.real(), x2.real()), + # paddle.logical_and(x1.imag(), x2.imag()), + # ) + return paddle.logical_and(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_and(x1, x2) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.tan(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) - return paddle.tan(x) +def logical_not( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + return paddle.logical_and( + paddle.logical_not(x.real()), paddle.logical_not(x.imag()) + ) + return paddle.logical_not(x.astype("float32")) + return paddle.logical_not(x) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def atan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.atan(x.astype("float32")).astype(ret_dtype) - if x.dtype in [paddle.complex64, paddle.complex128]: - atanh_iz = paddle_backend.atanh(paddle.complex(-x.imag(), x.real())) - return paddle.complex(atanh_iz.imag(), -atanh_iz.real()) - return paddle.atan(x) +def logical_or( + x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x1): + return paddle.logical_or( + paddle.logical_or(x1.real(), x2.real()), + paddle.logical_or(x1.imag(), x2.imag()), + ) + return paddle.logical_or(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_or(x1, x2) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def atan2( +def logical_xor( x1: paddle.Tensor, x2: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.atan2(x1, x2).astype(ret_dtype) + if ret_dtype in [paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128]: + # this logic works well when both inputs are complex but when one of them + # is casted from real to complex, the imaginary part is zero which messes + # with the XOR logic + # if paddle.is_complex(x1): + # return paddle.logical_xor( + # paddle.logical_xor(x1.real(), x2.real()), + # paddle.logical_xor(x1.imag(), x2.imag()), + # ) + return paddle.logical_xor(x1.astype("float32"), x2.astype("float32")) + return paddle.logical_xor(x1, x2) -def log(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ +def maximum( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], + /, + *, + use_where: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [ paddle.int8, paddle.int16, - paddle.int32, - paddle.int64, paddle.uint8, paddle.float16, paddle.complex64, paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x): - return paddle.complex(paddle.log(paddle.abs(x)), paddle.angle(x)) - return paddle.log(x.astype("float32")).astype(x.dtype) - return paddle.log(x) - - -def exp(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - return paddle.exp(x) - if paddle.is_complex(x): - return paddle.multiply( - paddle.exp(x.real()), - paddle.complex(paddle.cos(x.imag()), paddle.sin(x.imag())), - ) - return paddle_backend.pow(math.e, x).astype(x.dtype) - - -def exp2( - x: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.pow(2, x) + if paddle.is_complex(x1): + use_where = True + else: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if use_where: + return paddle_backend.where( + paddle_backend.greater_equal(x1, x2), x1, x2 + ).astype(ret_dtype) + return paddle.maximum(x1, x2).astype(ret_dtype) -def subtract( +def minimum( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - alpha: Optional[Union[int, float]] = None, + use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [paddle.int8, paddle.uint8, paddle.float16, paddle.bool]: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if alpha not in (1, None): - x2 = paddle_backend.multiply(x2, alpha) - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.subtract(x1, x2).astype(ret_dtype) + if x1.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x1): + use_where = True + else: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if use_where: + return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( + ret_dtype + ) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def remainder( + return paddle.minimum(x1, x2).astype(ret_dtype) + + +def multiply( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - modulus: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if not modulus: - res = paddle_backend.divide(x1, x2) - res_floored = paddle_backend.where( - paddle_backend.greater_equal(res, 0.0), - paddle_backend.floor(res), - paddle_backend.ceil(res), - ) - diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) - return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) - if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: x1, x2 = x1.astype("float32"), x2.astype("float32") - return paddle.remainder(x1, x2).astype(ret_dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, - backend_version, -) -def atanh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.int32, - paddle.int64, - paddle.uint8, - paddle.float16, - ]: - ret_dtype = x.dtype - return paddle.atanh(x.astype("float32")).astype(ret_dtype) - if paddle.is_complex(x): - return 0.5 * (paddle_backend.log(1 + x) - paddle_backend.log(1 - x)) - return paddle.atanh(x) + return paddle.multiply(x1, x2).astype(ret_dtype) -def bitwise_right_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], +def nan_to_num( + x: paddle.Tensor, /, *, + copy: Optional[bool] = True, + nan: Optional[Union[float, int]] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") / 2 ** x2.astype("float64")).astype( - ret_dtype - ) + with ivy.ArrayMode(False): + if ivy.is_int_dtype(x): + if posinf is None: + posinf = ivy.iinfo(x).max + if neginf is None: + neginf = ivy.iinfo(x).min + elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): + if posinf is None: + posinf = ivy.finfo(x).max + if neginf is None: + neginf = ivy.finfo(x).min + ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret > 0), + paddle.to_tensor(posinf, dtype=x.dtype), + ret, + ) + ret = ivy.where( + ivy.logical_and(ivy.isinf(ret), ret < 0), + paddle.to_tensor(neginf, dtype=x.dtype), + ret, + ) + if copy: + return ret.clone() + else: + x = ret + return x -def bitwise_left_shift( - x1: Union[int, bool, paddle.Tensor], - x2: Union[int, bool, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, +def negative( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - return paddle.floor(x1.astype("float64") * 2 ** x2.astype("float64")).astype( - ret_dtype - ) - - -# Extra # -# ------# - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, - backend_version, -) -def erf(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - # TODO: add support for complex x, supported in scipy only atm - if x.dtype in [paddle.int8, paddle.int16, paddle.int32, paddle.int64, paddle.uint8]: - return paddle.erf(x.astype("float32")).astype(x.dtype) - return paddle.erf(x) + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + if x.dtype == paddle.bool: + return paddle.logical_not(x) + return paddle.neg(x) -def minimum( +def not_equal( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - if x1.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.logical_not(paddle_backend.equal(x1, x2)) - if use_where: - return paddle_backend.where(paddle_backend.less_equal(x1, x2), x1, x2).astype( - ret_dtype - ) - return paddle.minimum(x1, x2).astype(ret_dtype) +def positive( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if not isinstance(x, paddle.Tensor): + x = paddle.to_tensor( + x, dtype=ivy.default_dtype(item=x, as_native=True) + ).squeeze() + return x.clone() -def maximum( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version +) +def pow( x1: Union[float, paddle.Tensor], x2: Union[float, paddle.Tensor], /, *, - use_where: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: x1, x2, ret_dtype = _elementwise_helper(x1, x2) @@ -1243,35 +1078,18 @@ def maximum( paddle.int16, paddle.uint8, paddle.float16, - paddle.complex64, - paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x1): - use_where = True - else: - x1, x2 = x1.astype("float32"), x2.astype("float32") - if use_where: - return paddle_backend.where( - paddle_backend.greater_equal(x1, x2), x1, x2 - ).astype(ret_dtype) - return paddle.maximum(x1, x2).astype(ret_dtype) - - -def reciprocal( - x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return paddle.reciprocal(x) - return paddle_backend.divide(1, x) - - -def deg2rad( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.dtype in [paddle.int32, paddle.int64, paddle.bool]: - return paddle.deg2rad(x.astype("float32")).astype(x.dtype) - return paddle.deg2rad(x) + return paddle.pow(x1.astype("float32"), x2.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x1): + # https://math.stackexchange.com/questions/476968/complex-power-of-a-complex-number + r = paddle.abs(x1) + theta = paddle.angle(x1) + res_mag = paddle.pow(r, x2.real()) / paddle.exp(x2.imag() * theta) + res_ang = paddle.log(r) * x2.imag() + theta * x2.real() + result = res_mag * paddle.complex(paddle.cos(res_ang), paddle.sin(res_ang)) + return result.astype(ret_dtype) + return paddle.pow(x1, x2) def rad2deg( @@ -1282,148 +1100,338 @@ def rad2deg( return paddle.rad2deg(x) -def trunc_divide( - x1: Union[float, paddle.Tensor], - x2: Union[float, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.trunc(paddle_backend.divide(x1, x2)) - - -def isreal( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if paddle.is_complex(x): - return paddle.logical_not(x.imag().astype(bool)) - else: - return paddle.ones_like(x, dtype="bool") +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "float32", + "float64", + "bool", + ) + } + }, + backend_version, +) +def real(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.real(x) -def fmod( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, +def reciprocal( + x: Union[float, paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - x1, x2, ret_dtype = _elementwise_helper(x1, x2) - res = paddle_backend.remainder(paddle_backend.abs(x1), paddle_backend.abs(x2)) - return paddle_backend.where(paddle_backend.less(x1, 0), -res, res) + if x.dtype in [paddle.float32, paddle.float64]: + return paddle.reciprocal(x) + return paddle_backend.divide(1, x) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128", "bool")}}, backend_version, ) -def lcm( - x1: paddle.Tensor, - x2: paddle.Tensor, +def remainder( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, + modulus: bool = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1_dtype = x1.dtype - x2_dtype = x2.dtype - if (x1_dtype, x2_dtype) == (paddle.int16, paddle.int16): - return paddle.cast( - paddle.lcm(paddle.cast(x1, paddle.int32), paddle.cast(x2, paddle.int32)), - paddle.int16, + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if not modulus: + res = paddle_backend.divide(x1, x2) + res_floored = paddle_backend.where( + paddle_backend.greater_equal(res, 0.0), + paddle_backend.floor(res), + paddle_backend.ceil(res), ) - elif x1_dtype != x2_dtype: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return paddle.lcm(x1, x2) + diff = paddle_backend.subtract(res, res_floored).astype(res.dtype) + return paddle_backend.round(paddle_backend.multiply(diff, x2)).astype(x1.dtype) + + if x1.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + return paddle.remainder(x1, x2).astype(ret_dtype) -def angle( - input: paddle.Tensor, +def round( + x: paddle.Tensor, /, *, decimals: int = 0, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + def _np_round(x, decimals): + # this is a logic to mimic np.round behaviour + # which rounds odd numbers up and even numbers down at limits like 0.5 + eps = 1e-6 * paddle.sign(x) + + # check if the integer is even or odd + candidate_ints = paddle_backend.remainder(paddle_backend.trunc(x), 2.0).astype( + bool + ) + # check if the fraction is exactly half + candidate_fractions = paddle_backend.equal( + paddle_backend.abs(paddle_backend.subtract(x, paddle_backend.trunc(x))), + 0.5, + ) + x = paddle_backend.where( + paddle.logical_and(~candidate_ints, candidate_fractions), + x - eps, + x, + ) + factor = paddle_backend.pow(10.0, decimals).astype(x.dtype) + factor_denom = ivy.where(ivy.isinf(x), 1.0, factor) + return paddle_backend.divide( + paddle.round(paddle_backend.multiply(x, factor)), factor_denom + ) + + if x.dtype not in [paddle.float32, paddle.float64]: + if paddle.is_complex(x): + return paddle.complex( + _np_round(x.real(), decimals), _np_round(x.imag(), decimals) + ) + return _np_round(x.astype("float32"), decimals).astype(x.dtype) + return _np_round(x, decimals).astype(x.dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def sign( + x: paddle.Tensor, /, *, - deg: Optional[bool] = None, + np_variant: Optional[bool] = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - result = paddle.angle(input) - if deg: - result = paddle.rad2deg(result) - return result + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + paddle.bfloat16, + paddle.bool, + ]: + return paddle.sgn(x.astype("float32")).astype(x.dtype) + return paddle.sgn(x) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8")}}, backend_version + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, ) -def gcd( - x1: Union[paddle.Tensor, int, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], +def sin(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.sin(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sin(re) * paddle.cosh(im), paddle.cos(re) * paddle.sinh(im) + ) + return paddle.sin(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def sinh(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + ret_dtype = x.dtype + return paddle.sinh(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + re = x.real() + im = x.imag() + return paddle.complex( + paddle.sinh(re) * paddle.cos(im), paddle.cosh(re) * paddle.sin(im) + ) + return paddle.sinh(x) + + +def sqrt(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + """Calculate the square root with type handling.""" + + if paddle.is_complex(x): + angle = paddle.angle(x) + return paddle.complex( + paddle.cos(angle / 2), paddle.sin(angle / 2) + ) * paddle.sqrt(paddle.abs(x)) + + if x.dtype in {paddle.float32, paddle.float64}: + return paddle.sqrt(x) + + intermediate_dtype, output_dtype = _determine_sqrt_dtype_cast(x.dtype) + if intermediate_dtype: + result = paddle.sqrt(x.astype(intermediate_dtype)) + return result.astype(output_dtype) + + raise ValueError(f"Unsupported data type for sqrt: {x.dtype}") + + +def square( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + return paddle.square(x) + if paddle.is_complex(x): + return paddle.complex( + paddle.square(paddle.real(x)) - paddle.square(paddle.imag(x)), + 2.0 * paddle.real(x) * paddle.imag(x), + ) + return paddle_backend.pow(x, 2).astype(x.dtype) + + +def subtract( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.gcd(x1, x2) + x1, x2, ret_dtype = _elementwise_helper(x1, x2) + if x1.dtype in [paddle.int8, paddle.uint8, paddle.float16, paddle.bool]: + x1, x2 = x1.astype("float32"), x2.astype("float32") + if alpha not in (1, None): + x2 = paddle_backend.multiply(x2, alpha) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return paddle.subtract(x1, x2).astype(ret_dtype) @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "float32", - "float64", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, backend_version, ) -def imag( - val: paddle.Tensor, +def tan(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + ret_dtype = x.dtype + return paddle.tan(x.astype("float32")).astype(ret_dtype) + if paddle.is_complex(x): + tanh_ix = paddle_backend.tanh(paddle.complex(-x.imag(), x.real())) + return paddle.complex(tanh_ix.imag(), -tanh_ix.real()) + return paddle.tan(x) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("bool", "bfloat16")}}, + backend_version, +) +def tanh( + x: paddle.Tensor, /, *, complex_mode="jax", out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + paddle.uint8, + paddle.float16, + ]: + return paddle.tanh(x.astype("float32")).astype(x.dtype) + if paddle.is_complex(x): + tanh_a = paddle.tanh(paddle.real(x)) + tan_b = paddle.tan(paddle.imag(x)) + return (tanh_a + 1j * tan_b) / (1 + 1j * (tanh_a * tan_b)) + return paddle.tanh(x) + + +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32")}, + backend_version, +) +def trapz( + y: paddle.Tensor, /, *, + x: Optional[paddle.Tensor] = None, + dx: Optional[float] = 1.0, + axis: Optional[int] = -1, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.imag(val) + if x is None: + d = dx + else: + if x.ndim == 1: + d = paddle.diff(x) + # reshape to correct shape + shape = [1] * y.ndim + shape[axis] = d.shape[0] + d = d.reshape(shape) + else: + d = paddle.diff(x, axis=axis) + + slice1 = [slice(None)] * y.ndim + slice2 = [slice(None)] * y.ndim + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) -def nan_to_num( - x: paddle.Tensor, + with ivy.ArrayMode(False): + if y.shape[axis] < 2: + return ivy.zeros_like(ivy.squeeze(y, axis=axis)) + ret = ivy.sum( + ivy.divide( + ivy.multiply( + d, + ivy.add( + ivy.get_item(y, tuple(slice1)), ivy.get_item(y, tuple(slice2)) + ), + ), + 2.0, + ), + axis=axis, + ) + + return ret + + +def trunc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex(paddle.trunc(x.real()), paddle.trunc(x.imag())) + return paddle.trunc(x.astype("float32")).astype(x.dtype) + return paddle.trunc(x) + + +def trunc_divide( + x1: Union[float, paddle.Tensor], + x2: Union[float, paddle.Tensor], /, *, - copy: Optional[bool] = True, - nan: Optional[Union[float, int]] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - with ivy.ArrayMode(False): - if ivy.is_int_dtype(x): - if posinf is None: - posinf = ivy.iinfo(x).max - if neginf is None: - neginf = ivy.iinfo(x).min - elif ivy.is_float_dtype(x) or ivy.is_complex_dtype(x): - if posinf is None: - posinf = ivy.finfo(x).max - if neginf is None: - neginf = ivy.finfo(x).min - ret = ivy.where(ivy.isnan(x), paddle.to_tensor(nan, dtype=x.dtype), x) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret > 0), - paddle.to_tensor(posinf, dtype=x.dtype), - ret, - ) - ret = ivy.where( - ivy.logical_and(ivy.isinf(ret), ret < 0), - paddle.to_tensor(neginf, dtype=x.dtype), - ret, - ) - if copy: - return ret.clone() - else: - x = ret - return x + return paddle_backend.trunc(paddle_backend.divide(x1, x2)) diff --git a/ivy/functional/backends/paddle/experimental/activations.py b/ivy/functional/backends/paddle/experimental/activations.py index 81bc6bdf25cb7..069bf40aeac5b 100644 --- a/ivy/functional/backends/paddle/experimental/activations.py +++ b/ivy/functional/backends/paddle/experimental/activations.py @@ -9,6 +9,24 @@ from . import backend_version +def elu( + x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.elu(x, alpha=alpha) + + if paddle.is_complex(x): + ret = ( + paddle_backend.where( + paddle_backend.greater(x, 0), + x, + paddle_backend.multiply(alpha, paddle_backend.expm1(x)), + ), + ) + return ret + return F.elu(x.cast("float32"), alpha).cast(x.dtype) + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -33,28 +51,6 @@ def logit(x: paddle.Tensor, /, *, eps: Optional[float] = None, out=None): ).cast(x.dtype) -def thresholded_relu( - x: paddle.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.thresholded_relu(x, threshold=threshold) - return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( - x.dtype - ) - - -def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - if x.dtype in [paddle.float32, paddle.float64]: - return F.relu6(x) - if paddle.is_complex(x): - return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) - return F.relu6(x.cast("float32")).cast(x.dtype) - - def logsigmoid( input: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: @@ -69,6 +65,14 @@ def logsigmoid( return F.log_sigmoid(input.cast("float32")).cast(input.dtype) +def relu6(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + if x.dtype in [paddle.float32, paddle.float64]: + return F.relu6(x) + if paddle.is_complex(x): + return paddle.complex(F.relu6(x.real()), F.relu6(x.imag())) + return F.relu6(x.cast("float32")).cast(x.dtype) + + def selu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: return F.selu(x) @@ -95,19 +99,15 @@ def silu(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle. return F.silu(x.cast("float32")).cast(x.dtype) -def elu( - x: paddle.Tensor, /, *, alpha: float = 1.0, out: Optional[paddle.Tensor] = None +def thresholded_relu( + x: paddle.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = 0, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if x.dtype in [paddle.float32, paddle.float64]: - return F.elu(x, alpha=alpha) - - if paddle.is_complex(x): - ret = ( - paddle_backend.where( - paddle_backend.greater(x, 0), - x, - paddle_backend.multiply(alpha, paddle_backend.expm1(x)), - ), - ) - return ret - return F.elu(x.cast("float32"), alpha).cast(x.dtype) + return F.thresholded_relu(x, threshold=threshold) + return paddle_backend.where(paddle_backend.greater(x, threshold), x, 0).cast( + x.dtype + ) diff --git a/ivy/functional/backends/paddle/experimental/creation.py b/ivy/functional/backends/paddle/experimental/creation.py index 04d15e026a29d..09ab86a05324d 100644 --- a/ivy/functional/backends/paddle/experimental/creation.py +++ b/ivy/functional/backends/paddle/experimental/creation.py @@ -15,6 +15,11 @@ import ivy from .. import backend_version + +# --- Helpers --- # +# --------------- # + + # noinspection PyProtectedMember # Helpers for calculating Window Functions # ---------------------------------------- @@ -29,39 +34,28 @@ def _kaiser_window(window_length, beta): ) / paddle_backend.i0(beta) -# Array API Standard # -# -------------------# +# --- Main --- # +# ------------ # -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +def blackman_window( + size: int, + /, *, + periodic: Optional[bool] = True, dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if window_length < 2: - return paddle.ones([window_length], dtype=dtype) - if periodic is False: - return _kaiser_window(window_length, beta).cast(dtype) + if size < 2: + return paddle.ones([size], dtype=dtype) + if periodic: + count = paddle.arange(size) / size else: - return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) - - -def vorbis_window( - window_length: paddle.Tensor, - *, - dtype: Optional[paddle.dtype] = paddle.float32, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if window_length == 0: - return paddle.to_tensor([], dtype=dtype) - i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) - pi = paddle.full(shape=i.shape, fill_value=math.pi) - return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( - dtype - ) + count = paddle.linspace(start=0, stop=size, num=size) + return ( + (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) + + (0.08 * paddle.cos(2 * math.pi * 2 * count)) + ).cast(dtype) def hann_window( @@ -81,6 +75,26 @@ def hann_window( return (0.5 - 0.5 * paddle.cos(2 * math.pi * count)).cast(dtype) +# Array API Standard # +# -------------------# + + +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if window_length < 2: + return paddle.ones([window_length], dtype=dtype) + if periodic is False: + return _kaiser_window(window_length, beta).cast(dtype) + else: + return _kaiser_window(window_length + 1, beta)[:-1].cast(dtype) + + def tril_indices( n_rows: int, n_cols: Optional[int] = None, @@ -101,6 +115,32 @@ def tril_indices( ) +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "complex", + ) + } + }, + backend_version, +) +def trilu( + x: paddle.Tensor, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if upper: + return paddle.triu(x=x, diagonal=k) + return paddle.tril(x=x, diagonal=k) + + @with_supported_dtypes( {"2.4.2 and below": ("float64", "float32", "int32", "int64")}, backend_version, @@ -134,26 +174,6 @@ def unsorted_segment_min( return res -def blackman_window( - size: int, - /, - *, - periodic: Optional[bool] = True, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if size < 2: - return paddle.ones([size], dtype=dtype) - if periodic: - count = paddle.arange(size) / size - else: - count = paddle.linspace(start=0, stop=size, num=size) - return ( - (0.42 - 0.5 * paddle.cos(2 * math.pi * count)) - + (0.08 * paddle.cos(2 * math.pi * 2 * count)) - ).cast(dtype) - - def unsorted_segment_sum( data: paddle.Tensor, segment_ids: paddle.Tensor, @@ -188,27 +208,16 @@ def unsorted_segment_sum( return res -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "complex", - ) - } - }, - backend_version, -) -def trilu( - x: paddle.Tensor, - /, +def vorbis_window( + window_length: paddle.Tensor, *, - k: int = 0, - upper: bool = True, + dtype: Optional[paddle.dtype] = paddle.float32, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if upper: - return paddle.triu(x=x, diagonal=k) - return paddle.tril(x=x, diagonal=k) + if window_length == 0: + return paddle.to_tensor([], dtype=dtype) + i = paddle_backend.arange(1, window_length * 2, 2, device=ivy.default_device()) + pi = paddle.full(shape=i.shape, fill_value=math.pi) + return paddle.sin((pi / 2) * (paddle.sin(pi * i / (window_length * 2)) ** 2)).cast( + dtype + ) diff --git a/ivy/functional/backends/paddle/experimental/elementwise.py b/ivy/functional/backends/paddle/experimental/elementwise.py index 8fe58b9e1519a..95b5be21b6aad 100644 --- a/ivy/functional/backends/paddle/experimental/elementwise.py +++ b/ivy/functional/backends/paddle/experimental/elementwise.py @@ -17,79 +17,102 @@ from .. import backend_version -@with_supported_dtypes( - {"2.5.1 and below": ("float32", "float64")}, - backend_version, -) -def lgamma( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - return paddle.lgamma(x) +_BERNOULLI_COEFS = [ + 12, + -720, + 30240, + -1209600, + 47900160, + -1307674368000 / 691, + 74724249600, + -10670622842880000 / 3617, + 5109094217170944000 / 43867, + -802857662698291200000 / 174611, + 14101100039391805440000 / 77683, + -1693824136731743669452800000 / 236364091, + 186134520519971831808000000 / 657931, + -37893265687455865519472640000000 / 3392780147, + 759790291646040068357842010112000000 / 1723168255201, + -134196726836183700385281186201600000000 / 7709321041217, +] -@with_supported_dtypes( - {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, - backend_version, -) -def fmax( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x1.dtype != x2.dtype: - x1, x2 = promote_types_of_inputs(x1, x2) - return paddle.fmax(x1, x2) +# --- Helpers --- # +# --------------- # -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) - return paddle.divide(paddle.sin(y), y) +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim -def float_power( - x1: Union[paddle.Tensor, float, list, tuple], - x2: Union[paddle.Tensor, float, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1 = paddle.cast(x1, dtype="float64") - x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power - return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis -def frexp( - x: Union[paddle.Tensor, Number], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException +def _np_ndim(x): + return ivy.array(x).ndim -def ldexp( - x1: Union[paddle.Tensor, Number], - x2: Union[paddle.Tensor, Number], +# --- Main --- # +# ------------ # + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def allclose( + x1: paddle.Tensor, + x2: paddle.Tensor, /, *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - out_dtype = x1.dtype - x1, x2 = promote_types_of_inputs(x1, x2) - with ivy.ArrayMode(False): - if ivy.any(ivy.less(x2, 0)): - pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 - neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 - ret = ivy.multiply(ivy.pow(2, pos_exp), x1) - ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) - else: - ret = ivy.multiply(ivy.pow(2, x2), x1) - return ivy.astype(ret, out_dtype, copy=False) +) -> bool: + return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) + + +@with_supported_dtypes( + { + "2.5.1 and below": ( + "complex64", + "complex128", + "float32", + "float64", + "int32", + "int64", + ) + }, + backend_version, +) +def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + return paddle.conj(x) def copysign( @@ -107,38 +130,17 @@ def copysign( return result -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version -) -def nansum( - x: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[paddle.dtype] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) - if result.shape == [1]: - result = paddle.fluid.layers.squeeze(result, [0]) - return result - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def isclose( +def count_nonzero( a: paddle.Tensor, - b: paddle.Tensor, /, *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, + axis: Optional[Union[int, list, tuple]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + with ivy.ArrayMode(False): + return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) def diff( @@ -167,185 +169,69 @@ def _tensor(val): ) -def signbit( - x: Union[paddle.Tensor, float, int, list, tuple], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.less( - paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 - ) - - -def hypot( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -@with_unsupported_device_and_dtypes( +@with_supported_dtypes( { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def allclose( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - rtol: Optional[float] = 1e-05, - atol: Optional[float] = 1e-08, - equal_nan: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> bool: - return paddle.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan).squeeze(0) - - -def fix( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - return ivy.trunc(x) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version -) -def nextafter( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - with ivy.ArrayMode(False): - eps = ivy.finfo(x1.dtype).eps - return ivy.where( - ivy.equal(x1, x2), - x2, - ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), + "2.5.0 and below": ( + "float32", + "float64", ) - - -_BERNOULLI_COEFS = [ - 12, - -720, - 30240, - -1209600, - 47900160, - -1307674368000 / 691, - 74724249600, - -10670622842880000 / 3617, - 5109094217170944000 / 43867, - -802857662698291200000 / 174611, - 14101100039391805440000 / 77683, - -1693824136731743669452800000 / 236364091, - 186134520519971831808000000 / 657931, - -37893265687455865519472640000000 / 3392780147, - 759790291646040068357842010112000000 / 1723168255201, - -134196726836183700385281186201600000000 / 7709321041217, -] - - -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "int32", - "int64", - "uint8", - "uint16", - "float16", - "bool", - ) - } }, - backend_version, -) -def zeta( - x: paddle.Tensor, - q: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - s, a = ivy.promote_types_of_inputs(x, q) - s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) - N = M = ( - paddle.to_tensor(8.0, dtype="float32") - if q.dtype == paddle.float32 - else paddle.to_tensor(8.0, dtype="float64") - ) - assert M <= len(_BERNOULLI_COEFS) - k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) - S = paddle.sum((a_ + k) ** -s_, -1) - Q = ivy.divide((q + N) ** (1 - x), x - 1) - T0 = (q + N) ** -x - m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) - s_over_a = (s_ + m) / (a_ + N) - s_over_a = ivy.where( - s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a - ) - T1 = paddle.cumprod(s_over_a, -1)[..., ::2] - # t=np.array(T1) - T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) - coefs = paddle.unsqueeze( - paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), - tuple(range(a.ndim)), - ) - T1 = T1 / coefs - T = T0 * (0.5 + paddle.sum(T1, -1)) - ans = S + Q + T - mask = x < 1 - ans[mask] = ivy.nan - return ans + backend_version, +) +def digamma( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.digamma(x) -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim +def fix( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + with ivy.ArrayMode(False): + return ivy.trunc(x) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +def float_power( + x1: Union[paddle.Tensor, float, list, tuple], + x2: Union[paddle.Tensor, float, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1 = paddle.cast(x1, dtype="float64") + x2 = paddle.cast(x2, dtype="float64") # Compute the element-wise power + return paddle.cast(paddle.pow(x1, x2), dtype=paddle.float64) -def _np_ndim(x): - return ivy.array(x).ndim +@with_supported_dtypes( + {"2.5.1 and below": ("float64", "float32", "int32", "int64")}, + backend_version, +) +def fmax( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x1.dtype != x2.dtype: + x1, x2 = promote_types_of_inputs(x1, x2) + return paddle.fmax(x1, x2) + + +def frexp( + x: Union[paddle.Tensor, Number], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException @with_supported_dtypes( @@ -578,47 +464,60 @@ def gradient( return outvals -def xlogy( - x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +def hypot( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - x, y, ret_dtype = _elementwise_helper(x, y) - with ivy.ArrayMode(False): - x_ok = ivy.not_equal(x, 0.0) - safe_x = ivy.where(x_ok, x, 1.0) - safe_y = ivy.where(x_ok, y, 1.0) - return ivy.where( - x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) - ).cast(ret_dtype) + raise IvyNotImplementedException() -def count_nonzero( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def isclose( a: paddle.Tensor, + b: paddle.Tensor, + /, + *, + rtol: Optional[float] = 1e-05, + atol: Optional[float] = 1e-08, + equal_nan: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + +def ldexp( + x1: Union[paddle.Tensor, Number], + x2: Union[paddle.Tensor, Number], /, *, - axis: Optional[Union[int, list, tuple]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: + out_dtype = x1.dtype + x1, x2 = promote_types_of_inputs(x1, x2) with ivy.ArrayMode(False): - return ivy.sum(ivy.not_equal(a, 0), axis=axis, keepdims=keepdims, dtype=dtype) + if ivy.any(ivy.less(x2, 0)): + pos_exp = ivy.greater_equal(x2, 0).astype(x2.dtype) * x2 + neg_exp = ivy.less(x2, 0).astype(x2.dtype) * x2 + ret = ivy.multiply(ivy.pow(2, pos_exp), x1) + ret = ivy.divide(ret, ivy.pow(2, -neg_exp)) + else: + ret = ivy.multiply(ivy.pow(2, x2), x1) + return ivy.astype(ret, out_dtype, copy=False) @with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - "float32", - "float64", - "int32", - "int64", - ) - }, + {"2.5.1 and below": ("float32", "float64")}, backend_version, ) -def conj(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: - return paddle.conj(x) +def lgamma( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + return paddle.lgamma(x) def modf( @@ -628,19 +527,128 @@ def modf( return paddle.modf(x, out=out) -@with_supported_dtypes( - { - "2.5.0 and below": ( - "float32", - "float64", +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "float16")}}, backend_version +) +def nansum( + x: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[paddle.dtype] = None, + keepdims: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + result = paddle.nansum(x, axis=axis, dtype=dtype, keepdim=keepdims) + if result.shape == [1]: + result = paddle.fluid.layers.squeeze(result, [0]) + return result + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def nextafter( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + with ivy.ArrayMode(False): + eps = ivy.finfo(x1.dtype).eps + return ivy.where( + ivy.equal(x1, x2), + x2, + ivy.where(ivy.greater(x2, x1), ivy.add(x1, eps), ivy.subtract(x1, eps)), ) + + +def signbit( + x: Union[paddle.Tensor, float, int, list, tuple], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.less( + paddle_backend.where(x.astype(bool), x, paddle_backend.divide(1.0, x)), 0.0 + ) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version +) +def sinc(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.Tensor: + y = ivy.pi * paddle.where(x == 0, paddle.to_tensor(1.0e-20, dtype=x.dtype), x) + return paddle.divide(paddle.sin(y), y) + + +def xlogy( + x: paddle.Tensor, y: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + x, y, ret_dtype = _elementwise_helper(x, y) + with ivy.ArrayMode(False): + x_ok = ivy.not_equal(x, 0.0) + safe_x = ivy.where(x_ok, x, 1.0) + safe_y = ivy.where(x_ok, y, 1.0) + return ivy.where( + x_ok, ivy.multiply(safe_x, ivy.log(safe_y)), ivy.zeros_like(x) + ).cast(ret_dtype) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "float16", + "bool", + ) + } }, backend_version, ) -def digamma( +def zeta( x: paddle.Tensor, + q: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.digamma(x) + with ivy.ArrayMode(False): + s, a = ivy.promote_types_of_inputs(x, q) + s_, a_ = paddle.unsqueeze(x, -1), paddle.unsqueeze(q, -1) + N = M = ( + paddle.to_tensor(8.0, dtype="float32") + if q.dtype == paddle.float32 + else paddle.to_tensor(8.0, dtype="float64") + ) + assert M <= len(_BERNOULLI_COEFS) + k = paddle.unsqueeze(ivy.arange(N, dtype=q.dtype), tuple(range(q.ndim))) + S = paddle.sum((a_ + k) ** -s_, -1) + Q = ivy.divide((q + N) ** (1 - x), x - 1) + T0 = (q + N) ** -x + m = paddle.unsqueeze(ivy.arange(2 * M, dtype=s.dtype), tuple(range(s.ndim))) + s_over_a = (s_ + m) / (a_ + N) + s_over_a = ivy.where( + s_over_a == 0, paddle.ones_like(s_over_a) * 1e-20, s_over_a + ) + T1 = paddle.cumprod(s_over_a, -1)[..., ::2] + # t=np.array(T1) + T1 = paddle.clip(T1, max=ivy.finfo(T1.dtype).max) + coefs = paddle.unsqueeze( + paddle.to_tensor(_BERNOULLI_COEFS[: T1.shape[-1]], dtype=T1.dtype), + tuple(range(a.ndim)), + ) + T1 = T1 / coefs + T = T0 * (0.5 + paddle.sum(T1, -1)) + ans = S + Q + T + mask = x < 1 + ans[mask] = ivy.nan + return ans diff --git a/ivy/functional/backends/paddle/experimental/layers.py b/ivy/functional/backends/paddle/experimental/layers.py index fafd86275beb3..4fdba46985ad9 100644 --- a/ivy/functional/backends/paddle/experimental/layers.py +++ b/ivy/functional/backends/paddle/experimental/layers.py @@ -13,6 +13,11 @@ ) from .. import backend_version + +# --- Helpers --- # +# --------------- # + + # local @@ -26,6 +31,255 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling +# --- Main --- # +# ------------ # + + +def adaptive_max_pool2d( + input: paddle.Tensor, output_size: Union[Sequence[int], int] +) -> paddle.Tensor: + squeeze = input.ndim == 3 + x = paddle.unsqueeze(input, axis=0) if squeeze else input + ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) + return paddle.squeeze(ret, axis=0) if squeeze else ret + + +def avg_pool1d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool2d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, + /, + *, + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def avg_pool3d( + x: paddle.Tensor, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, + /, + *, + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def dct( + x: paddle.Tensor, + /, + *, + type: Optional[Literal[1, 2, 3, 4]] = 2, + n: Optional[int] = None, + axis: Optional[int] = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout1d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 3 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout2d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 4 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("bfloat16", "float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def dropout3d( + x: paddle.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + axis = data_format.index("C") - 5 + x.ndim + return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) + + +def embedding( + weights: paddle.Tensor, + indices: paddle.Tensor, + /, + *, + max_norm: Optional[int] = None, + out=None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def fft( + x: paddle.Tensor, + dim: int, + /, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if not isinstance(dim, int): + raise IvyValueError(f"Expecting instead of {type(dim)}") + + if n is None: + n = x.shape[dim] + + if dim < -x.ndim or dim >= x.ndim: + raise IvyValueError( + f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" + ) + + if not isinstance(n, int): + raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") + + if n <= 1: + raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") + + valid_norm_modes = ["backward", "ortho", "forward"] + if norm not in valid_norm_modes: + raise IvyValueError( + f"Unrecognized normalization mode {norm}, expecting one of" + f" {valid_norm_modes}" + ) + + if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: + x = x.cast(paddle.complex128) + else: + x = x.cast(paddle.complex64) + + return paddle.fft.fft(x, n, dim, norm=norm) + + +@with_supported_dtypes( + { + "2.5.1 and below": ( + "complex64", + "complex128", + ) + }, + backend_version, +) +def fft2( + x: paddle.Tensor, + *, + dim: Optional[Union[int, Tuple[int]]] = None, + norm: Optional[str] = "backward", + s: Optional[Union[int, Tuple[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + res = paddle.fft.fft2(x, s, dim, norm) + return res.astype("complex128") + + +def ifft( + x: paddle.Tensor, + dim: int, + *, + norm: Optional[str] = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() + + +def ifftn( + x: paddle.Tensor, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, + *, + norm: Optional[str] = "backward", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.fft.ifftn(x, s, axes, norm) + + +def interpolate( + x: paddle.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +): + raise IvyNotImplementedException() + + @with_supported_device_and_dtypes( { "2.5.1 and below": { @@ -237,230 +491,6 @@ def max_pool3d( return res -def avg_pool1d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool2d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def avg_pool3d( - x: paddle.Tensor, - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def dct( - x: paddle.Tensor, - /, - *, - type: Optional[Literal[1, 2, 3, 4]] = 2, - n: Optional[int] = None, - axis: Optional[int] = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def fft( - x: paddle.Tensor, - dim: int, - /, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if not isinstance(dim, int): - raise IvyValueError(f"Expecting instead of {type(dim)}") - - if n is None: - n = x.shape[dim] - - if dim < -x.ndim or dim >= x.ndim: - raise IvyValueError( - f"Invalid dim {dim}, expecting a value ranging from {-x.ndim} to {x.ndim-1}" - ) - - if not isinstance(n, int): - raise TypeError(f"Expecting int type for 'n', instead of {type(n)}") - - if n <= 1: - raise IvyValueError(f"Invalid number of data points {n}, expecting more than 1") - - valid_norm_modes = ["backward", "ortho", "forward"] - if norm not in valid_norm_modes: - raise IvyValueError( - f"Unrecognized normalization mode {norm}, expecting one of" - f" {valid_norm_modes}" - ) - - if x.dtype in [paddle.int64, paddle.float64, paddle.complex128]: - x = x.cast(paddle.complex128) - else: - x = x.cast(paddle.complex64) - - return paddle.fft.fft(x, n, dim, norm=norm) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout1d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 3 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout2d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 4 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("bfloat16", "float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def dropout3d( - x: paddle.Tensor, - prob: float, - /, - *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axis = data_format.index("C") - 5 + x.ndim - return paddle.nn.functional.dropout(x, p=prob, axis=axis, training=training) - - -def ifft( - x: paddle.Tensor, - dim: int, - *, - norm: Optional[str] = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def embedding( - weights: paddle.Tensor, - indices: paddle.Tensor, - /, - *, - max_norm: Optional[int] = None, - out=None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - -def interpolate( - x: paddle.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Optional[Literal["linear", "bilinear", "trilinear"]] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -): - raise IvyNotImplementedException() - - -def adaptive_max_pool2d( - input: paddle.Tensor, output_size: Union[Sequence[int], int] -) -> paddle.Tensor: - squeeze = input.ndim == 3 - x = paddle.unsqueeze(input, axis=0) if squeeze else input - ret = paddle.nn.functional.adaptive_max_pool2d(x, output_size) - return paddle.squeeze(ret, axis=0) if squeeze else ret - - -def ifftn( - x: paddle.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.fft.ifftn(x, s, axes, norm) - - @with_unsupported_dtypes( {"2.5.1 and below": ("bfloat16", "float16", "complex64", "complex128", "bool")}, backend_version, @@ -475,24 +505,3 @@ def rfftn( ) -> paddle.Tensor: result = paddle.fft.rfftn(x, s, axes, norm) return result.astype("complex128") - - -@with_supported_dtypes( - { - "2.5.1 and below": ( - "complex64", - "complex128", - ) - }, - backend_version, -) -def fft2( - x: paddle.Tensor, - *, - dim: Optional[Union[int, Tuple[int]]] = None, - norm: Optional[str] = "backward", - s: Optional[Union[int, Tuple[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - res = paddle.fft.fft2(x, s, dim, norm) - return res.astype("complex128") diff --git a/ivy/functional/backends/paddle/experimental/linear_algebra.py b/ivy/functional/backends/paddle/experimental/linear_algebra.py index 00f9e583ad2e7..22e5447d528c0 100644 --- a/ivy/functional/backends/paddle/experimental/linear_algebra.py +++ b/ivy/functional/backends/paddle/experimental/linear_algebra.py @@ -12,6 +12,26 @@ from .. import backend_version +def adjoint( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + _check_valid_dimension_size(x) + return paddle.moveaxis(x, -2, -1).conj() + + +def cond( + x: paddle.Tensor, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[paddle.Tensor] = None, +) -> Any: + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) @@ -45,28 +65,14 @@ def diagflat( )(diag) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version -) -def kron( +def dot( a: paddle.Tensor, b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.kron(a, b) - - -def matrix_exp( - x: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # TODO: this is elementwise exp, should be changed to matrix exp ASAP - # return paddle.exp(x) - raise IvyNotImplementedException() + return paddle.dot(a, b, out=out) def eig( @@ -79,24 +85,17 @@ def eigvals(x: paddle.Tensor, /) -> paddle.Tensor: return paddle.linalg.eig(x)[0] -def adjoint( - x: paddle.Tensor, +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16")}}, backend_version +) +def kron( + a: paddle.Tensor, + b: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - _check_valid_dimension_size(x) - return paddle.moveaxis(x, -2, -1).conj() - - -def cond( - x: paddle.Tensor, - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[paddle.Tensor] = None, -) -> Any: - raise IvyNotImplementedException() + return paddle.kron(a, b) def lu_factor( @@ -109,17 +108,15 @@ def lu_factor( raise IvyNotImplementedException() -def dot( - a: paddle.Tensor, - b: paddle.Tensor, +def matrix_exp( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.dot(a, b, out=out) - - -dot.support_native_out = True + # TODO: this is elementwise exp, should be changed to matrix exp ASAP + # return paddle.exp(x) + raise IvyNotImplementedException() @with_supported_device_and_dtypes( @@ -145,3 +142,6 @@ def multi_dot( out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: return paddle.linalg.multi_dot(x) + + +dot.support_native_out = True diff --git a/ivy/functional/backends/paddle/experimental/losses.py b/ivy/functional/backends/paddle/experimental/losses.py index 70441eb4655f3..54f155697c070 100644 --- a/ivy/functional/backends/paddle/experimental/losses.py +++ b/ivy/functional/backends/paddle/experimental/losses.py @@ -26,20 +26,21 @@ }, backend_version, ) -def l1_loss( +def huber_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, ) -> paddle.Tensor: - return F.l1_loss(input, target, reduction=reduction) + return paddle.fluid.layers.huber_loss(input, target, delta=delta) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( + "float16", "int8", "int16", "int32", @@ -53,24 +54,20 @@ def l1_loss( }, backend_version, ) -def smooth_l1_loss( +def l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - beta: Optional[float] = 1.0, reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return paddle.nn.functional.smooth_l1_loss( - input, target, reduction=reduction, beta=beta - ) + return F.l1_loss(input, target, reduction=reduction) @with_unsupported_device_and_dtypes( { "2.5.1 and below": { "cpu": ( - "float16", "int8", "int16", "int32", @@ -84,11 +81,14 @@ def smooth_l1_loss( }, backend_version, ) -def huber_loss( +def smooth_l1_loss( input: paddle.Tensor, target: paddle.Tensor, /, *, - delta: Optional[float] = 1.0, + beta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", ) -> paddle.Tensor: - return paddle.fluid.layers.huber_loss(input, target, delta=delta) + return paddle.nn.functional.smooth_l1_loss( + input, target, reduction=reduction, beta=beta + ) diff --git a/ivy/functional/backends/paddle/experimental/manipulation.py b/ivy/functional/backends/paddle/experimental/manipulation.py index 59e781c937591..c195b7894d40f 100644 --- a/ivy/functional/backends/paddle/experimental/manipulation.py +++ b/ivy/functional/backends/paddle/experimental/manipulation.py @@ -44,7 +44,6 @@ -3.04682672343198398683e-1, 6.76795274409476084995e-1, ] - _i0B = [ -7.23318048787475395456e-18, -4.83050448594418207126e-18, @@ -74,197 +73,174 @@ ] -def moveaxis( - a: paddle.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if isinstance(source, tuple): - source = list(source) - if isinstance(destination, tuple): - source = list(destination) - if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: - return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) - return paddle.moveaxis(a, source, destination) +def atleast_1d( + *arys: paddle.Tensor, copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 1: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=0)) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res + + +def atleast_2d( + *arys: paddle.Tensor, copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim < 2: + with ivy.ArrayMode(False): + res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) + else: + res.append(ary) + if len(res) == 1: + return res[0] + return res @with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, + {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, ) -def heaviside( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle.heaviside(x1, x2) +def atleast_3d( + *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[paddle.Tensor]: + res = [] + for ary in arys: + ary = ivy.array(ary, copy=copy).data + if ary.ndim == 0: + result = ary.reshape((1, 1, 1)) + elif ary.ndim == 1: + result = ary[None, :, None] + elif ary.ndim == 2: + result = ary[:, :, None] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return res -def flipud( - m: paddle.Tensor, +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + def _broadcast_shape(s1, s2): + len_1 = len(s1) + len_2 = len(s2) + if len_1 == 0: + return () if len_2 == 0 else s2 + elif len_1 != 0 and len_2 == 0: + return s1 + else: + return paddle.broadcast_shape(s1, s2) + + if len(shapes) == 0: + raise ValueError("shapes=[] must be non-empty") + elif len(shapes) == 1: + return shapes[0] + result = _broadcast_shape(shapes[0], shapes[1]) + for i in range(2, len(shapes)): + result = _broadcast_shape(result, shapes[i]) + # paddle outputs -1 if the output dimension is 0 + result = [0 if dim == -1 else dim for dim in result] + return tuple(result) + + +def concat_from_sequence( + input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) - return paddle.flip(m, axis=0) + with ivy.ArrayMode(False): + if new_axis == 0: + return ivy.concat(input_sequence, axis=axis) + elif new_axis == 1: + return ivy.stack(input_sequence, axis=axis) -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int16", "float16")}}, - backend_version, -) -def vstack( - arrays: Sequence[paddle.Tensor], +def dsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], /, *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=0) - else: - return ivy.stack(arrays, axis=0) + copy: Optional[bool] = None, +) -> List[paddle.Tensor]: + if ary.ndim < 3: + raise ivy.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def dstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - if arrays[0].ndim >= 2: - return ivy.concat(arrays, axis=1) + arrays = ivy.atleast_2d(*arrays) + if not isinstance(arrays, list): + arrays = [arrays] + if arrays[0].ndim < 3: + return ivy.stack(arrays, axis=-1) else: - return ivy.concat(arrays, axis=0) + return ivy.concat(arrays, axis=2) -@with_supported_device_and_dtypes( - { - "2.5.1 and above": { - "cpu": ( - "bool", - "int32", - "int64", - "float32", - "float64", - ), - "gpu": ("float16",), - }, - }, - backend_version, -) -def rot90( - m: paddle.Tensor, +def expand( + x: paddle.Tensor, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: Optional[int] = 1, - axes: Optional[Tuple[int, int]] = (0, 1), out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) - return paddle.rot90(m, k=k, axes=axes) + return paddle_backend.broadcast_to(x, shape) @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, + {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version ) -def top_k( - x: paddle.Tensor, - k: int, - /, - *, - axis: int = -1, - largest: Optional[bool] = True, - sorted: bool = True, - out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] - ) - with ivy.ArrayMode(False): - indices = ivy.argsort(x, axis=axis, descending=largest) - indices = paddle.index_select(indices, paddle.arange(end=k), axis) - if not sorted: - indices = paddle.sort(indices, axis=axis) - val = ivy.take_along_axis(x, indices, axis) - return topk_res(val, indices) - - -def fliplr( - m: paddle.Tensor, - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: - return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) - return paddle.flip(m, axis=1) - - -def i0( - x: paddle.Tensor, +def fill_diagonal( + a: paddle.Tensor, + v: Union[int, float], /, *, - out: Optional[paddle.Tensor] = None, + wrap: bool = False, ) -> paddle.Tensor: - def _i0_1(x): - return paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), - ) - - def _i0_2(x): - return paddle_backend.divide( - paddle_backend.multiply( - paddle_backend.exp(x), - _chbevl( - paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B - ), - ), - paddle_backend.sqrt(x), - ) - - def _chbevl(x, vals): - b0 = vals[0] - b1 = 0.0 - - for i in range(1, len(vals)): - b2 = b1 - b1 = b0 - b0 = paddle_backend.add( - paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] - ) - return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) + shape = a.shape + max_end = paddle.prod(paddle.to_tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() + end = max_end if end > max_end else end + a = paddle.reshape(a, (-1,)) + w = paddle.zeros(a.shape, dtype=bool) + ins = paddle.arange(0, max_end) + steps = paddle.arange(0, end, step) - x = paddle_backend.abs(x) - return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) + for i in steps: + i = ins == i + w = paddle.logical_or(w, i) + v = paddle.to_tensor(v, dtype=a.dtype) + a = paddle.where(w, v, a) + a = paddle.reshape(a, shape) + return a def flatten( @@ -306,105 +282,165 @@ def _flatten(x, start_dim, end_dim): return _flatten(x, start_dim, end_dim) -def vsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], +def fliplr( + m: paddle.Tensor, /, *, copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 2: - raise ivy.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=1).cast(m.dtype) + return paddle.flip(m, axis=1) -def dsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Sequence[int], paddle.Tensor], +def flipud( + m: paddle.Tensor, /, *, copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim < 3: - raise ivy.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.flip(m.cast("float32"), axis=0).cast(m.dtype) + return paddle.flip(m, axis=0) -def atleast_1d( - *arys: paddle.Tensor, copy: Optional[bool] = None +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def heaviside( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle.heaviside(x1, x2) + + +def hsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Tuple[int, ...]], + /, + *, + copy: Optional[bool] = None, ) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 1: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=0)) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res + if ary.ndim == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def dstack( +def hstack( arrays: Sequence[paddle.Tensor], /, *, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: with ivy.ArrayMode(False): - arrays = ivy.atleast_2d(*arrays) - if not isinstance(arrays, list): - arrays = [arrays] - if arrays[0].ndim < 3: - return ivy.stack(arrays, axis=-1) + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=1) else: - return ivy.concat(arrays, axis=2) + return ivy.concat(arrays, axis=0) -def atleast_2d( - *arys: paddle.Tensor, copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim < 2: - with ivy.ArrayMode(False): - res.append(ivy.expand_dims(ary, axis=list(range(2 - ary.ndim)))) - else: - res.append(ary) - if len(res) == 1: - return res[0] - return res +def i0( + x: paddle.Tensor, + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + def _i0_1(x): + return paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl(paddle_backend.subtract(paddle_backend.divide(x, 2.0), 2.0), _i0A), + ) + + def _i0_2(x): + return paddle_backend.divide( + paddle_backend.multiply( + paddle_backend.exp(x), + _chbevl( + paddle_backend.subtract(paddle_backend.divide(32.0, x), 2.0), _i0B + ), + ), + paddle_backend.sqrt(x), + ) + def _chbevl(x, vals): + b0 = vals[0] + b1 = 0.0 -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("float16",)}}, + for i in range(1, len(vals)): + b2 = b1 + b1 = b0 + b0 = paddle_backend.add( + paddle_backend.subtract(paddle_backend.multiply(x, b1), b2), vals[i] + ) + return paddle_backend.multiply(0.5, paddle_backend.subtract(b0, b2)) + + x = paddle_backend.abs(x) + return paddle_backend.where(paddle_backend.less_equal(x, 8.0), _i0_1(x), _i0_2(x)) + + +def moveaxis( + a: paddle.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if isinstance(source, tuple): + source = list(source) + if isinstance(destination, tuple): + source = list(destination) + if a.dtype in [paddle.int8, paddle.int16, paddle.uint8]: + return paddle.moveaxis(a.cast("float32"), source, destination).cast(a.dtype) + return paddle.moveaxis(a, source, destination) + + +@with_supported_device_and_dtypes( + { + "2.5.1 and above": { + "cpu": ( + "bool", + "int32", + "int64", + "float32", + "float64", + ), + "gpu": ("float16",), + }, + }, backend_version, ) -def atleast_3d( - *arys: Union[paddle.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[paddle.Tensor]: - res = [] - for ary in arys: - ary = ivy.array(ary, copy=copy).data - if ary.ndim == 0: - result = ary.reshape((1, 1, 1)) - elif ary.ndim == 1: - result = ary[None, :, None] - elif ary.ndim == 2: - result = ary[:, :, None] - else: - result = ary - res.append(result) - if len(res) == 1: - return res[0] - else: - return res +def rot90( + m: paddle.Tensor, + /, + *, + copy: Optional[bool] = None, + k: Optional[int] = 1, + axes: Optional[Tuple[int, int]] = (0, 1), + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if (k % 4) and m.dtype in [paddle.int8, paddle.int16, paddle.uint8, paddle.float16]: + return paddle.rot90(m.cast("float32"), k=k, axes=axes).cast(m.dtype) + return paddle.rot90(m, k=k, axes=axes) @with_unsupported_device_and_dtypes( @@ -484,65 +520,31 @@ def take_along_axis( return paddle.take_along_axis(arr, indices, axis) -def hsplit( - ary: paddle.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[paddle.Tensor]: - if ary.ndim == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - def _broadcast_shape(s1, s2): - len_1 = len(s1) - len_2 = len(s2) - if len_1 == 0: - return () if len_2 == 0 else s2 - elif len_1 != 0 and len_2 == 0: - return s1 - else: - return paddle.broadcast_shape(s1, s2) - - if len(shapes) == 0: - raise ValueError("shapes=[] must be non-empty") - elif len(shapes) == 1: - return shapes[0] - result = _broadcast_shape(shapes[0], shapes[1]) - for i in range(2, len(shapes)): - result = _broadcast_shape(result, shapes[i]) - # paddle outputs -1 if the output dimension is 0 - result = [0 if dim == -1 else dim for dim in result] - return tuple(result) - - -def expand( +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def top_k( x: paddle.Tensor, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.broadcast_to(x, shape) - - -def concat_from_sequence( - input_sequence: Union[Tuple[paddle.Tensor], List[paddle.Tensor]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: + axis: int = -1, + largest: Optional[bool] = True, + sorted: bool = True, + out: Optional[Tuple[paddle.Tensor, paddle.Tensor]] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", paddle.Tensor), ("indices", paddle.Tensor)] + ) with ivy.ArrayMode(False): - if new_axis == 0: - return ivy.concat(input_sequence, axis=axis) - elif new_axis == 1: - return ivy.stack(input_sequence, axis=axis) + indices = ivy.argsort(x, axis=axis, descending=largest) + indices = paddle.index_select(indices, paddle.arange(end=k), axis) + if not sorted: + indices = paddle.sort(indices, axis=axis) + val = ivy.take_along_axis(x, indices, axis) + return topk_res(val, indices) @with_unsupported_device_and_dtypes( @@ -608,35 +610,32 @@ def unique_consecutive( ) +def vsplit( + ary: paddle.Tensor, + indices_or_sections: Union[int, Sequence[int], paddle.Tensor], + /, + *, + copy: Optional[bool] = None, +) -> List[paddle.Tensor]: + if ary.ndim < 2: + raise ivy.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, copy=copy, num_or_size_splits=indices_or_sections, axis=0) + + @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "int16", "uint8", "float16")}}, backend_version + {"2.5.1 and below": {"cpu": ("int16", "float16")}}, + backend_version, ) -def fill_diagonal( - a: paddle.Tensor, - v: Union[int, float], +def vstack( + arrays: Sequence[paddle.Tensor], /, *, - wrap: bool = False, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - shape = a.shape - max_end = paddle.prod(paddle.to_tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (paddle.cumprod(paddle.to_tensor(shape[:-1]), dim=0)).sum() - end = max_end if end > max_end else end - a = paddle.reshape(a, (-1,)) - w = paddle.zeros(a.shape, dtype=bool) - ins = paddle.arange(0, max_end) - steps = paddle.arange(0, end, step) - - for i in steps: - i = ins == i - w = paddle.logical_or(w, i) - v = paddle.to_tensor(v, dtype=a.dtype) - a = paddle.where(w, v, a) - a = paddle.reshape(a, shape) - return a + with ivy.ArrayMode(False): + if arrays[0].ndim >= 2: + return ivy.concat(arrays, axis=0) + else: + return ivy.stack(arrays, axis=0) diff --git a/ivy/functional/backends/paddle/experimental/norms.py b/ivy/functional/backends/paddle/experimental/norms.py index c1fe25a638596..afce1abca3d81 100644 --- a/ivy/functional/backends/paddle/experimental/norms.py +++ b/ivy/functional/backends/paddle/experimental/norms.py @@ -97,11 +97,27 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( - (x.ndim > 1 and x.ndim < 6) - and (scale is not None and scale.ndim == 1) - and (offset is not None and offset.ndim == 1) -) +def instance_norm( + x: paddle.Tensor, + mean: paddle.Tensor, + variance: paddle.Tensor, + /, + *, + scale: Optional[paddle.Tensor] = None, + offset: Optional[paddle.Tensor] = None, + training: Optional[bool] = False, + eps: Optional[float] = 1e-5, + momentum: Optional[float] = 1e-1, + data_format: Optional[str] = "NSC", + out: Optional[ + Tuple[ + paddle.Tensor, + paddle.Tensor, + paddle.Tensor, + ] + ] = None, +) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: + raise IvyNotImplementedException() def l1_normalize( @@ -132,30 +148,14 @@ def l2_normalize( raise IvyNotImplementedException() -def instance_norm( - x: paddle.Tensor, - mean: paddle.Tensor, - variance: paddle.Tensor, - /, - *, - scale: Optional[paddle.Tensor] = None, - offset: Optional[paddle.Tensor] = None, - training: Optional[bool] = False, - eps: Optional[float] = 1e-5, - momentum: Optional[float] = 1e-1, - data_format: Optional[str] = "NSC", - out: Optional[ - Tuple[ - paddle.Tensor, - paddle.Tensor, - paddle.Tensor, - ] - ] = None, -) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor,]: - raise IvyNotImplementedException() - - def lp_normalize( x: paddle.Tensor, /, *, p: float = 2, axis: int = None, out: paddle.Tensor = None ) -> paddle.Tensor: raise IvyNotImplementedException() + + +batch_norm.partial_mixed_handler = lambda x, *args, scale, offset, **kwargs: ( + (x.ndim > 1 and x.ndim < 6) + and (scale is not None and scale.ndim == 1) + and (offset is not None and offset.ndim == 1) +) diff --git a/ivy/functional/backends/paddle/experimental/random.py b/ivy/functional/backends/paddle/experimental/random.py index 542f665e4bb10..1c95f96b39179 100644 --- a/ivy/functional/backends/paddle/experimental/random.py +++ b/ivy/functional/backends/paddle/experimental/random.py @@ -12,6 +12,70 @@ from paddle.device import core from ivy import with_supported_device_and_dtypes + +# bernoulli +@with_supported_device_and_dtypes( + { + "2.5.0 and above": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + }, + "2.4.2 and below": { + "cpu": ( + "float32", + "float64", + ), + "gpu": ("float16", "float32", "float64"), + }, + }, + backend_version, +) +def bernoulli( + probs: Union[float, paddle.Tensor], + *, + logits: Union[float, paddle.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: core.Place, + dtype: paddle.dtype, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + if probs is not None: + probs = probs + elif logits is not None: + probs = ivy.softmax(logits) + probs = paddle.cast(probs, dtype) + probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs + probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) + sample = paddle.bernoulli(probs) + return to_device(sample, device) + + +# beta +def beta( + alpha: Union[float, paddle.Tensor], + beta: Union[float, paddle.Tensor], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + device: core.Place = None, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if seed is not None: + paddle.seed(seed) + shape = _check_bounds_and_get_shape(alpha, beta, shape) + dtype = paddle.float32 if dtype is None else dtype + beta = paddle.cast(beta, alpha.dtype) + dist = paddle.distribution.Beta(alpha, beta) + sample = dist.sample(shape) + sample = paddle.cast(sample, dtype) + return to_device(sample, device) if device is not None else sample + + # dirichlet @@ -51,29 +115,6 @@ def dirichlet( return res -# beta -def beta( - alpha: Union[float, paddle.Tensor], - beta: Union[float, paddle.Tensor], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - device: core.Place = None, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - shape = _check_bounds_and_get_shape(alpha, beta, shape) - dtype = paddle.float32 if dtype is None else dtype - beta = paddle.cast(beta, alpha.dtype) - dist = paddle.distribution.Beta(alpha, beta) - sample = dist.sample(shape) - sample = paddle.cast(sample, dtype) - return to_device(sample, device) if device is not None else sample - - def gamma( alpha: Union[float, paddle.Tensor], beta: Union[float, paddle.Tensor], @@ -99,43 +140,3 @@ def poisson( out: Optional[paddle.Tensor] = None, ): raise IvyNotImplementedException() - - -# bernoulli -@with_supported_device_and_dtypes( - { - "2.5.0 and above": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - }, - "2.4.2 and below": { - "cpu": ( - "float32", - "float64", - ), - "gpu": ("float16", "float32", "float64"), - }, - }, - backend_version, -) -def bernoulli( - probs: Union[float, paddle.Tensor], - *, - logits: Union[float, paddle.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: core.Place, - dtype: paddle.dtype, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if seed is not None: - paddle.seed(seed) - if probs is not None: - probs = probs - elif logits is not None: - probs = ivy.softmax(logits) - probs = paddle.cast(probs, dtype) - probs = paddle.unsqueeze(probs, 0) if len(probs.shape) == 0 else probs - probs = paddle.maximum(probs, paddle.full_like(probs, 1e-6)) - sample = paddle.bernoulli(probs) - return to_device(sample, device) diff --git a/ivy/functional/backends/paddle/experimental/statistical.py b/ivy/functional/backends/paddle/experimental/statistical.py index 0cc4e40739733..8047578332080 100644 --- a/ivy/functional/backends/paddle/experimental/statistical.py +++ b/ivy/functional/backends/paddle/experimental/statistical.py @@ -11,120 +11,120 @@ from . import backend_version -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "complex64", - "complex128", - "bool", - ) - } - }, - backend_version, -) -def median( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # keepdims is set to True because in versions up to 2.5.1 - # there was a problem when the axis was defined and it was the - # only axis in the tensor so it needs to be handled manually - - ret_dtype = input.dtype - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(input): - ret = paddle.complex( - paddle.median(input.real(), axis=axis, keepdim=True), - paddle.median(input.imag(), axis=axis, keepdim=True), - ) - else: - ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) - else: - ret = paddle.median(input, axis=axis, keepdim=True) - if not keepdims: - ret = paddle_backend.squeeze(ret, axis=axis) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == input.ndim: - axis = None - if (input.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) +# --- Helpers --- # +# --------------- # -def nanmean( - a: paddle.Tensor, - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - ret_dtype = dtype if dtype is not None else a.dtype - a = a.cast( - ret_dtype - ) # this is necessary to match other FWs behaviour which cast before calculation - if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: - if paddle.is_complex(a): - ret = paddle.complex( - paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), - paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), - ) +def __find_cummax( + x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None +) -> Tuple[paddle.Tensor, paddle.Tensor]: + indices = [] + values = [] + x_dtype = x.dtype if dtype is None else dtype + if ( + isinstance(x.tolist()[0], list) + and len(x[0].shape) >= 1 + and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) + ): + if axis >= 1: + if not isinstance(x, list): + x = x.tolist() + for ret1 in x: + value, indice = __find_cummax( + paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype + ) + indices.append(indice) + values.append(value) else: - ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] else: - ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) - - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == a.ndim: - axis = None - if (a.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) - + if not isinstance(x, list): + x = x.tolist() + n = 0 + for idx, y in enumerate(x): + if x[n] > y: + values.append(x[n]) + elif x[n] <= y or idx == 0: + n = idx + values.append(y) + indices.append(n) -def _validate_quantile(q): - if isinstance(q, float): - q = paddle.to_tensor(q) - if q.ndim == 1 and q.size < 10: - for i in range(q.size): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(x, paddle.Tensor): + return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( + indices, dtype="int64" + ) else: - if not (paddle.all(0 <= q) and paddle.all(q <= 1)): - return False - return True + return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] - if len(axis) == 0: - raise ValueError("Axis can't be empty!") + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) + else: + indices.append((lst, tuple(prefix))) + return indices - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis +def _compute_quantile_wrapper( + x, + q, + axis=None, + keepdims=False, + interpolation="linear", +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation not in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + return _handle_axis( + x, + q, + _quantile, + keepdims=keepdims, + axis=axis, + interpolation=interpolation, + ) def _handle_axis(a, q, fn, keepdims=False, axis=None, interpolation="nearest"): @@ -216,154 +216,39 @@ def _quantile(a, q, axis=None, interpolation="nearest"): return out.astype(ret_dtype) -def _compute_quantile_wrapper( - x, - q, - axis=None, - keepdims=False, - interpolation="linear", -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation not in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) - return _handle_axis( - x, - q, - _quantile, - keepdims=keepdims, - axis=axis, - interpolation=interpolation, - ) +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + if len(axis) == 0: + raise ValueError("Axis can't be empty!") -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bfloat16", - "complex64", - "complex128", - ) - } - }, - backend_version, -) -def quantile( - a: paddle.Tensor, - q: Union[paddle.Tensor, float], - /, - *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: Optional[bool] = False, - interpolation: Optional[str] = "linear", - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - x=a, - q=q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) - + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") -def corrcoef( - x: paddle.Tensor, - /, - *, - y: Optional[paddle.Tensor] = None, - rowvar: Optional[bool] = True, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis -def histogram( - a: paddle.Tensor, - /, - *, - bins: Optional[Union[int, paddle.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[paddle.Tensor] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[paddle.Tensor] = None, - density: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> Tuple[paddle.Tensor]: - if range is None: - min_range = 0 - max_range = 0 +def _validate_quantile(q): + if isinstance(q, float): + q = paddle.to_tensor(q) + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False else: - min_range = range[0] - max_range = range[1] - return paddle.histogram(a, bins=bins, min=min_range, max=max_range) - - -def nanmedian( - input: paddle.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: Optional[bool] = False, - dtype: Optional[paddle.dtype] = None, - overwrite_input: Optional[bool] = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: - if dtype is None: - dtype = input.dtype - input = input.cast("float32") - paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) - + if not (paddle.all(0 <= q) and paddle.all(q <= 1)): + return False + return True -@with_unsupported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "int8", - "int16", - "uint8", - "float16", - "bool", - ) - } - }, - backend_version, -) -def unravel_index( - indices: paddle.Tensor, - shape: Tuple[int], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if indices.ndim == 0: - indices = indices.unsqueeze(0) - coord = [] - indices = indices - for dim in reversed(shape): - coord.append((indices % dim).astype("int32")) - indices = paddle.floor(indices / dim) - return tuple(reversed(coord)) +# --- Main --- # +# ------------ # @with_unsupported_device_and_dtypes( @@ -397,34 +282,15 @@ def bincount( ) -def igamma( - a: paddle.Tensor, +def corrcoef( + x: paddle.Tensor, /, *, - x: paddle.Tensor, + y: Optional[paddle.Tensor] = None, + rowvar: Optional[bool] = True, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - results = [] - ret_dtype = a.dtype if out is None else out.dtype - if paddle.float16 in [a.dtype, x.dtype]: - a = a.astype("float32") - x = x.astype("float32") - - for ai, xi in zip(a.flatten(), x.flatten()): - ai = ai.astype("float64") - xi = xi.astype("float64") - - def integrand(t): - return paddle.exp(-t) * paddle.pow(t, ai - 1) - - intervals = paddle.linspace(0, xi, 10001).astype("float64") - interval_width = xi / 10000 - values = integrand(intervals) - integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) - result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) - results.append(result) - - return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) + raise IvyNotImplementedException() def cov( @@ -536,88 +402,6 @@ def cummax( return ivy.flip(x, axis=axis), ivy.flip(indices, axis=axis) -def __find_cummax( - x: paddle.Tensor, axis: int = 0, dtype: Optional[paddle.dtype] = None -) -> Tuple[paddle.Tensor, paddle.Tensor]: - indices = [] - values = [] - x_dtype = x.dtype if dtype is None else dtype - if ( - isinstance(x.tolist()[0], list) - and len(x[0].shape) >= 1 - and (isinstance(x[0], paddle.Tensor) or isinstance(x[0], ivy.Array)) - ): - if axis >= 1: - if not isinstance(x, list): - x = x.tolist() - for ret1 in x: - value, indice = __find_cummax( - paddle.to_tensor(ret1, dtype=x_dtype), axis=axis - 1, dtype=x_dtype - ) - indices.append(indice) - values.append(value) - else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - else: - if not isinstance(x, list): - x = x.tolist() - n = 0 - for idx, y in enumerate(x): - if x[n] > y: - values.append(x[n]) - elif x[n] <= y or idx == 0: - n = idx - values.append(y) - indices.append(n) - - if isinstance(x, paddle.Tensor): - return paddle.to_tensor(values, dtype=x.dtype), paddle.to_tensor( - indices, dtype="int64" - ) - else: - return ivy.array(values, dtype=x_dtype), ivy.array(indices, dtype="int64") - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16")}}, backend_version, @@ -645,3 +429,227 @@ def cummin( if reverse: cummin_x = paddle.flip(cummin_x, axis=[axis]) return cummin_x.cast(dtype) + + +def histogram( + a: paddle.Tensor, + /, + *, + bins: Optional[Union[int, paddle.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[paddle.Tensor] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[paddle.Tensor] = None, + density: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> Tuple[paddle.Tensor]: + if range is None: + min_range = 0 + max_range = 0 + else: + min_range = range[0] + max_range = range[1] + return paddle.histogram(a, bins=bins, min=min_range, max=max_range) + + +def igamma( + a: paddle.Tensor, + /, + *, + x: paddle.Tensor, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + results = [] + ret_dtype = a.dtype if out is None else out.dtype + if paddle.float16 in [a.dtype, x.dtype]: + a = a.astype("float32") + x = x.astype("float32") + + for ai, xi in zip(a.flatten(), x.flatten()): + ai = ai.astype("float64") + xi = xi.astype("float64") + + def integrand(t): + return paddle.exp(-t) * paddle.pow(t, ai - 1) + + intervals = paddle.linspace(0, xi, 10001).astype("float64") + interval_width = xi / 10000 + values = integrand(intervals) + integral = paddle.multiply((values[:-1] + values[1:]) / 2, interval_width) + result = paddle.divide(paddle.sum(integral), paddle.exp(paddle.lgamma(ai))) + results.append(result) + + return paddle.to_tensor(results, dtype=ret_dtype).reshape(a.shape) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "complex64", + "complex128", + "bool", + ) + } + }, + backend_version, +) +def median( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # keepdims is set to True because in versions up to 2.5.1 + # there was a problem when the axis was defined and it was the + # only axis in the tensor so it needs to be handled manually + + ret_dtype = input.dtype + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(input): + ret = paddle.complex( + paddle.median(input.real(), axis=axis, keepdim=True), + paddle.median(input.imag(), axis=axis, keepdim=True), + ) + else: + ret = paddle.median(input.cast("float32"), axis=axis, keepdim=True) + else: + ret = paddle.median(input, axis=axis, keepdim=True) + if not keepdims: + ret = paddle_backend.squeeze(ret, axis=axis) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == input.ndim: + axis = None + if (input.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + +def nanmean( + a: paddle.Tensor, + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + ret_dtype = dtype if dtype is not None else a.dtype + a = a.cast( + ret_dtype + ) # this is necessary to match other FWs behaviour which cast before calculation + if a.dtype not in [paddle.int64, paddle.float32, paddle.float64]: + if paddle.is_complex(a): + ret = paddle.complex( + paddle.nanmean(a.real(), axis=axis, keepdim=keepdims), + paddle.nanmean(a.imag(), axis=axis, keepdim=keepdims), + ) + else: + ret = paddle.nanmean(a.cast("float32"), axis=axis, keepdim=keepdims) + else: + ret = paddle.nanmean(a, axis=axis, keepdim=keepdims) + + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == a.ndim: + axis = None + if (a.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + +def nanmedian( + input: paddle.Tensor, + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[paddle.dtype] = None, + overwrite_input: Optional[bool] = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if input.dtype not in [paddle.int32, paddle.int64, paddle.float32, paddle.float64]: + if dtype is None: + dtype = input.dtype + input = input.cast("float32") + paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + return paddle.nanmedian(x=input, axis=axis, keepdim=keepdims).cast(dtype) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bfloat16", + "complex64", + "complex128", + ) + } + }, + backend_version, +) +def quantile( + a: paddle.Tensor, + q: Union[paddle.Tensor, float], + /, + *, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: Optional[bool] = False, + interpolation: Optional[str] = "linear", + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + x=a, + q=q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) + + +@with_unsupported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "int8", + "int16", + "uint8", + "float16", + "bool", + ) + } + }, + backend_version, +) +def unravel_index( + indices: paddle.Tensor, + shape: Tuple[int], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if indices.ndim == 0: + indices = indices.unsqueeze(0) + coord = [] + indices = indices + for dim in reversed(shape): + coord.append((indices % dim).astype("int32")) + indices = paddle.floor(indices / dim) + + return tuple(reversed(coord)) diff --git a/ivy/functional/backends/paddle/general.py b/ivy/functional/backends/paddle/general.py index f4258c4f452b9..31ad5024f28f7 100644 --- a/ivy/functional/backends/paddle/general.py +++ b/ivy/functional/backends/paddle/general.py @@ -13,24 +13,8 @@ from ivy.utils.exceptions import _check_inplace_update_support -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, paddle.Tensor): - if exclusive and not x.stop_gradient: - return False - return True - return False - - -def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: - return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) - - -def container_types(): - return [] - - -def current_backend_str() -> str: - return "paddle" +# --- Helpers --- # +# --------------- # def _check_query(query): @@ -45,59 +29,20 @@ def _check_query(query): ) -def get_item( - x: paddle.Tensor, - /, - query: Union[paddle.Tensor, Tuple], - *, - copy: bool = None, -) -> paddle.Tensor: - dtype = x.dtype - if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: - ret = x.cast("float32").__getitem__(query).cast(dtype) - elif dtype in [paddle.complex64, paddle.complex128]: - ret = paddle.complex( - x.real().__getitem__(query), - x.imag().__getitem__(query), - ) - else: - ret = x.__getitem__(query) - return ret - - -get_item.partial_mixed_handler = ( - lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape -) +# --- Main --- # +# ------------ # -def to_numpy( - x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif paddle.is_tensor(x): - if copy: - return np.array(x) - else: - return np.asarray(x) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") +def array_equal(x0: paddle.Tensor, x1: paddle.Tensor, /) -> bool: + return bool(paddle_backend.all(paddle_backend.equal(x0, x1))) -def to_scalar(x: paddle.Tensor, /) -> Number: - if isinstance(x, (Number, complex)): - return x - return x.item() +def container_types(): + return [] -def to_list(x: paddle.Tensor, /) -> list: - return x.tolist() +def current_backend_str() -> str: + return "paddle" def gather( @@ -287,6 +232,26 @@ def gather_nd( return out +def get_item( + x: paddle.Tensor, + /, + query: Union[paddle.Tensor, Tuple], + *, + copy: bool = None, +) -> paddle.Tensor: + dtype = x.dtype + if dtype in [paddle.int8, paddle.int16, paddle.float16, paddle.bfloat16]: + ret = x.cast("float32").__getitem__(query).cast(dtype) + elif dtype in [paddle.complex64, paddle.complex128]: + ret = paddle.complex( + x.real().__getitem__(query), + x.imag().__getitem__(query), + ) + else: + ret = x.__getitem__(query) + return ret + + def get_num_dims( x: paddle.Tensor, /, *, as_array: bool = False ) -> Union[paddle.Tensor, int]: @@ -356,6 +321,48 @@ def inplace_variables_supported(): return False +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, paddle.Tensor): + if exclusive and not x.stop_gradient: + return False + return True + return False + + +def isin( + elements: paddle.Tensor, + test_elements: paddle.Tensor, + /, + *, + assume_unique: Optional[bool] = False, + invert: Optional[bool] = False, +) -> paddle.Tensor: + input_shape = elements.shape + if elements.ndim == 0: + elements = paddle_backend.expand_dims(elements, axis=0) + if test_elements.ndim == 0: + test_elements = paddle_backend.expand_dims(test_elements, axis=0) + if not assume_unique: + test_elements = paddle_backend.unique_values(test_elements) + + elements = elements.reshape([-1]) + test_elements = test_elements.reshape([-1]) + + output = paddle_backend.any( + paddle_backend.equal( + paddle_backend.expand_dims(elements, axis=-1), test_elements + ), + axis=-1, + ) + return paddle_backend.logical_xor( + paddle_backend.reshape(output, input_shape), invert + ) + + +def itemsize(x: paddle.Tensor) -> int: + return x.element_size() + + def multiprocessing(context=None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -512,6 +519,36 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: paddle.Tensor, /) -> list: + return x.tolist() + + +def to_numpy( + x: Union[paddle.Tensor, List[paddle.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif paddle.is_tensor(x): + if copy: + return np.array(x) + else: + return np.asarray(x) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a Paddle Tensor.") + + +def to_scalar(x: paddle.Tensor, /) -> Number: + if isinstance(x, (Number, complex)): + return x + return x.item() + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -595,35 +632,6 @@ def _vmap(*args, **kwargs): return _vmap -def isin( - elements: paddle.Tensor, - test_elements: paddle.Tensor, - /, - *, - assume_unique: Optional[bool] = False, - invert: Optional[bool] = False, -) -> paddle.Tensor: - input_shape = elements.shape - if elements.ndim == 0: - elements = paddle_backend.expand_dims(elements, axis=0) - if test_elements.ndim == 0: - test_elements = paddle_backend.expand_dims(test_elements, axis=0) - if not assume_unique: - test_elements = paddle_backend.unique_values(test_elements) - - elements = elements.reshape([-1]) - test_elements = test_elements.reshape([-1]) - - output = paddle_backend.any( - paddle_backend.equal( - paddle_backend.expand_dims(elements, axis=-1), test_elements - ), - axis=-1, - ) - return paddle_backend.logical_xor( - paddle_backend.reshape(output, input_shape), invert - ) - - -def itemsize(x: paddle.Tensor) -> int: - return x.element_size() +get_item.partial_mixed_handler = ( + lambda x, query, **kwargs: _check_query(query) and 0 not in x.shape +) diff --git a/ivy/functional/backends/paddle/gradients.py b/ivy/functional/backends/paddle/gradients.py index 9faf1f8d2dca0..f7a7fee303be4 100644 --- a/ivy/functional/backends/paddle/gradients.py +++ b/ivy/functional/backends/paddle/gradients.py @@ -20,24 +20,45 @@ ) -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = x.astype(ivy.default_float_dtype()) - if not x.is_leaf: - ret = x.detach() - ret.stop_gradient = False +# --- Helpers --- # +# --------------- # + + +def _get_jac_one_arg_fn(grad_fn, xs, out_idx): + nested_indices = iter(ivy.all_nested_indices(xs)) + + def one_arg_fn(x): + idx = next(nested_indices) + new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x + ret = grad_fn(new_xs) + for i in out_idx: + ret = ret[i] return ret - ret = paddle_backend.copy_array(x) - ret.stop_gradient = False - return ret + return one_arg_fn -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, paddle.Tensor) and not x.stop_gradient +def _get_one_out_fn(grad_fn, xs, fn_ret): + out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) -def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: - return x.value() + def one_out_fn(o): + out_idx = next(out_nested_indices) + out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape + one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) + jacobian = ivy.nested_map( + xs, + lambda x: jacobian_to_ivy( + paddle.incubate.autograd.Jacobian( + one_arg_fn, ivy.to_native(x.expand_dims()) + ), + x.shape, + out_shape, + ), + shallow=False, + ) + return jacobian + + return one_out_fn def _grad_func(y, xs, retain_grads): @@ -103,6 +124,10 @@ def grad_(x): return grads +# --- Main --- # +# ------------ # + + @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -159,100 +184,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = paddle.grad(y, x, allow_unused=True)[0] - grad = grad if grad is not None else paddle.zeros_like(x) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def stop_gradient( - x: Optional[paddle.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[paddle.Tensor] = None, -): - is_var = is_variable(x) - x.stop_gradient = True - if is_var and preserve_type: - return variable(x) - return x - - -def _get_jac_one_arg_fn(grad_fn, xs, out_idx): - nested_indices = iter(ivy.all_nested_indices(xs)) - - def one_arg_fn(x): - idx = next(nested_indices) - new_xs = ivy.set_nest_at_index(xs, idx, x, shallow=False) if idx else x - ret = grad_fn(new_xs) - for i in out_idx: - ret = ret[i] - return ret - - return one_arg_fn - - -def _get_one_out_fn(grad_fn, xs, fn_ret): - out_nested_indices = iter(ivy.all_nested_indices(fn_ret)) - - def one_out_fn(o): - out_idx = next(out_nested_indices) - out_shape = ivy.index_nest(grad_fn(xs), out_idx).shape - one_arg_fn = _get_jac_one_arg_fn(grad_fn, xs, out_idx) - jacobian = ivy.nested_map( - xs, - lambda x: jacobian_to_ivy( - paddle.incubate.autograd.Jacobian( - one_arg_fn, ivy.to_native(x.expand_dims()) - ), - x.shape, - out_shape, - ), - shallow=False, - ) - return jacobian - - return one_out_fn - - -def jacobian_to_ivy(jacobian, in_shape, out_shape): - jac_ivy = ivy.to_ivy(jacobian[:]) - jac_shape = out_shape + in_shape - jac_reshaped = jac_ivy.reshape(jac_shape) - return jac_reshaped - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(xs): - fn_ret = grad_fn(xs) - one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) - jacobian = ivy.nested_map(fn_ret, one_out_fn) - return jacobian - - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -290,5 +221,82 @@ def _inner(x): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, paddle.Tensor) and not x.stop_gradient + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(xs): + fn_ret = grad_fn(xs) + one_out_fn = _get_one_out_fn(grad_fn, xs, fn_ret) + jacobian = ivy.nested_map(fn_ret, one_out_fn) + return jacobian + + return callback_fn + + +def jacobian_to_ivy(jacobian, in_shape, out_shape): + jac_ivy = ivy.to_ivy(jacobian[:]) + jac_shape = out_shape + in_shape + jac_reshaped = jac_ivy.reshape(jac_shape) + return jac_reshaped + + +def stop_gradient( + x: Optional[paddle.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[paddle.Tensor] = None, +): + is_var = is_variable(x) + x.stop_gradient = True + if is_var and preserve_type: + return variable(x) + return x + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = paddle.grad(y, x, allow_unused=True)[0] + grad = grad if grad is not None else paddle.zeros_like(x) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = x.astype(ivy.default_float_dtype()) + if not x.is_leaf: + ret = x.detach() + ret.stop_gradient = False + return ret + ret = paddle_backend.copy_array(x) + ret.stop_gradient = False + return ret + + +def variable_data(x: paddle.Tensor, /) -> paddle.Tensor: + return x.value() + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/paddle/layers.py b/ivy/functional/backends/paddle/layers.py index f70c374bc70de..462d2b96f23e8 100644 --- a/ivy/functional/backends/paddle/layers.py +++ b/ivy/functional/backends/paddle/layers.py @@ -19,8 +19,8 @@ from . import backend_version -def _is_list_or_tuple(inp): - return isinstance(inp, (list, tuple)) +# --- Helpers --- # +# --------------- # def _convert_to_list(value, n, name="padding", _type=int): @@ -37,6 +37,27 @@ def _convert_to_list(value, n, name="padding", _type=int): return value_list +def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): + if filter_format == "channel_first": + filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) + + # adding dilation in input + x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations + for i in range(dims): + if x_dilations[i] > 1: + h = x.shape[1 + i] + new_height = h + (h - 1) * (x_dilations[i] - 1) + h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] + x = paddle_backend.swapaxes(x, 1 + i, -1) + x = paddle.matmul(x, h) + x = paddle_backend.swapaxes(x, -1, 1 + i) + return x, filters + + +def _is_list_or_tuple(inp): + return isinstance(inp, (list, tuple)) + + def _pad_before_conv(x, filters, strides, padding, dims, dilations, data_format): dilations = _convert_to_list(dilations, dims, "dilations") strides = _convert_to_list(strides, dims, "strides") @@ -122,21 +143,8 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding -def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): - if filter_format == "channel_first": - filters = paddle.transpose(filters, (*range(2, dims + 2), 1, 0)) - - # adding dilation in input - x_dilations = [x_dilations] * dims if isinstance(x_dilations, int) else x_dilations - for i in range(dims): - if x_dilations[i] > 1: - h = x.shape[1 + i] - new_height = h + (h - 1) * (x_dilations[i] - 1) - h = paddle.eye(new_height, dtype=x.dtype)[:: x_dilations[i]] - x = paddle_backend.swapaxes(x, 1 + i, -1) - x = paddle.matmul(x, h) - x = paddle_backend.swapaxes(x, -1, 1 + i) - return x, filters +# --- Main --- # +# ------------ # def conv1d( @@ -259,21 +267,6 @@ def conv2d_transpose( return res -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: paddle.Tensor, - filters: paddle.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: Optional[str] = "NHWC", - dilations: Optional[Union[int, Tuple[int, int]]] = 1, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("float16",)}}, backend_version, @@ -494,3 +487,18 @@ def conv_general_transpose( if data_format == "channel_last": res = res.transpose(0, *range(2, dims + 2), 1) return res + + +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: paddle.Tensor, + filters: paddle.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: Optional[str] = "NHWC", + dilations: Optional[Union[int, Tuple[int, int]]] = 1, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() diff --git a/ivy/functional/backends/paddle/linear_algebra.py b/ivy/functional/backends/paddle/linear_algebra.py index 8a99ab77110fc..61e2f5744c1dc 100644 --- a/ivy/functional/backends/paddle/linear_algebra.py +++ b/ivy/functional/backends/paddle/linear_algebra.py @@ -108,6 +108,33 @@ def det(x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None) -> paddle.T return ret +# Extra # +# ----- # + + +def diag( + x: paddle.Tensor, + /, + *, + k: int = 0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) + ) + return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) + return paddle.diag(x, offset=k) + + def diagonal( x: paddle.Tensor, /, @@ -136,6 +163,16 @@ def diagonal( return paddle.diagonal(x, offset=offset, axis1=axis1, axis2=axis2) +def eig( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> Tuple[paddle.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] + ) + eigenvalues, eigenvectors = paddle.linalg.eig(x) + return result_tuple(eigenvalues, eigenvectors) + + def eigh( x: paddle.Tensor, /, @@ -318,16 +355,6 @@ def matrix_norm( return ret -def eig( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> Tuple[paddle.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", paddle.Tensor), ("eigenvectors", paddle.Tensor)] - ) - eigenvalues, eigenvectors = paddle.linalg.eig(x) - return result_tuple(eigenvalues, eigenvectors) - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -423,18 +450,6 @@ def pinv( return paddle.linalg.pinv(x, rcond=rtol) -def tensorsolve( - x1: paddle.Tensor, - x2: paddle.Tensor, - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - # Implemented as a composite function in ivy.functional.ivy.linear_algebra - raise IvyNotImplementedException() - - @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, @@ -560,6 +575,18 @@ def tensordot( return ret.squeeze().cast(ret_dtype) if x1.ndim == axes else ret.cast(ret_dtype) +def tensorsolve( + x1: paddle.Tensor, + x2: paddle.Tensor, + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # Implemented as a composite function in ivy.functional.ivy.linear_algebra + raise IvyNotImplementedException() + + @with_unsupported_device_and_dtypes( { "2.5.1 and below": { @@ -588,6 +615,28 @@ def trace( return ret.squeeze() if x.ndim <= 2 else ret +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, + backend_version, +) +def vander( + x: paddle.Tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + return paddle.pow( + paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), + paddle.arange(start, stop, step, dtype=x.dtype), + ) + + def vecdot( x1: paddle.Tensor, x2: paddle.Tensor, @@ -633,55 +682,6 @@ def vector_norm( ) -# Extra # -# ----- # - - -def diag( - x: paddle.Tensor, - /, - *, - k: int = 0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.diag(x.real(), offset=k), paddle.diag(x.imag(), offset=k) - ) - return paddle.diag(x.cast("float32"), offset=k).cast(x.dtype) - return paddle.diag(x, offset=k) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("uint8", "int8", "int16", "complex64", "complex128")}}, - backend_version, -) -def vander( - x: paddle.Tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - return paddle.pow( - paddle.moveaxis(paddle.unsqueeze(x, 0), 0, 1), - paddle.arange(start, stop, step, dtype=x.dtype), - ) - - @with_unsupported_dtypes( {"2.5.1 and below": ("unsigned", "int8", "int16", "float16")}, backend_version, diff --git a/ivy/functional/backends/paddle/manipulation.py b/ivy/functional/backends/paddle/manipulation.py index a52336b3c46a0..0902034d2362a 100644 --- a/ivy/functional/backends/paddle/manipulation.py +++ b/ivy/functional/backends/paddle/manipulation.py @@ -14,6 +14,33 @@ from ...ivy.manipulation import _calculate_out_shape +# --- Helpers --- # +# --------------- # + + +def _reshape_fortran_paddle(x, shape): + if len(x.shape) > 0: + x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) + return paddle_backend.permute_dims( + paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] + ) + + +# --- Main --- # +# ------------ # + + +def clip( + x: paddle.Tensor, + x_min: Union[Number, paddle.Tensor], + x_max: Union[Number, paddle.Tensor], + /, + *, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) + + # Array API Standard # # -------------------# @@ -54,6 +81,35 @@ def concat( return ret +def constant_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + paddings = [] + pad_width = list(pad_width) + for item in pad_width: + if len(item) != 2: + raise ivy.utils.exceptions.IvyException("Length of each item should be 2") + else: + paddings.append(item[0]) + paddings.append(item[1]) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.nn.functional.pad( + x.cast("float32"), pad=paddings, value=value + ).cast(x.dtype) + return paddle.nn.functional.pad(x=x, pad=paddings, value=value) + + def expand_dims( x: paddle.Tensor, /, @@ -97,12 +153,50 @@ def permute_dims( return paddle.transpose(x, axes) -def _reshape_fortran_paddle(x, shape): - if len(x.shape) > 0: - x = paddle_backend.permute_dims(x, list(reversed(range(x.ndim)))) - return paddle_backend.permute_dims( - paddle.reshape(x, shape[::-1]), list(range(len(shape)))[::-1] - ) +def repeat( + x: paddle.Tensor, + /, + repeats: Union[int, Iterable[int]], + *, + axis: int = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + # handle the case when repeats contains 0 as paddle doesn't support it + if (isinstance(repeats, Number) and repeats == 0) or ( + isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 + ): + if axis is None: + return paddle.to_tensor([], dtype=x.dtype) + else: + shape = x.shape + shape[axis] = 0 + return paddle.zeros(shape=shape).cast(x.dtype) + + if isinstance(repeats, paddle.Tensor) and repeats.size == 1: + repeats = repeats.item() + + if axis is not None: + axis = axis % x.ndim + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + return paddle.complex( + paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), + paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), + ) + + return paddle.repeat_interleave( + x.cast("float32"), repeats=repeats, axis=axis + ).cast(x.dtype) + + return paddle.repeat_interleave(x, repeats=repeats, axis=axis) def reshape( @@ -170,6 +264,62 @@ def roll( return paddle.roll(x, shift, axis) +# Extra # +# ------# + + +def split( + x: paddle.Tensor, + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, + axis: Optional[int] = 0, + with_remainder: Optional[bool] = False, +) -> List[paddle.Tensor]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + num_or_size_splits = x.shape[axis] + elif isinstance(num_or_size_splits, paddle.Tensor): + num_or_size_splits = num_or_size_splits.cast("int32") + num_or_size_splits = num_or_size_splits.tolist() + elif isinstance(num_or_size_splits, int): + num_chunks = x.shape[axis] // num_or_size_splits + remainder = x.shape[axis] % num_or_size_splits + if remainder != 0: + if with_remainder: + num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] + else: + raise ivy.utils.exceptions.IvyException( + "Split size is not compatible with input shape" + ) + + if isinstance(num_or_size_splits, (list, tuple)): + if sum(num_or_size_splits) < x.shape[axis]: + num_or_size_splits + type(num_or_size_splits)([-1]) + elif sum(num_or_size_splits) > x.shape[axis]: + raise ivy.utils.exceptions.IvyException( + "total split size is not compatible with input shape," + f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" + ) + + if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: + if paddle.is_complex(x): + imag_list = paddle.split(x.imag(), num_or_size_splits, axis) + real_list = paddle.split(x.real(), num_or_size_splits, axis) + return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] + ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) + return [tensor.cast(x.dtype) for tensor in ret] + return paddle.split(x, num_or_size_splits, axis) + + def squeeze( x: paddle.Tensor, /, @@ -241,106 +391,18 @@ def stack( return paddle.stack(arrays, axis=axis) -# Extra # -# ------# - - -def split( +def swapaxes( x: paddle.Tensor, + axis0: int, + axis1: int, /, *, copy: Optional[bool] = None, - num_or_size_splits: Optional[Union[int, List[int], paddle.Tensor]] = None, - axis: Optional[int] = 0, - with_remainder: Optional[bool] = False, -) -> List[paddle.Tensor]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - num_or_size_splits = x.shape[axis] - elif isinstance(num_or_size_splits, paddle.Tensor): - num_or_size_splits = num_or_size_splits.cast("int32") - num_or_size_splits = num_or_size_splits.tolist() - elif isinstance(num_or_size_splits, int): - num_chunks = x.shape[axis] // num_or_size_splits - remainder = x.shape[axis] % num_or_size_splits - if remainder != 0: - if with_remainder: - num_or_size_splits = [num_or_size_splits] * num_chunks + [remainder] - else: - raise ivy.utils.exceptions.IvyException( - "Split size is not compatible with input shape" - ) - - if isinstance(num_or_size_splits, (list, tuple)): - if sum(num_or_size_splits) < x.shape[axis]: - num_or_size_splits + type(num_or_size_splits)([-1]) - elif sum(num_or_size_splits) > x.shape[axis]: - raise ivy.utils.exceptions.IvyException( - "total split size is not compatible with input shape," - f" got {sum(num_or_size_splits)} which is more than x.shape[axis]" - ) - - if x.dtype in [paddle.int16, paddle.complex64, paddle.complex128]: - if paddle.is_complex(x): - imag_list = paddle.split(x.imag(), num_or_size_splits, axis) - real_list = paddle.split(x.real(), num_or_size_splits, axis) - return [paddle.complex(a, b) for a, b in zip(real_list, imag_list)] - ret = paddle.split(x.cast("int32"), num_or_size_splits, axis) - return [tensor.cast(x.dtype) for tensor in ret] - return paddle.split(x, num_or_size_splits, axis) - - -def repeat( - x: paddle.Tensor, - /, - repeats: Union[int, Iterable[int]], - *, - axis: int = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - # handle the case when repeats contains 0 as paddle doesn't support it - if (isinstance(repeats, Number) and repeats == 0) or ( - isinstance(repeats, paddle.Tensor) and repeats.size == 1 and repeats.item() == 0 - ): - if axis is None: - return paddle.to_tensor([], dtype=x.dtype) - else: - shape = x.shape - shape[axis] = 0 - return paddle.zeros(shape=shape).cast(x.dtype) - - if isinstance(repeats, paddle.Tensor) and repeats.size == 1: - repeats = repeats.item() - - if axis is not None: - axis = axis % x.ndim - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - paddle.bool, - ]: - if paddle.is_complex(x): - return paddle.complex( - paddle.repeat_interleave(x.real(), repeats=repeats, axis=axis), - paddle.repeat_interleave(x.imag(), repeats=repeats, axis=axis), - ) - - return paddle.repeat_interleave( - x.cast("float32"), repeats=repeats, axis=axis - ).cast(x.dtype) - - return paddle.repeat_interleave(x, repeats=repeats, axis=axis) + axes = [x for x in range(x.ndim)] + axes[axis0], axes[axis1] = axes[axis1], axes[axis0] + return paddle_backend.permute_dims(x, axes) def tile( @@ -383,70 +445,6 @@ def tile( return paddle.tile(x, repeats) -def constant_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - value: Number = 0.0, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - paddings = [] - pad_width = list(pad_width) - for item in pad_width: - if len(item) != 2: - raise ivy.utils.exceptions.IvyException("Length of each item should be 2") - else: - paddings.append(item[0]) - paddings.append(item[1]) - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.nn.functional.pad( - x.cast("float32"), pad=paddings, value=value - ).cast(x.dtype) - return paddle.nn.functional.pad(x=x, pad=paddings, value=value) - - -def zero_pad( - x: paddle.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[paddle.Tensor] = None, -): - return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) - - -def swapaxes( - x: paddle.Tensor, - axis0: int, - axis1: int, - /, - *, - copy: Optional[bool] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - axes = [x for x in range(x.ndim)] - axes[axis0], axes[axis1] = axes[axis1], axes[axis0] - return paddle_backend.permute_dims(x, axes) - - -def clip( - x: paddle.Tensor, - x_min: Union[Number, paddle.Tensor], - x_max: Union[Number, paddle.Tensor], - /, - *, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - return paddle_backend.minimum(paddle_backend.maximum(x, x_min), x_max) - - def unstack( x: paddle.Tensor, /, @@ -482,3 +480,13 @@ def unstack( if keepdims: return [paddle_backend.expand_dims(r, axis=axis) for r in ret] return ret + + +def zero_pad( + x: paddle.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[paddle.Tensor] = None, +): + return paddle_backend.constant_pad(x, pad_width=pad_width, value=0) diff --git a/ivy/functional/backends/paddle/random.py b/ivy/functional/backends/paddle/random.py index 3ca640ccf531b..85d02e5b69146 100644 --- a/ivy/functional/backends/paddle/random.py +++ b/ivy/functional/backends/paddle/random.py @@ -19,39 +19,69 @@ ) from . import backend_version -# Extra # -# ------# + +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ( + "float32", + "float64", + ) + } + }, + backend_version, +) +def multinomial( + population_size: int, + num_samples: int, + /, + *, + batch_size: int = 1, + probs: Optional[paddle.Tensor] = None, + replace: bool = True, + device: Place, + seed: Optional[int] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + if probs is None: + probs = paddle.ones((batch_size, num_samples)) / population_size + probs = paddle.cast(probs, paddle.float32) + if seed: + paddle.seed(seed) + x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) + return x @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def random_uniform( +def randint( + low: Union[int, paddle.Tensor], + high: Union[int, paddle.Tensor], + /, *, - low: Union[float, paddle.Tensor] = 0.0, - high: Union[float, paddle.Tensor] = 1.0, - shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: paddle.dtype, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, device: Place, - seed=None, + dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, + seed: Optional[int] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) + _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - # Set range and seed - rng = high - low + range = high - low if seed: _ = paddle.seed(seed) - random_base = paddle.uniform(shape, min=0.0, max=1.0) - return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( - dtype + _retval = paddle.cast( + paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype ) + return _retval if shape else _retval.squeeze(axis=0) @with_unsupported_device_and_dtypes( @@ -80,68 +110,39 @@ def random_normal( return paddle.normal(mean, std).cast(dtype) -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ( - "float32", - "float64", - ) - } - }, - backend_version, -) -def multinomial( - population_size: int, - num_samples: int, - /, - *, - batch_size: int = 1, - probs: Optional[paddle.Tensor] = None, - replace: bool = True, - device: Place, - seed: Optional[int] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - if probs is None: - probs = paddle.ones((batch_size, num_samples)) / population_size - probs = paddle.cast(probs, paddle.float32) - if seed: - paddle.seed(seed) - x = paddle.multinomial(probs, num_samples=num_samples, replacement=replace) - return x +# Extra # +# ------# @with_unsupported_device_and_dtypes( {"2.5.1 and below": {"cpu": ("int8",)}}, backend_version, ) -def randint( - low: Union[int, paddle.Tensor], - high: Union[int, paddle.Tensor], - /, +def random_uniform( *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + low: Union[float, paddle.Tensor] = 0.0, + high: Union[float, paddle.Tensor] = 1.0, + shape: Optional[Union[paddle.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: paddle.dtype, device: Place, - dtype: Optional[Union[paddle.dtype, ivy.Dtype]] = None, - seed: Optional[int] = None, + seed=None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: if not dtype: dtype = ivy.default_int_dtype() dtype = ivy.as_native_dtype(dtype) - _randint_check_dtype_and_bound(low, high, dtype) low = paddle.cast(low, "float32") if isinstance(low, paddle.Tensor) else low high = paddle.cast(high, "float32") if isinstance(high, paddle.Tensor) else high shape = _check_bounds_and_get_shape(low, high, shape).shape - range = high - low + # Set range and seed + rng = high - low if seed: _ = paddle.seed(seed) + random_base = paddle.uniform(shape, min=0.0, max=1.0) - _retval = paddle.cast( - paddle.uniform(shape or [1], min=0.0, max=1.0) * range + low, dtype + return paddle_backend.add(paddle_backend.multiply(random_base, rng), low).cast( + dtype ) - return _retval if shape else _retval.squeeze(axis=0) def seed(*, seed_value: int = 0) -> None: diff --git a/ivy/functional/backends/paddle/searching.py b/ivy/functional/backends/paddle/searching.py index 9519e03c23916..36fd90f771b40 100644 --- a/ivy/functional/backends/paddle/searching.py +++ b/ivy/functional/backends/paddle/searching.py @@ -94,6 +94,31 @@ def argmin( return ret.astype(dtype) +# Extra # +# ----- # + + +def argwhere( + x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None +) -> paddle.Tensor: + if x.ndim == 0: + return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") + if x.dtype in [ + paddle.int8, + paddle.uint8, + paddle.float16, + paddle.complex64, + paddle.complex128, + ]: + if paddle.is_complex(x): + real_idx = paddle.nonzero(x.real()) + imag_idx = paddle.nonzero(x.imag()) + idx = paddle.concat([real_idx, imag_idx], axis=0) + return paddle.unique(idx, axis=0) + return paddle.nonzero(x.cast("float32")) + return paddle.nonzero(x) + + def nonzero( x: paddle.Tensor, /, @@ -175,28 +200,3 @@ def where( result = paddle.where(condition, x1, x2) return result.squeeze().cast(ret_dtype) if scalar_out else result.cast(ret_dtype) - - -# Extra # -# ----- # - - -def argwhere( - x: paddle.Tensor, /, *, out: Optional[paddle.Tensor] = None -) -> paddle.Tensor: - if x.ndim == 0: - return paddle.zeros(shape=[int(bool(x.item())), 0], dtype="int64") - if x.dtype in [ - paddle.int8, - paddle.uint8, - paddle.float16, - paddle.complex64, - paddle.complex128, - ]: - if paddle.is_complex(x): - real_idx = paddle.nonzero(x.real()) - imag_idx = paddle.nonzero(x.imag()) - idx = paddle.concat([real_idx, imag_idx], axis=0) - return paddle.unique(idx, axis=0) - return paddle.nonzero(x.cast("float32")) - return paddle.nonzero(x) diff --git a/ivy/functional/backends/paddle/sorting.py b/ivy/functional/backends/paddle/sorting.py index 6a1d8f38f10d8..b12e326d19c8c 100644 --- a/ivy/functional/backends/paddle/sorting.py +++ b/ivy/functional/backends/paddle/sorting.py @@ -33,29 +33,13 @@ def argsort( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, backend_version, ) -def sort( - x: paddle.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[paddle.Tensor] = None, +def msort( + a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None ) -> paddle.Tensor: - if x.dtype in [ - paddle.int8, - paddle.int16, - paddle.uint8, - paddle.float16, - paddle.bool, - ]: - return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( - x.dtype - ) - return paddle.sort(x, axis=axis, descending=descending) + return paddle.sort(a, axis=0) @with_unsupported_device_and_dtypes( @@ -115,10 +99,26 @@ def searchsorted( @with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("int8", "uint8", "int16", "float16", "complex")}}, + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, backend_version, ) -def msort( - a: Union[paddle.Tensor, list, tuple], /, *, out: Optional[paddle.Tensor] = None +def sort( + x: paddle.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - return paddle.sort(a, axis=0) + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bool, + ]: + return paddle.sort(x.cast("float32"), axis=axis, descending=descending).cast( + x.dtype + ) + return paddle.sort(x, axis=axis, descending=descending) diff --git a/ivy/functional/backends/paddle/statistical.py b/ivy/functional/backends/paddle/statistical.py index 4d0297b1b7aac..b03ff9d943fcd 100644 --- a/ivy/functional/backends/paddle/statistical.py +++ b/ivy/functional/backends/paddle/statistical.py @@ -1,7 +1,3 @@ -# global - -torch_scatter = None - from typing import Union, Optional, Sequence import paddle @@ -16,45 +12,163 @@ # local from . import backend_version -# Array API Standard # -# -------------------# +# global +torch_scatter = None -def min( + +# --- Helpers --- # +# --------------- # + + +def _std(x, axis, correction, keepdim): + u = paddle_backend.mean(x, axis=axis, keepdims=True) + out = paddle_backend.sum( + paddle_backend.pow(paddle_backend.subtract(x, u), 2), + axis=axis, + keepdims=keepdim, + ) + num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() + num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() + n = num_elm_out / num_elm_in + out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) + if correction: + n = paddle_backend.sqrt( + paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) + ) + out = paddle_backend.multiply(out, n) + return out + + +# --- Main --- # +# ------------ # + + +# Extra # +# ----- # +@with_supported_device_and_dtypes( + { + "2.5.1 and below": { + "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") + } + }, + backend_version, +) +def cumprod( x: paddle.Tensor, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[paddle.dtype] = None, out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: - ret_dtype = x.dtype - if x.dtype in [ + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ + paddle.uint8, paddle.int8, paddle.int16, + paddle.float16, + ]: + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumprod(x, dim=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.ones_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumprod(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) + else: + x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +@with_unsupported_device_and_dtypes( + {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, + backend_version, +) +def cumsum( + x: paddle.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[paddle.dtype] = None, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + dtype = dtype if dtype is not None else x.dtype + x = paddle.cast(x, dtype) + if ivy.as_native_dtype(dtype) in [ paddle.uint8, + paddle.int8, paddle.float16, - paddle.bfloat16, - paddle.complex64, - paddle.complex128, paddle.bool, ]: - if paddle.is_complex(x): - real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) - imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) - ret = paddle.complex(real, imag) - else: - ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) + x = paddle.cast(x, "float32") + if not (exclusive or reverse): + return paddle.cumsum(x, axis=axis).cast(dtype) + elif exclusive and reverse: + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle_backend.swapaxes(x, axis, -1) + return paddle_backend.flip(x, axis=(axis,)).cast(dtype) + elif exclusive: + x = paddle_backend.swapaxes(x, axis, -1) + x = paddle_backend.concat( + [ + paddle.zeros_like( + paddle_backend.get_item(x, (..., slice(-1, None, None))) + ), + paddle_backend.get_item(x, (..., slice(None, -1, None))), + ], + axis=-1, + ) + x = paddle.cumsum(x, -1) + return paddle_backend.swapaxes(x, axis, -1).cast(dtype) else: - ret = paddle.amin(x, axis=axis, keepdim=keepdims) - # The following code is to simulate other frameworks - # output shapes behaviour since min output dim is 1 in paddle - if isinstance(axis, Sequence): - if len(axis) == x.ndim: - axis = None - if (x.ndim == 1 or axis is None) and not keepdims: - ret = ret.squeeze() - return ret.astype(ret_dtype) + x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) + return paddle_backend.flip(x, axis=axis).cast(dtype) + + +def einsum( + equation: str, + *operands: paddle.Tensor, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + raise IvyNotImplementedException() def max( @@ -128,6 +242,47 @@ def mean( return ret.astype(ret_dtype) +# Array API Standard # +# -------------------# + + +def min( + x: paddle.Tensor, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + ret_dtype = x.dtype + if x.dtype in [ + paddle.int8, + paddle.int16, + paddle.uint8, + paddle.float16, + paddle.bfloat16, + paddle.complex64, + paddle.complex128, + paddle.bool, + ]: + if paddle.is_complex(x): + real = paddle.amin(x.real(), axis=axis, keepdim=keepdims) + imag = paddle.amin(x.imag(), axis=axis, keepdim=keepdims) + ret = paddle.complex(real, imag) + else: + ret = paddle.amin(x.cast("float32"), axis=axis, keepdim=keepdims) + else: + ret = paddle.amin(x, axis=axis, keepdim=keepdims) + # The following code is to simulate other frameworks + # output shapes behaviour since min output dim is 1 in paddle + if isinstance(axis, Sequence): + if len(axis) == x.ndim: + axis = None + if (x.ndim == 1 or axis is None) and not keepdims: + ret = ret.squeeze() + return ret.astype(ret_dtype) + + def prod( x: paddle.Tensor, /, @@ -150,25 +305,6 @@ def prod( return ret -def _std(x, axis, correction, keepdim): - u = paddle_backend.mean(x, axis=axis, keepdims=True) - out = paddle_backend.sum( - paddle_backend.pow(paddle_backend.subtract(x, u), 2), - axis=axis, - keepdims=keepdim, - ) - num_elm_in = paddle.prod(paddle.to_tensor(x.shape)).item() - num_elm_out = paddle.prod(paddle.to_tensor(out.shape)).item() - n = num_elm_out / num_elm_in - out = paddle_backend.sqrt(paddle_backend.multiply(out, n)) - if correction: - n = paddle_backend.sqrt( - paddle_backend.divide(num_elm_in, (num_elm_in - correction * num_elm_out)) - ) - out = paddle_backend.multiply(out, n) - return out - - def std( x: paddle.Tensor, /, @@ -217,130 +353,3 @@ def var( ) -> paddle.Tensor: ret = paddle_backend.pow(_std(x, axis, correction, keepdims), 2).cast(x.dtype) return ret - - -# Extra # -# ----- # -@with_supported_device_and_dtypes( - { - "2.5.1 and below": { - "cpu": ("int32", "int64", "float64", "complex128", "float32", "complex64") - } - }, - backend_version, -) -def cumprod( - x: paddle.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ - paddle.uint8, - paddle.int8, - paddle.int16, - paddle.float16, - ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumprod(x, dim=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.ones_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumprod(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) - else: - x = paddle.cumprod(paddle_backend.flip(x, axis=(axis,)), dim=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -@with_unsupported_device_and_dtypes( - {"2.5.1 and below": {"cpu": ("complex64", "complex128")}}, - backend_version, -) -def cumsum( - x: paddle.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[paddle.dtype] = None, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - dtype = dtype if dtype is not None else x.dtype - x = paddle.cast(x, dtype) - if ivy.as_native_dtype(dtype) in [ - paddle.uint8, - paddle.int8, - paddle.float16, - paddle.bool, - ]: - x = paddle.cast(x, "float32") - if not (exclusive or reverse): - return paddle.cumsum(x, axis=axis).cast(dtype) - elif exclusive and reverse: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle_backend.swapaxes(x, axis, -1) - return paddle_backend.flip(x, axis=(axis,)).cast(dtype) - elif exclusive: - x = paddle_backend.swapaxes(x, axis, -1) - x = paddle_backend.concat( - [ - paddle.zeros_like( - paddle_backend.get_item(x, (..., slice(-1, None, None))) - ), - paddle_backend.get_item(x, (..., slice(None, -1, None))), - ], - axis=-1, - ) - x = paddle.cumsum(x, -1) - return paddle_backend.swapaxes(x, axis, -1).cast(dtype) - else: - x = paddle.cumsum(paddle_backend.flip(x, axis=(axis,)), axis=axis) - return paddle_backend.flip(x, axis=axis).cast(dtype) - - -def einsum( - equation: str, - *operands: paddle.Tensor, - out: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - raise IvyNotImplementedException() diff --git a/ivy/functional/backends/tensorflow/activations.py b/ivy/functional/backends/tensorflow/activations.py index de43f95f9a8f7..719d6b8d6cccf 100644 --- a/ivy/functional/backends/tensorflow/activations.py +++ b/ivy/functional/backends/tensorflow/activations.py @@ -30,6 +30,11 @@ def gelu( return tf.nn.gelu(x, approximate) +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return x * tf.nn.relu6(x + 3) / 6 + + def leaky_relu( x: Tensor, /, @@ -41,6 +46,23 @@ def leaky_relu( return tf.nn.leaky_relu(x, alpha) +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def log_softmax( + x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None +): + return tf.nn.log_softmax(x, axis) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def mish( + x: Tensor, + /, + *, + out: Optional[Tensor] = None, +) -> Tensor: + return x * tf.math.tanh(tf.math.softplus(x)) + + def relu(x: Tensor, /, *, complex_mode="jax", out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu(x) @@ -89,25 +111,3 @@ def softplus( if threshold is not None: return tf.where(x_beta > threshold, x, res) return res - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def log_softmax( - x: Tensor, /, *, axis: Optional[int] = None, out: Optional[Tensor] = None -): - return tf.nn.log_softmax(x, axis) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def mish( - x: Tensor, - /, - *, - out: Optional[Tensor] = None, -) -> Tensor: - return x * tf.math.tanh(tf.math.softplus(x)) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def hardswish(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return x * tf.nn.relu6(x + 3) / 6 diff --git a/ivy/functional/backends/tensorflow/control_flow_ops.py b/ivy/functional/backends/tensorflow/control_flow_ops.py index 9b6f42c456d75..1f600eebbe861 100644 --- a/ivy/functional/backends/tensorflow/control_flow_ops.py +++ b/ivy/functional/backends/tensorflow/control_flow_ops.py @@ -1,42 +1,20 @@ import tensorflow as tf -# def if_exp(cond, if_true, if_false, expr_repr): -# def true_fn(): -# return if_true() -# -# def false_fn(): -# return if_false() -# -# return tf.cond(cond, true_fn, false_fn) +# --- Helpers --- # +# --------------- # -def if_else(cond, body_fn, orelse_fn, vars): - # back-compatibility - if isinstance(cond, bool): - v = cond - cond = lambda *_: v - cond = bool(cond(*vars)) - # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) - # use pythonic placeholder until the graph compiler supports callable arguments - - if cond: - return body_fn(*vars) - else: - return orelse_fn(*vars) +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) -def while_loop(test_fn, body_fn, vars): - def body_fn_wrapper(*loop_vars): - return body_fn(*loop_vars) - - def test_fn_wrapper(*loop_vars): - return test_fn(*loop_vars) +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} - if not vars: - vars = (0,) - return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) +# --- Main --- # +# ------------ # def for_loop( @@ -70,9 +48,40 @@ def empty_function(*args): return _dict_to_tuple(vars_dict) -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} +# def if_exp(cond, if_true, if_false, expr_repr): +# def true_fn(): +# return if_true() +# +# def false_fn(): +# return if_false() +# +# return tf.cond(cond, true_fn, false_fn) -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) +def if_else(cond, body_fn, orelse_fn, vars): + # back-compatibility + if isinstance(cond, bool): + v = cond + cond = lambda *_: v + cond = bool(cond(*vars)) + # return tf.cond(cond, lambda: body_fn(*vars), lambda: orelse_fn(*vars)) + + # use pythonic placeholder until the graph compiler supports callable arguments + + if cond: + return body_fn(*vars) + else: + return orelse_fn(*vars) + + +def while_loop(test_fn, body_fn, vars): + def body_fn_wrapper(*loop_vars): + return body_fn(*loop_vars) + + def test_fn_wrapper(*loop_vars): + return test_fn(*loop_vars) + + if not vars: + vars = (0,) + + return tf.while_loop(test_fn_wrapper, body_fn_wrapper, loop_vars=vars) diff --git a/ivy/functional/backends/tensorflow/creation.py b/ivy/functional/backends/tensorflow/creation.py index 25ce085eb766d..9db7dc926d4a9 100644 --- a/ivy/functional/backends/tensorflow/creation.py +++ b/ivy/functional/backends/tensorflow/creation.py @@ -20,6 +20,18 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +# --- Main --- # +# ------------ # + + # Array API Standard # # -------------------# @@ -97,6 +109,17 @@ def asarray( return tf.identity(ret) if copy else ret +def copy_array( + x: Union[tf.Tensor, tf.Variable], + *, + to_ivy_array: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if to_ivy_array: + return ivy.to_ivy(tf.identity(x)) + return tf.identity(x) + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -181,6 +204,27 @@ def from_dlpack( return tf.experimental.dlpack.from_dlpack(dlcapsule) +@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) +def frombuffer( + buffer: bytes, + dtype: Optional[tf.DType] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(buffer, bytearray): + buffer = bytes(buffer) + ret = tf.io.decode_raw(buffer, dtype) + dtype = tf.dtypes.as_dtype(dtype) + if offset > 0: + offset = int(offset / dtype.size) + if count > -1: + ret = ret[offset : offset + count] + else: + ret = ret[offset:] + + return ret + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -210,10 +254,6 @@ def full_like( return tf.experimental.numpy.full_like(x, fill_value, dtype=dtype) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - def linspace( start: Union[tf.Tensor, tf.Variable, float], stop: Union[tf.Tensor, tf.Variable, float], @@ -265,6 +305,23 @@ def meshgrid( return res +def one_hot( + indices: Union[tf.Tensor, tf.Variable], + depth: int, + /, + *, + on_value: Optional[Number] = None, + off_value: Optional[Number] = None, + axis: Optional[int] = None, + dtype: Optional[tf.DType] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.one_hot( + indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype + ) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -310,6 +367,29 @@ def triu( return tf.experimental.numpy.triu(x, k) +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: str, +) -> Tuple[Union[tf.Tensor, tf.Variable]]: + n_cols = n_rows if n_cols is None else n_cols + + if n_rows < 0 or n_cols < 0: + n_rows, n_cols = 0, 0 + + ret = [[], []] + + for i in range(0, min(n_rows, n_cols - k), 1): + for j in range(max(0, k + i), n_cols, 1): + ret[0].append(i) + ret[1].append(j) + + return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) + + def zeros( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -336,75 +416,3 @@ def zeros_like( array = asarray - - -def copy_array( - x: Union[tf.Tensor, tf.Variable], - *, - to_ivy_array: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if to_ivy_array: - return ivy.to_ivy(tf.identity(x)) - return tf.identity(x) - - -def one_hot( - indices: Union[tf.Tensor, tf.Variable], - depth: int, - /, - *, - on_value: Optional[Number] = None, - off_value: Optional[Number] = None, - axis: Optional[int] = None, - dtype: Optional[tf.DType] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.one_hot( - indices, depth, on_value=on_value, off_value=off_value, axis=axis, dtype=dtype - ) - - -@with_unsupported_dtypes({"2.13.0 and below": ("uint32", "uint64")}, backend_version) -def frombuffer( - buffer: bytes, - dtype: Optional[tf.DType] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(buffer, bytearray): - buffer = bytes(buffer) - ret = tf.io.decode_raw(buffer, dtype) - dtype = tf.dtypes.as_dtype(dtype) - if offset > 0: - offset = int(offset / dtype.size) - if count > -1: - ret = ret[offset : offset + count] - else: - ret = ret[offset:] - - return ret - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: str, -) -> Tuple[Union[tf.Tensor, tf.Variable]]: - n_cols = n_rows if n_cols is None else n_cols - - if n_rows < 0 or n_cols < 0: - n_rows, n_cols = 0, 0 - - ret = [[], []] - - for i in range(0, min(n_rows, n_cols - k), 1): - for j in range(max(0, k + i), n_cols, 1): - ret[0].append(i) - ret[1].append(j) - - return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) diff --git a/ivy/functional/backends/tensorflow/data_type.py b/ivy/functional/backends/tensorflow/data_type.py index 819a8fc381d75..d6100609a38d8 100644 --- a/ivy/functional/backends/tensorflow/data_type.py +++ b/ivy/functional/backends/tensorflow/data_type.py @@ -28,7 +28,6 @@ tf.complex128: "complex128", tf.bool: "bool", } - native_dtype_dict = { "int8": tf.int8, "int16": tf.int16, @@ -91,93 +90,6 @@ def __repr__(self): ) -# Array API Standard # -# -------------------# - - -def astype( - x: Union[tf.Tensor, tf.Variable], - dtype: Union[DType, str], - /, - *, - copy: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return tf.experimental.numpy.copy(x) if copy else x - return tf.cast(x, dtype) - - -def broadcast_arrays( - *arrays: Union[tf.Tensor, tf.Variable], -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(arrays) > 1: - try: - desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - if len(arrays) > 2: - for i in range(2, len(arrays)): - try: - desired_shape = tf.broadcast_dynamic_shape( - desired_shape, arrays[i].shape - ) - except tf.errors.InvalidArgumentError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - else: - return [arrays[0]] - result = [] - for tensor in arrays: - result.append(tf.broadcast_to(tensor, desired_shape)) - - return result - - -def broadcast_to( - x: Union[tf.Tensor, tf.Variable], - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if tf.rank(x) > len(shape): - return tf.broadcast_to(tf.reshape(x, -1), shape) - return tf.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - if ivy.as_native_dtype(type) == tf.bfloat16: - return Finfo(Bfloat16Finfo()) - return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: - if isinstance(type, (tf.Tensor, np.ndarray)): - type = type.dtype - return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def result_type( - *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], -) -> ivy.Dtype: - if len(arrays_and_dtypes) <= 1: - return tf.experimental.numpy.result_type(arrays_and_dtypes) - - result = tf.experimental.numpy.result_type( - arrays_and_dtypes[0], arrays_and_dtypes[1] - ) - for i in range(2, len(arrays_and_dtypes)): - result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) - return as_ivy_dtype(result) - - # Extra # # ------# @@ -247,6 +159,62 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: Union[tf.Tensor, tf.Variable], + dtype: Union[DType, str], + /, + *, + copy: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return tf.experimental.numpy.copy(x) if copy else x + return tf.cast(x, dtype) + + +def broadcast_arrays( + *arrays: Union[tf.Tensor, tf.Variable], +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(arrays) > 1: + try: + desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + if len(arrays) > 2: + for i in range(2, len(arrays)): + try: + desired_shape = tf.broadcast_dynamic_shape( + desired_shape, arrays[i].shape + ) + except tf.errors.InvalidArgumentError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + else: + return [arrays[0]] + result = [] + for tensor in arrays: + result.append(tf.broadcast_to(tensor, desired_shape)) + + return result + + +def broadcast_to( + x: Union[tf.Tensor, tf.Variable], + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if tf.rank(x) > len(shape): + return tf.broadcast_to(tf.reshape(x, -1), shape) + return tf.broadcast_to(x, shape) + + def dtype( x: Union[tf.Tensor, tf.Variable, np.ndarray], *, as_native: bool = False ) -> ivy.Dtype: @@ -269,6 +237,22 @@ def dtype_bits(dtype_in: Union[tf.DType, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> Finfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + if ivy.as_native_dtype(type) == tf.bfloat16: + return Finfo(Bfloat16Finfo()) + return Finfo(tf.experimental.numpy.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[DType, str, tf.Tensor, tf.Variable, np.ndarray], /) -> np.iinfo: + if isinstance(type, (tf.Tensor, np.ndarray)): + type = type.dtype + return tf.experimental.numpy.iinfo(ivy.as_ivy_dtype(type)) + + def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -278,6 +262,16 @@ def is_native_dtype(dtype_in: Union[tf.DType, str], /) -> bool: return False -# ToDo: -# 1. result_type: Add support for bfloat16 with int16 -# 2. can_cast : Add support for complex64, complex128 +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def result_type( + *arrays_and_dtypes: Union[tf.Tensor, tf.Variable, tf.DType], +) -> ivy.Dtype: + if len(arrays_and_dtypes) <= 1: + return tf.experimental.numpy.result_type(arrays_and_dtypes) + + result = tf.experimental.numpy.result_type( + arrays_and_dtypes[0], arrays_and_dtypes[1] + ) + for i in range(2, len(arrays_and_dtypes)): + result = tf.experimental.numpy.result_type(result, arrays_and_dtypes[i]) + return as_ivy_dtype(result) diff --git a/ivy/functional/backends/tensorflow/device.py b/ivy/functional/backends/tensorflow/device.py index d0872e5e2da9d..ca7c545ecba6e 100644 --- a/ivy/functional/backends/tensorflow/device.py +++ b/ivy/functional/backends/tensorflow/device.py @@ -4,9 +4,6 @@ Collection of TensorFlow general functions, wrapped to fit Ivy syntax and signature. """ - -# global -_round = round import tensorflow as tf from typing import Union, Optional @@ -14,6 +11,33 @@ import ivy from ivy.functional.ivy.device import Profiler as BaseProfiler +# global +_round = round + + +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._options = tf.profiler.experimental.ProfilerOptions( + host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 + ) + + def start(self): + tf.profiler.experimental.start(self._save_dir, options=self._options) + + def stop(self): + tf.profiler.experimental.stop() + + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() + + +# --- Helpers --- # +# --------------- # + def _same_device(dev_a, dev_b): if dev_a is None or dev_b is None: @@ -23,34 +47,8 @@ def _same_device(dev_a, dev_b): ) -def dev( - x: Union[tf.Tensor, tf.Variable], - /, - *, - as_native: bool = False, -) -> Union[ivy.Device, str]: - dv = x.device - if as_native: - return dv - return as_ivy_dev(dv) - - -def to_device( - x: Union[tf.Tensor, tf.Variable], - device: str, - /, - *, - stream: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if device is None: - return x - device = as_native_dev(device) - current_dev = dev(x) - if not _same_device(current_dev, device): - with tf.device("/" + device.upper()): - return tf.identity(x) - return x +# --- Main --- # +# ------------ # def as_ivy_dev(device: str, /): @@ -79,47 +77,57 @@ def clear_cached_mem_on_dev(device: str, /): return None -def num_gpus() -> int: - return len(tf.config.list_physical_devices("GPU")) +def dev( + x: Union[tf.Tensor, tf.Variable], + /, + *, + as_native: bool = False, +) -> Union[ivy.Device, str]: + dv = x.device + if as_native: + return dv + return as_ivy_dev(dv) def gpu_is_available() -> bool: return len(tf.config.list_physical_devices("GPU")) > 0 -def tpu_is_available() -> bool: - try: - resolver = tf.distribute.cluster_resolver.TPUClusterResolver() - tf.config.experimental_connect_to_cluster(resolver) - tf.tpu.experimental.initialize_tpu_system(resolver) - tf.config.list_logical_devices("TPU") - tf.distribute.experimental.TPUStrategy(resolver) - return True - except ValueError: - return False - - def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): default_device = ivy.default_device(device_shifting_dev, as_native=True) with tf.device(default_device): return fn(*args, **kwargs) -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._options = tf.profiler.experimental.ProfilerOptions( - host_tracer_level=3, python_tracer_level=1, device_tracer_level=1 - ) +def num_gpus() -> int: + return len(tf.config.list_physical_devices("GPU")) - def start(self): - tf.profiler.experimental.start(self._save_dir, options=self._options) - def stop(self): - tf.profiler.experimental.stop() +def to_device( + x: Union[tf.Tensor, tf.Variable], + device: str, + /, + *, + stream: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if device is None: + return x + device = as_native_dev(device) + current_dev = dev(x) + if not _same_device(current_dev, device): + with tf.device("/" + device.upper()): + return tf.identity(x) + return x - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +def tpu_is_available() -> bool: + try: + resolver = tf.distribute.cluster_resolver.TPUClusterResolver() + tf.config.experimental_connect_to_cluster(resolver) + tf.tpu.experimental.initialize_tpu_system(resolver) + tf.config.list_logical_devices("TPU") + tf.distribute.experimental.TPUStrategy(resolver) + return True + except ValueError: + return False diff --git a/ivy/functional/backends/tensorflow/elementwise.py b/ivy/functional/backends/tensorflow/elementwise.py index 1e060e8e8cae3..d7b27fad44bc9 100644 --- a/ivy/functional/backends/tensorflow/elementwise.py +++ b/ivy/functional/backends/tensorflow/elementwise.py @@ -56,6 +56,32 @@ def add( return tf.add(x1, x2) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def angle( + input: Union[tf.Tensor, tf.Variable], + /, + *, + deg: Optional[bool] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if deg: + return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) + else: + return tf.math.angle(input, name=None) + + def asin( x: Union[tf.Tensor, tf.Variable], /, @@ -218,6 +244,16 @@ def cosh( return tf.cosh(x) +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def deg2rad( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.deg2rad(x) + + def divide( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -245,6 +281,20 @@ def equal( return tf.math.equal(x1, x2) +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def erf( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.erf(x) + + def exp( x: Union[tf.Tensor, tf.Variable], /, @@ -313,6 +363,37 @@ def fmin( return ret +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, + backend_version, +) +def fmod( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + # tf.math.floormod returns wrong results + res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) + return tf.where(x1 < 0, -res, res) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def gcd( + x1: Union[tf.Tensor, tf.Variable, int, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = promote_types_of_inputs(x1, x2) + return tf.experimental.numpy.gcd(x1, x2) + + @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def greater( x1: Union[float, tf.Tensor, tf.Variable], @@ -337,6 +418,28 @@ def greater_equal( return tf.math.greater_equal(x1, x2) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "bfloat16", + "int32", + ) + }, + backend_version, +) +def imag( + val: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.imag(val, name=None) + + def isfinite( x: Union[tf.Tensor, tf.Variable], /, @@ -387,6 +490,15 @@ def isnan( return tf.math.is_nan(x) +def isreal( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.isreal(x) + + @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) def lcm( x1: Union[tf.Tensor, tf.Variable], @@ -471,16 +583,6 @@ def logaddexp( return tf.experimental.numpy.logaddexp(x1, x2) -@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) -def real( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.real(x) - - @with_unsupported_dtypes( { "2.13.0 and below": ( @@ -555,6 +657,64 @@ def logical_xor( return tf.math.logical_xor(tf.cast(x1, tf.bool), tf.cast(x2, tf.bool)) +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def maximum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.maximum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "complex", + ) + }, + backend_version, +) +def minimum( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + use_where: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + dtype = x1.dtype + if use_where: + return tf.math.minimum(x1, x2) + x1 = tf.cast(x1, tf.float64) + x2 = tf.cast(x2, tf.float64) + return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) + + def multiply( x1: Union[float, tf.Tensor, tf.Variable], x2: Union[float, tf.Tensor, tf.Variable], @@ -566,6 +726,31 @@ def multiply( return tf.math.multiply(x1, x2) +def nan_to_num( + x: Union[tf.Tensor, tf.Variable], + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + posinf = posinf if posinf is not None else x.dtype.max + neginf = neginf if neginf is not None else x.dtype.min + posinf = tf.constant(posinf, x.dtype) + neginf = tf.constant(neginf, x.dtype) + nan = tf.constant(nan, x.dtype) + ret = tf.where(tf.math.is_nan(x), nan, x) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) + ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) + if copy: + return ret + else: + x = ret + return x + + @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) @@ -623,6 +808,49 @@ def pow( return tf.experimental.numpy.power(x1, x2) +def rad2deg( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.rad2deg(x) + + +@with_unsupported_dtypes({"2.13.0 and below": ("float16",)}, backend_version) +def real( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.real(x) + + +@with_unsupported_dtypes( + { + "2.13.0 and below": ( + "uint8", + "uint16", + "uint32", + "uint64", + "int8", + "int16", + "int32", + "int64", + ) + }, + backend_version, +) +def reciprocal( + x: Union[float, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.reciprocal(x) + + @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def remainder( x1: Union[float, tf.Tensor, tf.Variable], @@ -785,232 +1013,4 @@ def trunc( return ret -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def erf( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.erf(x) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def maximum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.maximum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 + tf.math.abs(x1 - x2)) / 2, dtype=dtype) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "complex", - ) - }, - backend_version, -) -def minimum( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - use_where: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - dtype = x1.dtype - if use_where: - return tf.math.minimum(x1, x2) - x1 = tf.cast(x1, tf.float64) - x2 = tf.cast(x2, tf.float64) - return tf.cast((x1 + x2 - tf.math.abs(x1 - x2)) / 2, dtype) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "int8", - "int16", - "int32", - "int64", - ) - }, - backend_version, -) -def reciprocal( - x: Union[float, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.reciprocal(x) - - -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def deg2rad( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.deg2rad(x) - - -def rad2deg( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rad2deg(x) - - -def isreal( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isreal(x) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64", "complex", "bool")}, - backend_version, -) -def fmod( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - # tf.math.floormod returns wrong results - res = tf.experimental.numpy.remainder(tf.math.abs(x1), tf.math.abs(x2)) - return tf.where(x1 < 0, -res, res) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def gcd( - x1: Union[tf.Tensor, tf.Variable, int, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - return tf.experimental.numpy.gcd(x1, x2) - - gcd.support_native_out = False - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def angle( - input: Union[tf.Tensor, tf.Variable], - /, - *, - deg: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if deg: - return tf.math.angle(input, name=None) * (180 / tf.experimental.numpy.pi) - else: - return tf.math.angle(input, name=None) - - -@with_unsupported_dtypes( - { - "2.13.0 and below": ( - "uint8", - "uint16", - "uint32", - "uint64", - "bfloat16", - "int32", - ) - }, - backend_version, -) -def imag( - val: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.imag(val, name=None) - - -def nan_to_num( - x: Union[tf.Tensor, tf.Variable], - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - posinf = posinf if posinf is not None else x.dtype.max - neginf = neginf if neginf is not None else x.dtype.min - posinf = tf.constant(posinf, x.dtype) - neginf = tf.constant(neginf, x.dtype) - nan = tf.constant(nan, x.dtype) - ret = tf.where(tf.math.is_nan(x), nan, x) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret > 0), posinf, ret) - ret = tf.where(tf.math.logical_and(tf.math.is_inf(ret), ret < 0), neginf, ret) - if copy: - return ret - else: - x = ret - return x diff --git a/ivy/functional/backends/tensorflow/experimental/activations.py b/ivy/functional/backends/tensorflow/experimental/activations.py index f798d7b907d98..0f5126a6678ff 100644 --- a/ivy/functional/backends/tensorflow/experimental/activations.py +++ b/ivy/functional/backends/tensorflow/experimental/activations.py @@ -10,6 +10,15 @@ from . import backend_version +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: + alpha = tf.cast(alpha, x.dtype) + ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) + + def logit( x: Union[tf.Tensor, tf.Variable], /, @@ -25,16 +34,9 @@ def logit( return tf.cast(tf.math.log(x / (1 - x)), x_dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) -def thresholded_relu( - x: Tensor, - /, - *, - threshold: Union[int, float] = 0, - out: Optional[Tensor] = None, -) -> Tensor: - threshold = tf.cast(threshold, x.dtype) - return tf.cast(tf.where(x > threshold, x, 0), x.dtype) +@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) +def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: + return tf.math.log_sigmoid(input) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -42,11 +44,6 @@ def relu6(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: return tf.nn.relu6(x) -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def logsigmoid(input: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: - return tf.math.log_sigmoid(input) - - @with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) def selu(x: Tensor, /, *, out: Optional[Tensor] = None) -> Tensor: ret = tf.nn.selu(x) @@ -68,10 +65,13 @@ def silu( return ivy.astype(ret, x.dtype) -@with_supported_dtypes({"2.13.0 and below": ("float",)}, backend_version) -def elu(x: Tensor, /, *, alpha: float = 1.0, out: Optional[Tensor] = None) -> Tensor: - alpha = tf.cast(alpha, x.dtype) - ret = tf.cast(tf.where(x > 0, x, tf.multiply(alpha, tf.math.expm1(x))), x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) +@with_unsupported_dtypes({"2.13.0 and below": ("complex", "bool")}, backend_version) +def thresholded_relu( + x: Tensor, + /, + *, + threshold: Union[int, float] = 0, + out: Optional[Tensor] = None, +) -> Tensor: + threshold = tf.cast(threshold, x.dtype) + return tf.cast(tf.where(x > threshold, x, 0), x.dtype) diff --git a/ivy/functional/backends/tensorflow/experimental/creation.py b/ivy/functional/backends/tensorflow/experimental/creation.py index 96cd0c15b84be..4462e9d61977e 100644 --- a/ivy/functional/backends/tensorflow/experimental/creation.py +++ b/ivy/functional/backends/tensorflow/experimental/creation.py @@ -8,28 +8,40 @@ from .. import backend_version -# Array API Standard # -# -------------------# +def blackman_window( + size: int, + /, + *, + periodic: bool = True, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if size < 2: + return tf.ones([size], dtype=tf.result_type(size, 0.0)) + if periodic: + count = tf.arange(size) / size + else: + count = tf.linspace(start=0, stop=size, num=size) + return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( + 0.08 * tf.cos(2 * tf.pi * 2 * count) + ) -@with_unsupported_device_and_dtypes( - {"2.13.0 and below": {"cpu": ("bfloat16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, + +def hann_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if window_length < 2: - return tf.ones([window_length], dtype=dtype) - if periodic is False: - return tf.signal.kaiser_window(window_length, beta, dtype=dtype) + if size < 2: + return tf.ones([size], dtype=dtype) + if periodic: + return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] else: - return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] + return tf.signal.hann_window(size, periodic=False, dtype=dtype) def kaiser_bessel_derived_window( @@ -42,29 +54,28 @@ def kaiser_bessel_derived_window( return tf.signal.kaiser_bessel_derived_window(window_length, beta, dtype) -def vorbis_window( - window_length: Union[tf.Tensor, tf.Variable], - *, - dtype: tf.DType = tf.dtypes.float32, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) +# Array API Standard # +# -------------------# -def hann_window( - size: int, - /, - *, +@with_unsupported_device_and_dtypes( + {"2.13.0 and below": {"cpu": ("bfloat16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, periodic: bool = True, + beta: float = 12.0, + *, dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=dtype) - if periodic: - return tf.signal.hann_window(size + 1, periodic=False, dtype=dtype)[:-1] + if window_length < 2: + return tf.ones([window_length], dtype=dtype) + if periodic is False: + return tf.signal.kaiser_window(window_length, beta, dtype=dtype) else: - return tf.signal.hann_window(size, periodic=False, dtype=dtype) + return tf.signal.kaiser_window(window_length + 1, beta, dtype=dtype)[:-1] def tril_indices( @@ -90,6 +101,20 @@ def tril_indices( return tuple(tf.convert_to_tensor(ret, dtype=tf.int64)) +@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) +def trilu( + x: Union[tf.Tensor, tf.Variable], + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if upper: + return tf.experimental.numpy.triu(x, k) + return tf.experimental.numpy.tril(x, k) + + def unsorted_segment_min( data: tf.Tensor, segment_ids: tf.Tensor, @@ -98,26 +123,6 @@ def unsorted_segment_min( return tf.math.unsorted_segment_min(data, segment_ids, num_segments) -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if size < 2: - return tf.ones([size], dtype=tf.result_type(size, 0.0)) - if periodic: - count = tf.arange(size) / size - else: - count = tf.linspace(start=0, stop=size, num=size) - - return (0.42 - 0.5 * tf.cos(2 * tf.pi * count)) + ( - 0.08 * tf.cos(2 * tf.pi * 2 * count) - ) - - def unsorted_segment_sum( data: tf.Tensor, segment_ids: tf.Tensor, @@ -126,15 +131,10 @@ def unsorted_segment_sum( return tf.math.unsorted_segment_sum(data, segment_ids, num_segments) -@with_unsupported_dtypes({"2.13.0 and below": ("bool",)}, backend_version) -def trilu( - x: Union[tf.Tensor, tf.Variable], - /, +def vorbis_window( + window_length: Union[tf.Tensor, tf.Variable], *, - k: int = 0, - upper: bool = True, + dtype: tf.DType = tf.dtypes.float32, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if upper: - return tf.experimental.numpy.triu(x, k) - return tf.experimental.numpy.tril(x, k) + return tf.signal.vorbis_window(window_length, dtype=dtype, name=None) diff --git a/ivy/functional/backends/tensorflow/experimental/elementwise.py b/ivy/functional/backends/tensorflow/experimental/elementwise.py index 96dba1c88a4dd..487ed04c624b4 100644 --- a/ivy/functional/backends/tensorflow/experimental/elementwise.py +++ b/ivy/functional/backends/tensorflow/experimental/elementwise.py @@ -11,61 +11,54 @@ from .. import backend_version -@with_supported_dtypes( - {"2.13.0 and below": ("float16", "float32", "float64")}, - backend_version, -) -def lgamma( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.math.lgamma(x) +# --- Helpers --- # +# --------------- # -def sinc( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x = ivy.pi * x - return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) +def _normalize_axis_index(ax: int, ndim: int) -> int: + if ax >= ndim or ax < -ndim: + raise ValueError("axis index is out of range") + return (ax + ndim) % ndim -@with_supported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version -) -def fmax( +def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) + if len(set(axis)) != len(axis): + raise ValueError("repeated axis") + return axis + + +# --- Main --- # +# ------------ # + + +def allclose( x1: Union[tf.Tensor, tf.Variable], x2: Union[tf.Tensor, tf.Variable], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = promote_types_of_inputs(x1, x2) - x1 = tf.where(tf.math.is_nan(x1), x2, x1) - x2 = tf.where(tf.math.is_nan(x2), x1, x2) - return tf.experimental.numpy.maximum(x1, x2) +) -> bool: + return tf.experimental.numpy.allclose( + x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan + ) -@with_unsupported_dtypes( - {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version -) -def float_power( - x1: Union[tf.Tensor, tf.Variable, float, list, tuple], - x2: Union[tf.Tensor, tf.Variable, float, list, tuple], +def conj( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): - out_dtype = tf.complex128 - else: - out_dtype = tf.float64 - return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) + return tf.math.conj(x) def copysign( @@ -103,67 +96,35 @@ def count_nonzero( ) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def nansum( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[tf.DType] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) - - -def isclose( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan - ) - - -def signbit( - x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], +@with_unsupported_dtypes( + {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version +) +def diff( + x: Union[tf.Tensor, tf.Variable, list, tuple], /, *, + n: int = 1, + axis: int = -1, + prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, + append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.signbit(x) + if n == 0: + return x + if prepend is not None: + x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) + if append is not None: + x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) + return tf.experimental.numpy.diff(x, n=n, axis=axis) -def hypot( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], +def digamma( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) - - -def allclose( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> bool: - return tf.experimental.numpy.allclose( - x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan - ) + return tf.math.digamma(x) @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) @@ -176,74 +137,56 @@ def fix( return tf.cast(tf.where(x > 0, tf.math.floor(x), tf.math.ceil(x)), x.dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) -def nextafter( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.nextafter(x1, x2) - - @with_unsupported_dtypes( {"2.13.0 and below": ("uint8", "uint16", "uint32", "uint64")}, backend_version ) -def diff( - x: Union[tf.Tensor, tf.Variable, list, tuple], +def float_power( + x1: Union[tf.Tensor, tf.Variable, float, list, tuple], + x2: Union[tf.Tensor, tf.Variable, float, list, tuple], /, *, - n: int = 1, - axis: int = -1, - prepend: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, - append: Optional[Union[tf.Tensor, tf.Variable, int, float, list, tuple]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if n == 0: - return x - if prepend is not None: - x = tf.experimental.numpy.append(prepend, x, axis=axis if axis != -1 else None) - if append is not None: - x = tf.experimental.numpy.append(x, append, axis=axis if axis != -1 else None) - return tf.experimental.numpy.diff(x, n=n, axis=axis) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)): + out_dtype = tf.complex128 + else: + out_dtype = tf.float64 + return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype) @with_supported_dtypes( - { - "2.13.0 and below": ( - "float32", - "float64", - ) - }, - backend_version, + {"2.13.0 and below": ("bfloat16", "float16", "float32", "float64")}, backend_version ) -def zeta( - x: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, tf.Variable], +def fmax( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.zeta(x, q) - - -def _normalize_axis_index(ax: int, ndim: int) -> int: - if ax >= ndim or ax < -ndim: - raise ValueError("axis index is out of range") - return (ax + ndim) % ndim + x1, x2 = promote_types_of_inputs(x1, x2) + x1 = tf.where(tf.math.is_nan(x1), x2, x1) + x2 = tf.where(tf.math.is_nan(x2), x1, x2) + return tf.experimental.numpy.maximum(x1, x2) -def _normalize_axis_tuple(axis: Union[int, list, tuple], ndim: int) -> Tuple[int, ...]: - if type(axis) not in (tuple, list): - try: - axis = [operator.index(axis)] - except TypeError: - pass - axis = tuple([_normalize_axis_index(ax, ndim) for ax in axis]) - if len(set(axis)) != len(axis): - raise ValueError("repeated axis") - return axis +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) +def frexp( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[ + Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] + ] = None, +) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: + e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) + e = tf.cast(e, x.dtype) + while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): + e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) + m = x / tf.math.pow(2, e) + e = tf.cast(e, tf.int32) + return m, e def gradient( @@ -428,36 +371,29 @@ def gradient( return outvals -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float16", - "float32", - "float64", - "complex64", - "complex128", - ) - }, - backend_version, -) -def xlogy( - x: Union[tf.Tensor, tf.Variable], - y: Union[tf.Tensor, tf.Variable], +def hypot( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - x, y = promote_types_of_inputs(x, y) - return tf.math.xlogy(x, y) + return tf.math.sqrt(tf.math.square(x1) + tf.math.square(x2)) -def conj( - x: Union[tf.Tensor, tf.Variable], +def isclose( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.conj(x) + return tf.experimental.numpy.isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan + ) @with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) @@ -481,22 +417,17 @@ def ldexp( return tf.cast(ret, out_dtype) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned",)}, backend_version) -def frexp( +@with_supported_dtypes( + {"2.13.0 and below": ("float16", "float32", "float64")}, + backend_version, +) +def lgamma( x: Union[tf.Tensor, tf.Variable], /, *, - out: Optional[ - Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]] - ] = None, -) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]: - e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype)) - e = tf.cast(e, x.dtype) - while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1): - e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype) - m = x / tf.math.pow(2, e) - e = tf.cast(e, tf.int32) - return m, e + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.lgamma(x) def modf( @@ -508,10 +439,87 @@ def modf( return tf.math.modf(x) -def digamma( +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def nansum( x: Union[tf.Tensor, tf.Variable], /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[tf.DType] = None, + keepdims: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.digamma(x) + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nansum(x, axis=axis, dtype=dtype, keepdims=keepdims) + + +@with_unsupported_dtypes({"2.13.0 and below": ("bflaot16", "float16")}, backend_version) +def nextafter( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.nextafter(x1, x2) + + +def signbit( + x: Union[tf.Tensor, tf.Variable, float, int, list, tuple], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.signbit(x) + + +def sinc( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x = ivy.pi * x + return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float16", + "float32", + "float64", + "complex64", + "complex128", + ) + }, + backend_version, +) +def xlogy( + x: Union[tf.Tensor, tf.Variable], + y: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + x, y = promote_types_of_inputs(x, y) + return tf.math.xlogy(x, y) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float32", + "float64", + ) + }, + backend_version, +) +def zeta( + x: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.math.zeta(x, q) diff --git a/ivy/functional/backends/tensorflow/experimental/layers.py b/ivy/functional/backends/tensorflow/experimental/layers.py index 94c60963871b5..2a528a74cdf3f 100644 --- a/ivy/functional/backends/tensorflow/experimental/layers.py +++ b/ivy/functional/backends/tensorflow/experimental/layers.py @@ -20,6 +20,10 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode, _get_size +# --- Helpers --- # +# --------------- # + + def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_last"): # Determine depth pooling kernel, strides, depth_pooling = _depth_max_pooling_helper( @@ -30,229 +34,71 @@ def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_ return x, kernel, strides, depth_pooling -def max_pool1d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) +def _fft2_helper(x, shape, axes): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) - if data_format == "NCW": - x = tf.transpose(x, (0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor ) - if not depth_pooling: - new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - padding = [(pad_w // 2, pad_w - pad_w // 2)] + shape = shape_initialization(shape, axes, x) - if ceil_mode: - padding[0] = _padding_ceil_mode( - x.shape[1], new_kernel[0], padding[0], strides[0] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) + rank = rank_initialization(axes) - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - if depth_pooling: - res = tf.transpose(res, (0, 2, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCW": - return tf.transpose(res, (0, 2, 1)) - return res + perm = get_perm(input_rank_tensor, axes) + x = transpose_x(x, perm, perform_transpose) -def max_pool2d( - x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + x = fft2_operations(x, rank) - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) + x = transpose_x(x, tf.argsort(perm), perform_transpose) - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - if not depth_pooling: - new_kernel = [ - kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) - ] - if isinstance(padding, str): - pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) - padding = [ - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] + return x - x_shape = x.shape[1:-1] - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) +def _fft2_norm( + x: Union[tf.Tensor, tf.Variable], + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", +): + n = tf.constant(s[0] * s[1], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.sqrt(n) + elif norm == "forward": + return x / n else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def max_pool3d( +def _fft_norm( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + dim: int, /, *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NCDHW": - x = tf.transpose(x, (0, 2, 3, 4, 1)) - kernel = ( - [kernel[i] for i in [0, 2, 3, 4, 1]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 2, 3, 4, 1]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 2, 3, 4, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_last" - ) - - if not depth_pooling: - x_shape = x.shape[1:-1] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - padding = [ - (pad_d // 2, pad_d - pad_d // 2), - (pad_h // 2, pad_h - pad_h // 2), - (pad_w // 2, pad_w - pad_w // 2), - ] - - if ceil_mode: - for i in range(dims): - padding[i] = _padding_ceil_mode( - x_shape[i], new_kernel[i], padding[i], strides[i] - ) - padding = [(0, 0)] + list(padding) + [(0, 0)] - x = tf.pad(x, padding, constant_values=-math.inf) + norm: str = "backward", +): + n = tf.constant(x.shape[dim], dtype=x.dtype) + if norm == "backward": + return x + elif norm == "ortho": + return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) + elif norm == "forward": + return x / tf.cast(n, x.dtype) else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - - if depth_pooling: - res = tf.transpose(res, (0, 2, 3, 4, 1)) - # converting minimum value to -inf because tensorflow clips -inf to minimum value - res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) - if data_format == "NCDHW": - return tf.transpose(res, (0, 4, 1, 2, 3)) - return res + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): @@ -280,13 +126,123 @@ def _handle_manual_pad_avg_pool(x, kernel, strides, padding, ceil_mode, dims): return padding, pad_specific, c -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) -def avg_pool1d( +def _ifft_norm( x: Union[tf.Tensor, tf.Variable], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, + dim: int, + *, + norm: str = "backward", +): + n = x.shape[dim] + if norm == "backward": + return x + elif norm == "ortho": + return x * math.sqrt(n) + elif norm == "forward": + return x * n + else: + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + + +def _ifftn_helper(x, shape, axes, norm): + x = fft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + norm_factor = norm_initialization(norm, shape, x) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = ifft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +def _rfftn_helper(x, shape, axes, norm): + x = rfft_input_validation(tf.convert_to_tensor(x)) + input_shape = x.shape + input_rank_tensor = tf.rank(x) + + shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) + + axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) + + perform_padding, perform_transpose = perform_actions_initialization( + shape, axes, input_shape, input_rank_tensor + ) + + shape = shape_initialization(shape, axes, x) + + rank = rank_initialization(axes) + + norm_factor = norm_initialization(norm, shape, x) + + x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + + perm = get_perm(input_rank_tensor, axes) + + x = transpose_x(x, perm, perform_transpose) + + x = rfft_operations(x, rank, norm_factor) + + x = transpose_x(x, tf.argsort(perm), perform_transpose) + + x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + + return x + + +def _right_pad_or_crop(tensor, shape): + input_shape = tf.shape(tensor) + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + with tf.control_dependencies( + [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] + ): + shape = tf.identity(shape) + shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) + + pad_sizes = tf.math.maximum(shape - input_shape, 0) + pad_sizes = tf.expand_dims(pad_sizes, -1) + pad_sizes = tf.concat( + [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 + ) + tensor = tf.pad(tensor, pad_sizes, constant_values=0) + + crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) + tensor = tf.slice(tensor, crop_tensor, shape) + return tensor + + +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "float64")}, backend_version) +def avg_pool1d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, *, data_format: str = "NWC", count_include_pad: bool = False, @@ -547,31 +503,15 @@ def avg_pool3d( return res -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version -) -def pool( - x: Union[tf.Tensor, tf.Variable], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, - *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.nn.pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ) +def axes_initialization(shape, axes, input_shape, input_rank_tensor): + if axes is None: + axes = ( + tf.range(-tf.size(input_shape), 0) + if shape is None + else tf.range(-tf.size(shape), 0) + ) + axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) + return axes @with_supported_dtypes({"2.13.0 and below": ("float32", "float64")}, backend_version) @@ -599,112 +539,6 @@ def dct( return dct_out -def idct( - x: Union[tf.Tensor, tf.Variable], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] - return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) - - -def _fft_norm( - x: Union[tf.Tensor, tf.Variable], - dim: int, - /, - *, - norm: str = "backward", -): - n = tf.constant(x.shape[dim], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype) - elif norm == "forward": - return x / tf.cast(n, x.dtype) - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -def _ifft_norm( - x: Union[tf.Tensor, tf.Variable], - dim: int, - *, - norm: str = "backward", -): - n = x.shape[dim] - if norm == "backward": - return x - elif norm == "ortho": - return x * math.sqrt(n) - elif norm == "forward": - return x * n - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft( - x: Union[tf.Tensor, tf.Variable], - dim: int, - /, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" - ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" - ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.shape[dim] != n: - s = list(x.shape) - if s[dim] > n: - index = [slice(None)] * len(s) - index[dim] = slice(0, n) - x = x[tuple(index)] - del index - else: - s[dim] = n - s[dim] - z = tf.zeros(s, x.dtype) - x = tf.concat([x, z], dim) - del s - operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" - if dim != -1 or dim != len(x.shape) - 1: - permute = [i for i in range(len(x.shape))] - permute[dim], permute[-1] = permute[-1], permute[dim] - x = tf.transpose(x, permute) - ret = tf.signal.fft(x, operation_name) - ret = tf.transpose(ret, permute) - del permute - else: - ret = tf.signal.fft(x, operation_name) - ret = _fft_norm(ret, dim, norm=norm) - return ret - - def dropout( x: Union[tf.Tensor, tf.Variable], prob: float, @@ -791,10 +625,181 @@ def dropout3d( return res -def ifft( - x: Union[tf.Tensor, tf.Variable], - dim: int, - *, +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def embedding( + weights: Union[tf.Tensor, tf.Variable], + indices: Union[tf.Tensor, tf.Variable], + /, + *, + max_norm: Optional[float] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft( + x: Union[tf.Tensor, tf.Variable], + dim: int, + /, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if x.shape[dim] != n: + s = list(x.shape) + if s[dim] > n: + index = [slice(None)] * len(s) + index[dim] = slice(0, n) + x = x[tuple(index)] + del index + else: + s[dim] = n - s[dim] + z = tf.zeros(s, x.dtype) + x = tf.concat([x, z], dim) + del s + operation_name = f"{n} points FFT at dim {dim} with {norm} normalization" + if dim != -1 or dim != len(x.shape) - 1: + permute = [i for i in range(len(x.shape))] + permute[dim], permute[-1] = permute[-1], permute[dim] + x = tf.transpose(x, permute) + ret = tf.signal.fft(x, operation_name) + ret = tf.transpose(ret, permute) + del permute + else: + ret = tf.signal.fft(x, operation_name) + ret = _fft_norm(ret, dim, norm=norm) + return ret + + +@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def fft2( + x: Union[tf.Tensor, tf.Variable], + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if len(x.shape) > 2: + result = _fft2_helper(x, s, dim) + else: + x_new = trans_x_to_s(x, s, dim) + x_complex = tf.cast(x_new, tf.complex128) + result = tf.signal.fft2d(x_complex) + + result = _fft2_norm(result, s, dim, norm) + if x.dtype == tf.complex64: + result = tf.cast(result, dtype=tf.complex128) + return result + + +def fft2_operations(x, rank): + if x.shape.rank == 1: + x = tf.signal.fft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} + ) + else: + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.fft(x), + 1: lambda: tf.signal.fft2d(x), + 2: lambda: tf.signal.fft3d(x), + }, + ) + return x + + +# --- IFFTN --- # +def fft_input_validation(x): + if not x.dtype.is_complex: + raise TypeError( + "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( + x.dtype + ) + ) + return x + + +def get_perm(input_rank_tensor, axes): + all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) + perm = tf.concat( + [ + tf.boolean_mask( + all_dims, + tf.foldl( + lambda acc, elem: tf.math.logical_and( + acc, tf.math.not_equal(all_dims, elem) + ), + axes, + initializer=tf.fill(all_dims.shape, True), + ), + ), + axes, + ], + 0, + ) + return perm + + +def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): + if perform_padding: + pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) + pad_shape = tf.tensor_scatter_nd_update( + pad_shape, tf.expand_dims(axes, -1), shape + ) + x = _right_pad_or_crop(x, pad_shape) + return x + + +def idct( + x: Union[tf.Tensor, tf.Variable], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> tf.Tensor: + inverse_type = {1: 1, 2: 3, 3: 2, 4: 4}[type] + return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) + + +def ifft( + x: Union[tf.Tensor, tf.Variable], + dim: int, + *, norm: str = "backward", n: Union[int, Tuple[int]] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, @@ -846,19 +851,41 @@ def ifft( return ret -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def embedding( - weights: Union[tf.Tensor, tf.Variable], - indices: Union[tf.Tensor, tf.Variable], - /, +def ifft_operations(x, rank, norm_factor): + if x.shape.rank == 1: + x = tf.signal.ifft(x) + elif x.shape.rank == 2: + x = tf.switch_case( + rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} + ) + else: + x = tf.switch_case( + rank - 1, + { + 0: lambda: tf.signal.ifft(x), + 1: lambda: tf.signal.ifft2d(x), + 2: lambda: tf.signal.ifft3d(x), + }, + ) + x = x * norm_factor + return x + + +def ifftn( + x: Union[tf.Tensor, tf.Variable], + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - max_norm: Optional[float] = None, + norm: Optional[str] = "backward", out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return tf.nn.embedding_lookup(weights, indices, max_norm=max_norm) + result = _ifftn_helper(x, s, axes, norm) + + if out is not None: + out = result + return out + else: + return result def interpolate( @@ -919,417 +946,292 @@ def interpolate( return ret -interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 - (not align_corners and (len(x.shape) - 2) < 2) - and mode not in ["nearest", "area", "bicubic", "nd"] -) - - -def _fft2_norm( - x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", -): - n = tf.constant(s[0] * s[1], dtype=x.dtype) - if norm == "backward": - return x - elif norm == "ortho": - return x / tf.sqrt(n) - elif norm == "forward": - return x / n - else: - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - - -def trans_x_to_s( +def max_pool1d( x: Union[tf.Tensor, tf.Variable], - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - """Change the shape of the input array x to the desired output shape s.""" - if x.dtype != tf.complex128 and x.dtype != tf.complex64: - x = tf.cast(x, tf.float32) - x_shape = x.shape - if dim == (-1, -2) or dim == (1, 0): - s = (s[1], s[0]) - if s[0] >= x_shape[0] and s[1] >= x_shape[1]: - paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) - x_new = tf.pad(x, paddings=paddings) - elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): - x_new = x[: s[0], : s[1]] - if s[0] != x_new.shape[0]: - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif s[1] != x_new.shape[1]: - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], 1) - elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[0] - x_new.shape[0] - z = tf.zeros((size, s[1]), dtype=x.dtype) - x_new = tf.concat([x_new, z], 0) - elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): - x_new = x[: s[0], : s[1]] - size = s[1] - x_new.shape[1] - z = tf.zeros((s[0], size), dtype=x.dtype) - x_new = tf.concat([x_new, z], axis=1) - else: - x_new = x[: s[0], : s[1]] - return x_new - + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -def fft2_operations(x, rank): - if x.shape.rank == 1: - x = tf.signal.fft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.fft(x), 1: lambda: tf.signal.fft2d(x)} + if data_format == "NCW": + x = tf.transpose(x, (0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.fft(x), - 1: lambda: tf.signal.fft2d(x), - 2: lambda: tf.signal.fft3d(x), - }, + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - return x - -def _fft2_helper(x, shape, axes): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = fft2_operations(x, rank) + if not depth_pooling: + new_kernel = [kernel[0] + (kernel[0] - 1) * (dilation[0] - 1)] + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + padding = [(pad_w // 2, pad_w - pad_w // 2)] - x = transpose_x(x, tf.argsort(perm), perform_transpose) + if ceil_mode: + padding[0] = _padding_ceil_mode( + x.shape[1], new_kernel[0], padding[0], strides[0] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - return x + if depth_pooling: + res = tf.transpose(res, (0, 2, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCW": + return tf.transpose(res, (0, 2, 1)) + return res -@with_supported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def fft2( +def max_pool2d( x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if len(x.shape) > 2: - result = _fft2_helper(x, s, dim) - else: - x_new = trans_x_to_s(x, s, dim) - x_complex = tf.cast(x_new, tf.complex128) - result = tf.signal.fft2d(x_complex) - - result = _fft2_norm(result, s, dim, norm) - if x.dtype == tf.complex64: - result = tf.cast(result, dtype=tf.complex128) - return result - + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -# --- IFFTN --- # -def fft_input_validation(x): - if not x.dtype.is_complex: - raise TypeError( - "Invalid FFT input: `x` must be of a complex dtype. Received: {}".format( - x.dtype - ) + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 1]] if len(kernel) == (dims + 2) else kernel ) - return x - - -def shape_and_axes_validation(shape, axes, input_rank_tensor): - if shape is not None: - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - checks_shape = [ - tf.debugging.assert_less_equal( - tf.size(shape), - input_rank_tensor, - message=( - "Argument `shape` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(shape), - ) - ] - with tf.control_dependencies(checks_shape): - shape = tf.identity(shape) - - if axes is not None: - axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) - checks_axes = [ - tf.debugging.assert_less_equal( - tf.size(axes), - input_rank_tensor, - message=( - "Argument `axes` cannot have length greater than the rank of `x`. " - "Received: {}" - ).format(axes), - ), - tf.debugging.assert_less( - axes, - input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - tf.debugging.assert_greater_equal( - axes, - -input_rank_tensor, - message=( - "Argument `axes` contains invalid indices. Received: {}" - ).format(axes), - ), - ] - with tf.control_dependencies(checks_axes): - axes = tf.identity(axes) - - if shape is not None and axes is not None: - checks_shape_axes = [ - tf.debugging.assert_equal( - tf.size(shape), - tf.size(axes), - message=( - "Arguments `shape` and `axes` must have equal length. " - "Received: {}, {}" - ).format(shape, axes), - ) - ] - with tf.control_dependencies(checks_shape_axes): - shape, axes = tf.identity_n([shape, axes]) - - return shape, axes - - -def axes_initialization(shape, axes, input_shape, input_rank_tensor): - if axes is None: - axes = ( - tf.range(-tf.size(input_shape), 0) - if shape is None - else tf.range(-tf.size(shape), 0) + strides = ( + [strides[i] for i in [0, 2, 3, 1]] + if len(strides) == (dims + 2) + else strides ) - axes = tf.where(tf.math.less(axes, 0), axes + input_rank_tensor, axes) - return axes - - -def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): - perform_padding = shape is not None - perform_transpose = tf.math.logical_not( - tf.math.reduce_all( - tf.math.equal( - axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) - ) + padding = ( + [padding[i] for i in [0, 2, 3, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - ) - return perform_padding, perform_transpose - - -def shape_initialization(shape, axes, x): - if shape is None: - shape = tf.gather(tf.shape(x), axes, axis=0) - return shape + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" + ) -def rank_initialization(axes): - rank = tf.size(axes) - with tf.control_dependencies( - [ - tf.debugging.assert_less_equal( - rank, 3, message="N-D FFT supported only up to 3-D." - ) + if not depth_pooling: + new_kernel = [ + kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(dims) ] - ): - rank = tf.identity(rank) - - return rank + if isinstance(padding, str): + pad_h = _handle_padding(x.shape[1], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x.shape[2], strides[1], new_kernel[1], padding) + padding = [ + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] + x_shape = x.shape[1:-1] -def norm_initialization(norm, shape, x): - if norm == "backward": - norm_factor = tf.constant(1, x.dtype) - elif norm == "forward" or norm == "ortho": - norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) - if norm == "ortho": - norm_factor = tf.math.sqrt(norm_factor) - return norm_factor + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) -def get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor): - if perform_padding: - pad_shape = -tf.ones([input_rank_tensor], dtype=tf.int32) - pad_shape = tf.tensor_scatter_nd_update( - pad_shape, tf.expand_dims(axes, -1), shape - ) - x = _right_pad_or_crop(x, pad_shape) - return x + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res -def get_perm(input_rank_tensor, axes): - all_dims = tf.range(input_rank_tensor, dtype=tf.dtypes.int32) - perm = tf.concat( - [ - tf.boolean_mask( - all_dims, - tf.foldl( - lambda acc, elem: tf.math.logical_and( - acc, tf.math.not_equal(all_dims, elem) - ), - axes, - initializer=tf.fill(all_dims.shape, True), - ), - ), - axes, - ], - 0, +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def max_pool3d( + x: Union[tf.Tensor, tf.Variable], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims ) - return perm - -def ifft_operations(x, rank, norm_factor): - if x.shape.rank == 1: - x = tf.signal.ifft(x) - elif x.shape.rank == 2: - x = tf.switch_case( - rank - 1, {0: lambda: tf.signal.ifft(x), 1: lambda: tf.signal.ifft2d(x)} + if data_format == "NCDHW": + x = tf.transpose(x, (0, 2, 3, 4, 1)) + kernel = ( + [kernel[i] for i in [0, 2, 3, 4, 1]] + if len(kernel) == (dims + 2) + else kernel ) - else: - x = tf.switch_case( - rank - 1, - { - 0: lambda: tf.signal.ifft(x), - 1: lambda: tf.signal.ifft2d(x), - 2: lambda: tf.signal.ifft3d(x), - }, + strides = ( + [strides[i] for i in [0, 2, 3, 4, 1]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 2, 3, 4, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - x = x * norm_factor - return x - - -def transpose_x(x, perm, perform_transpose): - x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) - return x - - -def static_output_shape(input_shape, shape, axes): - output_shape = input_shape.as_list() - if shape is not None: - if axes is None: - axes = list(range(-len(shape), 0)) - if isinstance(shape, tf.Tensor): - if isinstance(axes, tf.Tensor): - output_shape = [None] * len(output_shape) - else: - for ax in axes: - output_shape[ax] = None - else: - for idx, ax in enumerate(axes): - output_shape[ax] = shape[idx] - return tf.TensorShape(output_shape) - - -def _right_pad_or_crop(tensor, shape): - input_shape = tf.shape(tensor) - shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) - with tf.control_dependencies( - [tf.debugging.assert_less_equal(tf.size(shape), tf.size(input_shape))] - ): - shape = tf.identity(shape) - shape = tf.concat([input_shape[: tf.size(input_shape) - tf.size(shape)], shape], 0) - - pad_sizes = tf.math.maximum(shape - input_shape, 0) - pad_sizes = tf.expand_dims(pad_sizes, -1) - pad_sizes = tf.concat( - [tf.zeros(pad_sizes.shape, dtype=tf.dtypes.int32), pad_sizes], -1 - ) - tensor = tf.pad(tensor, pad_sizes, constant_values=0) - - crop_tensor = tf.zeros(shape.shape, dtype=tf.dtypes.int32) - tensor = tf.slice(tensor, crop_tensor, shape) - return tensor - - -def _ifftn_helper(x, shape, axes, norm): - x = fft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_last" ) - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) + if not depth_pooling: + x_shape = x.shape[1:-1] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + padding = [ + (pad_d // 2, pad_d - pad_d // 2), + (pad_h // 2, pad_h - pad_h // 2), + (pad_w // 2, pad_w - pad_w // 2), + ] - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) + if ceil_mode: + for i in range(dims): + padding[i] = _padding_ceil_mode( + x_shape[i], new_kernel[i], padding[i], strides[i] + ) + padding = [(0, 0)] + list(padding) + [(0, 0)] + x = tf.pad(x, padding, constant_values=-math.inf) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) - perm = get_perm(input_rank_tensor, axes) + res = tf.nn.pool(x, kernel, "MAX", strides, "VALID", dilations=dilation) - x = transpose_x(x, perm, perform_transpose) + if depth_pooling: + res = tf.transpose(res, (0, 2, 3, 4, 1)) + # converting minimum value to -inf because tensorflow clips -inf to minimum value + res = tf.where(res <= ivy.finfo(res.dtype).min, -math.inf, res) + if data_format == "NCDHW": + return tf.transpose(res, (0, 4, 1, 2, 3)) + return res - x = ifft_operations(x, rank, norm_factor) - x = transpose_x(x, tf.argsort(perm), perform_transpose) +def norm_initialization(norm, shape, x): + if norm == "backward": + norm_factor = tf.constant(1, x.dtype) + elif norm == "forward" or norm == "ortho": + norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype) + if norm == "ortho": + norm_factor = tf.math.sqrt(norm_factor) + return norm_factor - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - return x +def perform_actions_initialization(shape, axes, input_shape, input_rank_tensor): + perform_padding = shape is not None + perform_transpose = tf.math.logical_not( + tf.math.reduce_all( + tf.math.equal( + axes, tf.range(input_rank_tensor - tf.size(axes), input_rank_tensor) + ) + ) + ) + return perform_padding, perform_transpose -def ifftn( +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float64", "float16")}, backend_version +) +def pool( x: Union[tf.Tensor, tf.Variable], - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, *, - norm: Optional[str] = "backward", + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - result = _ifftn_helper(x, s, axes, norm) + return tf.nn.pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ) - if out is not None: - out = result - return out - else: - return result +def rank_initialization(axes): + rank = tf.size(axes) + with tf.control_dependencies( + [ + tf.debugging.assert_less_equal( + rank, 3, message="N-D FFT supported only up to 3-D." + ) + ] + ): + rank = tf.identity(rank) -""" -RFFTN Function -""" + return rank def rfft_input_validation(x): @@ -1364,40 +1266,6 @@ def rfft_operations(x, rank, norm_factor): return x -def _rfftn_helper(x, shape, axes, norm): - x = rfft_input_validation(tf.convert_to_tensor(x)) - input_shape = x.shape - input_rank_tensor = tf.rank(x) - - shape_, axes_ = shape_and_axes_validation(shape, axes, input_rank_tensor) - - axes = axes_initialization(shape, axes, input_shape, input_rank_tensor) - - perform_padding, perform_transpose = perform_actions_initialization( - shape, axes, input_shape, input_rank_tensor - ) - - shape = shape_initialization(shape, axes, x) - - rank = rank_initialization(axes) - - norm_factor = norm_initialization(norm, shape, x) - - x = get_x_after_pad_or_crop(x, shape, axes, perform_padding, input_rank_tensor) - - perm = get_perm(input_rank_tensor, axes) - - x = transpose_x(x, perm, perform_transpose) - - x = rfft_operations(x, rank, norm_factor) - - x = transpose_x(x, tf.argsort(perm), perform_transpose) - - x = tf.ensure_shape(x, static_output_shape(input_shape, shape_, axes_)) - - return x - - @with_supported_device_and_dtypes( { "2.5.0 and above": { @@ -1427,3 +1295,139 @@ def rfftn( else: # return result return tf.cast(result, tf.complex128) + + +def shape_and_axes_validation(shape, axes, input_rank_tensor): + if shape is not None: + shape = tf.convert_to_tensor(shape, dtype=tf.dtypes.int32) + checks_shape = [ + tf.debugging.assert_less_equal( + tf.size(shape), + input_rank_tensor, + message=( + "Argument `shape` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(shape), + ) + ] + with tf.control_dependencies(checks_shape): + shape = tf.identity(shape) + + if axes is not None: + axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32) + checks_axes = [ + tf.debugging.assert_less_equal( + tf.size(axes), + input_rank_tensor, + message=( + "Argument `axes` cannot have length greater than the rank of `x`. " + "Received: {}" + ).format(axes), + ), + tf.debugging.assert_less( + axes, + input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + tf.debugging.assert_greater_equal( + axes, + -input_rank_tensor, + message=( + "Argument `axes` contains invalid indices. Received: {}" + ).format(axes), + ), + ] + with tf.control_dependencies(checks_axes): + axes = tf.identity(axes) + + if shape is not None and axes is not None: + checks_shape_axes = [ + tf.debugging.assert_equal( + tf.size(shape), + tf.size(axes), + message=( + "Arguments `shape` and `axes` must have equal length. " + "Received: {}, {}" + ).format(shape, axes), + ) + ] + with tf.control_dependencies(checks_shape_axes): + shape, axes = tf.identity_n([shape, axes]) + + return shape, axes + + +def shape_initialization(shape, axes, x): + if shape is None: + shape = tf.gather(tf.shape(x), axes, axis=0) + return shape + + +def static_output_shape(input_shape, shape, axes): + output_shape = input_shape.as_list() + if shape is not None: + if axes is None: + axes = list(range(-len(shape), 0)) + if isinstance(shape, tf.Tensor): + if isinstance(axes, tf.Tensor): + output_shape = [None] * len(output_shape) + else: + for ax in axes: + output_shape[ax] = None + else: + for idx, ax in enumerate(axes): + output_shape[ax] = shape[idx] + return tf.TensorShape(output_shape) + + +def trans_x_to_s( + x: Union[tf.Tensor, tf.Variable], + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), +) -> Union[tf.Tensor, tf.Variable]: + """Change the shape of the input array x to the desired output shape s.""" + if x.dtype != tf.complex128 and x.dtype != tf.complex64: + x = tf.cast(x, tf.float32) + x_shape = x.shape + if dim == (-1, -2) or dim == (1, 0): + s = (s[1], s[0]) + if s[0] >= x_shape[0] and s[1] >= x_shape[1]: + paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]]) + x_new = tf.pad(x, paddings=paddings) + elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape): + x_new = x[: s[0], : s[1]] + if s[0] != x_new.shape[0]: + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif s[1] != x_new.shape[1]: + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], 1) + elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[0] - x_new.shape[0] + z = tf.zeros((size, s[1]), dtype=x.dtype) + x_new = tf.concat([x_new, z], 0) + elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape): + x_new = x[: s[0], : s[1]] + size = s[1] - x_new.shape[1] + z = tf.zeros((s[0], size), dtype=x.dtype) + x_new = tf.concat([x_new, z], axis=1) + else: + x_new = x[: s[0], : s[1]] + return x_new + + +def transpose_x(x, perm, perform_transpose): + x = tf.cond(perform_transpose, lambda: tf.transpose(x, perm=perm), lambda: x) + return x + + +interpolate.partial_mixed_handler = lambda x, *args, mode="linear", scale_factor=None, recompute_scale_factor=None, align_corners=None, **kwargs: ( # noqa: E501 + (not align_corners and (len(x.shape) - 2) < 2) + and mode not in ["nearest", "area", "bicubic", "nd"] +) +"""RFFTN Function.""" diff --git a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py index c457e4ceb0ac0..ff132130891d3 100644 --- a/ivy/functional/backends/tensorflow/experimental/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/experimental/linear_algebra.py @@ -11,33 +11,56 @@ from .. import backend_version -@with_unsupported_dtypes( - {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version -) -def eigh_tridiagonal( - alpha: Union[tf.Tensor, tf.Variable], - beta: Union[tf.Tensor, tf.Variable], +def adjoint( + x: Union[tf.Tensor, tf.Variable], /, *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] - ] = None, - tol: Optional[float] = None, -) -> Union[ - tf.Tensor, - tf.Variable, - Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], -]: - return tf.linalg.eigh_tridiagonal( - alpha, - beta, - eigvals_only=eigvals_only, - select=select, - select_range=select_range, - tol=tol, - ) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + _check_valid_dimension_size(x) + return tf.linalg.adjoint(x) + + +@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( + x: Union[tf.Tensor, tf.Variable], + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + svd = tf.linalg.svd(x, compute_uv=False) + if len(x.shape) >= 3: + ax = len(x.shape) // 2 + elif len(x.shape) >= 3 and p == -1: + ax = [-1, -2] + else: + ax = None + if p is None or p == 2: + k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) + elif p == "nuc": + svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) + k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) + elif p == "fro": + k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] + ) + elif p < 0: + if p == -1: + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=0), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) + elif p == -2: + k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) + elif p == -float("inf"): + k = tf.reduce_min( + tf.reduce_sum(tf.abs(x), axis=1), axis=ax + ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) + else: + k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( + tf.linalg.inv(x), ord=p, axis=[-2, -1] + ) + return k def diagflat( @@ -75,23 +98,14 @@ def diagflat( return ret -def kron( - a: Union[tf.Tensor, tf.Variable], - b: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.kron(a, b) - - -def matrix_exp( - x: Union[tf.Tensor, tf.Variable], +def dot( + a: tf.Tensor, + b: tf.Tensor, /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.linalg.expm(x) + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + return tf.tensordot(a, b, axes=1) def eig( @@ -105,6 +119,35 @@ def eig( return tf.linalg.eig(x) +@with_unsupported_dtypes( + {"2.13.0 and below": ("int", "float16", "bfloat16")}, backend_version +) +def eigh_tridiagonal( + alpha: Union[tf.Tensor, tf.Variable], + beta: Union[tf.Tensor, tf.Variable], + /, + *, + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], tf.Tensor, tf.Variable] + ] = None, + tol: Optional[float] = None, +) -> Union[ + tf.Tensor, + tf.Variable, + Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]], +]: + return tf.linalg.eigh_tridiagonal( + alpha, + beta, + eigvals_only=eigvals_only, + select=select, + select_range=select_range, + tol=tol, + ) + + def eigvals( x: Union[tf.Tensor, tf.Variable], /, @@ -114,14 +157,33 @@ def eigvals( return tf.linalg.eigvals(x) -def adjoint( +def kron( + a: Union[tf.Tensor, tf.Variable], + b: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.kron(a, b) + + +def lu_factor( + x: Union[tf.Tensor, tf.Variable], + /, + *, + pivot: Optional[bool] = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + raise IvyNotImplementedException() + + +def matrix_exp( x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - _check_valid_dimension_size(x) - return tf.linalg.adjoint(x) + return tf.linalg.expm(x) @with_supported_dtypes( @@ -151,66 +213,4 @@ def multi_dot( return dot_out -@with_unsupported_dtypes({"1.25.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( - x: Union[tf.Tensor, tf.Variable], - /, - *, - p: Optional[Union[None, int, str]] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - svd = tf.linalg.svd(x, compute_uv=False) - if len(x.shape) >= 3: - ax = len(x.shape) // 2 - elif len(x.shape) >= 3 and p == -1: - ax = [-1, -2] - else: - ax = None - if p is None or p == 2: - k = tf.reduce_max(svd, axis=ax) / tf.reduce_min(svd, axis=ax) - elif p == "nuc": - svd_inv = tf.linalg.svd(tf.linalg.inv(x), compute_uv=False) - k = tf.reduce_sum(svd, axis=ax) * tf.reduce_sum(svd_inv, axis=ax) - elif p == "fro": - k = tf.norm(x, ord="euclidean", axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord="euclidean", axis=[-2, -1] - ) - elif p < 0: - if p == -1: - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=0), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=0), axis=ax) - elif p == -2: - k = tf.reduce_min(svd, axis=ax) / tf.reduce_max(svd, axis=ax) - elif p == -float("inf"): - k = tf.reduce_min( - tf.reduce_sum(tf.abs(x), axis=1), axis=ax - ) * tf.reduce_min(tf.reduce_sum(tf.abs(tf.linalg.inv(x)), axis=1), axis=ax) - else: - k = tf.norm(x, ord=p, axis=[-2, -1]) * tf.norm( - tf.linalg.inv(x), ord=p, axis=[-2, -1] - ) - return k - - -def lu_factor( - x: Union[tf.Tensor, tf.Variable], - /, - *, - pivot: Optional[bool] = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - raise IvyNotImplementedException() - - -def dot( - a: tf.Tensor, - b: tf.Tensor, - /, - *, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - return tf.tensordot(a, b, axes=1) - - dot.support_native_out = True diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index 7e0a806ee1e19..bf0be94c43b24 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -10,99 +10,126 @@ import ivy -def moveaxis( - a: Union[tf.Tensor, tf.Variable], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, +def atleast_1d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.moveaxis(a, source, destination) +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_1d(*arys) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def heaviside( - x1: Union[tf.Tensor, tf.Variable], - x2: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) +def atleast_2d( + *arys: Union[tf.Tensor, tf.Variable], + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_2d(*arys) -def flipud( - m: Union[tf.Tensor, tf.Variable], +def atleast_3d( + *arys: Union[tf.Tensor, tf.Variable, bool, Number], + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + return tf.experimental.numpy.atleast_3d(*arys) + + +def broadcast_shapes( + *shapes: Union[List[int], List[Tuple]], +) -> Tuple[int, ...]: + if len(shapes) > 1: + desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) + if len(shapes) > 2: + for i in range(2, len(shapes)): + desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) + else: + return [shapes[0]] + return tuple(desired_shape.numpy().tolist()) + + +def concat_from_sequence( + input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], /, *, - copy: Optional[bool] = None, + new_axis: int = 0, + axis: int = 0, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.flipud(m) + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + highest_dtype = input_sequence[0].dtype + for i in input_sequence: + highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) + if new_axis == 0: + ret = tf.concat(input_sequence, axis=axis) + return ret + elif new_axis == 1: + ret = tf.stack(input_sequence, axis=axis) + return ret -def vstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + +def dsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vstack(arrays) + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) < 3: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def dstack( arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.hstack(arrays) + return tf.experimental.numpy.dstack(arrays) -def rot90( - m: Union[tf.Tensor, tf.Variable], +def expand( + x: Union[tf.Tensor, tf.Variable], + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Union[tf.Tensor, tf.Variable] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.rot90(m, k, axes) + shape = list(shape) + for i, dim in enumerate(shape): + if dim < 0: + shape[i] = x.shape.num_elements() / tf.reduce_prod( + [s for s in shape if s > 0] + ) + return tf.broadcast_to(x, shape) -@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) -def top_k( - x: tf.Tensor, - k: int, +def fill_diagonal( + a: tf.Tensor, + v: Union[int, float], /, *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, -) -> Tuple[tf.Tensor, tf.Tensor]: - k = min(k, x.shape[axis]) - if not largest: - indices = tf.experimental.numpy.argsort(x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) + wrap: bool = False, +): + shape = tf.shape(a) + max_end = tf.math.reduce_prod(shape) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] else: - indices = tf.experimental.numpy.argsort(-x, axis=axis) - indices = tf.experimental.numpy.take( - indices, tf.experimental.numpy.arange(k), axis=axis - ) - indices = tf.dtypes.cast(indices, tf.int32) - if not sorted: - indices = tf.sort(indices, axis=axis) - topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) - val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) - indices = tf.dtypes.cast(indices, tf.int64) - return topk_res(val, indices) + step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) + a = tf.reshape(a, (-1,)) + end = min(end, max_end) + indices = [[i] for i in range(0, end, step)] + ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) + a = tf.tensor_scatter_nd_update(a, indices, ups) + a = tf.reshape(a, shape) + return a def fliplr( @@ -115,72 +142,80 @@ def fliplr( return tf.experimental.numpy.fliplr(m) -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) -def i0( - x: Union[tf.Tensor, tf.Variable], +def flipud( + m: Union[tf.Tensor, tf.Variable], /, *, + copy: Optional[bool] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.math.bessel_i0(x, name=None) + return tf.experimental.numpy.flipud(m) -def vsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def heaviside( + x1: Union[tf.Tensor, tf.Variable], + x2: Union[tf.Tensor, tf.Variable], /, *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.cast(tf.experimental.numpy.heaviside(x1, x2), x1.dtype) -def dsplit( +def hsplit( ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, ) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) < 3: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def atleast_1d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_1d(*arys) +def hstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.hstack(arrays) -def dstack( - arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) +def i0( + x: Union[tf.Tensor, tf.Variable], /, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.dstack(arrays) + return tf.math.bessel_i0(x, name=None) -def atleast_2d( - *arys: Union[tf.Tensor, tf.Variable], +def moveaxis( + a: Union[tf.Tensor, tf.Variable], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_2d(*arys) + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.moveaxis(a, source, destination) -def atleast_3d( - *arys: Union[tf.Tensor, tf.Variable, bool, Number], +def rot90( + m: Union[tf.Tensor, tf.Variable], + /, + *, copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - return tf.experimental.numpy.atleast_3d(*arys) + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Union[tf.Tensor, tf.Variable] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.rot90(m, k, axes) def take_along_axis( @@ -228,69 +263,36 @@ def take_along_axis( return tf.experimental.numpy.take_along_axis(arr, indices, axis) -def hsplit( - ary: Union[tf.Tensor, tf.Variable], - indices_or_sections: Union[int, Tuple[int, ...]], +@with_unsupported_dtypes({"2.13.0 and below": ("unsigned", "complex")}, backend_version) +def top_k( + x: tf.Tensor, + k: int, /, *, - copy: Optional[bool] = None, -) -> List[Union[tf.Tensor, tf.Variable]]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -def broadcast_shapes( - *shapes: Union[List[int], List[Tuple]], -) -> Tuple[int, ...]: - if len(shapes) > 1: - desired_shape = tf.broadcast_dynamic_shape(shapes[0], shapes[1]) - if len(shapes) > 2: - for i in range(2, len(shapes)): - desired_shape = tf.broadcast_dynamic_shape(desired_shape, shapes[i]) + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[tf.Tensor, tf.Tensor]] = None, +) -> Tuple[tf.Tensor, tf.Tensor]: + k = min(k, x.shape[axis]) + if not largest: + indices = tf.experimental.numpy.argsort(x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) else: - return [shapes[0]] - return tuple(desired_shape.numpy().tolist()) - - -def expand( - x: Union[tf.Tensor, tf.Variable], - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - shape = list(shape) - for i, dim in enumerate(shape): - if dim < 0: - shape[i] = x.shape.num_elements() / tf.reduce_prod( - [s for s in shape if s > 0] - ) - return tf.broadcast_to(x, shape) - - -def concat_from_sequence( - input_sequence: Union[Tuple[tf.Tensor], List[tf.Tensor]], - /, - *, - new_axis: int = 0, - axis: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - highest_dtype = input_sequence[0].dtype - for i in input_sequence: - highest_dtype = ivy.as_native_dtype(ivy.promote_types(highest_dtype, i.dtype)) - - if new_axis == 0: - ret = tf.concat(input_sequence, axis=axis) - return ret - elif new_axis == 1: - ret = tf.stack(input_sequence, axis=axis) - return ret + indices = tf.experimental.numpy.argsort(-x, axis=axis) + indices = tf.experimental.numpy.take( + indices, tf.experimental.numpy.arange(k), axis=axis + ) + indices = tf.dtypes.cast(indices, tf.int32) + if not sorted: + indices = tf.sort(indices, axis=axis) + topk_res = NamedTuple("top_k", [("values", tf.Tensor), ("indices", tf.Tensor)]) + val = tf.experimental.numpy.take_along_axis(x, indices, axis=axis) + indices = tf.dtypes.cast(indices, tf.int64) + return topk_res(val, indices) def unique_consecutive( @@ -346,26 +348,24 @@ def unique_consecutive( ) -def fill_diagonal( - a: tf.Tensor, - v: Union[int, float], +def vsplit( + ary: Union[tf.Tensor, tf.Variable], + indices_or_sections: Union[int, Sequence[int], tf.Tensor, tf.Variable], /, *, - wrap: bool = False, -): - shape = tf.shape(a) - max_end = tf.math.reduce_prod(shape) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + tf.reduce_sum(tf.math.cumprod(shape[:-1])) - a = tf.reshape(a, (-1,)) - end = min(end, max_end) - indices = [[i] for i in range(0, end, step)] - ups = tf.convert_to_tensor([v] * len(indices), dtype=a.dtype) - a = tf.tensor_scatter_nd_update(a, indices, ups) - a = tf.reshape(a, shape) - return a + copy: Optional[bool] = None, +) -> List[Union[tf.Tensor, tf.Variable]]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vstack(arrays) diff --git a/ivy/functional/backends/tensorflow/experimental/norms.py b/ivy/functional/backends/tensorflow/experimental/norms.py index 5fc38b60283af..22f00921bcf96 100644 --- a/ivy/functional/backends/tensorflow/experimental/norms.py +++ b/ivy/functional/backends/tensorflow/experimental/norms.py @@ -4,30 +4,6 @@ from . import backend_version -def l1_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - -def l2_normalize( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[int] = None, - out: Optional[tf.Tensor] = None, -) -> tf.Tensor: - denorm = tf.norm(x, axis=axis, keepdims=True) - denorm = tf.math.maximum(denorm, 1e-12) - return tf.math.divide(x, denorm) - - @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) def batch_norm( x: Union[tf.Tensor, tf.Variable], @@ -153,6 +129,30 @@ def instance_norm( ) +def l1_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, ord=1, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + +def l2_normalize( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[int] = None, + out: Optional[tf.Tensor] = None, +) -> tf.Tensor: + denorm = tf.norm(x, axis=axis, keepdims=True) + denorm = tf.math.maximum(denorm, 1e-12) + return tf.math.divide(x, denorm) + + def lp_normalize( x: Union[tf.Tensor, tf.Variable], /, diff --git a/ivy/functional/backends/tensorflow/experimental/random.py b/ivy/functional/backends/tensorflow/experimental/random.py index a5bbf97481be0..fc31d6a1c8d3f 100644 --- a/ivy/functional/backends/tensorflow/experimental/random.py +++ b/ivy/functional/backends/tensorflow/experimental/random.py @@ -15,6 +15,51 @@ ) +def bernoulli( + probs: Union[float, tf.Tensor, tf.Variable], + *, + logits: Union[float, tf.Tensor, tf.Variable] = None, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: str, + dtype: DType, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if seed is not None: + tf.random.set_seed(seed) + if logits is not None: + logits = tf.cast(logits, dtype) + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + probs = tf.cast(probs, dtype) + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return tfp.distributions.Bernoulli( + logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True + ).sample(shape, seed) + + +def beta( + alpha: Union[float, tf.Tensor, tf.Variable], + beta: Union[float, tf.Tensor, tf.Variable], + /, + *, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + device: Optional[str] = None, + dtype: Optional[Union[DType, ivy.Dtype]] = None, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if not dtype: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + shape = _check_bounds_and_get_shape(alpha, beta, shape).shape + alpha = tf.cast(alpha, dtype) + beta = tf.cast(beta, dtype) + return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) + + # dirichlet @with_unsupported_dtypes( { @@ -54,26 +99,6 @@ def dirichlet( ) -def beta( - alpha: Union[float, tf.Tensor, tf.Variable], - beta: Union[float, tf.Tensor, tf.Variable], - /, - *, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: Optional[str] = None, - dtype: Optional[Union[DType, ivy.Dtype]] = None, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if not dtype: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - shape = _check_bounds_and_get_shape(alpha, beta, shape).shape - alpha = tf.cast(alpha, dtype) - beta = tf.cast(beta, dtype) - return tfp.distributions.Beta(alpha, beta).sample(shape, seed=seed) - - def gamma( alpha: Union[float, tf.Tensor, tf.Variable], beta: Union[float, tf.Tensor, tf.Variable], @@ -117,28 +142,3 @@ def poisson( if tf.reduce_any(lam < 0): return tf.where(lam < 0, fill_value, ret) return ret - - -def bernoulli( - probs: Union[float, tf.Tensor, tf.Variable], - *, - logits: Union[float, tf.Tensor, tf.Variable] = None, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - device: str, - dtype: DType, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if seed is not None: - tf.random.set_seed(seed) - if logits is not None: - logits = tf.cast(logits, dtype) - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - probs = tf.cast(probs, dtype) - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return tfp.distributions.Bernoulli( - logits=logits, probs=probs, dtype=dtype, allow_nan_stats=True - ).sample(shape, seed) diff --git a/ivy/functional/backends/tensorflow/experimental/statistical.py b/ivy/functional/backends/tensorflow/experimental/statistical.py index 9cd91437a5f4a..2a5feef312527 100644 --- a/ivy/functional/backends/tensorflow/experimental/statistical.py +++ b/ivy/functional/backends/tensorflow/experimental/statistical.py @@ -14,202 +14,79 @@ from copy import deepcopy -def histogram( - a: tf.Tensor, - /, - *, - bins: Optional[Union[int, tf.Tensor]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[tf.DType] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[tf.Tensor] = None, - density: Optional[bool] = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Tuple[tf.Tensor]: - min_a = tf.reduce_min(a) - max_a = tf.reduce_max(a) - if isinstance(bins, tf.Tensor) and range: - raise ivy.exceptions.IvyException( - "Must choose between specifying bins and range or bin edges directly" - ) - if range: - if isinstance(bins, int): - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - elif isinstance(bins, int): - range = (min_a, max_a) - bins = tf.cast( - tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype - ) - if tf.shape(bins)[0] < 2: - raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") - if min_a < bins[0] and not extend_lower_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_lower_interval to deal with this." - ) - if max_a > bins[-1] and not extend_upper_interval: - raise ivy.exceptions.IvyException( - "Values of x outside of the intervals cause errors in tensorflow backend. " - "Consider using extend_upper_interval to deal with this." - ) - ret = tfp.stats.histogram( - x=a, - edges=bins, - axis=axis, - weights=weights, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - name="histogram", - ) - if density: - pass - # TODO: Tensorflow native dtype argument is not working - if dtype: - ret = tf.cast(ret, dtype) - bins = tf.cast(bins, dtype) - # TODO: weird error when returning bins: return ret, bins - return ret +# --- Helpers --- # +# --------------- # -@with_supported_dtypes( - { - "2.13.0 and below": ( - "float", - "complex", +def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: + values, indices = [], [] + if ( + isinstance(x[0], tf.Tensor) + and isinstance(x[0].numpy().tolist(), list) + and len(x[0].numpy().tolist()) >= 1 + ): + if axis >= 1: + for ret1 in x: + value, indice = __find_cummax(ret1, axis=axis - 1) + indices.append(indice) + values.append(value) + else: + x_list = x.numpy() + z_list = __get_index(x_list.tolist()) + indices, values, n1 = x_list.copy(), x_list.copy(), {} + indices.fill(0) + values.fill(0) + z_list = sorted(z_list, key=lambda i: i[1]) + for y, y_index in z_list: + multi_index = y_index + if tuple(multi_index[1:]) not in n1: + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + elif ( + y + >= x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + ): + n1[tuple(multi_index[1:])] = multi_index[0] + indices[y_index] = multi_index[0] + values[y_index] = y + else: + indices[y_index] = n1[tuple(multi_index[1:])] + values[y_index] = x_list[ + tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) + ] + else: + x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) + values, indices = tf.scan( + lambda a, b: ( + a + if a > b + or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 + else b + ), + (x, x_indices), ) - }, - backend_version, -) -def median( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tfp.stats.percentile( - input, - 50.0, - axis=axis, - interpolation="midpoint", - keepdims=keepdims, + + return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( + tf.convert_to_tensor(indices), dtype=tf.int64 ) -def nanmean( - a: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - np_math_ops.enable_numpy_methods_on_tensor() - return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) - +def __get_index(lst, indices=None, prefix=None): + if indices is None: + indices = [] + if prefix is None: + prefix = [] -def _validate_quantile(q): - if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: - for i in range(tf.size(q)): - if not (0.0 <= q[i] <= 1.0): - return False + if isinstance(lst, list): + for i, sub_lst in enumerate(lst): + sub_indices = prefix + [i] + __get_index(sub_lst, indices, sub_indices) else: - if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): - return False - return True - - -def to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis - - -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = tf.experimental.numpy.ndim(a) - axis_arg = deepcopy(axis) - if axis is not None: - axis = to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) - - for i, s in enumerate(sorted(keep)): - a = tf.experimental.numpy.moveaxis(a, s, i) - a = tf.reshape( - a, - [ - *a.shape[:nkeep], - -1, - ], - ) - axis_arg = -1 - - ret = fn(a, q, axis=axis_arg) - - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] - - return ret - - -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if tf.experimental.numpy.ndim(q) > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = tf.reshape(a, [-1]) - elif axis != 0: - a = tf.experimental.numpy.moveaxis(a, axis, 0) - axis = 0 - - n = a.shape[axis] - - indices = q * (n - 1) - - a = tf.sort(a, axis) - - indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) - indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) - - weights = indices - tf.cast(indices_below, dtype=ret_dtype) - - indices_below = tf.clip_by_value(indices_below, 0, n - 1) - indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) - tensor_upper = tf.gather(a, indices_upper, axis=axis) - tensor_below = tf.gather(a, indices_below, axis=axis) - - pred = weights <= 0.5 - out = tf.where(pred, tensor_below, tensor_upper) - - return tf.cast(out, ret_dtype) + indices.append((lst, tuple(prefix))) + return indices def _compute_quantile_wrapper( @@ -247,50 +124,39 @@ def _compute_quantile_wrapper( ) -def quantile( - a: Union[tf.Tensor, tf.Variable], - q: Union[tf.Tensor, float], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - interpolation: str = "linear", - keepdims: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, - axis=axis, - keepdims=keepdims, - interpolation=interpolation, - ) +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = tf.experimental.numpy.ndim(a) + axis_arg = deepcopy(axis) + if axis is not None: + axis = to_positive_axis(axis, nd) + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) -def corrcoef( - x: tf.Tensor, - /, - *, - y: tf.Tensor, - rowvar: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> tf.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = tf.concat([x, y], axis=axis) + for i, s in enumerate(sorted(keep)): + a = tf.experimental.numpy.moveaxis(a, s, i) + a = tf.reshape( + a, + [ + *a.shape[:nkeep], + -1, + ], + ) + axis_arg = -1 - if rowvar: - mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) - cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) - else: - mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) - cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) + ret = fn(a, q, axis=axis_arg) - cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) - cor = cov2_t @ cov_t @ cov2_t - return cor + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret def _nanmedian_helper(input, axis=None, keepdims=False): @@ -452,22 +318,52 @@ def _nanmedian_helper(input, axis=None, keepdims=False): return result -def nanmedian( - input: Union[tf.Tensor, tf.Variable], - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if overwrite_input: - copied_input = tf.identity(input) - return _nanmedian_helper(copied_input, axis, keepdims) +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if tf.experimental.numpy.ndim(q) > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = tf.reshape(a, [-1]) + elif axis != 0: + a = tf.experimental.numpy.moveaxis(a, axis, 0) + axis = 0 + + n = a.shape[axis] + + indices = q * (n - 1) + a = tf.sort(a, axis) + + indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32) + indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32) + + weights = indices - tf.cast(indices_below, dtype=ret_dtype) + + indices_below = tf.clip_by_value(indices_below, 0, n - 1) + indices_upper = tf.clip_by_value(indices_upper, 0, n - 1) + tensor_upper = tf.gather(a, indices_upper, axis=axis) + tensor_below = tf.gather(a, indices_below, axis=axis) + + pred = weights <= 0.5 + out = tf.where(pred, tensor_below, tensor_upper) + + return tf.cast(out, ret_dtype) + + +def _validate_quantile(q): + if tf.experimental.numpy.ndim(q) == 1 and tf.size(q) < 10: + for i in range(tf.size(q)): + if not (0.0 <= q[i] <= 1.0): + return False else: - result = _nanmedian_helper(input, axis, keepdims) - return result + if not (tf.math.reduce_all(0 <= q) and tf.math.reduce_all(q <= 1)): + return False + return True + + +# --- Main --- # +# ------------ # @with_supported_device_and_dtypes( @@ -505,19 +401,30 @@ def bincount( ) -@with_supported_device_and_dtypes( - { - "2.13.0 and below": { - "cpu": ("float32", "float64"), - "gpu": ("bfloat16", "float16", "float32", "float64"), - } - }, - backend_version, -) -def igamma( - a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None +def corrcoef( + x: tf.Tensor, + /, + *, + y: tf.Tensor, + rowvar: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> tf.Tensor: - return tf.math.igamma(a, x) + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = tf.concat([x, y], axis=axis) + + if rowvar: + mean_t = tf.reduce_mean(xarr, axis=1, keepdims=True) + cov_t = ((xarr - mean_t) @ tf.transpose(xarr - mean_t)) / (x.shape[1] - 1) + else: + mean_t = tf.reduce_mean(xarr, axis=0, keepdims=True) + cov_t = (tf.transpose(xarr - mean_t) @ (xarr - mean_t)) / (x.shape[1] - 1) + + cov2_t = tf.linalg.diag(1 / tf.sqrt(tf.linalg.diag_part(cov_t))) + cor = cov2_t @ cov_t @ cov2_t + return cor @with_unsupported_dtypes({"2.13.0 and below": ("float16", "bfloat16")}, backend_version) @@ -680,77 +587,6 @@ def cummax( return __find_cummax(x, axis=axis) -def __find_cummax(x: tf.Tensor, axis: int = 0) -> Tuple[tf.Tensor, tf.Tensor]: - values, indices = [], [] - if ( - isinstance(x[0], tf.Tensor) - and isinstance(x[0].numpy().tolist(), list) - and len(x[0].numpy().tolist()) >= 1 - ): - if axis >= 1: - for ret1 in x: - value, indice = __find_cummax(ret1, axis=axis - 1) - indices.append(indice) - values.append(value) - else: - x_list = x.numpy() - z_list = __get_index(x_list.tolist()) - indices, values, n1 = x_list.copy(), x_list.copy(), {} - indices.fill(0) - values.fill(0) - z_list = sorted(z_list, key=lambda i: i[1]) - for y, y_index in z_list: - multi_index = y_index - if tuple(multi_index[1:]) not in n1: - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - elif ( - y - >= x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - ): - n1[tuple(multi_index[1:])] = multi_index[0] - indices[y_index] = multi_index[0] - values[y_index] = y - else: - indices[y_index] = n1[tuple(multi_index[1:])] - values[y_index] = x_list[ - tuple([n1[tuple(multi_index[1:])]] + list(multi_index[1:])) - ] - else: - x_indices = tf.convert_to_tensor(list(range(0, x.shape[0])), dtype=x.dtype) - values, indices = tf.scan( - lambda a, b: ( - a - if a > b - or tf.experimental.numpy.where(x[0].numpy() == b[0].numpy()) == 0 - else b - ), - (x, x_indices), - ) - - return tf.convert_to_tensor(values, dtype=x.dtype), tf.cast( - tf.convert_to_tensor(indices), dtype=tf.int64 - ) - - -def __get_index(lst, indices=None, prefix=None): - if indices is None: - indices = [] - if prefix is None: - prefix = [] - - if isinstance(lst, list): - for i, sub_lst in enumerate(lst): - sub_indices = prefix + [i] - __get_index(sub_lst, indices, sub_indices) - else: - indices.append((lst, tuple(prefix))) - return indices - - @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "complex")}, backend_version, @@ -781,3 +617,175 @@ def cummin( return cummin_x else: return tf.cast(cummin_x, dtype) + + +def histogram( + a: tf.Tensor, + /, + *, + bins: Optional[Union[int, tf.Tensor]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[tf.DType] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[tf.Tensor] = None, + density: Optional[bool] = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Tuple[tf.Tensor]: + min_a = tf.reduce_min(a) + max_a = tf.reduce_max(a) + if isinstance(bins, tf.Tensor) and range: + raise ivy.exceptions.IvyException( + "Must choose between specifying bins and range or bin edges directly" + ) + if range: + if isinstance(bins, int): + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + elif isinstance(bins, int): + range = (min_a, max_a) + bins = tf.cast( + tf.linspace(start=range[0], stop=range[1], num=bins + 1), dtype=a.dtype + ) + if tf.shape(bins)[0] < 2: + raise ivy.exceptions.IvyException("bins must have at least 1 bin (size > 1)") + if min_a < bins[0] and not extend_lower_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_lower_interval to deal with this." + ) + if max_a > bins[-1] and not extend_upper_interval: + raise ivy.exceptions.IvyException( + "Values of x outside of the intervals cause errors in tensorflow backend. " + "Consider using extend_upper_interval to deal with this." + ) + ret = tfp.stats.histogram( + x=a, + edges=bins, + axis=axis, + weights=weights, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + name="histogram", + ) + if density: + pass + # TODO: Tensorflow native dtype argument is not working + if dtype: + ret = tf.cast(ret, dtype) + bins = tf.cast(bins, dtype) + # TODO: weird error when returning bins: return ret, bins + return ret + + +@with_supported_device_and_dtypes( + { + "2.13.0 and below": { + "cpu": ("float32", "float64"), + "gpu": ("bfloat16", "float16", "float32", "float64"), + } + }, + backend_version, +) +def igamma( + a: tf.Tensor, /, *, x: tf.Tensor, out: Optional[tf.Tensor] = None +) -> tf.Tensor: + return tf.math.igamma(a, x) + + +@with_supported_dtypes( + { + "2.13.0 and below": ( + "float", + "complex", + ) + }, + backend_version, +) +def median( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tfp.stats.percentile( + input, + 50.0, + axis=axis, + interpolation="midpoint", + keepdims=keepdims, + ) + + +def nanmean( + a: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + np_math_ops.enable_numpy_methods_on_tensor() + return tf.experimental.numpy.nanmean(a, axis=axis, keepdims=keepdims, dtype=dtype) + + +def nanmedian( + input: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if overwrite_input: + copied_input = tf.identity(input) + return _nanmedian_helper(copied_input, axis, keepdims) + + else: + result = _nanmedian_helper(input, axis, keepdims) + return result + + +def quantile( + a: Union[tf.Tensor, tf.Variable], + q: Union[tf.Tensor, float], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + interpolation: str = "linear", + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + ) + + +def to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis diff --git a/ivy/functional/backends/tensorflow/general.py b/ivy/functional/backends/tensorflow/general.py index ba239549e3fdd..e6052c0b5e49e 100644 --- a/ivy/functional/backends/tensorflow/general.py +++ b/ivy/functional/backends/tensorflow/general.py @@ -23,12 +23,26 @@ _round = round -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, (tf.Tensor, tf.Variable)): - if exclusive and isinstance(x, tf.Variable): - return False - return True - return False +# --- Helpers --- # +# --------------- # + + +def _check_query(query): + return not isinstance(query, list) and ( + not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) + ) + + +def _update_view(view, base): + for fn, args, kwargs, index in view._manipulation_stack: + base = ivy.__dict__[fn](base, *args, **kwargs) + base = base[index] if ivy.exists(index) else base + view.data = base.data + return view + + +# --- Main --- # +# ------------ # def array_equal( @@ -48,58 +62,6 @@ def current_backend_str() -> str: return "tensorflow" -def _check_query(query): - return not isinstance(query, list) and ( - not (ivy.is_array(query) and ivy.is_bool_dtype(query) ^ bool(query.ndim > 0)) - ) - - -def get_item( - x: Union[tf.Tensor, tf.Variable], - /, - query: Union[tf.Tensor, tf.Variable, Tuple], - *, - copy: bool = None, -) -> Union[tf.Tensor, tf.Variable]: - return x.__getitem__(query) - - -get_item.partial_mixed_handler = lambda x, query, **kwargs: ( - all(_check_query(i) for i in query) - if isinstance(query, tuple) - else _check_query(query) -) - - -def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: - # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions - if ( - ivy.is_array(x) - and get_num_dims(x) == 0 - and ivy.as_native_dtype(x.dtype) is tf.bfloat16 - ): - x = tf.expand_dims(x, 0) - if copy: - return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) - else: - return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) - if copy: - return np.array(tf.convert_to_tensor(x)) - else: - return np.asarray(tf.convert_to_tensor(x)) - - -def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: - ret = to_numpy(x).item() - if x.dtype == tf.bfloat16: - return float(ret) - return ret - - -def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: - return x.numpy().tolist() - - def gather( params: Union[tf.Tensor, tf.Variable], indices: Union[tf.Tensor, tf.Variable], @@ -115,38 +77,6 @@ def gather( return tf.gather(params, indices, axis=axis, batch_dims=batch_dims) -def gather_nd_helper(params, indices): - indices_shape = tf.shape(indices) - params_shape = tf.shape(params) - num_index_dims = indices_shape[-1] - result_dim_sizes_list = [ - tf.math.reduce_prod(params_shape[i + 1 :]) for i in range(len(params_shape) - 1) - ] + [1] - result_dim_sizes = tf.convert_to_tensor(result_dim_sizes_list, dtype=indices.dtype) - implicit_indices_factor = result_dim_sizes[num_index_dims - 1] - flat_params = tf.reshape(params, (-1,)) - new_shape = [1] * (len(indices_shape) - 1) + [num_index_dims] - indices_scales = tf.reshape(result_dim_sizes[0:num_index_dims], new_shape) - indices_for_flat_tiled = tf.reshape( - tf.reduce_sum(indices * indices_scales, -1, keepdims=True), (-1, 1) - ) - indices_for_flat_tiled = tf.repeat( - indices_for_flat_tiled, implicit_indices_factor, axis=1 - ) - implicit_indices = tf.repeat( - tf.expand_dims(tf.range(implicit_indices_factor), 0), - indices_for_flat_tiled.shape[0], - axis=0, - ) - indices_for_flat = indices_for_flat_tiled + implicit_indices - flat_indices_for_flat = tf.reshape(indices_for_flat, (-1,)) - flat_gather = tf.gather(flat_params, flat_indices_for_flat) - res = tf.reshape( - flat_gather, tf.concat([indices_shape[:-1], params_shape[num_index_dims:]], 0) - ) - return res - - def gather_nd( params: Union[tf.Tensor, tf.Variable], indices: Union[tf.Tensor, tf.Variable], @@ -184,6 +114,48 @@ def gather_nd( return result +def gather_nd_helper(params, indices): + indices_shape = tf.shape(indices) + params_shape = tf.shape(params) + num_index_dims = indices_shape[-1] + result_dim_sizes_list = [ + tf.math.reduce_prod(params_shape[i + 1 :]) for i in range(len(params_shape) - 1) + ] + [1] + result_dim_sizes = tf.convert_to_tensor(result_dim_sizes_list, dtype=indices.dtype) + implicit_indices_factor = result_dim_sizes[num_index_dims - 1] + flat_params = tf.reshape(params, (-1,)) + new_shape = [1] * (len(indices_shape) - 1) + [num_index_dims] + indices_scales = tf.reshape(result_dim_sizes[0:num_index_dims], new_shape) + indices_for_flat_tiled = tf.reshape( + tf.reduce_sum(indices * indices_scales, -1, keepdims=True), (-1, 1) + ) + indices_for_flat_tiled = tf.repeat( + indices_for_flat_tiled, implicit_indices_factor, axis=1 + ) + implicit_indices = tf.repeat( + tf.expand_dims(tf.range(implicit_indices_factor), 0), + indices_for_flat_tiled.shape[0], + axis=0, + ) + indices_for_flat = indices_for_flat_tiled + implicit_indices + flat_indices_for_flat = tf.reshape(indices_for_flat, (-1,)) + flat_gather = tf.gather(flat_params, flat_indices_for_flat) + res = tf.reshape( + flat_gather, tf.concat([indices_shape[:-1], params_shape[num_index_dims:]], 0) + ) + return res + + +def get_item( + x: Union[tf.Tensor, tf.Variable], + /, + query: Union[tf.Tensor, tf.Variable, Tuple], + *, + copy: bool = None, +) -> Union[tf.Tensor, tf.Variable]: + return x.__getitem__(query) + + def get_num_dims(x, /, *, as_array=False): return ( tf.cast(tf.shape(tf.shape(x))[0], tf.int64) @@ -286,18 +258,49 @@ def inplace_update( return val -def _update_view(view, base): - for fn, args, kwargs, index in view._manipulation_stack: - base = ivy.__dict__[fn](base, *args, **kwargs) - base = base[index] if ivy.exists(index) else base - view.data = base.data - return view - - def inplace_variables_supported(): return True +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, (tf.Tensor, tf.Variable)): + if exclusive and isinstance(x, tf.Variable): + return False + return True + return False + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def isin( + elements: tf.Tensor, + test_elements: tf.Tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> tf.Tensor: + input_shape = elements.shape + + if tf.rank(elements) == 0: + elements = tf.reshape(elements, [1]) + if tf.rank(test_elements) == 0: + test_elements = tf.reshape(test_elements, [1]) + if not assume_unique: + test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] + + elements = tf.reshape(elements, [-1]) + test_elements = tf.reshape(test_elements, [-1]) + + output = tf.reduce_any( + tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 + ) + return tf.reshape(output, input_shape) ^ invert + + +def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: + return x.dtype.size + + def multiprocessing(context: Optional[str] = None): return ( _multiprocessing if context is None else _multiprocessing.get_context(context) @@ -347,9 +350,6 @@ def scatter_flat( return res -scatter_flat.support_native_out = True - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def scatter_nd( indices: Union[tf.Tensor, tf.Variable], @@ -406,9 +406,6 @@ def scatter_nd( return res -scatter_nd.support_native_out = True - - def shape( x: Union[tf.Tensor, tf.Variable], /, @@ -421,6 +418,35 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: Union[tf.Tensor, tf.Variable], /) -> list: + return x.numpy().tolist() + + +def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray: + # TensorFlow fails to convert bfloat16 tensor when it has 0 dimensions + if ( + ivy.is_array(x) + and get_num_dims(x) == 0 + and ivy.as_native_dtype(x.dtype) is tf.bfloat16 + ): + x = tf.expand_dims(x, 0) + if copy: + return np.squeeze(np.array(tf.convert_to_tensor(x)), 0) + else: + return np.squeeze(np.asarray(tf.convert_to_tensor(x)), 0) + if copy: + return np.array(tf.convert_to_tensor(x)) + else: + return np.asarray(tf.convert_to_tensor(x)) + + +def to_scalar(x: Union[tf.Tensor, tf.Variable], /) -> Number: + ret = to_numpy(x).item() + if x.dtype == tf.bfloat16: + return float(ret) + return ret + + def vmap( func: Callable, in_axes: Union[int, Sequence[int], Sequence[None]] = 0, @@ -507,32 +533,10 @@ def _vmap(*args, **kwargs): return _vmap -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def isin( - elements: tf.Tensor, - test_elements: tf.Tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> tf.Tensor: - input_shape = elements.shape - - if tf.rank(elements) == 0: - elements = tf.reshape(elements, [1]) - if tf.rank(test_elements) == 0: - test_elements = tf.reshape(test_elements, [1]) - if not assume_unique: - test_elements = tf.unique(tf.reshape(test_elements, [-1]))[0] - - elements = tf.reshape(elements, [-1]) - test_elements = tf.reshape(test_elements, [-1]) - - output = tf.reduce_any( - tf.equal(tf.expand_dims(elements, -1), test_elements), axis=-1 - ) - return tf.reshape(output, input_shape) ^ invert - - -def itemsize(x: Union[tf.Tensor, tf.Variable]) -> int: - return x.dtype.size +get_item.partial_mixed_handler = lambda x, query, **kwargs: ( + all(_check_query(i) for i in query) + if isinstance(query, tuple) + else _check_query(query) +) +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True diff --git a/ivy/functional/backends/tensorflow/gradients.py b/ivy/functional/backends/tensorflow/gradients.py index 3eb4207d2df36..6968f2f245ac8 100644 --- a/ivy/functional/backends/tensorflow/gradients.py +++ b/ivy/functional/backends/tensorflow/gradients.py @@ -21,17 +21,8 @@ ) -def variable(x, /): - with tf.device(ivy.dev(x, as_native=True)): - return tf.Variable(x, trainable=True) - - -def is_variable(x, /, *, exclusive=False): - return isinstance(x, tf.Variable) - - -def variable_data(x: tf.Variable, /) -> tf.Variable: - return x.value() +# --- Helpers --- # +# --------------- # def _grad_func(y, xs, xs_required, tape): @@ -63,6 +54,10 @@ def _grad_func(y, xs, xs_required, tape): return grads +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: Union[tf.Tensor, tf.Variable], @@ -117,81 +112,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - def grad_fn(xs): - grads = ivy.nested_map( - xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False - ) - with tf.GradientTape(watch_accessed_variables=False) as tape: - xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) - tape.watch(xs) - y = func(xs) - y = y.to_native(y) - grads_ = tape.gradient(y, xs) - grads_ = ivy.nested_map( - grads_, - lambda x: ivy.to_ivy(x), - include_derived=True, - ) - grads_ = ivy.to_ivy(grads_) - grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) - grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) - xs = ivy.to_ivy(xs) - if isinstance(xs, ivy.Array): - grads = grads_ - else: - ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) - y = ivy.to_ivy(y) - return y, grads - - return grad_fn - - -def stop_gradient( - x: Union[tf.Tensor, tf.Variable], - /, - *, - preserve_type: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - is_var = is_variable(x) - x = tf.stop_gradient(x) - if is_var and preserve_type: - return variable(x) - return x - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - - def callback_fn(x_in): - with tf.GradientTape(persistent=True) as tape: - ivy.nested_map(x_in, ivy.copy_array) - x_in = ivy.to_native(x_in, nested=True) - tape.watch(x_in) - y = grad_fn(x_in) - - # Deal with multiple outputs - if not isinstance(y, ivy.NativeArray): - jacobian = ivy.nested_map( - y, - lambda yi: ivy.to_ivy( - tape.jacobian(yi, x_in, unconnected_gradients="zero"), - nested=True, - ), - include_derived=True, - ) - else: - jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) - return jacobian - - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -240,5 +160,93 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive=False): + return isinstance(x, tf.Variable) + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + + def callback_fn(x_in): + with tf.GradientTape(persistent=True) as tape: + ivy.nested_map(x_in, ivy.copy_array) + x_in = ivy.to_native(x_in, nested=True) + tape.watch(x_in) + y = grad_fn(x_in) + + # Deal with multiple outputs + if not isinstance(y, ivy.NativeArray): + jacobian = ivy.nested_map( + y, + lambda yi: ivy.to_ivy( + tape.jacobian(yi, x_in, unconnected_gradients="zero"), + nested=True, + ), + include_derived=True, + ) + else: + jacobian = ivy.to_ivy(tape.jacobian(y, x_in)) + return jacobian + + return callback_fn + + +def stop_gradient( + x: Union[tf.Tensor, tf.Variable], + /, + *, + preserve_type: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + is_var = is_variable(x) + x = tf.stop_gradient(x) + if is_var and preserve_type: + return variable(x) + return x + + +def value_and_grad(func): + def grad_fn(xs): + grads = ivy.nested_map( + xs, lambda x: ivy.zeros_like(x), include_derived=True, shallow=False + ) + with tf.GradientTape(watch_accessed_variables=False) as tape: + xs = ivy.nested_map(xs, lambda x: ivy.to_native(x), include_derived=True) + tape.watch(xs) + y = func(xs) + y = y.to_native(y) + grads_ = tape.gradient(y, xs) + grads_ = ivy.nested_map( + grads_, + lambda x: ivy.to_ivy(x), + include_derived=True, + ) + grads_ = ivy.to_ivy(grads_) + grad_idxs = ivy.nested_argwhere(grads_, lambda x: ivy.is_ivy_array(x)) + grad_array_vals = list(ivy.multi_index_nest(grads_, grad_idxs)) + xs = ivy.to_ivy(xs) + if isinstance(xs, ivy.Array): + grads = grads_ + else: + ivy.set_nest_at_indices(grads, grad_idxs, grad_array_vals) + y = ivy.to_ivy(y) + return y, grads + + return grad_fn + + +def variable(x, /): + with tf.device(ivy.dev(x, as_native=True)): + return tf.Variable(x, trainable=True) + + +def variable_data(x: tf.Variable, /) -> tf.Variable: + return x.value() + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/tensorflow/layers.py b/ivy/functional/backends/tensorflow/layers.py index c3070846e21cb..d361fdd9827dc 100644 --- a/ivy/functional/backends/tensorflow/layers.py +++ b/ivy/functional/backends/tensorflow/layers.py @@ -17,6 +17,10 @@ ) +# --- Helpers --- # +# --------------- # + + def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): if filter_format == "channel_first": filters = tf.transpose(filters, (*range(2, dims + 2), 1, 0)) @@ -33,6 +37,24 @@ def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): return x, filters +def _output_shape( + x_shape, filter_shape, output_shape, strides, padding, dims, dilations +): + dilations = [dilations] * dims if isinstance(dilations, int) else dilations + strides = [strides] * dims if isinstance(strides, int) else strides + if output_shape is None: + out_shape = [ + _deconv_length( + x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] + ) + for i in range(dims) + ] + output_shape = [x_shape[0], *out_shape, filter_shape[-2]] + elif len(output_shape) == dims: + output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] + return output_shape + + def _pad_before_conv(x, filters, strides, padding, dims, dilations): dilations = [dilations] * dims if isinstance(dilations, int) else dilations strides = [strides] * dims if isinstance(strides, int) else strides @@ -64,22 +86,8 @@ def _pad_before_conv(x, filters, strides, padding, dims, dilations): ) -def _output_shape( - x_shape, filter_shape, output_shape, strides, padding, dims, dilations -): - dilations = [dilations] * dims if isinstance(dilations, int) else dilations - strides = [strides] * dims if isinstance(strides, int) else strides - if output_shape is None: - out_shape = [ - _deconv_length( - x_shape[i + 1], strides[i], filter_shape[i], padding, dilations[i] - ) - for i in range(dims) - ] - output_shape = [x_shape[0], *out_shape, filter_shape[-2]] - elif len(output_shape) == dims: - output_shape = [x_shape[0]] + output_shape + [filter_shape[-2]] - return output_shape +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) @@ -205,32 +213,6 @@ def conv2d_transpose( return res -@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) -def depthwise_conv2d( - x: Union[tf.Tensor, tf.Variable], - filters: Union[tf.Tensor, tf.Variable], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NCHW": - x = tf.transpose(x, (0, 2, 3, 1)) - if tf.rank(filters) == 3: - filters = tf.expand_dims(filters, -1) - x = _pad_before_conv(x, filters, strides, padding, 2, dilations) - strides = [1, strides[0], strides[1], 1] - res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) - if data_format == "NCHW": - return tf.transpose(res, (0, 3, 1, 2)) - return res - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) def conv3d( x: Union[tf.Tensor, tf.Variable], @@ -491,3 +473,29 @@ def conv_general_transpose( if data_format == "channel_first": res = tf.transpose(res, (0, dims + 1, *range(1, dims + 1))) return res + + +@with_unsupported_dtypes({"2.13.0 and below": ("bfloat16", "complex")}, backend_version) +def depthwise_conv2d( + x: Union[tf.Tensor, tf.Variable], + filters: Union[tf.Tensor, tf.Variable], + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NCHW": + x = tf.transpose(x, (0, 2, 3, 1)) + if tf.rank(filters) == 3: + filters = tf.expand_dims(filters, -1) + x = _pad_before_conv(x, filters, strides, padding, 2, dilations) + strides = [1, strides[0], strides[1], 1] + res = tf.nn.depthwise_conv2d(x, filters, strides, "VALID", "NHWC", dilations) + if data_format == "NCHW": + return tf.transpose(res, (0, 3, 1, 2)) + return res diff --git a/ivy/functional/backends/tensorflow/linear_algebra.py b/ivy/functional/backends/tensorflow/linear_algebra.py index 5281feade12b5..28c91cb2bcac4 100644 --- a/ivy/functional/backends/tensorflow/linear_algebra.py +++ b/ivy/functional/backends/tensorflow/linear_algebra.py @@ -79,6 +79,21 @@ def det( return tf.linalg.det(x) +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def diag( + x: Union[tf.Tensor, tf.Variable], + /, + *, + k: int = 0, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.diag(x, k=k) + + @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def diagonal( x: Union[tf.Tensor, tf.Variable], @@ -614,6 +629,21 @@ def trace( return tf.experimental.numpy.trace(x, offset=offset, axis1=axis1, axis2=axis2) +@with_unsupported_dtypes( + {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, + backend_version, +) +def vander( + x: Union[tf.Tensor, tf.Variable], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.experimental.numpy.vander(x, N=N, increasing=increasing) + + @with_unsupported_dtypes( {"2.13.0 and below": ("bfloat16", "float16", "complex")}, backend_version, @@ -667,36 +697,6 @@ def vector_norm( return tf.reduce_sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord) -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def diag( - x: Union[tf.Tensor, tf.Variable], - /, - *, - k: int = 0, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.diag(x, k=k) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("bfloat16", "float16", "complex", "unsigned")}, - backend_version, -) -def vander( - x: Union[tf.Tensor, tf.Variable], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.experimental.numpy.vander(x, N=N, increasing=increasing) - - @with_unsupported_dtypes( { "2.13.0 and below": ( diff --git a/ivy/functional/backends/tensorflow/manipulation.py b/ivy/functional/backends/tensorflow/manipulation.py index 41800bca81d4f..a26c08e2c8abe 100644 --- a/ivy/functional/backends/tensorflow/manipulation.py +++ b/ivy/functional/backends/tensorflow/manipulation.py @@ -15,12 +15,47 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _reshape_fortran_tf(x, shape): if len(x.shape) > 0: x = tf.transpose(x) return tf.transpose(tf.reshape(x, shape[::-1])) +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def clip( + x: Union[tf.Tensor, tf.Variable], + x_min: Union[Number, tf.Tensor, tf.Variable], + x_max: Union[Number, tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): + promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) + promoted_type = ivy.as_native_dtype( + ivy.promote_types(promoted_type, x_max.dtype) + ) + x = tf.cast(x, promoted_type) + x_min = tf.cast(x_min, promoted_type) + x_max = tf.cast(x_max, promoted_type) + if tf.size(x) == 0: + ret = x + elif x.dtype == tf.bool: + ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) + ret = tf.cast(ret, x.dtype) + else: + ret = tf.clip_by_value(x, x_min, x_max) + return ret + + # Array API Standard # # -------------------# @@ -54,6 +89,14 @@ def concat( raise ivy.utils.exceptions.IvyIndexError(error) +def constant_pad( + x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None +): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width, constant_values=value) + + def expand_dims( x: Union[tf.Tensor, tf.Variable], /, @@ -106,6 +149,18 @@ def permute_dims( return tf.transpose(x, perm=axes) +@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) +def repeat( + x: Union[tf.Tensor, tf.Variable], + /, + repeats: Union[int, List[int]], + *, + axis: int = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + return tf.repeat(x, repeats, axis) + + def reshape( x: Union[tf.Tensor, tf.Variable], /, @@ -148,6 +203,47 @@ def roll( return ret +# Extra # +# ------# + + +def split( + x: Union[tf.Tensor, tf.Variable], + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> Union[tf.Tensor, tf.Variable]: + if x.shape == (): + if num_or_size_splits is not None and num_or_size_splits != 1: + raise ivy.utils.exceptions.IvyException( + "input array had no shape, but num_sections specified was {}".format( + num_or_size_splits + ) + ) + return [x] + if num_or_size_splits is None: + dim_size = tf.shape(x)[axis] + num_or_size_splits = int(dim_size) + if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): + num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) + num_or_size_splits = num_or_size_splits.numpy().tolist() + elif isinstance(num_or_size_splits, int) and with_remainder: + num_chunks = x.shape[axis] / num_or_size_splits + num_chunks_int = math.floor(num_chunks) + remainder = num_chunks - num_chunks_int + if remainder != 0: + num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ + int(remainder * num_or_size_splits) + ] + + return tf.split(x, num_or_size_splits, axis) + + def squeeze( x: Union[tf.Tensor, tf.Variable], /, @@ -202,57 +298,25 @@ def stack( raise ivy.utils.exceptions.IvyIndexError(e) -# Extra # -# ------# - - -def split( - x: Union[tf.Tensor, tf.Variable], +def swapaxes( + x, + axis0, + axis1, /, *, copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], Union[tf.Tensor, tf.Variable]] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> Union[tf.Tensor, tf.Variable]: - if x.shape == (): - if num_or_size_splits is not None and num_or_size_splits != 1: - raise ivy.utils.exceptions.IvyException( - "input array had no shape, but num_sections specified was {}".format( - num_or_size_splits - ) - ) - return [x] - if num_or_size_splits is None: - dim_size = tf.shape(x)[axis] - num_or_size_splits = int(dim_size) - if isinstance(num_or_size_splits, (tf.Tensor, tf.Variable)): - num_or_size_splits = tf.cast(num_or_size_splits, tf.int32) - num_or_size_splits = num_or_size_splits.numpy().tolist() - elif isinstance(num_or_size_splits, int) and with_remainder: - num_chunks = x.shape[axis] / num_or_size_splits - num_chunks_int = math.floor(num_chunks) - remainder = num_chunks - num_chunks_int - if remainder != 0: - num_or_size_splits = [num_or_size_splits] * num_chunks_int + [ - int(remainder * num_or_size_splits) - ] - - return tf.split(x, num_or_size_splits, axis) - - -@with_supported_dtypes({"2.13.0 and below": ("int32", "int64")}, backend_version) -def repeat( - x: Union[tf.Tensor, tf.Variable], - /, - repeats: Union[int, List[int]], - *, - axis: int = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - return tf.repeat(x, repeats, axis) +): + x_shape = x.shape + num_dims = len(x_shape) + axis0 %= num_dims + axis1 %= num_dims + config = list(range(num_dims)) + config.pop(axis0) + config.insert(axis0, axis1) + config.pop(axis1) + config.insert(axis1, axis0) + return tf.transpose(x, config) @with_unsupported_dtypes( @@ -293,68 +357,6 @@ def tile( return tf.tile(x, repeats) -def constant_pad( - x, /, pad_width, *, value=0, out: Optional[Union[tf.Tensor, tf.Variable]] = None -): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width, constant_values=value) - - -def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): - if x.shape == (): - x = tf.reshape(x, (-1,)) - return tf.pad(x, pad_width) - - -def swapaxes( - x, - axis0, - axis1, - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -): - x_shape = x.shape - num_dims = len(x_shape) - axis0 %= num_dims - axis1 %= num_dims - config = list(range(num_dims)) - config.pop(axis0) - config.insert(axis0, axis1) - config.pop(axis1) - config.insert(axis1, axis0) - return tf.transpose(x, config) - - -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def clip( - x: Union[tf.Tensor, tf.Variable], - x_min: Union[Number, tf.Tensor, tf.Variable], - x_max: Union[Number, tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if hasattr(x_min, "dtype") and hasattr(x_max, "dtype"): - promoted_type = ivy.as_native_dtype(ivy.promote_types(x.dtype, x_min.dtype)) - promoted_type = ivy.as_native_dtype( - ivy.promote_types(promoted_type, x_max.dtype) - ) - x = tf.cast(x, promoted_type) - x_min = tf.cast(x_min, promoted_type) - x_max = tf.cast(x_max, promoted_type) - if tf.size(x) == 0: - ret = x - elif x.dtype == tf.bool: - ret = tf.clip_by_value(tf.cast(x, tf.float16), x_min, x_max) - ret = tf.cast(ret, x.dtype) - else: - ret = tf.clip_by_value(x, x_min, x_max) - return ret - - def unstack( x: Union[tf.Tensor, tf.Variable], /, @@ -369,3 +371,9 @@ def unstack( if keepdims: return [tf.expand_dims(r, axis) for r in ret] return ret + + +def zero_pad(x, /, pad_width, *, out: Optional[Union[tf.Tensor, tf.Variable]] = None): + if x.shape == (): + x = tf.reshape(x, (-1,)) + return tf.pad(x, pad_width) diff --git a/ivy/functional/backends/tensorflow/random.py b/ivy/functional/backends/tensorflow/random.py index b0f3b97a3d877..ae0f6b4a26c31 100644 --- a/ivy/functional/backends/tensorflow/random.py +++ b/ivy/functional/backends/tensorflow/random.py @@ -22,47 +22,6 @@ from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, tf.Tensor, tf.Variable] = 0.0, - high: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, - dtype: DType, - device: str, - seed: Optional[int] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - shape = _check_bounds_and_get_shape(low, high, shape).shape - low = tf.cast(low, dtype) - high = tf.cast(high, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) - - -def random_normal( - *, - mean: Union[float, tf.Tensor, tf.Variable] = 0.0, - std: Union[float, tf.Tensor, tf.Variable] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: DType, - seed: Optional[int] = None, - device: str, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - mean = tf.cast(mean, dtype) - std = tf.cast(std, dtype) - if seed: - tf.random.set_seed(seed) - return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) - - @with_unsupported_dtypes({"2.13.0 and below": ("bfloat16",)}, backend_version) def multinomial( population_size: int, @@ -142,6 +101,47 @@ def randint( return tf.cast(tf.random.uniform(shape, low, high, "float32", seed=seed), dtype) +def random_normal( + *, + mean: Union[float, tf.Tensor, tf.Variable] = 0.0, + std: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: DType, + seed: Optional[int] = None, + device: str, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + mean = tf.cast(mean, dtype) + std = tf.cast(std, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.normal(shape, mean, std, dtype=dtype, seed=seed) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, tf.Tensor, tf.Variable] = 0.0, + high: Union[float, tf.Tensor, tf.Variable] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int], tf.Tensor]] = None, + dtype: DType, + device: str, + seed: Optional[int] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + shape = _check_bounds_and_get_shape(low, high, shape).shape + low = tf.cast(low, dtype) + high = tf.cast(high, dtype) + if seed: + tf.random.set_seed(seed) + return tf.random.uniform(shape, low, high, dtype=dtype, seed=seed) + + def seed(*, seed_value: int = 0) -> None: tf.random.set_seed(seed_value) return diff --git a/ivy/functional/backends/tensorflow/searching.py b/ivy/functional/backends/tensorflow/searching.py index 047441cff9b08..13d461c8c54c5 100644 --- a/ivy/functional/backends/tensorflow/searching.py +++ b/ivy/functional/backends/tensorflow/searching.py @@ -77,6 +77,29 @@ def argmin( return tf.cast(ret, dtype) if dtype is not None else ret +# Extra # +# ----- # + + +def argwhere( + x: Union[tf.Tensor, tf.Variable], + /, + *, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if isinstance(x, tf.Variable): + x_ndim = x.shape.rank + else: + x_ndim = x.ndim + if x_ndim == 0: + return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") + where_x = tf.experimental.numpy.nonzero(x) + res = tf.experimental.numpy.concatenate( + [tf.expand_dims(item, -1) for item in where_x], -1 + ) + return res + + def nonzero( x: Union[tf.Tensor, tf.Variable], /, @@ -114,26 +137,3 @@ def where( ) -> Union[tf.Tensor, tf.Variable]: x1, x2 = ivy.promote_types_of_inputs(x1, x2) return tf.cast(tf.experimental.numpy.where(condition, x1, x2), x1.dtype) - - -# Extra # -# ----- # - - -def argwhere( - x: Union[tf.Tensor, tf.Variable], - /, - *, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - if isinstance(x, tf.Variable): - x_ndim = x.shape.rank - else: - x_ndim = x.ndim - if x_ndim == 0: - return tf.zeros(shape=[int(bool(x)), 0], dtype="int64") - where_x = tf.experimental.numpy.nonzero(x) - res = tf.experimental.numpy.concatenate( - [tf.expand_dims(item, -1) for item in where_x], -1 - ) - return res diff --git a/ivy/functional/backends/tensorflow/sorting.py b/ivy/functional/backends/tensorflow/sorting.py index 06403a1d358ca..2e1202a943ee0 100644 --- a/ivy/functional/backends/tensorflow/sorting.py +++ b/ivy/functional/backends/tensorflow/sorting.py @@ -27,29 +27,6 @@ def argsort( return tf.cast(ret, dtype=tf.int64) -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def sort( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - # TODO: handle stable sort when it's supported in tensorflow - # currently it supports only quicksort (unstable) - direction = "DESCENDING" if descending else "ASCENDING" - x = tf.convert_to_tensor(x) - is_bool = x.dtype.is_bool - if is_bool: - x = tf.cast(x, tf.int32) - ret = tf.sort(x, axis=axis, direction=direction) - if is_bool: - ret = tf.cast(ret, dtype=tf.bool) - return ret - - # msort @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) def msort( @@ -103,3 +80,26 @@ def searchsorted( if is_supported_int_ret_dtype: return tf.searchsorted(x, v, side=side, out_type=ret_dtype) return tf.cast(tf.searchsorted(x, v, side=side), ret_dtype) + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def sort( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + # TODO: handle stable sort when it's supported in tensorflow + # currently it supports only quicksort (unstable) + direction = "DESCENDING" if descending else "ASCENDING" + x = tf.convert_to_tensor(x) + is_bool = x.dtype.is_bool + if is_bool: + x = tf.cast(x, tf.int32) + ret = tf.sort(x, axis=axis, direction=direction) + if is_bool: + ret = tf.cast(ret, dtype=tf.bool) + return ret diff --git a/ivy/functional/backends/tensorflow/statistical.py b/ivy/functional/backends/tensorflow/statistical.py index 53fc8b5bb57e6..4ab04d803cd90 100644 --- a/ivy/functional/backends/tensorflow/statistical.py +++ b/ivy/functional/backends/tensorflow/statistical.py @@ -9,21 +9,82 @@ from . import backend_version from ivy.utils.einsum_parser import legalise_einsum_expr -# Array API Standard # -# -------------------# +# --- Helpers --- # +# --------------- # -@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) -def min( + +def _infer_dtype(dtype: tf.DType): + default_dtype = ivy.infer_default_dtype(dtype) + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return default_dtype + return dtype + + +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) +def cumprod( x: Union[tf.Tensor, tf.Variable], /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[tf.DType] = None, out: Optional[Union[tf.Tensor, tf.Variable]] = None, ) -> Union[tf.Tensor, tf.Variable]: - axis = tuple(axis) if isinstance(axis, list) else axis - return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumprod(x, axis, exclusive, reverse) + + +def cumsum( + x: Union[tf.Tensor, tf.Variable], + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if dtype is tf.bool: + dtype = ivy.default_int_dtype() + elif ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + else: + dtype = _infer_dtype(x.dtype) + dtype = ivy.as_native_dtype(dtype) + x = tf.cast(x, dtype) + return tf.math.cumsum(x, axis, exclusive, reverse) + + +@with_unsupported_dtypes( + {"2.13.0 and below": ("unsigned", "int8", "int16")}, + backend_version, +) +def einsum( + equation: str, + *operands: Union[tf.Tensor, tf.Variable], + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + dtype = _get_promoted_type_of_operands(operands) + equation = legalise_einsum_expr(*[equation, *operands]) + return tf.cast(tf.einsum(equation, *operands), dtype) @with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) @@ -51,11 +112,21 @@ def mean( return tf.math.reduce_mean(x, axis=axis, keepdims=keepdims) -def _infer_dtype(dtype: tf.DType): - default_dtype = ivy.infer_default_dtype(dtype) - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return default_dtype - return dtype +# Array API Standard # +# -------------------# + + +@with_unsupported_dtypes({"2.13.0 and below": ("complex",)}, backend_version) +def min( + x: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + axis = tuple(axis) if isinstance(axis, list) else axis + return tf.math.reduce_min(x, axis=axis, keepdims=keepdims) def prod( @@ -144,65 +215,3 @@ def var( * size / (size - correction) ) - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.13.0 and below": "bfloat16"}, backend_version) -def cumprod( - x: Union[tf.Tensor, tf.Variable], - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumprod(x, axis, exclusive, reverse) - - -def cumsum( - x: Union[tf.Tensor, tf.Variable], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[tf.DType] = None, - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if dtype is tf.bool: - dtype = ivy.default_int_dtype() - elif ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - else: - dtype = _infer_dtype(x.dtype) - dtype = ivy.as_native_dtype(dtype) - x = tf.cast(x, dtype) - return tf.math.cumsum(x, axis, exclusive, reverse) - - -@with_unsupported_dtypes( - {"2.13.0 and below": ("unsigned", "int8", "int16")}, - backend_version, -) -def einsum( - equation: str, - *operands: Union[tf.Tensor, tf.Variable], - out: Optional[Union[tf.Tensor, tf.Variable]] = None, -) -> Union[tf.Tensor, tf.Variable]: - dtype = _get_promoted_type_of_operands(operands) - equation = legalise_einsum_expr(*[equation, *operands]) - return tf.cast(tf.einsum(equation, *operands), dtype) diff --git a/ivy/functional/backends/torch/activations.py b/ivy/functional/backends/torch/activations.py index afb85ebe96f5d..82194664c14e1 100644 --- a/ivy/functional/backends/torch/activations.py +++ b/ivy/functional/backends/torch/activations.py @@ -17,25 +17,6 @@ from . import backend_version -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def relu( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.relu(x) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def leaky_relu( - x: torch.Tensor, - /, - *, - alpha: float = 0.2, - complex_mode="jax", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.leaky_relu(x, alpha) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def gelu( x: torch.Tensor, @@ -52,43 +33,31 @@ def gelu( return torch.nn.functional.gelu(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - if not ivy.is_array(x): - x = torch.tensor(x) - return torch.sigmoid(x, out=out) - - -sigmoid.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def softmax( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "complex", + "float16", + ) + }, + backend_version, +) +def hardswish( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if axis is None: - axis = -1 - return torch.nn.functional.softmax(x, axis) + return torch.nn.functional.hardswish(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def softplus( +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def leaky_relu( x: torch.Tensor, /, *, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, + alpha: float = 0.2, complex_mode="jax", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - kwargs = { - k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None - } - return torch.nn.functional.softplus(x, **kwargs) + return torch.nn.functional.leaky_relu(x, alpha) @with_unsupported_dtypes( @@ -125,16 +94,47 @@ def mish(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.mish(x) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "complex", - "float16", - ) - }, - backend_version, -) -def hardswish( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def relu( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.nn.functional.hardswish(x) + return torch.relu(x) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sigmoid(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + if not ivy.is_array(x): + x = torch.tensor(x) + return torch.sigmoid(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def softmax( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if axis is None: + axis = -1 + return torch.nn.functional.softmax(x, axis) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def softplus( + x: torch.Tensor, + /, + *, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, + complex_mode="jax", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + kwargs = { + k: v for k, v in {"beta": beta, "threshold": threshold}.items() if v is not None + } + return torch.nn.functional.softplus(x, **kwargs) + + +sigmoid.support_native_out = True diff --git a/ivy/functional/backends/torch/creation.py b/ivy/functional/backends/torch/creation.py index 7c5a41482e44c..8f1603b9c8f79 100644 --- a/ivy/functional/backends/torch/creation.py +++ b/ivy/functional/backends/torch/creation.py @@ -25,6 +25,10 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + # noinspection PyProtectedMember @@ -47,6 +51,28 @@ def _differentiable_linspace(start, stop, num, *, device, dtype=None): return res +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +def _stack_tensors(x, dtype): + if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): + for i, item in enumerate(x): + x[i] = _stack_tensors(item, dtype) + x = torch.stack(x) + else: + if isinstance(x, (list, tuple)): + if isinstance(x[0], torch.Tensor): + x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) + else: + x = torch.as_tensor(x, dtype=dtype) + return x + + +# --- Main --- # +# ------------ # + + @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def arange( start: float, @@ -78,23 +104,6 @@ def arange( return torch.arange(start, stop, step, dtype=dtype, device=device) -arange.support_native_out = True - - -def _stack_tensors(x, dtype): - if isinstance(x, (list, tuple)) and len(x) != 0 and isinstance(x[0], (list, tuple)): - for i, item in enumerate(x): - x[i] = _stack_tensors(item, dtype) - x = torch.stack(x) - else: - if isinstance(x, (list, tuple)): - if isinstance(x[0], torch.Tensor): - x = torch.stack([torch.as_tensor(i, dtype=dtype) for i in x]) - else: - x = torch.as_tensor(x, dtype=dtype) - return x - - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) @asarray_to_native_arrays_and_back @asarray_infer_device @@ -137,6 +146,17 @@ def asarray( return ret.clone().detach() if copy else ret +def copy_array( + x: torch.Tensor, + *, + to_ivy_array: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if to_ivy_array: + return ivy.to_ivy(x.clone()) + return x.clone() + + def empty( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -152,9 +172,6 @@ def empty( ) -empty.support_native_out = True - - def empty_like( x: torch.Tensor, /, @@ -225,14 +242,23 @@ def eye( return ret -eye.support_native_out = True - - def from_dlpack(x, /, *, out: Optional[torch.Tensor] = None): x = x.detach() if x.requires_grad else x return torch.utils.dlpack.from_dlpack(x) +def frombuffer( + buffer: bytes, + dtype: Optional[torch.dtype] = float, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> torch.Tensor: + buffer_copy = copy.deepcopy(buffer) + dtype = ivy.as_native_dtype(dtype) + + return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) + + def full( shape: Union[ivy.NativeShape, Sequence[int]], fill_value: Union[int, float, bool], @@ -254,9 +280,6 @@ def full( ) -full.support_native_out = True - - def full_like( x: torch.Tensor, /, @@ -270,10 +293,6 @@ def full_like( return torch.full_like(x, fill_value, dtype=dtype, device=device) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) - - @with_unsupported_device_and_dtypes( {"2.0.1 and below": {"cpu": ("float16",)}}, backend_version ) @@ -325,9 +344,6 @@ def linspace( return ans.to(dtype) -linspace.support_native_out = True - - def linspace_helper(start, stop, num, axis=None, *, dtype=None, device): num = num.detach().numpy().item() if isinstance(num, torch.Tensor) else num start_is_array = isinstance(start, torch.Tensor) @@ -428,6 +444,46 @@ def meshgrid( return res +def one_hot( + indices: torch.Tensor, + depth: int, + /, + *, + on_value: Optional[torch.Tensor] = None, + off_value: Optional[torch.Tensor] = None, + axis: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = torch.float32 + else: + if not on_none: + dtype = torch.tensor(on_value).dtype + elif not off_none: + dtype = torch.tensor(off_value).dtype + else: + dtype = ivy.as_native_dtype(dtype) + + on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) + off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) + + res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) + + if not on_none or not off_none: + res = torch.where(res == 1, on_value, off_value) + + if axis is not None: + res = torch.moveaxis(res, -1, axis) + + return res.to(device, dtype) + + def ones( shape: Union[ivy.NativeShape, Sequence[int]], *, @@ -438,18 +494,21 @@ def ones( return torch.ones(shape, dtype=dtype, device=device, out=out) -ones.support_native_out = True - - -def ones_like_v_0p4p0_and_above( +def ones_like_v_0p1p12_to_0p2p0( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.ones_like(x, dtype=dtype, device=device) +): + if len(x.shape) == 1: + for i in range(x.shape[0]): + x[i] = 1 + return x + for i in range(x.shape[0]): + x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) + return x def ones_like_v_0p3p0_to_0p3p1( @@ -463,21 +522,15 @@ def ones_like_v_0p3p0_to_0p3p1( return torch.ones_like(x, out=out) -def ones_like_v_0p1p12_to_0p2p0( +def ones_like_v_0p4p0_and_above( x: torch.Tensor, /, *, dtype: torch.dtype, device: torch.device, out: Optional[torch.Tensor] = None, -): - if len(x.shape) == 1: - for i in range(x.shape[0]): - x[i] = 1 - return x - for i in range(x.shape[0]): - x[i, :] = ones_like_v_0p1p12_to_0p2p0(x[i, :]) - return x +) -> torch.Tensor: + return torch.ones_like(x, dtype=dtype, device=device) def tril( @@ -486,16 +539,26 @@ def tril( return torch.tril(x, diagonal=k, out=out) -tril.support_native_out = True - - def triu( x: torch.Tensor, /, *, k: int = 0, out: Optional[torch.Tensor] = None ) -> torch.Tensor: return torch.triu(x, diagonal=k, out=out) -triu.support_native_out = True +def triu_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, + /, + *, + device: torch.device, +) -> Tuple[torch.Tensor]: + n_cols = n_rows if n_cols is None else n_cols + return tuple( + torch.triu_indices( + row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device + ) + ) def zeros( @@ -508,9 +571,6 @@ def zeros( return torch.zeros(shape, dtype=dtype, device=device, out=out) -zeros.support_native_out = True - - def zeros_like( x: torch.Tensor, /, @@ -527,82 +587,12 @@ def zeros_like( array = asarray - - -def copy_array( - x: torch.Tensor, - *, - to_ivy_array: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if to_ivy_array: - return ivy.to_ivy(x.clone()) - return x.clone() - - -def one_hot( - indices: torch.Tensor, - depth: int, - /, - *, - on_value: Optional[torch.Tensor] = None, - off_value: Optional[torch.Tensor] = None, - axis: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - on_none = on_value is None - off_none = off_value is None - - if dtype is None: - if on_none and off_none: - dtype = torch.float32 - else: - if not on_none: - dtype = torch.tensor(on_value).dtype - elif not off_none: - dtype = torch.tensor(off_value).dtype - else: - dtype = ivy.as_native_dtype(dtype) - - on_value = torch.tensor(1.0) if on_none else torch.tensor(on_value, dtype=dtype) - off_value = torch.tensor(0.0) if off_none else torch.tensor(off_value, dtype=dtype) - - res = torch.nn.functional.one_hot(indices.to(torch.int64), depth) - - if not on_none or not off_none: - res = torch.where(res == 1, on_value, off_value) - - if axis is not None: - res = torch.moveaxis(res, -1, axis) - - return res.to(device, dtype) - - -def frombuffer( - buffer: bytes, - dtype: Optional[torch.dtype] = float, - count: Optional[int] = -1, - offset: Optional[int] = 0, -) -> torch.Tensor: - buffer_copy = copy.deepcopy(buffer) - dtype = ivy.as_native_dtype(dtype) - - return torch.frombuffer(buffer_copy, dtype=dtype, count=count, offset=offset) - - -def triu_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, - /, - *, - device: torch.device, -) -> Tuple[torch.Tensor]: - n_cols = n_rows if n_cols is None else n_cols - return tuple( - torch.triu_indices( - row=n_rows, col=n_cols, offset=k, dtype=torch.int64, device=device - ) - ) +arange.support_native_out = True +empty.support_native_out = True +eye.support_native_out = True +full.support_native_out = True +linspace.support_native_out = True +ones.support_native_out = True +tril.support_native_out = True +triu.support_native_out = True +zeros.support_native_out = True diff --git a/ivy/functional/backends/torch/data_type.py b/ivy/functional/backends/torch/data_type.py index 05b4ccb37a4c4..67f5aeabdaf84 100644 --- a/ivy/functional/backends/torch/data_type.py +++ b/ivy/functional/backends/torch/data_type.py @@ -23,7 +23,6 @@ torch.complex128: "complex128", torch.bool: "bool", } - native_dtype_dict = { "int8": torch.int8, "int16": torch.int16, @@ -68,73 +67,6 @@ def smallest_normal(self): return self._torch_finfo.tiny -# Array API Standard # -# -------------------# - - -def astype( - x: torch.Tensor, - dtype: torch.dtype, - /, - *, - copy: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if x.dtype == dtype: - return x.clone() if copy else x - return x.to(dtype) - - -def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: - try: - return list(torch.broadcast_tensors(*arrays)) - except RuntimeError as e: - raise ivy.utils.exceptions.IvyBroadcastShapeError(e) - - -def broadcast_to( - x: torch.Tensor, - /, - shape: Union[ivy.NativeShape, Sequence[int]], - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) - if x.ndim > len(shape): - return torch.broadcast_to(x.reshape(-1), shape) - return torch.broadcast_to(x, shape) - - -@_handle_nestable_dtype_info -def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return Finfo(torch.finfo(ivy.as_native_dtype(type))) - - -@_handle_nestable_dtype_info -def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: - if isinstance(type, (torch.Tensor, np.ndarray)): - type = type.dtype - return torch.iinfo(ivy.as_native_dtype(type)) - - -def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: - input = [] - for val in arrays_and_dtypes: - torch_val = as_native_dtype(val) - if isinstance(torch_val, torch.dtype): - torch_val = torch.tensor(1, dtype=torch_val) - input.append(torch_val) - - result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) - - for i in range(2, len(input)): - result = torch.tensor(1, dtype=torch.result_type(result, input[i])) - return as_ivy_dtype(result.dtype) - - # Extra # # ------# @@ -204,6 +136,44 @@ def as_native_dtype( ) +# Array API Standard # +# -------------------# + + +def astype( + x: torch.Tensor, + dtype: torch.dtype, + /, + *, + copy: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if x.dtype == dtype: + return x.clone() if copy else x + return x.to(dtype) + + +def broadcast_arrays(*arrays: torch.Tensor) -> List[torch.Tensor]: + try: + return list(torch.broadcast_tensors(*arrays)) + except RuntimeError as e: + raise ivy.utils.exceptions.IvyBroadcastShapeError(e) + + +def broadcast_to( + x: torch.Tensor, + /, + shape: Union[ivy.NativeShape, Sequence[int]], + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_shapes_broadcastable(x.shape, shape) + if x.ndim > len(shape): + return torch.broadcast_to(x.reshape(-1), shape) + return torch.broadcast_to(x, shape) + + def dtype(x: Union[torch.tensor, np.ndarray], *, as_native: bool = False) -> ivy.Dtype: if as_native: return ivy.as_native_dtype(x.dtype) @@ -224,6 +194,20 @@ def dtype_bits(dtype_in: Union[torch.dtype, str, np.dtype], /) -> int: ) +@_handle_nestable_dtype_info +def finfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> Finfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return Finfo(torch.finfo(ivy.as_native_dtype(type))) + + +@_handle_nestable_dtype_info +def iinfo(type: Union[torch.dtype, str, torch.Tensor, np.ndarray], /) -> torch.iinfo: + if isinstance(type, (torch.Tensor, np.ndarray)): + type = type.dtype + return torch.iinfo(ivy.as_native_dtype(type)) + + def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: if not ivy.is_hashable_dtype(dtype_in): return False @@ -231,3 +215,18 @@ def is_native_dtype(dtype_in: Union[torch.dtype, str], /) -> bool: return True else: return False + + +def result_type(*arrays_and_dtypes: Union[torch.tensor, torch.dtype]) -> ivy.Dtype: + input = [] + for val in arrays_and_dtypes: + torch_val = as_native_dtype(val) + if isinstance(torch_val, torch.dtype): + torch_val = torch.tensor(1, dtype=torch_val) + input.append(torch_val) + + result = torch.tensor(1, dtype=torch.result_type(input[0], input[1])) + + for i in range(2, len(input)): + result = torch.tensor(1, dtype=torch.result_type(result, input[i])) + return as_ivy_dtype(result.dtype) diff --git a/ivy/functional/backends/torch/device.py b/ivy/functional/backends/torch/device.py index e7f768a7f5fe1..4a86dcf02f129 100644 --- a/ivy/functional/backends/torch/device.py +++ b/ivy/functional/backends/torch/device.py @@ -18,37 +18,26 @@ torch_scatter = None -# API # -# ----# +class Profiler(BaseProfiler): + def __init__(self, save_dir: str): + super(Profiler, self).__init__(save_dir) + self._prof = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True + ) -def dev( - x: torch.Tensor, /, *, as_native: bool = False -) -> Union[ivy.Device, torch.device]: - dv = x.device - if as_native: - if isinstance(dv, torch.device): - dv = dv.type - elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return torch.device(dv.replace("gpu", "mps")) - return torch.device(dv.replace("gpu", "cuda")) - return as_ivy_dev(dv) + def start(self): + self._prof.__enter__() + def stop(self): + self._prof.__exit__(None, None, None) + self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) -def to_device( - x: torch.Tensor, - device: torch.device, - /, - *, - stream: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if device is None: - return x - ret = x.to(as_native_dev(device)) - if isinstance(x, torch.nn.Parameter): - return torch.nn.Parameter(ret) - return ret + def __enter__(self): + self.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() def as_ivy_dev(device: torch.device, /): @@ -89,10 +78,21 @@ def clear_cached_mem_on_dev(device: Union[ivy.Device, torch.device], /) -> None: mps.empty_cache() -def num_gpus() -> int: - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return 1 - return torch.cuda.device_count() +# API # +# ----# + + +def dev( + x: torch.Tensor, /, *, as_native: bool = False +) -> Union[ivy.Device, torch.device]: + dv = x.device + if as_native: + if isinstance(dv, torch.device): + dv = dv.type + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return torch.device(dv.replace("gpu", "mps")) + return torch.device(dv.replace("gpu", "cuda")) + return as_ivy_dev(dv) def gpu_is_available() -> bool: @@ -103,13 +103,6 @@ def gpu_is_available() -> bool: return False -# noinspection PyUnresolvedReferences -def tpu_is_available() -> bool: - if importlib.util.find_spec("torch_xla") is not None: - return True - return False - - def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): args, kwargs, device_shifting_dev = _shift_native_arrays_on_default_device( *args, device_shifting_dev=device_shifting_dev, **kwargs @@ -121,22 +114,30 @@ def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): return fn(*args, **kwargs) -class Profiler(BaseProfiler): - def __init__(self, save_dir: str): - super(Profiler, self).__init__(save_dir) - self._prof = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True - ) +def num_gpus() -> int: + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return 1 + return torch.cuda.device_count() - def start(self): - self._prof.__enter__() - def stop(self): - self._prof.__exit__(None, None, None) - self._prof.export_chrome_trace(os.path.join(self._save_dir, "trace.json")) +def to_device( + x: torch.Tensor, + device: torch.device, + /, + *, + stream: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if device is None: + return x + ret = x.to(as_native_dev(device)) + if isinstance(x, torch.nn.Parameter): + return torch.nn.Parameter(ret) + return ret - def __enter__(self): - self.start() - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() +# noinspection PyUnresolvedReferences +def tpu_is_available() -> bool: + if importlib.util.find_spec("torch_xla") is not None: + return True + return False diff --git a/ivy/functional/backends/torch/elementwise.py b/ivy/functional/backends/torch/elementwise.py index a0a3f5c11e053..1d03e1b29cfcc 100644 --- a/ivy/functional/backends/torch/elementwise.py +++ b/ivy/functional/backends/torch/elementwise.py @@ -13,12 +13,49 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _cast_for_unary_op(x): if not isinstance(x, torch.Tensor): x = torch.tensor(x) return x +# --- Main --- # +# ------------ # + + +@handle_numpy_arrays_in_specific_backend +def abs( + x: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = _cast_for_unary_op(x) + if x.dtype is torch.bool: + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.abs(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.acos(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.acosh(x, out=out) + + @handle_numpy_arrays_in_specific_backend def add( x1: Union[float, torch.Tensor], @@ -34,48 +71,69 @@ def add( return torch.add(x1, x2, out=out) -add.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_xor( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def angle( + input: torch.Tensor, /, *, + deg: Optional[bool] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_xor(x1, x2, out=out) + if deg: + return torch.angle(input, out=out) * (180 / pi) + else: + return torch.angle(input, out=out) -bitwise_xor.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.asin(x, out=out) -def imag( - val: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if val.dtype not in (torch.complex64, torch.complex128): - ret = torch.imag(val.to(torch.complex64)) - return ret.to(val.dtype) - return torch.imag(val) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.asinh(x, out=out) -imag.support_native_out = False +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.atan(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version +) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) @handle_numpy_arrays_in_specific_backend -def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def atan2( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.atan2(x1, x2, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.expm1(x, out=out) + return torch.atanh(x, out=out) -expm1.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def bitwise_and( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_and(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -87,68 +145,49 @@ def bitwise_invert( return torch.bitwise_not(x, out=out) -bitwise_invert.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.isfinite(x) - - +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def isinf( - x: torch.Tensor, +def bitwise_left_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, - detect_positive: bool = True, - detect_negative: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - if detect_negative and detect_positive: - return torch.isinf(x) - elif detect_negative: - return torch.isneginf(x) - elif detect_positive: - return torch.isposinf(x) - return torch.full_like(x, False, dtype=torch.bool) + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_left_shift(x1, x2, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def bitwise_or( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.eq(x1, x2, out=out) - - -equal.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + return torch.bitwise_or(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less_equal( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def bitwise_right_shift( + x1: Union[int, bool, torch.Tensor], + x2: Union[int, bool, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.less_equal(x1, x2, out=out) - - -less_equal.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) + x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) + return torch.bitwise_right_shift(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_and( +def bitwise_xor( x1: Union[int, bool, torch.Tensor], x2: Union[int, bool, torch.Tensor], /, @@ -156,10 +195,7 @@ def bitwise_and( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_and(x1, x2, out=out) - - -bitwise_and.support_native_out = True + return torch.bitwise_xor(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @@ -173,135 +209,162 @@ def ceil(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.ceil(x, out=out) -ceil.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.cos(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.floor(x, out=out) + return torch.cosh(x, out=out) -floor.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.deg2rad(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmin( - x1: torch.Tensor, - x2: torch.Tensor, +@handle_numpy_arrays_in_specific_backend +def divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.fmin(x1, x2, out=None) - - -fmin.support_native_out = True + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + ret = torch.div(x1, x2) + if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): + ret = ivy.astype(ret, x1.dtype, copy=False) + else: + ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) + return ret -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def asin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.asin(x, out=out) +def equal( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.eq(x1, x2, out=out) -asin.support_native_out = True +# Extra # +# ------# -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def asinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.asinh(x, out=out) + return torch.erf(x, out=out) -asinh.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.exp(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def sign( - x: torch.Tensor, +def exp2( + x: Union[torch.Tensor, float, list, tuple], /, *, - np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - if "complex" in str(x.dtype): - if np_variant: - return torch.where( - x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j - ) - return torch.sgn(x, out=out) - return torch.sign(x, out=out) - - -sign.support_native_out = True + return torch.exp2(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def expm1(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.sqrt(x, out=out) - - -sqrt.support_native_out = True + return torch.expm1(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def cosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def floor(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.cosh(x, out=out) - - -cosh.support_native_out = True + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.floor(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log10(x, out=out) - - -log10.support_native_out = True +def floor_divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if ivy.exists(out): + if not ivy.is_float_dtype(out): + return ivy.inplace_update( + out, torch.floor(torch.div(x1, x2)).type(out.dtype) + ) + return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log2(x, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmin( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.fmin(x1, x2, out=None) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "complex")}, + backend_version, +) @handle_numpy_arrays_in_specific_backend -def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log1p(x, out=out) - - -log1p.support_native_out = True +def fmod( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmod(x1, x2, out=None) -@handle_numpy_arrays_in_specific_backend -def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.isnan(x) +def gcd( + x1: Union[torch.Tensor, int, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.gcd(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def less( +def greater( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -309,14 +372,12 @@ def less( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.lt(x1, x2, out=out) - - -less.support_native_out = True + return torch.greater(x1, x2, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def multiply( +def greater_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -324,56 +385,73 @@ def multiply( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.multiply(x1, x2, out=out) + return torch.greater_equal(x1, x2, out=out) -multiply.support_native_out = True +def imag( + val: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if val.dtype not in (torch.complex64, torch.complex128): + ret = torch.imag(val.to(torch.complex64)) + return ret.to(val.dtype) + return torch.imag(val) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def cos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def isfinite(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.cos(x, out=out) + return torch.isfinite(x) -cos.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def isinf( + x: torch.Tensor, + /, + *, + detect_positive: bool = True, + detect_negative: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = _cast_for_unary_op(x) + if detect_negative and detect_positive: + return torch.isinf(x) + elif detect_negative: + return torch.isneginf(x) + elif detect_positive: + return torch.isposinf(x) + return torch.full_like(x, False, dtype=torch.bool) @handle_numpy_arrays_in_specific_backend -def logical_not( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: +def isnan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.logical_not(x.type(torch.bool), out=out) + return torch.isnan(x) -logical_not.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.isreal(x) +@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], +def lcm( + x1: torch.Tensor, + x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2) - if ivy.is_float_dtype(x1.dtype) or ivy.is_complex_dtype(x1.dtype): - ret = ivy.astype(ret, x1.dtype, copy=False) - else: - ret = ivy.astype(ret, ivy.default_float_dtype(as_native=True), copy=False) - return ret - - -divide.support_native_out = True + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.lcm(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater( +def less( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -381,15 +459,12 @@ def greater( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater(x1, x2, out=out) - - -greater.support_native_out = True + return torch.lt(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def greater_equal( +def less_equal( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, @@ -397,46 +472,60 @@ def greater_equal( out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.greater_equal(x1, x2, out=out) - - -greater_equal.support_native_out = True + return torch.less_equal(x1, x2, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def acos(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: +def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.acos(x, out=out) + return torch.log(x, out=out) -acos.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log10(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log10(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def lcm( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.lcm(x1, x2, out=out) +def log1p(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log1p(x, out=out) -lcm.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def log2(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.log2(x, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def logical_xor( +def logaddexp( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.logaddexp(x1, x2, out=out) -logical_xor.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def logaddexp2( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = promote_types_of_inputs(x1, x2) + if not ivy.is_float_dtype(x1): + x1 = x1.type(ivy.default_float_dtype(as_native=True)) + x2 = x2.type(ivy.default_float_dtype(as_native=True)) + return torch.logaddexp2(x1, x2, out=out) @handle_numpy_arrays_in_specific_backend @@ -446,7 +535,12 @@ def logical_and( return torch.logical_and(x1.type(torch.bool), x2.type(torch.bool), out=out) -logical_and.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def logical_not( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.logical_not(x.type(torch.bool), out=out) @handle_numpy_arrays_in_specific_backend @@ -456,27 +550,72 @@ def logical_or( return torch.logical_or(x1.type(torch.bool), x2.type(torch.bool), out=out) -logical_or.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def acosh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.acosh(x, out=out) - - -acosh.support_native_out = True +def logical_xor( + x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + return torch.logical_xor(x1.type(torch.bool), x2.type(torch.bool), out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sin(x, out=out) +def maximum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 >= x2, x1, x2, out=out) + return torch.maximum(x1, x2, out=out) -sin.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def minimum( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + use_where: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + if use_where: + return torch.where(x1 <= x2, x1, x2, out=out) + return torch.minimum(x1, x2, out=out) + + +@handle_numpy_arrays_in_specific_backend +def multiply( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.multiply(x1, x2, out=out) + + +def nan_to_num( + x: torch.Tensor, + /, + *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if copy: + return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) + else: + x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) + return x @handle_numpy_arrays_in_specific_backend @@ -487,9 +626,6 @@ def negative( return torch.neg(x, out=out) -negative.support_native_out = True - - @handle_numpy_arrays_in_specific_backend def not_equal( x1: Union[float, torch.Tensor], @@ -502,56 +638,108 @@ def not_equal( return torch.not_equal(x1, x2, out=out) -not_equal.support_native_out = True +@handle_numpy_arrays_in_specific_backend +def positive( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.positive(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def tanh( - x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None +def pow( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tanh(x, out=out) + x1, x2 = ivy.promote_types_of_inputs(x1, x2) + return torch.pow(x1, x2, out=out) -tanh.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("complex64", "complex128")}, backend_version +) +@handle_numpy_arrays_in_specific_backend +def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.rad2deg(x, out=out) + + +def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.real(x) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def reciprocal( + x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.reciprocal(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend -def floor_divide( +def remainder( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, + modulus: bool = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if ivy.exists(out): - if not ivy.is_float_dtype(out): - return ivy.inplace_update( - out, torch.floor(torch.div(x1, x2)).type(out.dtype) - ) - return torch.floor(torch.div(x1, x2), out=out).type(x1.dtype) + if not modulus: + res = x1 / x2 + res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) + diff = res - res_floored + diff, x2 = ivy.promote_types_of_inputs(diff, x2) + if ivy.exists(out): + if out.dtype != x2.dtype: + return ivy.inplace_update( + out, torch.round(torch.mul(diff, x2)).to(out.dtype) + ) + return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) + return torch.remainder(x1, x2, out=out).to(x1.dtype) -floor_divide.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@handle_numpy_arrays_in_specific_backend +def round( + x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + if "int" in str(x.dtype): + if ivy.exists(out): + return ivy.inplace_update(out, x) + return x + return torch.round(x, decimals=decimals, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def bitwise_or( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], +def sign( + x: torch.Tensor, /, *, + np_variant: Optional[bool] = True, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_or(x1, x2, out=out) + x = _cast_for_unary_op(x) + if "complex" in str(x.dtype): + if np_variant: + return torch.where( + x.real != 0, torch.sign(x.real) + 0.0j, torch.sign(x.imag) + 0.0j + ) + return torch.sgn(x, out=out) + return torch.sign(x, out=out) -bitwise_or.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def sin(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sin(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -561,15 +749,11 @@ def sinh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.sinh(x, out=out) -sinh.support_native_out = True - - +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def positive( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: +def sqrt(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: x = _cast_for_unary_op(x) - return torch.positive(x) + return torch.sqrt(x, out=out) @handle_numpy_arrays_in_specific_backend @@ -578,37 +762,35 @@ def square(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.T return torch.square(x, out=out) -square.support_native_out = True - - @handle_numpy_arrays_in_specific_backend -def pow( +def subtract( x1: Union[float, torch.Tensor], x2: Union[float, torch.Tensor], /, *, + alpha: Optional[Union[int, float]] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.pow(x1, x2, out=out) + if alpha not in (1, None): + return torch.subtract(x1, x2, alpha=alpha, out=out) + return torch.subtract(x1, x2, out=out) -pow.support_native_out = True +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +@handle_numpy_arrays_in_specific_backend +def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.tan(x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def round( - x: torch.Tensor, /, *, decimals: int = 0, out: Optional[torch.Tensor] = None +def tanh( + x: torch.Tensor, /, *, complex_mode="jax", out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if "int" in str(x.dtype): - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.round(x, decimals=decimals, out=out) - - -round.support_native_out = True + x = _cast_for_unary_op(x) + return torch.tanh(x, out=out) def trapz( @@ -634,9 +816,6 @@ def trapz( return torch.trapezoid(y, x=x, dim=axis) -trapz.support_native_out = False - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) @handle_numpy_arrays_in_specific_backend def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: @@ -649,388 +828,87 @@ def trunc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return ret -trunc.support_native_out = True - - +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @handle_numpy_arrays_in_specific_backend -def abs( - x: Union[float, torch.Tensor], +def trunc_divide( + x1: Union[float, torch.Tensor], + x2: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x = _cast_for_unary_op(x) - if x.dtype is torch.bool: - if ivy.exists(out): - return ivy.inplace_update(out, x) - return x - return torch.abs(x, out=out) - - -abs.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def logaddexp( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.logaddexp(x1, x2, out=out) - - -logaddexp.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def logaddexp2( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - if not ivy.is_float_dtype(x1): - x1 = x1.type(ivy.default_float_dtype(as_native=True)) - x2 = x2.type(ivy.default_float_dtype(as_native=True)) - return torch.logaddexp2(x1, x2, out=out) - - -logaddexp2.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def tan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.tan(x, out=out) + ret = torch.div(x1, x2, rounding_mode="trunc") + if ivy.is_float_dtype(x1.dtype): + ret = ret.to(x1.dtype) + else: + ret = ret.to(ivy.default_float_dtype(as_native=True)) + return ret +add.support_native_out = True +bitwise_xor.support_native_out = True +imag.support_native_out = False +expm1.support_native_out = True +bitwise_invert.support_native_out = True +equal.support_native_out = True +less_equal.support_native_out = True +bitwise_and.support_native_out = True +ceil.support_native_out = True +floor.support_native_out = True +fmin.support_native_out = True +asin.support_native_out = True +asinh.support_native_out = True +sign.support_native_out = True +sqrt.support_native_out = True +cosh.support_native_out = True +log10.support_native_out = True +log1p.support_native_out = True +less.support_native_out = True +multiply.support_native_out = True +cos.support_native_out = True +logical_not.support_native_out = True +divide.support_native_out = True +greater.support_native_out = True +greater_equal.support_native_out = True +acos.support_native_out = True +lcm.support_native_out = True +logical_xor.support_native_out = True +logical_and.support_native_out = True +logical_or.support_native_out = True +acosh.support_native_out = True +sin.support_native_out = True +negative.support_native_out = True +not_equal.support_native_out = True +tanh.support_native_out = True +floor_divide.support_native_out = True +bitwise_or.support_native_out = True +sinh.support_native_out = True +square.support_native_out = True +pow.support_native_out = True +round.support_native_out = True +trapz.support_native_out = False +trunc.support_native_out = True +abs.support_native_out = True +logaddexp.support_native_out = True +logaddexp2.support_native_out = True tan.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def atan(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.atan(x, out=out) - - atan.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version -) # TODO Fixed in PyTorch 1.12.1 (this note excludes complex) -@handle_numpy_arrays_in_specific_backend -def atan2( - x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - return torch.atan2(x1, x2, out=out) - - atan2.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def log(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.log(x, out=out) - - log.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def exp(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.exp(x, out=out) - - exp.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def exp2( - x: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.exp2(x, out=out) - - exp2.support_native_out = True - - -@handle_numpy_arrays_in_specific_backend -def subtract( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - alpha: Optional[Union[int, float]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if alpha not in (1, None): - return torch.subtract(x1, x2, alpha=alpha, out=out) - return torch.subtract(x1, x2, out=out) - - subtract.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def remainder( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - modulus: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if not modulus: - res = x1 / x2 - res_floored = torch.where(res >= 0, torch.floor(res), torch.ceil(res)) - diff = res - res_floored - diff, x2 = ivy.promote_types_of_inputs(diff, x2) - if ivy.exists(out): - if out.dtype != x2.dtype: - return ivy.inplace_update( - out, torch.round(torch.mul(diff, x2)).to(out.dtype) - ) - return torch.round(torch.mul(diff, x2), out=out).to(x1.dtype) - return torch.remainder(x1, x2, out=out).to(x1.dtype) - - remainder.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def atanh(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.atanh(x, out=out) - - atanh.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_right_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - x2 = torch.clamp(x2, min=0, max=torch.iinfo(x2.dtype).bits - 1) - return torch.bitwise_right_shift(x1, x2, out=out) - - bitwise_right_shift.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def bitwise_left_shift( - x1: Union[int, bool, torch.Tensor], - x2: Union[int, bool, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2, array_api_promotion=True) - return torch.bitwise_left_shift(x1, x2, out=out) - - bitwise_left_shift.support_native_out = True - - -# Extra # -# ------# - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) -@handle_numpy_arrays_in_specific_backend -def erf(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.erf(x, out=out) - - erf.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def minimum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 <= x2, x1, x2, out=out) - return torch.minimum(x1, x2, out=out) - - minimum.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def maximum( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - use_where: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - if use_where: - return torch.where(x1 >= x2, x1, x2, out=out) - return torch.maximum(x1, x2, out=out) - - maximum.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def reciprocal( - x: Union[float, torch.Tensor], /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.reciprocal(x, out=out) - - reciprocal.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def deg2rad(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.deg2rad(x, out=out) - - deg2rad.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("complex64", "complex128")}, backend_version -) -@handle_numpy_arrays_in_specific_backend -def rad2deg(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.rad2deg(x, out=out) - - rad2deg.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -@handle_numpy_arrays_in_specific_backend -def trunc_divide( - x1: Union[float, torch.Tensor], - x2: Union[float, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = ivy.promote_types_of_inputs(x1, x2) - ret = torch.div(x1, x2, rounding_mode="trunc") - if ivy.is_float_dtype(x1.dtype): - ret = ret.to(x1.dtype) - else: - ret = ret.to(ivy.default_float_dtype(as_native=True)) - return ret - - -@handle_numpy_arrays_in_specific_backend -def isreal(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.isreal(x) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "complex")}, - backend_version, -) -@handle_numpy_arrays_in_specific_backend -def fmod( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmod(x1, x2, out=None) - - fmod.support_native_out = True - - -def gcd( - x1: Union[torch.Tensor, int, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.gcd(x1, x2, out=out) - - gcd.support_native_out = True - - -def angle( - input: torch.Tensor, - /, - *, - deg: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if deg: - return torch.angle(input, out=out) * (180 / pi) - else: - return torch.angle(input, out=out) - - angle.support_native_out = True - - -def nan_to_num( - x: torch.Tensor, - /, - *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if copy: - return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf, out=out) - else: - x = torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) - return x - - -def real(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.real(x) diff --git a/ivy/functional/backends/torch/experimental/activations.py b/ivy/functional/backends/torch/experimental/activations.py index 98d72ac526c45..008312f394795 100644 --- a/ivy/functional/backends/torch/experimental/activations.py +++ b/ivy/functional/backends/torch/experimental/activations.py @@ -10,6 +10,16 @@ from . import backend_version +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def elu( + x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + ret = torch.nn.functional.elu(x, alpha) + if ivy.exists(out): + return ivy.inplace_update(out, ret).astype(x.dtype) + return ivy.astype(ret, x.dtype) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def logit( x: torch.Tensor, @@ -21,15 +31,11 @@ def logit( return torch.logit(x, eps=eps, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) -def thresholded_relu( - x: torch.Tensor, - /, - *, - threshold: Optional[Union[int, float]] = None, - out: Optional[torch.Tensor] = None, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def logsigmoid( + input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - return torch.threshold(x, threshold=threshold, value=0) + return torch.nn.functional.logsigmoid(input) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -37,13 +43,6 @@ def relu6(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Te return torch.nn.functional.relu6(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def logsigmoid( - input: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.nn.functional.logsigmoid(input) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def selu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: ret = torch.nn.functional.selu(x) @@ -57,11 +56,12 @@ def silu(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Ten return torch.nn.functional.silu(x) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def elu( - x: torch.Tensor, /, *, alpha: float = 1.0, out: Optional[torch.Tensor] = None +@with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) +def thresholded_relu( + x: torch.Tensor, + /, + *, + threshold: Optional[Union[int, float]] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - ret = torch.nn.functional.elu(x, alpha) - if ivy.exists(out): - return ivy.inplace_update(out, ret).astype(x.dtype) - return ivy.astype(ret, x.dtype) + return torch.threshold(x, threshold=threshold, value=0) diff --git a/ivy/functional/backends/torch/experimental/creation.py b/ivy/functional/backends/torch/experimental/creation.py index a1b94a8a894cf..f6762fa46f08a 100644 --- a/ivy/functional/backends/torch/experimental/creation.py +++ b/ivy/functional/backends/torch/experimental/creation.py @@ -12,33 +12,20 @@ ) from .. import backend_version -# noinspection PyProtectedMember - - -# Array API Standard # -# -------------------# - -@with_unsupported_device_and_dtypes( - {"2.0.1 and below": {"cpu": ("float16",)}}, - backend_version, -) -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def blackman_window( + size: int, + /, *, + periodic: bool = True, dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.kaiser_window( - window_length, - periodic, - beta, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.blackman_window( + size, + periodic=periodic, dtype=dtype, - layout=torch.strided, - device=None, - requires_grad=False, ) @@ -64,29 +51,6 @@ def hamming_window( ) -def vorbis_window( - window_length: torch.tensor, - *, - dtype: torch.dtype = torch.float32, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.tensor( - [ - round( - math.sin( - (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) - ), - 8, - ) - for i in range(1, window_length * 2)[0::2] - ], - dtype=dtype, - ) - - -vorbis_window.support_native_out = False - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def hann_window( size: int, @@ -103,7 +67,34 @@ def hann_window( ) -hann_window.support_native_out = False +# noinspection PyProtectedMember + + +# Array API Standard # +# -------------------# + + +@with_unsupported_device_and_dtypes( + {"2.0.1 and below": {"cpu": ("float16",)}}, + backend_version, +) +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.kaiser_window( + window_length, + periodic, + beta, + dtype=dtype, + layout=torch.strided, + device=None, + requires_grad=False, + ) def tril_indices( @@ -126,6 +117,19 @@ def tril_indices( ) +def trilu( + x: torch.Tensor, + /, + *, + k: int = 0, + upper: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if upper: + return torch.triu(x, diagonal=k, out=out) + return torch.tril(x, diagonal=k, out=out) + + def unsorted_segment_min( data: torch.Tensor, segment_ids: torch.Tensor, @@ -152,25 +156,6 @@ def unsorted_segment_min( return res -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def blackman_window( - size: int, - /, - *, - periodic: bool = True, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.blackman_window( - size, - periodic=periodic, - dtype=dtype, - ) - - -blackman_window.support_native_out = False - - def unsorted_segment_sum( data: torch.Tensor, segment_ids: torch.Tensor, @@ -196,17 +181,27 @@ def unsorted_segment_sum( return res -def trilu( - x: torch.Tensor, - /, +def vorbis_window( + window_length: torch.tensor, *, - k: int = 0, - upper: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if upper: - return torch.triu(x, diagonal=k, out=out) - return torch.tril(x, diagonal=k, out=out) + dtype: torch.dtype = torch.float32, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.tensor( + [ + round( + math.sin( + (ivy.pi / 2) * (math.sin(ivy.pi * (i) / (window_length * 2)) ** 2) + ), + 8, + ) + for i in range(1, window_length * 2)[0::2] + ], + dtype=dtype, + ) +vorbis_window.support_native_out = False +hann_window.support_native_out = False +blackman_window.support_native_out = False trilu.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/elementwise.py b/ivy/functional/backends/torch/experimental/elementwise.py index dcb4e30dcf1dc..ab2d5a0c2125d 100644 --- a/ivy/functional/backends/torch/experimental/elementwise.py +++ b/ivy/functional/backends/torch/experimental/elementwise.py @@ -14,49 +14,69 @@ from .. import backend_version -@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) -def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - return torch.lgamma(x, out=out) +# --- Helpers --- # +# --------------- # -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def fmax( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.fmax(x1, x2, out=None) - +def _are_suitable_types_for_torch_lerp(input, end, weight): + suitable_types = [ + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ] -fmax.support_native_out = True + if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): + return False + else: + if input.dtype not in suitable_types or end.dtype not in suitable_types: + return False + if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): + return False + else: + if isinstance(weight, torch.Tensor): + if weight.dtype not in [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ]: + return False -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: - x = _cast_for_unary_op(x) - return torch.sinc(x, out=out) + return True -sinc.support_native_out = True +# --- Main --- # +# ------------ # -def float_power( - x1: Union[torch.Tensor, float, list, tuple], - x2: Union[torch.Tensor, float, list, tuple], +def allclose( + x1: torch.Tensor, + x2: torch.Tensor, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - # Native out is supported but with restrictions leading - # to failures hence letting ivy handle it. - x1, x2 = promote_types_of_inputs(x1, x2) - return torch.float_power(x1, x2, out=out) +) -> bool: + ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) + return torch.tensor(ret) -float_power.support_native_out = True +def conj( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + conj_x = torch.conj(x) + return torch.resolve_conj(input=conj_x) def copysign( @@ -73,9 +93,6 @@ def copysign( return torch.copysign(torch.as_tensor(x1), x2, out=out) -copysign.support_native_out = True - - def count_nonzero( a: torch.Tensor, /, @@ -107,42 +124,6 @@ def count_nonzero( return x -count_nonzero.support_native_out = False - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def nansum( - x: torch.Tensor, - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[torch.dtype] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) - - -nansum.support_native_out = False - - -def isclose( - a: torch.Tensor, - b: torch.Tensor, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - -isclose.support_native_out = False - - def diff( x: Union[torch.Tensor, list, tuple], /, @@ -167,41 +148,14 @@ def diff( return torch.diff(x, n=n, dim=axis, prepend=prepend, append=append) -def signbit( - x: Union[torch.Tensor, float, int, list, tuple], - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.signbit(x, out=out) - - -signbit.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def hypot( - x1: torch.Tensor, - x2: torch.Tensor, +def digamma( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.hypot(x1, x2) - - -def allclose( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[torch.Tensor] = None, -) -> bool: - ret = torch.allclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan) - return torch.tensor(ret) + return torch.special.digamma(x, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -214,39 +168,39 @@ def fix( return torch.fix(x, out=out) -fix.support_native_out = True +def float_power( + x1: Union[torch.Tensor, float, list, tuple], + x2: Union[torch.Tensor, float, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # Native out is supported but with restrictions leading + # to failures hence letting ivy handle it. + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.float_power(x1, x2, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def nextafter( +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def fmax( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nextafter(x1, x2) - - -nextafter.support_native_out = True + x1, x2 = promote_types_of_inputs(x1, x2) + return torch.fmax(x1, x2, out=None) -def zeta( +def frexp( x: torch.Tensor, - q: torch.Tensor, /, *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) - temp = torch.logical_and(temp, torch.le(q, 0)) - nan_indices = torch.logical_or(temp, torch.lt(x, 1)) - result = torch.special.zeta(x, q) - result.masked_fill_(nan_indices, float("nan")) - return result - - -zeta.support_native_out = False + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + mantissa, exponent = torch.frexp(x, out=out) + return mantissa, exponent def gradient( @@ -270,25 +224,28 @@ def gradient( return grad -@with_supported_dtypes( - {"2.0.1 and below": ("float16", "float32", "float64")}, - backend_version, -) -def xlogy( - x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None -) -> torch.tensor: - x, y = promote_types_of_inputs(x, y) - return torch.xlogy(x, y, out=out) +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def hypot( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.hypot(x1, x2) -def conj( - x: torch.Tensor, +def isclose( + a: torch.Tensor, + b: torch.Tensor, /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - conj_x = torch.conj(x) - return torch.resolve_conj(input=conj_x) + return torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) def ldexp( @@ -301,39 +258,6 @@ def ldexp( return torch.ldexp(x1, x2, out=out) -def _are_suitable_types_for_torch_lerp(input, end, weight): - suitable_types = [ - torch.int8, - torch.int16, - torch.int32, - torch.int64, - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ] - - if not isinstance(input, torch.Tensor) or not isinstance(end, torch.Tensor): - return False - else: - if input.dtype not in suitable_types or end.dtype not in suitable_types: - return False - - if not isinstance(weight, float) and not isinstance(weight, torch.Tensor): - return False - else: - if isinstance(weight, torch.Tensor): - if weight.dtype not in [ - torch.float16, - torch.bfloat16, - torch.float32, - torch.float64, - ]: - return False - - return True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def lerp( input: torch.Tensor, @@ -346,40 +270,100 @@ def lerp( return torch.lerp(input, end, weight, out=out) -lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( - _are_suitable_types_for_torch_lerp(input, end, weight) -) -lerp.support_native_out = True +@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, backend_version) +def lgamma(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + return torch.lgamma(x, out=out) -def frexp( +def modf( x: torch.Tensor, /, *, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - mantissa, exponent = torch.frexp(x, out=out) - return mantissa, exponent + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + modf_x = torch.modf(x) + return torch.resolve_modf(input=modf_x) -def modf( +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def nansum( x: torch.Tensor, /, *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[torch.dtype] = None, + keepdims: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - modf_x = torch.modf(x) - return torch.resolve_modf(input=modf_x) + dtype = ivy.as_native_dtype(dtype) + return torch.nansum(x, dim=axis, keepdim=keepdims, dtype=dtype) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def digamma( +def nextafter( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nextafter(x1, x2) + + +def signbit( + x: Union[torch.Tensor, float, int, list, tuple], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.signbit(x, out=out) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def sinc(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: + x = _cast_for_unary_op(x) + return torch.sinc(x, out=out) + + +@with_supported_dtypes( + {"2.0.1 and below": ("float16", "float32", "float64")}, + backend_version, +) +def xlogy( + x: torch.tensor, y: torch.tensor, /, *, out: Optional[torch.tensor] = None +) -> torch.tensor: + x, y = promote_types_of_inputs(x, y) + return torch.xlogy(x, y, out=out) + + +def zeta( x: torch.Tensor, + q: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.digamma(x, out=out) + temp = torch.logical_and(torch.ne(torch.remainder(x, 2), 0), torch.gt(x, 1)) + temp = torch.logical_and(temp, torch.le(q, 0)) + nan_indices = torch.logical_or(temp, torch.lt(x, 1)) + result = torch.special.zeta(x, q) + result.masked_fill_(nan_indices, float("nan")) + return result +fmax.support_native_out = True +sinc.support_native_out = True +float_power.support_native_out = True +copysign.support_native_out = True +count_nonzero.support_native_out = False +nansum.support_native_out = False +isclose.support_native_out = False +signbit.support_native_out = True +fix.support_native_out = True +nextafter.support_native_out = True +zeta.support_native_out = False +lerp.partial_mixed_handler = lambda input, end, weight, **kwargs: ( + _are_suitable_types_for_torch_lerp(input, end, weight) +) +lerp.support_native_out = True digamma.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/layers.py b/ivy/functional/backends/torch/experimental/layers.py index 9d35ad4f31232..0c26508af9aca 100644 --- a/ivy/functional/backends/torch/experimental/layers.py +++ b/ivy/functional/backends/torch/experimental/layers.py @@ -16,257 +16,34 @@ from ivy.functional.ivy.experimental.layers import _padding_ceil_mode -def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_first"): - # Determine depth pooling - kernel, strides, depth_pooling = _depth_max_pooling_helper( - x.shape, kernel, strides, dims=dims, data_format=data_format - ) - if depth_pooling: - x = torch.permute(x, (0, 2, 1, *range(3, dims + 2))) - return x, kernel, strides, depth_pooling - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def max_pool1d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 1 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NWC": - x = x.permute((0, 2, 1)) - kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel - strides = ( - [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides - ) - padding = ( - [padding[i] for i in [0, 2, 1]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) - - if not depth_pooling: - new_kernel = [dilation[0] * (kernel[0] - 1) + 1] - - if isinstance(padding, str): - pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2] - else: - pad_list = [item for sublist in padding for item in sublist] - - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) - - if depth_pooling: - res = torch.permute(res, (0, 2, 1)) - if data_format == "NWC": - res = res.permute((0, 2, 1)) - return res - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def max_pool2d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 2 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) - - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - kernel = ( - [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel - ) - strides = ( - [strides[i] for i in [0, 3, 1, 2]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 3, 1, 2]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - - # determine depth pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" - ) - - x_shape = list(x.shape[2:]) - if not depth_pooling: - new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] - - if isinstance(padding, str): - pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] - else: - # torch pad takes width padding first, then height padding - padding = (padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] +# --- Helpers --- # +# --------------- # - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = torch.permute(res, (0, 2, 1, 3)) - if data_format == "NHWC": - return res.permute(0, 2, 3, 1) - return res +def _add_ceil_pad_to_pad_list(num_pad, k, c): + return num_pad + (num_pad - ((k * num_pad) / (k - c))) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", +def _adjust_num_padded_values_to_ceil( + pad_specific, num_padded_values, x_shape, kernel, strides, dims +): + for i in range(dims): + pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] + _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) + num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( + num_padded_values[i][-1], kernel[i], c ) - }, - backend_version, -) -def max_pool3d( - x: torch.Tensor, - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dims = 3 - kernel, strides, padding, dilation = _validate_max_pool_params( - kernel, strides, padding, dilation, ceil_mode, dims=dims - ) + return num_padded_values - if data_format == "NDHWC": - x = x.permute(0, 4, 1, 2, 3) - kernel = ( - [kernel[i] for i in [0, 4, 1, 2, 3]] - if len(kernel) == (dims + 2) - else kernel - ) - strides = ( - [strides[i] for i in [0, 4, 1, 2, 3]] - if len(strides) == (dims + 2) - else strides - ) - padding = ( - [padding[i] for i in [0, 4, 1, 2, 3]] - if isinstance(padding, list) and len(padding) == (dims + 2) - else padding - ) - # Determine deptwise pooling - x, kernel, strides, depth_pooling = _determine_depth_max_pooling( - x, kernel, strides, dims, data_format="channel_first" +def _determine_depth_max_pooling(x, kernel, strides, dims, data_format="channel_first"): + # Determine depth pooling + kernel, strides, depth_pooling = _depth_max_pooling_helper( + x.shape, kernel, strides, dims=dims, data_format=data_format ) - - if not depth_pooling: - x_shape = x.shape[2:] - new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] - - if isinstance(padding, str): - pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) - pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) - pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) - pad_list = [ - pad_w // 2, - pad_w - pad_w // 2, - pad_h // 2, - pad_h - pad_h // 2, - pad_d // 2, - pad_d - pad_d // 2, - ] - else: - # torch pad takes width padding first, then height, then depth - padding = (padding[2], padding[1], padding[0]) - pad_list = [item for sublist in padding for item in sublist] - - x = torch.nn.functional.pad( - x, - pad_list, - value=float("-inf"), - ) - else: - if isinstance(padding, list) and any( - [item != 0 for sublist in padding for item in sublist] - ): - raise NotImplementedError( - "Nonzero explicit padding is not supported for depthwise max pooling" - ) - - res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) - if depth_pooling: - res = res.permute(0, 2, 1, 3, 4) - if data_format == "NDHWC": - res = res.permute(0, 2, 3, 4, 1) - return res - - -def _add_ceil_pad_to_pad_list(num_pad, k, c): - return num_pad + (num_pad - ((k * num_pad) / (k - c))) + x = torch.permute(x, (0, 2, 1, *range(3, dims + 2))) + return x, kernel, strides, depth_pooling def _get_specific_pad(x_shape, kernel, strides, padding, dims): @@ -293,6 +70,27 @@ def _get_specific_pad(x_shape, kernel, strides, padding, dims): return padding, pad_specific +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool1d(input, output_size): + return torch.nn.functional.adaptive_avg_pool1d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_avg_pool2d(input, output_size): + return torch.nn.functional.adaptive_avg_pool2d(input, output_size) + + +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def adaptive_max_pool2d( + input: torch.Tensor, output_size: Union[Sequence[int], int] +) -> torch.Tensor: + return torch.nn.functional.adaptive_max_pool2d(input, output_size) + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) def avg_pool1d( x: torch.Tensor, @@ -363,18 +161,6 @@ def avg_pool1d( return res -def _adjust_num_padded_values_to_ceil( - pad_specific, num_padded_values, x_shape, kernel, strides, dims -): - for i in range(dims): - pad = [pad_specific[i] // 2, pad_specific[i] - pad_specific[i] // 2] - _, c = _padding_ceil_mode(x_shape[i], kernel[i], pad, strides[i], True) - num_padded_values[i][-1] = _add_ceil_pad_to_pad_list( - num_padded_values[i][-1], kernel[i], c - ) - return num_padded_values - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -630,23 +416,222 @@ def dct( else: x = x * axis_dim_float - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(None, axis_dim) - dct_out = torch.real( - torch.fft.irfft( - scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis - ) - )[axis_idx] - return dct_out + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(None, axis_dim) + dct_out = torch.real( + torch.fft.irfft( + scale * torch.complex(x, real_zero), n=2 * axis_dim, axis=axis + ) + )[axis_idx] + return dct_out + + elif type == 4: + dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) + axis_idx = [slice(None)] * len(x.shape) + axis_idx[axis] = slice(1, None, 2) + dct_out = dct_2[axis_idx] + if norm == "ortho": + dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) + return dct_out + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def dropout( + x: torch.Tensor, + prob: float, + /, + *, + scale: bool = True, + dtype: torch.dtype = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + x = ivy.astype(x, dtype) if dtype else x + res = torch.nn.functional.dropout(x, prob, training=training) + res = torch.multiply(res, (1.0 - prob)) if not scale else res + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout1d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 3 + if data_format == "NWC": + perm = (0, 2, 1) if is_batched else (1, 0) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout1d(x, prob, training=training) + if data_format == "NWC": + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def dropout2d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 4 + if data_format == "NHWC": + perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout2d(x, prob, training=training) + if data_format == "NHWC": + perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def dropout3d( + x: torch.Tensor, + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_batched = len(x.shape) == 5 + if data_format == "NDHWC": + perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) + x = torch.permute(x, perm) + res = torch.nn.functional.dropout3d(x, prob, training=training) + if data_format == "NDHWC": + perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) + res = torch.permute(res, perm) + return res + + +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def embedding( + weights: torch.Tensor, + indices: torch.Tensor, + /, + *, + max_norm: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, + backend_version, +) +def fft( + x: torch.Tensor, + dim: int, + /, + *, + norm: str = "backward", + n: Union[int, Tuple[int]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not isinstance(dim, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(dim)}" + ) + if n is None: + n = x.shape[dim] + if n < -len(x.shape): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not isinstance(n, int): + raise ivy.utils.exceptions.IvyError( + f"Expecting instead of {type(n)}" + ) + if n <= 1: + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {n}, expecting more than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + if x.dtype in [torch.int64, torch.float64, torch.complex128]: + out_dtype = torch.complex128 + else: + out_dtype = torch.complex64 + return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) + - elif type == 4: - dct_2 = dct(x, type=2, n=2 * axis_dim, axis=axis, norm=None) - axis_idx = [slice(None)] * len(x.shape) - axis_idx[axis] = slice(1, None, 2) - dct_out = dct_2[axis_idx] - if norm == "ortho": - dct_out *= math.sqrt(0.5) * torch.rsqrt(axis_dim_float) - return dct_out +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def fft2( + x: torch.Tensor, + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if not all(isinstance(j, int) for j in dim): + raise ivy.utils.exceptions.IvyError( + f"Expecting {dim} to be a sequence of integers " + ) + if s is None: + s = (x.shape[dim[0]], x.shape[dim[1]]) + if all(j < -len(x.shape) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid dim {dim}, expecting ranging" + " from {-len(x.shape)} to {len(x.shape)-1} " + ) + if not all(isinstance(j, int) for j in s): + raise ivy.utils.exceptions.IvyError( + f"Expecting {s} to be a sequence of integers " + ) + if all(j <= 1 for j in s): + raise ivy.utils.exceptions.IvyError( + f"Invalid data points {s}, expecting s points larger than 1" + ) + if norm != "backward" and norm != "ortho" and norm != "forward": + raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") + return torch.tensor( + torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 + ) def idct( @@ -663,19 +648,9 @@ def idct( return dct(x, type=inverse_type, n=n, axis=axis, norm=norm, out=out) -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def fft( +def ifft( x: torch.Tensor, dim: int, - /, *, norm: str = "backward", n: Union[int, Tuple[int]] = None, @@ -702,89 +677,200 @@ def fft( ) if norm != "backward" and norm != "ortho" and norm != "forward": raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - if x.dtype in [torch.int64, torch.float64, torch.complex128]: - out_dtype = torch.complex128 - else: - out_dtype = torch.complex64 - return torch.fft.fft(x, n, dim, norm, out=out).to(dtype=out_dtype) + return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - ) - }, - backend_version, -) -def dropout( +def ifftn( x: torch.Tensor, - prob: float, - /, + s: Optional[Union[int, Tuple[int]]] = None, + axes: Optional[Union[int, Tuple[int]]] = None, *, - scale: bool = True, - dtype: torch.dtype = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, + norm: Optional[str] = "backward", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = ivy.astype(x, dtype) if dtype else x - res = torch.nn.functional.dropout(x, prob, training=training) - res = torch.multiply(res, (1.0 - prob)) if not scale else res - return res + return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) -dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( - kwargs.get("noise_shape") is None and kwargs.get("seed") is None -) +def interpolate( + x: torch.Tensor, + size: Union[Sequence[int], int], + /, + *, + mode: Literal[ + "linear", + "bilinear", + "trilinear", + "nearest", + "area", + "nearest_exact", + "tf_area", + "bicubic", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", + ] = "linear", + scale_factor: Optional[Union[Sequence[int], int]] = None, + recompute_scale_factor: Optional[bool] = None, + align_corners: Optional[bool] = None, + antialias: bool = False, + out: Optional[torch.Tensor] = None, +): + return torch.nn.functional.interpolate( + x, + size=size, + mode=mode, + align_corners=align_corners, + antialias=antialias, + scale_factor=scale_factor, + recompute_scale_factor=recompute_scale_factor, + ) -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def dropout1d( +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def max_pool1d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 3 + dims = 1 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NWC": - perm = (0, 2, 1) if is_batched else (1, 0) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout1d(x, prob, training=training) + x = x.permute((0, 2, 1)) + kernel = [kernel[i] for i in [0, 2, 1]] if len(kernel) == (dims + 2) else kernel + strides = ( + [strides[i] for i in [0, 2, 1]] if len(strides) == (dims + 2) else strides + ) + padding = ( + [padding[i] for i in [0, 2, 1]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + if not depth_pooling: + new_kernel = [dilation[0] * (kernel[0] - 1) + 1] + + if isinstance(padding, str): + pad_w = _handle_padding(x.shape[2], strides[0], new_kernel[0], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2] + else: + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool1d(x, kernel, strides, 0, dilation, ceil_mode) + + if depth_pooling: + res = torch.permute(res, (0, 2, 1)) if data_format == "NWC": - res = torch.permute(res, perm) + res = res.permute((0, 2, 1)) return res @with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, + { + "2.0.1 and below": ( + "float16", + "bfloat16", + ) + }, backend_version, ) -def dropout2d( +def max_pool2d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 4 + dims = 2 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) + if data_format == "NHWC": - perm = (0, 3, 1, 2) if is_batched else (2, 0, 1) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout2d(x, prob, training=training) + x = x.permute(0, 3, 1, 2) + kernel = ( + [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel + ) + strides = ( + [strides[i] for i in [0, 3, 1, 2]] + if len(strides) == (dims + 2) + else strides + ) + padding = ( + [padding[i] for i in [0, 3, 1, 2]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding + ) + + # determine depth pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" + ) + + x_shape = list(x.shape[2:]) + if not depth_pooling: + new_kernel = [kernel[i] + (kernel[i] - 1) * (dilation[i] - 1) for i in range(2)] + + if isinstance(padding, str): + pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] + else: + # torch pad takes width padding first, then height padding + padding = (padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] + + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), + ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + + res = torch.nn.functional.max_pool2d(x, kernel, strides, 0, dilation, ceil_mode) + if depth_pooling: + res = torch.permute(res, (0, 2, 1, 3)) if data_format == "NHWC": - perm = (0, 2, 3, 1) if is_batched else (1, 2, 0) - res = torch.permute(res, perm) + return res.permute(0, 2, 3, 1) return res @@ -797,184 +883,87 @@ def dropout2d( }, backend_version, ) -def dropout3d( +def max_pool3d( x: torch.Tensor, - prob: float, + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], /, *, - training: bool = True, data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - is_batched = len(x.shape) == 5 - if data_format == "NDHWC": - perm = (0, 4, 1, 2, 3) if is_batched else (3, 0, 1, 2) - x = torch.permute(x, perm) - res = torch.nn.functional.dropout3d(x, prob, training=training) - if data_format == "NDHWC": - perm = (0, 2, 3, 4, 1) if is_batched else (1, 2, 3, 0) - res = torch.permute(res, perm) - return res - + dims = 3 + kernel, strides, padding, dilation = _validate_max_pool_params( + kernel, strides, padding, dilation, ceil_mode, dims=dims + ) -def ifft( - x: torch.Tensor, - dim: int, - *, - norm: str = "backward", - n: Union[int, Tuple[int]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(dim, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(dim)}" - ) - if n is None: - n = x.shape[dim] - if n < -len(x.shape): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " + if data_format == "NDHWC": + x = x.permute(0, 4, 1, 2, 3) + kernel = ( + [kernel[i] for i in [0, 4, 1, 2, 3]] + if len(kernel) == (dims + 2) + else kernel ) - if not isinstance(n, int): - raise ivy.utils.exceptions.IvyError( - f"Expecting instead of {type(n)}" + strides = ( + [strides[i] for i in [0, 4, 1, 2, 3]] + if len(strides) == (dims + 2) + else strides ) - if n <= 1: - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {n}, expecting more than 1" + padding = ( + [padding[i] for i in [0, 4, 1, 2, 3]] + if isinstance(padding, list) and len(padding) == (dims + 2) + else padding ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.fft.ifft(x, n, dim, norm, out=out).resolve_conj() - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def embedding( - weights: torch.Tensor, - indices: torch.Tensor, - /, - *, - max_norm: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return torch.nn.functional.embedding(indices, weights, max_norm=max_norm) - - -embedding.support_native_out = False - -def interpolate( - x: torch.Tensor, - size: Union[Sequence[int], int], - /, - *, - mode: Literal[ - "linear", - "bilinear", - "trilinear", - "nearest", - "area", - "nearest_exact", - "tf_area", - "bicubic", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", - ] = "linear", - scale_factor: Optional[Union[Sequence[int], int]] = None, - recompute_scale_factor: Optional[bool] = None, - align_corners: Optional[bool] = None, - antialias: bool = False, - out: Optional[torch.Tensor] = None, -): - return torch.nn.functional.interpolate( - x, - size=size, - mode=mode, - align_corners=align_corners, - antialias=antialias, - scale_factor=scale_factor, - recompute_scale_factor=recompute_scale_factor, + # Determine deptwise pooling + x, kernel, strides, depth_pooling = _determine_depth_max_pooling( + x, kernel, strides, dims, data_format="channel_first" ) + if not depth_pooling: + x_shape = x.shape[2:] + new_kernel = [dilation[i] * (kernel[i] - 1) + 1 for i in range(dims)] -interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ - "tf_area", - "nd", - "bicubic_tensorflow", - "mitchellcubic", - "lanczos3", - "lanczos5", - "gaussian", -] - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_max_pool2d( - input: torch.Tensor, output_size: Union[Sequence[int], int] -) -> torch.Tensor: - return torch.nn.functional.adaptive_max_pool2d(input, output_size) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool1d(input, output_size): - return torch.nn.functional.adaptive_avg_pool1d(input, output_size) - - -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def adaptive_avg_pool2d(input, output_size): - return torch.nn.functional.adaptive_avg_pool2d(input, output_size) - + if isinstance(padding, str): + pad_d = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding) + pad_h = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding) + pad_w = _handle_padding(x_shape[2], strides[2], new_kernel[2], padding) + pad_list = [ + pad_w // 2, + pad_w - pad_w // 2, + pad_h // 2, + pad_h - pad_h // 2, + pad_d // 2, + pad_d - pad_d // 2, + ] + else: + # torch pad takes width padding first, then height, then depth + padding = (padding[2], padding[1], padding[0]) + pad_list = [item for sublist in padding for item in sublist] -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def fft2( - x: torch.Tensor, - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not all(isinstance(j, int) for j in dim): - raise ivy.utils.exceptions.IvyError( - f"Expecting {dim} to be a sequence of integers " - ) - if s is None: - s = (x.shape[dim[0]], x.shape[dim[1]]) - if all(j < -len(x.shape) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid dim {dim}, expecting ranging" - " from {-len(x.shape)} to {len(x.shape)-1} " - ) - if not all(isinstance(j, int) for j in s): - raise ivy.utils.exceptions.IvyError( - f"Expecting {s} to be a sequence of integers " - ) - if all(j <= 1 for j in s): - raise ivy.utils.exceptions.IvyError( - f"Invalid data points {s}, expecting s points larger than 1" + x = torch.nn.functional.pad( + x, + pad_list, + value=float("-inf"), ) - if norm != "backward" and norm != "ortho" and norm != "forward": - raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}") - return torch.tensor( - torch.fft.fft2(x, s, dim, norm, out=out), dtype=torch.complex128 - ) + else: + if isinstance(padding, list) and any( + [item != 0 for sublist in padding for item in sublist] + ): + raise NotImplementedError( + "Nonzero explicit padding is not supported for depthwise max pooling" + ) + res = torch.nn.functional.max_pool3d(x, kernel, strides, 0, dilation, ceil_mode) -def ifftn( - x: torch.Tensor, - s: Optional[Union[int, Tuple[int]]] = None, - axes: Optional[Union[int, Tuple[int]]] = None, - *, - norm: Optional[str] = "backward", - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, out=out) + if depth_pooling: + res = res.permute(0, 2, 1, 3, 4) + if data_format == "NDHWC": + res = res.permute(0, 2, 3, 4, 1) + return res @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -1010,3 +999,18 @@ def rfftn( return torch.tensor( torch.fft.rfftn(x, s, axes, norm=norm, out=out), dtype=torch.complex128 ) + + +dropout.partial_mixed_handler = lambda x, prob, **kwargs: ( + kwargs.get("noise_shape") is None and kwargs.get("seed") is None +) +embedding.support_native_out = False +interpolate.partial_mixed_handler = lambda *args, mode="linear", **kwargs: mode not in [ + "tf_area", + "nd", + "bicubic_tensorflow", + "mitchellcubic", + "lanczos3", + "lanczos5", + "gaussian", +] diff --git a/ivy/functional/backends/torch/experimental/linear_algebra.py b/ivy/functional/backends/torch/experimental/linear_algebra.py index 68e98bff30a60..8efd920016f44 100644 --- a/ivy/functional/backends/torch/experimental/linear_algebra.py +++ b/ivy/functional/backends/torch/experimental/linear_algebra.py @@ -12,6 +12,27 @@ from ivy.functional.ivy.experimental.linear_algebra import _check_valid_dimension_size +def adjoint( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_valid_dimension_size(x) + return torch.adjoint(x).resolve_conj() + + +@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) +def cond( + x: torch.Tensor, + /, + *, + p: Optional[Union[None, int, str]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.cond(x, p=p, out=out) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def diagflat( x: torch.Tensor, @@ -97,32 +118,14 @@ def diagflat( return ret -diagflat.support_native_out = False - - -def kron( +def dot( a: torch.Tensor, b: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.kron(a, b, out=out) - - -kron.support_native_out = True - - -def matrix_exp( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.linalg.matrix_exp(x) - - -matrix_exp.support_native_out = True + return torch.matmul(a, b) def eig( @@ -133,73 +136,56 @@ def eig( return torch.linalg.eig(x) -eig.support_native_out = False - - def eigvals(x: torch.Tensor, /) -> torch.Tensor: if not torch.is_complex(x): x = x.to(torch.complex128) return torch.linalg.eigvals(x) -eigvals.support_native_out = False - - -def adjoint( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_dimension_size(x) - return torch.adjoint(x).resolve_conj() - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def multi_dot( - x: Sequence[torch.Tensor], +def kron( + a: torch.Tensor, + b: torch.Tensor, /, *, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.multi_dot(x, out=out) - - -multi_dot.support_native_out = True +) -> torch.tensor: + return torch.kron(a, b, out=out) -@with_unsupported_dtypes({"2.0.0 and below": ("float16", "bfloat16")}, backend_version) -def cond( +def lu_factor( x: torch.Tensor, /, *, - p: Optional[Union[None, int, str]] = None, + pivot: Optional[bool] = True, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.cond(x, p=p, out=out) - - -cond.support_native_out = False +) -> Tuple[torch.Tensor]: + raise IvyNotImplementedException() -def lu_factor( +def matrix_exp( x: torch.Tensor, /, *, - pivot: Optional[bool] = True, out: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor]: - raise IvyNotImplementedException() +) -> torch.Tensor: + return torch.linalg.matrix_exp(x) -def dot( - a: torch.Tensor, - b: torch.Tensor, +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def multi_dot( + x: Sequence[torch.Tensor], /, *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.matmul(a, b) + return torch.linalg.multi_dot(x, out=out) +diagflat.support_native_out = False +kron.support_native_out = True +matrix_exp.support_native_out = True +eig.support_native_out = False +eigvals.support_native_out = False +multi_dot.support_native_out = True +cond.support_native_out = False dot.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/losses.py b/ivy/functional/backends/torch/experimental/losses.py index db3442685a678..e3b6e9485e543 100644 --- a/ivy/functional/backends/torch/experimental/losses.py +++ b/ivy/functional/backends/torch/experimental/losses.py @@ -3,6 +3,24 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, + backend_version, +) +def huber_loss( + input: torch.Tensor, + target: torch.Tensor, + /, + *, + reduction: Optional[str] = "mean", + delta: Optional[float] = 1.0, +) -> torch.Tensor: + return torch.nn.functional.huber_loss( + input, target, reduction=reduction, delta=delta + ) + + # Assuming ivy and backend_version are imported and defined properly @@ -52,20 +70,3 @@ def smooth_l1_loss( beta=beta, reduction=reduction, ) - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "int8", "int16", "int32", "int64", "bool")}, - backend_version, -) -def huber_loss( - input: torch.Tensor, - target: torch.Tensor, - /, - *, - reduction: Optional[str] = "mean", - delta: Optional[float] = 1.0, -) -> torch.Tensor: - return torch.nn.functional.huber_loss( - input, target, reduction=reduction, delta=delta - ) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index 8683957a72c97..3b344352dbb6e 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -10,63 +10,67 @@ import ivy -def moveaxis( - a: torch.Tensor, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.moveaxis(a, source, destination) +def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_1d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -moveaxis.support_native_out = False +def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: + transformed = torch.atleast_2d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -def heaviside( - x1: torch.tensor, - x2: torch.tensor, - /, - *, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.heaviside( - x1, - x2, - out=out, - ) +def atleast_3d( + *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None +) -> List[torch.Tensor]: + transformed = torch.atleast_3d(*arys) + if isinstance(transformed, tuple): + return list(transformed) + return transformed -heaviside.support_native_out = True +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: + return tuple(torch.broadcast_shapes(*shapes)) -def flipud( - m: torch.Tensor, +def concat_from_sequence( + input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], /, *, - copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.flipud(m) - - -flipud.support_native_out = False + new_axis: int = 0, + axis: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + is_tuple = type(input_sequence) is tuple + if is_tuple: + input_sequence = list(input_sequence) + if new_axis == 0: + ret = torch.cat(input_sequence, dim=axis) + return ret + elif new_axis == 1: + ret = torch.stack(input_sequence, dim=axis) + return ret -def vstack( - arrays: Sequence[torch.Tensor], +def dsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if not isinstance(arrays, tuple): - arrays = tuple(arrays) - return torch.vstack(arrays, out=None) + copy: Optional[bool] = None, +) -> List[torch.Tensor]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "dsplit only works on arrays of 3 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) -def hstack( +def dstack( arrays: Sequence[torch.Tensor], /, *, @@ -74,71 +78,49 @@ def hstack( ) -> torch.Tensor: if not isinstance(arrays, tuple): arrays = tuple(arrays) - return torch.hstack(arrays, out=None) + return torch.dstack(arrays, out=out) -def rot90( - m: torch.Tensor, +def expand( + x: torch.Tensor, + shape: Union[List[int], List[Tuple]], /, *, copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.rot90(m, k, axes) - - -def top_k( - x: torch.Tensor, - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - k = min(k, x.shape[axis]) - topk_res = NamedTuple( - "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] - ) - if not largest: - indices = torch.argsort(x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - else: - indices = torch.argsort(-x, dim=axis) - indices = torch.index_select(indices, axis, torch.arange(k)) - if not sorted: - indices = torch.sort(indices, dim=axis)[0] - val = torch.gather(x, axis, indices) - return topk_res(val, indices) - - -def fliplr( - m: torch.Tensor, - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - return torch.fliplr(m) - - -fliplr.support_native_out = False + return x.expand(shape) -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def i0( - x: torch.Tensor, +def fill_diagonal( + a: torch.Tensor, + v: Union[int, float], /, *, - out: Optional[torch.Tensor] = None, + wrap: bool = False, ) -> torch.Tensor: - return torch.i0(x, out=out) + shape = a.shape + max_end = torch.prod(torch.tensor(shape)) + end = max_end + if len(shape) == 2: + step = shape[1] + 1 + if not wrap: + end = shape[1] * shape[1] + else: + step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + end = max_end if end > max_end else end + a = torch.reshape(a, (-1,)) + w = torch.zeros(a.shape, dtype=bool).to(a.device) + ins = torch.arange(0, max_end).to(a.device) + steps = torch.arange(0, end, step).to(a.device) -i0.support_native_out = True + for i in steps: + i = ins == i + w = torch.logical_or(w, i) + a = torch.where(w, v, a) + a = torch.reshape(a, shape) + return a def flatten( @@ -154,47 +136,53 @@ def flatten( return torch.flatten(x, start_dim=start_dim, end_dim=end_dim) -flatten.partial_mixed_handler = ( - lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" -) +def fliplr( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.fliplr(m) -def vsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], +def flipud( + m: torch.Tensor, /, *, copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "vsplit only works on arrays of 2 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.flipud(m) -def dsplit( +def heaviside( + x1: torch.tensor, + x2: torch.tensor, + /, + *, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + return torch.heaviside( + x1, + x2, + out=out, + ) + + +def hsplit( ary: torch.Tensor, - indices_or_sections: Union[int, Sequence[int], torch.Tensor], + indices_or_sections: Union[int, Tuple[int, ...]], /, *, copy: Optional[bool] = None, ) -> List[torch.Tensor]: - if len(ary.shape) < 2: - raise ivy.utils.exceptions.IvyError( - "dsplit only works on arrays of 3 or more dimensions" - ) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=2) - - -def atleast_1d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_1d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed + if len(ary.shape) == 1: + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) -def dstack( +def hstack( arrays: Sequence[torch.Tensor], /, *, @@ -202,23 +190,41 @@ def dstack( ) -> torch.Tensor: if not isinstance(arrays, tuple): arrays = tuple(arrays) - return torch.dstack(arrays, out=out) + return torch.hstack(arrays, out=None) -def atleast_2d(*arys: torch.Tensor, copy: Optional[bool] = None) -> List[torch.Tensor]: - transformed = torch.atleast_2d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def i0( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.i0(x, out=out) -def atleast_3d( - *arys: Union[torch.Tensor, bool, Number], copy: Optional[bool] = None -) -> List[torch.Tensor]: - transformed = torch.atleast_3d(*arys) - if isinstance(transformed, tuple): - return list(transformed) - return transformed +def moveaxis( + a: torch.Tensor, + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], + /, + *, + copy: Optional[bool] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.moveaxis(a, source, destination) + + +def rot90( + m: torch.Tensor, + /, + *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.rot90(m, k, axes) @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -263,59 +269,30 @@ def take_along_axis( return torch.take_along_dim(arr, indices, axis, out=out) -def hsplit( - ary: torch.Tensor, - indices_or_sections: Union[int, Tuple[int, ...]], - /, - *, - copy: Optional[bool] = None, -) -> List[torch.Tensor]: - if len(ary.shape) == 1: - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) - return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=1) - - -take_along_axis.support_native_out = True - - -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: - return tuple(torch.broadcast_shapes(*shapes)) - - -broadcast_shapes.support_native_out = False - - -def expand( +def top_k( x: torch.Tensor, - shape: Union[List[int], List[Tuple]], - /, - *, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return x.expand(shape) - - -expand.support_native_out = False - - -def concat_from_sequence( - input_sequence: Union[Tuple[torch.Tensor], List[torch.Tensor]], + k: int, /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - is_tuple = type(input_sequence) is tuple - if is_tuple: - input_sequence = list(input_sequence) - if new_axis == 0: - ret = torch.cat(input_sequence, dim=axis) - return ret - elif new_axis == 1: - ret = torch.stack(input_sequence, dim=axis) - return ret + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + k = min(k, x.shape[axis]) + topk_res = NamedTuple( + "top_k", [("values", torch.Tensor), ("indices", torch.Tensor)] + ) + if not largest: + indices = torch.argsort(x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) + else: + indices = torch.argsort(-x, dim=axis) + indices = torch.index_select(indices, axis, torch.arange(k)) + if not sorted: + indices = torch.sort(indices, dim=axis)[0] + val = torch.gather(x, axis, indices) + return topk_res(val, indices) @with_unsupported_dtypes({"2.0.1 and below": ("complex", "float16")}, backend_version) @@ -342,32 +319,39 @@ def unique_consecutive( ) -def fill_diagonal( - a: torch.Tensor, - v: Union[int, float], +def vsplit( + ary: torch.Tensor, + indices_or_sections: Union[int, Sequence[int], torch.Tensor], /, *, - wrap: bool = False, + copy: Optional[bool] = None, +) -> List[torch.Tensor]: + if len(ary.shape) < 2: + raise ivy.utils.exceptions.IvyError( + "vsplit only works on arrays of 2 or more dimensions" + ) + return ivy.split(ary, num_or_size_splits=indices_or_sections, axis=0) + + +def vstack( + arrays: Sequence[torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - shape = a.shape - max_end = torch.prod(torch.tensor(shape)) - end = max_end - if len(shape) == 2: - step = shape[1] + 1 - if not wrap: - end = shape[1] * shape[1] - else: - step = 1 + (torch.cumprod(torch.tensor(shape[:-1]), 0)).sum() + if not isinstance(arrays, tuple): + arrays = tuple(arrays) + return torch.vstack(arrays, out=None) - end = max_end if end > max_end else end - a = torch.reshape(a, (-1,)) - w = torch.zeros(a.shape, dtype=bool).to(a.device) - ins = torch.arange(0, max_end).to(a.device) - steps = torch.arange(0, end, step).to(a.device) - for i in steps: - i = ins == i - w = torch.logical_or(w, i) - a = torch.where(w, v, a) - a = torch.reshape(a, shape) - return a +moveaxis.support_native_out = False +heaviside.support_native_out = True +flipud.support_native_out = False +fliplr.support_native_out = False +i0.support_native_out = True +flatten.partial_mixed_handler = ( + lambda *args, copy=None, start_dim=0, end_dim=1, order="C", **kwargs: order == "C" +) +take_along_axis.support_native_out = True +broadcast_shapes.support_native_out = False +expand.support_native_out = False diff --git a/ivy/functional/backends/torch/experimental/norms.py b/ivy/functional/backends/torch/experimental/norms.py index 5a3e0c31f764b..987280c69a0e9 100644 --- a/ivy/functional/backends/torch/experimental/norms.py +++ b/ivy/functional/backends/torch/experimental/norms.py @@ -5,33 +5,6 @@ from .. import backend_version -def l1_normalize( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) - - -l1_normalize.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def l2_normalize( - x: torch.Tensor, - /, - *, - axis: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) - - -l2_normalize.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) def batch_norm( x: torch.Tensor, @@ -73,15 +46,29 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def group_norm( + x: torch.Tensor, + num_groups: int = 1, + /, + *, + offset: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", + out: Optional[torch.Tensor] = None, +): + xdims = x.ndim + if data_format == "NSC": + x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) + xnormalized = torch.nn.functional.group_norm( + x, num_groups, weight=scale, bias=offset, eps=eps ) -) + + if data_format == "NSC": + xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) + + return xnormalized @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) @@ -127,40 +114,25 @@ def instance_norm( return xnormalized, runningmean, runningvariance -instance_norm.partial_mixed_handler = ( - lambda x, mean, variance, scale=None, offset=None, **kwargs: ( - x.ndim > 1 - and mean.ndim == 1 - and variance.ndim == 1 - and (scale is None or scale.ndim == 1) - and (offset is None or offset.ndim == 1) - ) -) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def group_norm( +def l1_normalize( x: torch.Tensor, - num_groups: int = 1, /, *, - offset: Optional[torch.Tensor] = None, - scale: Optional[torch.Tensor] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", + axis: Optional[int] = None, out: Optional[torch.Tensor] = None, -): - xdims = x.ndim - if data_format == "NSC": - x = torch.permute(x, dims=(0, xdims - 1, *range(1, xdims - 1))) - xnormalized = torch.nn.functional.group_norm( - x, num_groups, weight=scale, bias=offset, eps=eps - ) +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=1, dim=axis, out=out) - if data_format == "NSC": - xnormalized = torch.permute(xnormalized, dims=(0, *range(2, xdims), 1)) - return xnormalized +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def l2_normalize( + x: torch.Tensor, + /, + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.normalize(x, p=2, dim=axis, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) @@ -175,4 +147,24 @@ def lp_normalize( return torch.nn.functional.normalize(x, p=p, dim=axis, out=out) +l1_normalize.support_native_out = True +l2_normalize.support_native_out = True +batch_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) + ) +) +instance_norm.partial_mixed_handler = ( + lambda x, mean, variance, scale=None, offset=None, **kwargs: ( + x.ndim > 1 + and mean.ndim == 1 + and variance.ndim == 1 + and (scale is None or scale.ndim == 1) + and (offset is None or offset.ndim == 1) + ) +) lp_normalize.support_native_out = True diff --git a/ivy/functional/backends/torch/experimental/random.py b/ivy/functional/backends/torch/experimental/random.py index a65cee4f57431..b8e58426ed24b 100644 --- a/ivy/functional/backends/torch/experimental/random.py +++ b/ivy/functional/backends/torch/experimental/random.py @@ -12,23 +12,46 @@ ) -# dirichlet -@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) -def dirichlet( - alpha: Union[torch.tensor, float, Sequence[float]], - /, +# --- Helpers --- # +# --------------- # + + +def _poisson_with_neg_lam(lam, fill_value, device, dtype): + if torch.any(lam < 0): + pos_lam = torch.where(lam < 0, 0, lam) + ret = torch.poisson(pos_lam).type(dtype).to(device) + ret = torch.where(lam < 0, fill_value, ret) + else: + ret = torch.poisson(lam).type(dtype).to(device) + return ret + + +# --- Main --- # +# ------------ # + + +def bernoulli( + probs: Union[float, torch.Tensor], *, - size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - out: Optional[torch.Tensor] = None, + logits: Union[float, torch.Tensor] = None, + shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, + device: torch.device, + dtype: torch.dtype, seed: Optional[int] = None, - dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - size = size if size is not None else len(alpha) - if seed is not None: + if seed: torch.manual_seed(seed) - return torch.tensor( - torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), - dtype=dtype, + if logits is not None: + if not _check_shapes_broadcastable(shape, logits.shape): + shape = logits.shape + elif probs is not None: + if not _check_shapes_broadcastable(shape, probs.shape): + shape = probs.shape + return ( + torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) + .sample(shape) + .to(device, dtype) ) @@ -53,6 +76,26 @@ def beta( return ret +# dirichlet +@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) +def dirichlet( + alpha: Union[torch.tensor, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + out: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + size = size if size is not None else len(alpha) + if seed is not None: + torch.manual_seed(seed) + return torch.tensor( + torch.distributions.dirichlet.Dirichlet(alpha).rsample(sample_shape=size), + dtype=dtype, + ) + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def gamma( alpha: Union[float, torch.Tensor], @@ -74,16 +117,6 @@ def gamma( return ret -def _poisson_with_neg_lam(lam, fill_value, device, dtype): - if torch.any(lam < 0): - pos_lam = torch.where(lam < 0, 0, lam) - ret = torch.poisson(pos_lam).type(dtype).to(device) - ret = torch.where(lam < 0, fill_value, ret) - else: - ret = torch.poisson(lam).type(dtype).to(device) - return ret - - def poisson( lam: Union[float, torch.Tensor], *, @@ -104,28 +137,3 @@ def poisson( _check_shapes_broadcastable(lam.shape, list_shape) lam = torch.broadcast_to(lam, list_shape) return _poisson_with_neg_lam(lam, fill_value, device, dtype) - - -def bernoulli( - probs: Union[float, torch.Tensor], - *, - logits: Union[float, torch.Tensor] = None, - shape: Optional[Union[ivy.NativeArray, Sequence[int]]] = None, - device: torch.device, - dtype: torch.dtype, - seed: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if seed: - torch.manual_seed(seed) - if logits is not None: - if not _check_shapes_broadcastable(shape, logits.shape): - shape = logits.shape - elif probs is not None: - if not _check_shapes_broadcastable(shape, probs.shape): - shape = probs.shape - return ( - torch.distributions.bernoulli.Bernoulli(probs=probs, logits=logits) - .sample(shape) - .to(device, dtype) - ) diff --git a/ivy/functional/backends/torch/experimental/statistical.py b/ivy/functional/backends/torch/experimental/statistical.py index cc1d4d5fa05ce..78b4e43e4db7e 100644 --- a/ivy/functional/backends/torch/experimental/statistical.py +++ b/ivy/functional/backends/torch/experimental/statistical.py @@ -10,6 +10,333 @@ from copy import deepcopy +# --- Helpers --- # +# --------------- # + + +def _compute_quantile_wrapper( + x, q, axis=None, keepdims=False, interpolation="linear", out=None +): + if not _validate_quantile(q): + raise ValueError("Quantiles must be in the range [0, 1]") + if interpolation in [ + "linear", + "lower", + "higher", + "midpoint", + "nearest", + "nearest_jax", + ]: + if interpolation == "nearest_jax": + return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) + else: + return torch.quantile( + x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out + ) + else: + raise ValueError( + "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" + ) + + +def _handle_axis(a, q, fn, keepdims=False, axis=None): + nd = a.ndim + axis_arg = deepcopy(axis) + if axis is not None: + axis = _to_positive_axis(axis, nd) + + if len(axis) == 1: + axis_arg = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) + + for i, s in enumerate(sorted(keep)): + a = torch.moveaxis(a, s, i) + a = a.view( + [ + *a.shape[:nkeep], + -1, + ] + ) + axis_arg = -1 + + ret = fn(a, q, axis=axis_arg) + + if keepdims: + if axis is None: + index_ret = (None,) * nd + else: + index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) + ret = ret[(Ellipsis,) + index_ret] + + return ret + + +def _nanmedian(input, axis, keepdims): + dtype = input.dtype + temp = input.to(torch.float64) + num_dim = len(temp.size()) + keepdim_shape = list(temp.size()) + q = 0.5 + + axis = [axis] if isinstance(axis, int) else list(axis) + + for i in axis: + keepdim_shape[i] = 1 + axis = [num_dim + x if x < 0 else x for x in axis] + axis.sort() + dimension = len(temp.size()) + while len(axis) > 0: + axis1 = axis[0] + for axis2 in range(axis1 + 1, dimension): + temp = torch.transpose(temp, axis1, axis2) + axis1 = axis2 + axis = [x - 1 for x in axis] + axis.pop(0) + dimension = dimension - 1 + temp = torch.flatten(temp, start_dim=dimension - len(axis)) + ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") + if keepdims: + keepdim_shape = tuple(keepdim_shape) + ret = ret.reshape(keepdim_shape) + + if dtype in [torch.int32, torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=torch.float16) + else: + ret = torch.asarray(ret, dtype=torch.float32) + + return ret + + +def _quantile(a, q, axis=None): + ret_dtype = a.dtype + if isinstance(q, float): + q = torch.as_tensor(q) + if isinstance(q, torch.Tensor) and q.ndim > 1: + raise ValueError("q argument must be a scalar or 1-dimensional!") + if axis is None: + axis = 0 + a = a.flatten() + + n = a.shape[axis] + + indices = q * (n - 1) + + a = torch.sort(a, axis)[axis] + indices_below = torch.floor(indices).to(torch.int64) + indices_upper = torch.ceil(indices).to(torch.int64) + + weights = indices - indices_below.to(torch.float64) + + indices_below = torch.clip(indices_below, 0, n - 1) + indices_upper = torch.clip(indices_upper, 0, n - 1) + tensor_upper = torch.index_select(a, 0, indices_upper) + tensor_below = torch.index_select(a, 0, indices_below) + + pred = weights <= 0.5 + out = torch.where(pred, tensor_below, tensor_upper) + return out.to(ret_dtype) + + +def _to_positive_axis(axis, ndim): + if not isinstance(axis, (list, tuple)): + axis = [axis] + + if len(axis) == 0: + raise ValueError("Axis can't be empty!") + + if len(set(axis)) != len(axis): + raise ValueError("Duplicated axis!") + + for i in range(len(axis)): + if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): + raise ValueError("Axis must be int in range [-rank(x), rank(x))") + if axis[i] < 0: + axis[i] += ndim + return axis + + +def _validate_quantile(q): + if isinstance(q, float): + q = torch.as_tensor(q) + if q.ndim == 1 and torch.numel(q) < 10: + for i in range(torch.numel(q)): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (torch.all(0 <= q) and torch.all(q <= 1)): + return False + return True + + +# --- Main --- # +# ------------ # + + +def bincount( + x: torch.Tensor, + /, + *, + weights: Optional[torch.Tensor] = None, + minlength: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if weights is None: + ret = torch.bincount(x, minlength=minlength) + ret = ret.to(x.dtype) + else: + ret = torch.bincount(x, weights=weights, minlength=minlength) + ret = ret.to(weights.dtype) + return ret + + +def corrcoef( + x: torch.Tensor, + /, + *, + y: Optional[torch.Tensor] = None, + rowvar: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if y is None: + xarr = x + else: + axis = 0 if rowvar else 1 + xarr = torch.concat([x, y], dim=axis) + xarr = xarr.T if not rowvar else xarr + + return torch.corrcoef(xarr) + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def cov( + x1: torch.Tensor, + x2: torch.Tensor = None, + /, + *, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[torch.Tensor] = None, + aweights: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + # dtype casts separately + if fweights is not None: + fweights = fweights.type(torch.int64) + if aweights is not None: + aweights = aweights.type(torch.float64) + + if x1.dim() > 2: + raise ValueError("x1 has more than 2 dimensions") + + if x2 is not None: + if x2.dim() > 2: + raise ValueError("x2 has more than 2 dimensions") + + if ddof is None: + if bias == 0: + ddof = 1 + else: + ddof = 0 + + if dtype is None: + x1 = x1.type(torch.float64) + if x2 is not None: + x2 = x2.type(torch.float64) + else: + x1 = x1.type(dtype) + if x2 is not None: + x2 = x2.type(dtype) + + X = x1 + if not rowVar and len(x1.shape) != 1: + X = torch.t(x1) + + if x2 is not None: + if not rowVar and len(x2.shape) != 1: + x2 = torch.t(x2) + X = torch.vstack((X, x2)) + + return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, + backend_version, +) +def cummax( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if x.dtype in (torch.bool, torch.float16): + x = x.to(dtype=torch.float64) + elif x.dtype in (torch.int16, torch.int8, torch.uint8): + x = x.to(dtype=torch.int64) + elif x.dtype in (torch.complex64, torch.complex128): + x = x.real.to(dtype=torch.float64) + + if exclusive or reverse: + if exclusive and reverse: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + x1, x2 = torch.concat( + (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 + ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) + x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x1, x2 = torch.cummax(x, -1) + res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) + else: + x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) + res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) + return res1, res2 + + return torch.cummax(x, axis, out=out) + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cummin( + x: torch.Tensor, + /, + *, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + if not (reverse): + ret = torch.cummin(x, axis)[0] + else: + ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] + ret = torch.flip(ret, (axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret.to(dtype)) + return ret.to(dtype) + + @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -136,254 +463,54 @@ def histogram( return ret -histogram.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) -def median( - input: torch.Tensor, - /, - *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if isinstance(axis, tuple): - if len(axis) == 1: - axis = axis[0] - ret = quantile( - input, - 0.5, - axis=axis, - keepdims=keepdims, - interpolation="midpoint", - ) - if input.dtype in [torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif input.dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=input.dtype) - else: - ret = torch.asarray(ret, dtype=torch.float32) - return ret - - -median.support_native_out = False - - -def nanmean( +def igamma( a: torch.Tensor, /, *, - axis: Optional[Union[int, Tuple[int]]] = None, - keepdims: bool = False, - dtype: Optional[torch.dtype] = None, + x: torch.Tensor, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) - - -nanmean.support_native_out = True - - -def _validate_quantile(q): - if isinstance(q, float): - q = torch.as_tensor(q) - if q.ndim == 1 and torch.numel(q) < 10: - for i in range(torch.numel(q)): - if not (0.0 <= q[i] <= 1.0): - return False - else: - if not (torch.all(0 <= q) and torch.all(q <= 1)): - return False - return True - - -def _to_positive_axis(axis, ndim): - if not isinstance(axis, (list, tuple)): - axis = [axis] - - if len(axis) == 0: - raise ValueError("Axis can't be empty!") - - if len(set(axis)) != len(axis): - raise ValueError("Duplicated axis!") - - for i in range(len(axis)): - if not (isinstance(axis[i], int) and (ndim > axis[i] >= -ndim)): - raise ValueError("Axis must be int in range [-rank(x), rank(x))") - if axis[i] < 0: - axis[i] += ndim - return axis - - -def _handle_axis(a, q, fn, keepdims=False, axis=None): - nd = a.ndim - axis_arg = deepcopy(axis) - if axis is not None: - axis = _to_positive_axis(axis, nd) - - if len(axis) == 1: - axis_arg = axis[0] - else: - keep = set(range(nd)) - set(axis) - nkeep = len(keep) - - for i, s in enumerate(sorted(keep)): - a = torch.moveaxis(a, s, i) - a = a.view( - [ - *a.shape[:nkeep], - -1, - ] - ) - axis_arg = -1 - - ret = fn(a, q, axis=axis_arg) - - if keepdims: - if axis is None: - index_ret = (None,) * nd - else: - index_ret = tuple(None if i in axis else slice(None) for i in range(nd)) - ret = ret[(Ellipsis,) + index_ret] - - return ret - - -def _quantile(a, q, axis=None): - ret_dtype = a.dtype - if isinstance(q, float): - q = torch.as_tensor(q) - if isinstance(q, torch.Tensor) and q.ndim > 1: - raise ValueError("q argument must be a scalar or 1-dimensional!") - if axis is None: - axis = 0 - a = a.flatten() - - n = a.shape[axis] - - indices = q * (n - 1) - - a = torch.sort(a, axis)[axis] - indices_below = torch.floor(indices).to(torch.int64) - indices_upper = torch.ceil(indices).to(torch.int64) - - weights = indices - indices_below.to(torch.float64) - - indices_below = torch.clip(indices_below, 0, n - 1) - indices_upper = torch.clip(indices_upper, 0, n - 1) - tensor_upper = torch.index_select(a, 0, indices_upper) - tensor_below = torch.index_select(a, 0, indices_below) - - pred = weights <= 0.5 - out = torch.where(pred, tensor_below, tensor_upper) - return out.to(ret_dtype) - - -def _compute_quantile_wrapper( - x, q, axis=None, keepdims=False, interpolation="linear", out=None -): - if not _validate_quantile(q): - raise ValueError("Quantiles must be in the range [0, 1]") - if interpolation in [ - "linear", - "lower", - "higher", - "midpoint", - "nearest", - "nearest_jax", - ]: - if interpolation == "nearest_jax": - return _handle_axis(x, q, _quantile, keepdims=keepdims, axis=axis) - else: - return torch.quantile( - x, q, dim=axis, keepdim=keepdims, interpolation=interpolation, out=out - ) - else: - raise ValueError( - "Interpolation must be 'linear', 'lower', 'higher', 'midpoint' or 'nearest'" - ) + return torch.special.gammainc(a, x, out=out) -@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) -def quantile( - a: torch.Tensor, - q: Union[torch.Tensor, float], +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bool")}, backend_version) +def median( + input: torch.Tensor, /, *, - axis: Optional[Union[Sequence[int], int]] = None, + axis: Optional[Union[Tuple[int], int]] = None, keepdims: bool = False, - interpolation: str = "linear", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # added the nearest_jax mode to enable jax-like calculations for method="nearest" - return _compute_quantile_wrapper( - a, - q, + if isinstance(axis, tuple): + if len(axis) == 1: + axis = axis[0] + ret = quantile( + input, + 0.5, axis=axis, keepdims=keepdims, - interpolation=interpolation, - out=out, + interpolation="midpoint", ) + if input.dtype in [torch.int64, torch.float64]: + ret = torch.asarray(ret, dtype=torch.float64) + elif input.dtype in [torch.float16, torch.bfloat16]: + ret = torch.asarray(ret, dtype=input.dtype) + else: + ret = torch.asarray(ret, dtype=torch.float32) + return ret -quantile.support_native_out = True - - -def corrcoef( - x: torch.Tensor, +def nanmean( + a: torch.Tensor, /, *, - y: Optional[torch.Tensor] = None, - rowvar: bool = True, + axis: Optional[Union[int, Tuple[int]]] = None, + keepdims: bool = False, + dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if y is None: - xarr = x - else: - axis = 0 if rowvar else 1 - xarr = torch.concat([x, y], dim=axis) - xarr = xarr.T if not rowvar else xarr - - return torch.corrcoef(xarr) - - -def _nanmedian(input, axis, keepdims): - dtype = input.dtype - temp = input.to(torch.float64) - num_dim = len(temp.size()) - keepdim_shape = list(temp.size()) - q = 0.5 - - axis = [axis] if isinstance(axis, int) else list(axis) - - for i in axis: - keepdim_shape[i] = 1 - axis = [num_dim + x if x < 0 else x for x in axis] - axis.sort() - dimension = len(temp.size()) - while len(axis) > 0: - axis1 = axis[0] - for axis2 in range(axis1 + 1, dimension): - temp = torch.transpose(temp, axis1, axis2) - axis1 = axis2 - axis = [x - 1 for x in axis] - axis.pop(0) - dimension = dimension - 1 - temp = torch.flatten(temp, start_dim=dimension - len(axis)) - ret = torch.nanquantile(temp, q, dim=-1, keepdim=keepdims, interpolation="midpoint") - if keepdims: - keepdim_shape = tuple(keepdim_shape) - ret = ret.reshape(keepdim_shape) - - if dtype in [torch.int32, torch.int64, torch.float64]: - ret = torch.asarray(ret, dtype=torch.float64) - elif dtype in [torch.float16, torch.bfloat16]: - ret = torch.asarray(ret, dtype=torch.float16) - else: - ret = torch.asarray(ret, dtype=torch.float32) - - return ret + return torch.nanmean(a, dim=axis, keepdim=keepdims, dtype=dtype, out=out) @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) @@ -438,166 +565,33 @@ def nanmedian( return _nanmedian(input, axis, keepdims) -nanmedian.support_native_out = True - - -def bincount( - x: torch.Tensor, - /, - *, - weights: Optional[torch.Tensor] = None, - minlength: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if weights is None: - ret = torch.bincount(x, minlength=minlength) - ret = ret.to(x.dtype) - else: - ret = torch.bincount(x, weights=weights, minlength=minlength) - ret = ret.to(weights.dtype) - return ret - - -bincount.support_native_out = False - - -def igamma( +@with_unsupported_dtypes({"2.0.1 and below": ("bfloat16", "float16")}, backend_version) +def quantile( a: torch.Tensor, + q: Union[torch.Tensor, float], /, *, - x: torch.Tensor, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: bool = False, + interpolation: str = "linear", out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.special.gammainc(a, x, out=out) + # added the nearest_jax mode to enable jax-like calculations for method="nearest" + return _compute_quantile_wrapper( + a, + q, + axis=axis, + keepdims=keepdims, + interpolation=interpolation, + out=out, + ) +histogram.support_native_out = True +median.support_native_out = False +nanmean.support_native_out = True +quantile.support_native_out = True +nanmedian.support_native_out = True +bincount.support_native_out = False igamma.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def cov( - x1: torch.Tensor, - x2: torch.Tensor = None, - /, - *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[torch.Tensor] = None, - aweights: Optional[torch.Tensor] = None, - dtype: Optional[torch.dtype] = None, -) -> torch.Tensor: - # dtype casts separately - if fweights is not None: - fweights = fweights.type(torch.int64) - if aweights is not None: - aweights = aweights.type(torch.float64) - - if x1.dim() > 2: - raise ValueError("x1 has more than 2 dimensions") - - if x2 is not None: - if x2.dim() > 2: - raise ValueError("x2 has more than 2 dimensions") - - if ddof is None: - if bias == 0: - ddof = 1 - else: - ddof = 0 - - if dtype is None: - x1 = x1.type(torch.float64) - if x2 is not None: - x2 = x2.type(torch.float64) - else: - x1 = x1.type(dtype) - if x2 is not None: - x2 = x2.type(dtype) - - X = x1 - if not rowVar and len(x1.shape) != 1: - X = torch.t(x1) - - if x2 is not None: - if not rowVar and len(x2.shape) != 1: - x2 = torch.t(x2) - X = torch.vstack((X, x2)) - - return torch.cov(X, correction=ddof, fweights=fweights, aweights=aweights) - - cov.support_native_out = False - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("uint8", "bfloat16", "float16")}, - backend_version, -) -def cummax( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - if x.dtype in (torch.bool, torch.float16): - x = x.to(dtype=torch.float64) - elif x.dtype in (torch.int16, torch.int8, torch.uint8): - x = x.to(dtype=torch.int64) - elif x.dtype in (torch.complex64, torch.complex128): - x = x.real.to(dtype=torch.float64) - - if exclusive or reverse: - if exclusive and reverse: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - x1, x2 = torch.concat( - (torch.zeros_like(x1[..., -1:]), x1[..., :-1]), -1 - ), torch.concat((torch.zeros_like(x2[..., -1:]), x2[..., :-1]), -1) - x1, x2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x1, x2 = torch.cummax(x, -1) - res1, res2 = torch.transpose(x1, axis, -1), torch.transpose(x2, axis, -1) - else: - x1, x2 = torch.cummax(torch.flip(x, dims=(axis,)), axis) - res1, res2 = torch.flip(x1, dims=(axis,)), torch.flip(x2, dims=(axis,)) - return res1, res2 - - return torch.cummax(x, axis, out=out) - - -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cummin( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - if not (reverse): - ret = torch.cummin(x, axis)[0] - else: - ret = torch.cummin(torch.flip(x, dims=(axis,)), axis)[0] - ret = torch.flip(ret, (axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret.to(dtype)) - return ret.to(dtype) diff --git a/ivy/functional/backends/torch/general.py b/ivy/functional/backends/torch/general.py index f8cd14960cddb..c75cc289b5bc7 100644 --- a/ivy/functional/backends/torch/general.py +++ b/ivy/functional/backends/torch/general.py @@ -4,11 +4,6 @@ from numbers import Number from operator import mul from typing import Optional, Union, Sequence, Callable, List, Tuple - -try: - import functorch -except ImportError: - functorch = () # for torch 1.10.1 import numpy as np import torch @@ -18,9 +13,18 @@ from . import backend_version, is_variable from ...ivy.general import _broadcast_to +try: + import functorch +except ImportError: + functorch = () # for torch 1.10.1 + torch_scatter = None +# --- Helpers --- # +# --------------- # + + def _parse_index(indices, ndims): ind = list() for so in indices: @@ -44,12 +48,8 @@ def _parse_index(indices, ndims): return ind -def is_native_array(x, /, *, exclusive=False): - if isinstance(x, torch.Tensor): - if exclusive and x.requires_grad: - return False - return True - return False +# --- Main --- # +# ------------ # @with_unsupported_dtypes({"2.0.1 and below": ("complex", "bfloat16")}, backend_version) @@ -66,111 +66,6 @@ def current_backend_str() -> str: return "torch" -def neg_step(query): - return ( - not isinstance(query, (int, bool)) - and not ivy.is_array(query) - and query is not None - and query is not Ellipsis - and ( - (isinstance(query, slice) and query.step is not None and query.step < 0) - or ( - not isinstance(query, slice) - and any( - isinstance(q, slice) and q.step is not None and q.step < 0 - for q in query - ) - ) - ) - ) - - -def get_item( - x: torch.Tensor, - /, - query: Union[torch.Tensor, Tuple], - *, - copy: bool = None, -) -> torch.Tensor: - return x.__getitem__(query) - - -get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) - - -def set_item( - x: torch.Tensor, - query: Union[torch.Tensor, Tuple], - val: torch.Tensor, - /, - *, - copy: Optional[bool] = False, -) -> torch.Tensor: - if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: - val = val.to(x.dtype) - if copy: - x = x.clone() - x.__setitem__(query, val) - return x - - -set_item.partial_mixed_handler = ( - lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad -) - - -def to_numpy( - x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True -) -> Union[np.ndarray, List[np.ndarray]]: - if isinstance(x, (float, int, bool)): - return x - elif isinstance(x, np.ndarray): - if copy: - return x.copy() - else: - return x - elif torch.is_tensor(x): - x = x.resolve_neg().resolve_conj() - if copy: - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16") - return x.detach().cpu().numpy() - else: - raise ivy.utils.exceptions.IvyException( - "Overwriting the same address is not supported for torch." - ) - elif isinstance(x, list): - return [ivy.to_numpy(u) for u in x] - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - -def to_scalar(x: torch.Tensor, /) -> Number: - if isinstance(x, (float, int)): - return x - return x.item() - - -def to_list(x: torch.Tensor, /) -> list: - if isinstance(x, np.ndarray): - return x.tolist() - elif torch.is_tensor(x): - if x.dtype is torch.bfloat16: - default_dtype = ivy.default_float_dtype(as_native=True) - if default_dtype is torch.bfloat16: - x = x.to(torch.float32) - else: - x = x.to(default_dtype) - return x.detach().cpu().numpy().astype("bfloat16").tolist() - else: - return x.detach().cpu().numpy().tolist() - raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") - - def gather( params: torch.Tensor, indices: torch.Tensor, @@ -207,6 +102,36 @@ def gather( return result +def gather_nd( + params: torch.Tensor, + indices: torch.Tensor, + /, + *, + batch_dims: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) + batch_dims = batch_dims % len(params.shape) + result = [] + if batch_dims == 0: + result = gather_nd_helper(params, indices) + else: + for b in range(batch_dims): + if b == 0: + zip_list = [(p, i) for p, i in zip(params, indices)] + else: + zip_list = [ + (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z + ] + for z in zip_list: + p, i = z + r = gather_nd_helper(p, i) + result.append(r) + result = torch.stack(result) + result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) + return result + + def gather_nd_helper(params, indices): indices_shape = indices.shape params_shape = params.shape @@ -237,34 +162,14 @@ def gather_nd_helper(params, indices): return res -def gather_nd( - params: torch.Tensor, - indices: torch.Tensor, +def get_item( + x: torch.Tensor, /, + query: Union[torch.Tensor, Tuple], *, - batch_dims: int = 0, - out: Optional[torch.Tensor] = None, + copy: bool = None, ) -> torch.Tensor: - ivy.utils.assertions.check_gather_nd_input_valid(params, indices, batch_dims) - batch_dims = batch_dims % len(params.shape) - result = [] - if batch_dims == 0: - result = gather_nd_helper(params, indices) - else: - for b in range(batch_dims): - if b == 0: - zip_list = [(p, i) for p, i in zip(params, indices)] - else: - zip_list = [ - (p, i) for z in [zip(p1, i1) for p1, i1 in zip_list] for p, i in z - ] - for z in zip_list: - p, i = z - r = gather_nd_helper(p, i) - result.append(r) - result = torch.stack(result) - result = result.reshape([*params.shape[0:batch_dims], *result.shape[1:]]) - return result + return x.__getitem__(query) def get_num_dims( @@ -338,6 +243,37 @@ def inplace_variables_supported(): return True +def is_native_array(x, /, *, exclusive=False): + if isinstance(x, torch.Tensor): + if exclusive and x.requires_grad: + return False + return True + return False + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version +) +def isin( + elements: torch.tensor, + test_elements: torch.tensor, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> torch.tensor: + return torch.isin( + elements, + test_elements, + assume_unique=assume_unique, + invert=invert, + ) + + +def itemsize(x: torch.tensor) -> int: + return x.element_size() + + def multiprocessing(context: Optional[str] = None): import torch.multiprocessing @@ -346,6 +282,25 @@ def multiprocessing(context: Optional[str] = None): return torch.multiprocessing.get_context(context) +def neg_step(query): + return ( + not isinstance(query, (int, bool)) + and not ivy.is_array(query) + and query is not None + and query is not Ellipsis + and ( + (isinstance(query, slice) and query.step is not None and query.step < 0) + or ( + not isinstance(query, slice) + and any( + isinstance(q, slice) and q.step is not None and q.step < 0 + for q in query + ) + ) + ) + ) + + @with_unsupported_dtypes( { "2.0.1 and below": ("bfloat16",), @@ -395,9 +350,6 @@ def scatter_flat( return res -scatter_flat.support_native_out = True - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -492,7 +444,20 @@ def scatter_nd( return res -scatter_nd.support_native_out = True +def set_item( + x: torch.Tensor, + query: Union[torch.Tensor, Tuple], + val: torch.Tensor, + /, + *, + copy: Optional[bool] = False, +) -> torch.Tensor: + if hasattr(x, "dtype") and hasattr(val, "dtype") and x.dtype != val.dtype: + val = val.to(x.dtype) + if copy: + x = x.clone() + x.__setitem__(query, val) + return x def shape( @@ -507,6 +472,58 @@ def shape( return ivy.Shape(x.shape) +def to_list(x: torch.Tensor, /) -> list: + if isinstance(x, np.ndarray): + return x.tolist() + elif torch.is_tensor(x): + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16").tolist() + else: + return x.detach().cpu().numpy().tolist() + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + +def to_numpy( + x: Union[torch.Tensor, List[torch.Tensor]], /, *, copy: bool = True +) -> Union[np.ndarray, List[np.ndarray]]: + if isinstance(x, (float, int, bool)): + return x + elif isinstance(x, np.ndarray): + if copy: + return x.copy() + else: + return x + elif torch.is_tensor(x): + x = x.resolve_neg().resolve_conj() + if copy: + if x.dtype is torch.bfloat16: + default_dtype = ivy.default_float_dtype(as_native=True) + if default_dtype is torch.bfloat16: + x = x.to(torch.float32) + else: + x = x.to(default_dtype) + return x.detach().cpu().numpy().astype("bfloat16") + return x.detach().cpu().numpy() + else: + raise ivy.utils.exceptions.IvyException( + "Overwriting the same address is not supported for torch." + ) + elif isinstance(x, list): + return [ivy.to_numpy(u) for u in x] + raise ivy.utils.exceptions.IvyException("Expected a pytorch tensor.") + + +def to_scalar(x: torch.Tensor, /) -> Number: + if isinstance(x, (float, int)): + return x + return x.item() + + @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def vmap( func: Callable, @@ -523,27 +540,10 @@ def _vmap(*args): return _vmap -@with_unsupported_dtypes( - {"2.0.1 and below": ("bfloat16", "float16", "complex", "bool")}, backend_version +get_item.partial_mixed_handler = lambda x, query, **kwargs: not neg_step(query) +set_item.partial_mixed_handler = ( + lambda x, query, val, **kwargs: not neg_step(query) and not x.requires_grad ) -def isin( - elements: torch.tensor, - test_elements: torch.tensor, - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> torch.tensor: - return torch.isin( - elements, - test_elements, - assume_unique=assume_unique, - invert=invert, - ) - - +scatter_flat.support_native_out = True +scatter_nd.support_native_out = True isin.support_native_out = True - - -def itemsize(x: torch.tensor) -> int: - return x.element_size() diff --git a/ivy/functional/backends/torch/gradients.py b/ivy/functional/backends/torch/gradients.py index 012e8aafe4a17..ad88bc84f0d3d 100644 --- a/ivy/functional/backends/torch/gradients.py +++ b/ivy/functional/backends/torch/gradients.py @@ -19,20 +19,8 @@ ) -def variable(x, /): - if ivy.is_int_dtype(x.dtype): - x = ivy.astype(x, ivy.default_float_dtype()).to_native() - if not x.is_leaf: - return x.detach().requires_grad_() - return x.clone().requires_grad_() - - -def is_variable(x, /, *, exclusive: bool = False): - return isinstance(x, torch.Tensor) and x.requires_grad - - -def variable_data(x: torch.Tensor, /) -> torch.Tensor: - return x.data +# --- Helpers --- # +# --------------- # def _grad_func(y, xs, retain_grads): @@ -92,6 +80,10 @@ def grad_(x): return grads +# --- Main --- # +# ------------ # + + def execute_with_gradients( func, xs: torch.Tensor, @@ -140,61 +132,6 @@ def execute_with_gradients( return _process_func_ret_and_grads(func_ret, grads, retain_grads) -def value_and_grad(func): - grad_fn = lambda xs: ivy.to_native(func(xs)) - - def callback_fn(xs): - y = grad_fn(xs) - - def autograd_fn(x): - x = ivy.to_native(x) - grad = torch.autograd.grad(y, x, allow_unused=True)[0] - grad = ( - grad - if grad is not None - else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) - ) - grad = ivy.to_ivy(grad) - return grad - - grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) - y = ivy.to_ivy(y) - return y, grads - - return callback_fn - - -def stop_gradient( - x: Optional[torch.Tensor], - /, - *, - preserve_type: bool = True, - out: Optional[torch.Tensor] = None, -): - if is_variable(x) and preserve_type: - if x.grad_fn: - x = x.detach() - x.requires_grad = True - elif x.grad: - x.grad.data.zero_() - return x - return x.detach() - - -def jac(func: Callable): - grad_fn = lambda x_in: ivy.to_native( - func(ivy.to_ivy(x_in, nested=True)), - nested=True, - include_derived=True, - ) - callback_fn = lambda x_in: ivy.to_ivy( - torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), - nested=True, - include_derived=True, - ) - return callback_fn - - def grad(f, argnums=0): if grad.nth == 0: grad.f_original = f @@ -257,5 +194,76 @@ def _inner(*args, **kwargs): return _nth_derivative(grad.nth) +def is_variable(x, /, *, exclusive: bool = False): + return isinstance(x, torch.Tensor) and x.requires_grad + + +def jac(func: Callable): + grad_fn = lambda x_in: ivy.to_native( + func(ivy.to_ivy(x_in, nested=True)), + nested=True, + include_derived=True, + ) + callback_fn = lambda x_in: ivy.to_ivy( + torch.func.jacfwd(grad_fn)((ivy.to_native(x_in, nested=True))), + nested=True, + include_derived=True, + ) + return callback_fn + + +def stop_gradient( + x: Optional[torch.Tensor], + /, + *, + preserve_type: bool = True, + out: Optional[torch.Tensor] = None, +): + if is_variable(x) and preserve_type: + if x.grad_fn: + x = x.detach() + x.requires_grad = True + elif x.grad: + x.grad.data.zero_() + return x + return x.detach() + + +def value_and_grad(func): + grad_fn = lambda xs: ivy.to_native(func(xs)) + + def callback_fn(xs): + y = grad_fn(xs) + + def autograd_fn(x): + x = ivy.to_native(x) + grad = torch.autograd.grad(y, x, allow_unused=True)[0] + grad = ( + grad + if grad is not None + else ivy.to_native(ivy.zeros_like(ivy.to_ivy(x))) + ) + grad = ivy.to_ivy(grad) + return grad + + grads = ivy.nested_map(xs, autograd_fn, include_derived=True, shallow=False) + y = ivy.to_ivy(y) + return y, grads + + return callback_fn + + +def variable(x, /): + if ivy.is_int_dtype(x.dtype): + x = ivy.astype(x, ivy.default_float_dtype()).to_native() + if not x.is_leaf: + return x.detach().requires_grad_() + return x.clone().requires_grad_() + + +def variable_data(x: torch.Tensor, /) -> torch.Tensor: + return x.data + + grad.f_original = None grad.nth = 0 diff --git a/ivy/functional/backends/torch/layers.py b/ivy/functional/backends/torch/layers.py index 0ae11f08e2e9b..ecb3f3f29d9bd 100644 --- a/ivy/functional/backends/torch/layers.py +++ b/ivy/functional/backends/torch/layers.py @@ -11,22 +11,8 @@ from ivy.functional.ivy.layers import _handle_padding, _deconv_length -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16", "bfloat16", "complex")}, - backend_version, -) -def linear( - x: torch.Tensor, - weight: torch.Tensor, - /, - *, - bias: Optional[torch.Tensor] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.nn.functional.linear(x, weight, bias) - - -linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 +# --- Helpers --- # +# --------------- # def _ff_xd_before_conv(x, filters, dims, filter_format, x_dilations): @@ -126,6 +112,10 @@ def _pad_before_conv_tranpose( return not_valid_pad, padding_list, output_padding +# --- Main --- # +# ------------ # + + @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version, @@ -288,48 +278,6 @@ def conv2d_transpose( return res -@with_unsupported_dtypes( - { - "2.0.1 and below": ( - "float16", - "bfloat16", - "complex", - ) - }, - backend_version, -) -# noinspection PyUnresolvedReferences -def depthwise_conv2d( - x: torch.Tensor, - filters: torch.Tensor, - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - strides = [strides] * 2 if isinstance(strides, int) else strides - dilations = [dilations] * 2 if isinstance(dilations, int) else dilations - if data_format == "NHWC": - x = x.permute(0, 3, 1, 2) - filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters - filters = torch.unsqueeze(filters, -1) - dims_in = filters.shape[-2] - filters = filters.permute(2, 3, 0, 1) - x, padding = _pad_before_conv( - x, filters, strides, padding, 2, dilations, "channel_first" - ) - # noinspection PyArgumentEqualDefault - res = torch.nn.functional.conv2d( - x, filters, None, strides, padding, dilations, dims_in - ) - if data_format == "NHWC": - return res.permute(0, 2, 3, 1) - return res - - @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "complex")}, backend_version ) @@ -544,3 +492,63 @@ def conv_general_transpose( if data_format == "channel_last": res = res.permute(0, *range(2, dims + 2), 1) return res + + +@with_unsupported_dtypes( + { + "2.0.1 and below": ( + "float16", + "bfloat16", + "complex", + ) + }, + backend_version, +) +# noinspection PyUnresolvedReferences +def depthwise_conv2d( + x: torch.Tensor, + filters: torch.Tensor, + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + strides = [strides] * 2 if isinstance(strides, int) else strides + dilations = [dilations] * 2 if isinstance(dilations, int) else dilations + if data_format == "NHWC": + x = x.permute(0, 3, 1, 2) + filters = ivy.squeeze(filters, 3).to_native() if filters.ndim == 4 else filters + filters = torch.unsqueeze(filters, -1) + dims_in = filters.shape[-2] + filters = filters.permute(2, 3, 0, 1) + x, padding = _pad_before_conv( + x, filters, strides, padding, 2, dilations, "channel_first" + ) + # noinspection PyArgumentEqualDefault + res = torch.nn.functional.conv2d( + x, filters, None, strides, padding, dilations, dims_in + ) + if data_format == "NHWC": + return res.permute(0, 2, 3, 1) + return res + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16", "bfloat16", "complex")}, + backend_version, +) +def linear( + x: torch.Tensor, + weight: torch.Tensor, + /, + *, + bias: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.nn.functional.linear(x, weight, bias) + + +linear.partial_mixed_handler = lambda x, weight, **kwargs: weight.ndim == 2 diff --git a/ivy/functional/backends/torch/linear_algebra.py b/ivy/functional/backends/torch/linear_algebra.py index 950c2ff6fd295..2f9a760617db2 100644 --- a/ivy/functional/backends/torch/linear_algebra.py +++ b/ivy/functional/backends/torch/linear_algebra.py @@ -40,9 +40,6 @@ def cholesky( return ret -cholesky.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "complex")}, backend_version) def cross( x1: torch.Tensor, @@ -68,15 +65,24 @@ def cross( ) -cross.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def det(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch.Tensor: return torch.linalg.det(x, out=out) -det.support_native_out = True +# Extra # +# ----- # + + +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def diag( + x: torch.Tensor, + /, + *, + k: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.tensor: + return torch.diag(x, diagonal=k) def diagonal( @@ -91,6 +97,17 @@ def diagonal( return torch.diagonal(x, offset=offset, dim1=axis1, dim2=axis2) +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def eig( + x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + result_tuple = NamedTuple( + "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] + ) + eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) + return result_tuple(eigenvalues, eigenvectors) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -102,9 +119,6 @@ def eigh( return result_tuple(eigenvalues, eigenvectors) -eigh.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def eigvalsh( x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None @@ -112,9 +126,6 @@ def eigvalsh( return torch.linalg.eigvalsh(x, UPLO=UPLO, out=out) -eigvalsh.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, backend_version) def inner( x1: torch.Tensor, x2: torch.Tensor, /, *, out: Optional[torch.Tensor] = None @@ -132,9 +143,6 @@ def inner( return torch.inner(x1, x2, out=out) -inner.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def inv( x: torch.Tensor, @@ -157,9 +165,6 @@ def inv( return ret -inv.support_native_out = True - - @with_unsupported_dtypes( {"2.0.1 and below": ("float16", "bfloat16", "bool")}, backend_version ) @@ -191,9 +196,6 @@ def matmul( return torch.matmul(x1, x2, out=out) -matmul.support_native_out = True - - @with_supported_dtypes({"2.0.1 and below": ("float", "complex")}, backend_version) def matrix_norm( x: torch.Tensor, @@ -207,23 +209,6 @@ def matrix_norm( return torch.linalg.matrix_norm(x, ord=ord, dim=axis, keepdim=keepdims, out=out) -matrix_norm.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def eig( - x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor]: - result_tuple = NamedTuple( - "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)] - ) - eigenvalues, eigenvectors = torch.linalg.eig(x, out=out) - return result_tuple(eigenvalues, eigenvectors) - - -eig.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_power( x: torch.Tensor, n: int, /, *, out: Optional[torch.Tensor] = None @@ -231,9 +216,6 @@ def matrix_power( return torch.linalg.matrix_power(x, n, out=out) -matrix_power.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def matrix_rank( x: torch.Tensor, @@ -271,9 +253,6 @@ def matrix_rank( return ret -matrix_rank.support_native_out = True - - def matrix_transpose( x: torch.Tensor, /, *, conjugate: bool = False, out: Optional[torch.Tensor] = None ) -> torch.Tensor: @@ -294,9 +273,6 @@ def outer( return torch.outer(x1, x2, out=out) -outer.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def pinv( x: torch.Tensor, @@ -310,21 +286,6 @@ def pinv( return torch.linalg.pinv(x, rtol, out=out) -pinv.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def tensorsolve( - x1: torch.Tensor, - x2: torch.Tensor, - /, - *, - axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.linalg.tensorsolve(x1, x2, dims=axes) - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def qr( x: torch.Tensor, @@ -359,9 +320,6 @@ def slogdet( return results(sign, logabsdet) -slogdet.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def solve( x1: torch.Tensor, @@ -420,9 +378,6 @@ def svdvals(x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None) -> torch. return torch.linalg.svdvals(x, out=out) -svdvals.support_native_out = True - - # ToDo: re-add int32 support once # (https://github.com/pytorch/pytorch/issues/84530) is fixed @with_unsupported_dtypes({"2.0.1 and below": ("int32",)}, backend_version) @@ -448,6 +403,18 @@ def tensordot( return ret +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def tensorsolve( + x1: torch.Tensor, + x2: torch.Tensor, + /, + *, + axes: Optional[Union[int, Tuple[List[int], List[int]]]] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.linalg.tensorsolve(x1, x2, dims=axes) + + @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) def trace( x: torch.Tensor, @@ -465,6 +432,30 @@ def trace( return ret +@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) +def vander( + x: torch.tensor, + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[torch.tensor] = None, +) -> torch.tensor: + # torch.vander hasn't been used as it produces 0 gradients + N = ivy.default(N, x.shape[-1]) + start, stop, step = N - 1, -1, -1 + if increasing: + start, stop, step = 0, N, 1 + ret = torch.pow( + torch.transpose(torch.unsqueeze(x, 0), 0, 1), + torch.arange(start, stop, step), + out=out, + ) + if ret.dtype != x.dtype: + return ret.to(x.dtype) + return ret + + def vecdot( x1: torch.Tensor, x2: torch.Tensor, @@ -485,9 +476,6 @@ def vecdot( return torch.tensordot(x1, x2, dims=([axis], [axis])).to(dtype) -vecdot.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("integer",)}, backend_version) def vector_norm( x: torch.Tensor, @@ -507,48 +495,6 @@ def vector_norm( return torch.linalg.vector_norm(x, ord, axis, keepdims, out=out) -vector_norm.support_native_out = True - - -# Extra # -# ----- # - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def diag( - x: torch.Tensor, - /, - *, - k: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.tensor: - return torch.diag(x, diagonal=k) - - -@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version) -def vander( - x: torch.tensor, - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[torch.tensor] = None, -) -> torch.tensor: - # torch.vander hasn't been used as it produces 0 gradients - N = ivy.default(N, x.shape[-1]) - start, stop, step = N - 1, -1, -1 - if increasing: - start, stop, step = 0, N, 1 - ret = torch.pow( - torch.transpose(torch.unsqueeze(x, 0), 0, 1), - torch.arange(start, stop, step), - out=out, - ) - if ret.dtype != x.dtype: - return ret.to(x.dtype) - return ret - - @with_unsupported_dtypes( { "2.0.1 and below": ( @@ -578,4 +524,22 @@ def vector_to_skew_symmetric_matrix( return torch.cat((row1, row2, row3), -2, out=out) +cholesky.support_native_out = True +cross.support_native_out = True +det.support_native_out = True +eigh.support_native_out = True +eigvalsh.support_native_out = True +inner.support_native_out = True +inv.support_native_out = True +matmul.support_native_out = True +matrix_norm.support_native_out = True +eig.support_native_out = True +matrix_power.support_native_out = True +matrix_rank.support_native_out = True +outer.support_native_out = True +pinv.support_native_out = True +slogdet.support_native_out = True +svdvals.support_native_out = True +vecdot.support_native_out = True +vector_norm.support_native_out = True vector_to_skew_symmetric_matrix.support_native_out = True diff --git a/ivy/functional/backends/torch/manipulation.py b/ivy/functional/backends/torch/manipulation.py index fe8b37b4dcafd..1b5efcec1cf06 100644 --- a/ivy/functional/backends/torch/manipulation.py +++ b/ivy/functional/backends/torch/manipulation.py @@ -14,12 +14,42 @@ from . import backend_version +# --- Helpers --- # +# --------------- # + + def _reshape_fortran_torch(x, shape): if len(x.shape) > 0: x = x.permute(*reversed(range(len(x.shape)))) return x.reshape(shape[::-1]).permute(list(range(len(shape)))[::-1]) +# --- Main --- # +# ------------ # + + +@with_unsupported_dtypes( + {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version +) +def clip( + x: torch.Tensor, + x_min: Union[Number, torch.Tensor], + x_max: Union[Number, torch.Tensor], + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if hasattr(x_min, "dtype"): + x_min = torch.asarray(x_min, device=x.device) + x_max = torch.asarray(x_max, device=x.device) + promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) + promoted_type = torch.promote_types(promoted_type, x.dtype) + x_min = x_min.to(promoted_type) + x_max = x_max.to(promoted_type) + x = x.to(promoted_type) + return torch.clamp(x, x_min, x_max, out=out) + + # Array API Standard # # -------------------# @@ -43,7 +73,26 @@ def concat( return torch.cat(xs, dim=axis, out=out) -concat.support_native_out = True +def constant_pad( + x: torch.Tensor, + /, + pad_width: List[List[int]], + *, + value: Number = 0.0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if 0 in x.shape: + new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] + return torch.ones(new_shape, dtype=x.dtype) * value + if x.shape == (): + x = x.unsqueeze(0) + if isinstance(pad_width, torch.Tensor): + pad_width = pad_width.detach().cpu().numpy().tolist() + pad_width_flat: List[int] = list() + for pad_width_sec in reversed(pad_width): + for item in pad_width_sec: + pad_width_flat.append(item) + return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) def expand_dims( @@ -93,6 +142,23 @@ def permute_dims( return torch.permute(x, axes) +@with_unsupported_dtypes( + {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version +) +def repeat( + x: torch.Tensor, + /, + repeats: Union[int, Iterable[int]], + *, + axis: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if len(x.shape) == 0 and axis in [0, -1]: + axis = None + repeats = torch.tensor(repeats) + return torch.repeat_interleave(x, repeats, axis) + + def reshape( x: torch.Tensor, /, @@ -132,69 +198,6 @@ def roll( return torch.roll(x, shift, axis) -def squeeze( - x: torch.Tensor, - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - copy: Optional[bool] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if isinstance(axis, int): - if x.size(dim=axis) > 1: - raise ValueError( - "Expected dimension of size [{}, {}], but found " - "dimension size {}".format(-x.dim(), x.dim(), axis) - ) - if x.shape[axis] != 1: - raise ivy.utils.exceptions.IvyException( - f"Expected size of axis to be 1 but was {x.shape[axis]}" - ) - return torch.squeeze(x, axis) - if axis is None: - if copy: - newarr = torch.clone(x) - return torch.squeeze(newarr) - return torch.squeeze(x) - newarr = torch.clone(x) - if isinstance(axis, tuple): - axis = list(axis) - normalise_axis = [ - (len(x.shape) - abs(element)) if element < 0 else element for element in axis - ] - normalise_axis.sort() - axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] - dim = x.dim() - for i in axis_updated_after_squeeze: - shape = x.shape[i] - if shape > 1 and (shape < -dim or dim <= shape): - raise ValueError( - "Expected dimension of size [{}, {}], " - "but found dimension size {}".format(-dim, dim, shape) - ) - else: - if copy: - newarr = torch.squeeze(newarr, i) - else: - x = torch.squeeze(x, i) - if copy: - return newarr - return x - - -def stack( - arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], - /, - *, - axis: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.stack(arrays, axis, out=out) - - -stack.support_native_out = True - - # Extra # # ------# @@ -244,61 +247,64 @@ def split( return list(torch.split(x, num_or_size_splits, axis)) -@with_unsupported_dtypes( - {"2.0.1 and below": ("int8", "int16", "uint8")}, backend_version -) -def repeat( +def squeeze( x: torch.Tensor, /, - repeats: Union[int, Iterable[int]], *, - axis: Optional[int] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + copy: Optional[bool] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if len(x.shape) == 0 and axis in [0, -1]: - axis = None - repeats = torch.tensor(repeats) - return torch.repeat_interleave(x, repeats, axis) - - -def tile( - x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - if isinstance(repeats, torch.Tensor): - repeats = repeats.detach().cpu().numpy().tolist() - return x.repeat(repeats) + if isinstance(axis, int): + if x.size(dim=axis) > 1: + raise ValueError( + "Expected dimension of size [{}, {}], but found " + "dimension size {}".format(-x.dim(), x.dim(), axis) + ) + if x.shape[axis] != 1: + raise ivy.utils.exceptions.IvyException( + f"Expected size of axis to be 1 but was {x.shape[axis]}" + ) + return torch.squeeze(x, axis) + if axis is None: + if copy: + newarr = torch.clone(x) + return torch.squeeze(newarr) + return torch.squeeze(x) + newarr = torch.clone(x) + if isinstance(axis, tuple): + axis = list(axis) + normalise_axis = [ + (len(x.shape) - abs(element)) if element < 0 else element for element in axis + ] + normalise_axis.sort() + axis_updated_after_squeeze = [dim - key for (key, dim) in enumerate(normalise_axis)] + dim = x.dim() + for i in axis_updated_after_squeeze: + shape = x.shape[i] + if shape > 1 and (shape < -dim or dim <= shape): + raise ValueError( + "Expected dimension of size [{}, {}], " + "but found dimension size {}".format(-dim, dim, shape) + ) + else: + if copy: + newarr = torch.squeeze(newarr, i) + else: + x = torch.squeeze(x, i) + if copy: + return newarr + return x -def constant_pad( - x: torch.Tensor, +def stack( + arrays: Union[Tuple[torch.Tensor], List[torch.Tensor]], /, - pad_width: List[List[int]], *, - value: Number = 0.0, + axis: int = 0, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if 0 in x.shape: - new_shape = [s + sum(pad_width[i]) for i, s in enumerate(x.shape)] - return torch.ones(new_shape, dtype=x.dtype) * value - if x.shape == (): - x = x.unsqueeze(0) - if isinstance(pad_width, torch.Tensor): - pad_width = pad_width.detach().cpu().numpy().tolist() - pad_width_flat: List[int] = list() - for pad_width_sec in reversed(pad_width): - for item in pad_width_sec: - pad_width_flat.append(item) - return torch.nn.functional.pad(x, pad_width_flat, mode="constant", value=value) - - -def zero_pad( - x: torch.Tensor, - /, - pad_width: List[List[int]], - *, - out: Optional[torch.Tensor] = None, -): - return constant_pad(x, pad_width, value=0.0) + return torch.stack(arrays, axis, out=out) def swapaxes( @@ -313,29 +319,12 @@ def swapaxes( return torch.transpose(x, axis0, axis1) -@with_unsupported_dtypes( - {"2.0.1 and below": ("bool", "float16", "complex")}, backend_version -) -def clip( - x: torch.Tensor, - x_min: Union[Number, torch.Tensor], - x_max: Union[Number, torch.Tensor], - /, - *, - out: Optional[torch.Tensor] = None, +def tile( + x: torch.Tensor, /, repeats: Sequence[int], *, out: Optional[torch.Tensor] = None ) -> torch.Tensor: - if hasattr(x_min, "dtype"): - x_min = torch.asarray(x_min, device=x.device) - x_max = torch.asarray(x_max, device=x.device) - promoted_type = torch.promote_types(x_min.dtype, x_max.dtype) - promoted_type = torch.promote_types(promoted_type, x.dtype) - x_min = x_min.to(promoted_type) - x_max = x_max.to(promoted_type) - x = x.to(promoted_type) - return torch.clamp(x, x_min, x_max, out=out) - - -clip.support_native_out = True + if isinstance(repeats, torch.Tensor): + repeats = repeats.detach().cpu().numpy().tolist() + return x.repeat(repeats) def unstack( @@ -355,3 +344,18 @@ def unstack( if keepdims: return [r.unsqueeze(axis) for r in ret] return ret + + +def zero_pad( + x: torch.Tensor, + /, + pad_width: List[List[int]], + *, + out: Optional[torch.Tensor] = None, +): + return constant_pad(x, pad_width, value=0.0) + + +concat.support_native_out = True +stack.support_native_out = True +clip.support_native_out = True diff --git a/ivy/functional/backends/torch/random.py b/ivy/functional/backends/torch/random.py index 55495e074b067..929ec4048a34d 100644 --- a/ivy/functional/backends/torch/random.py +++ b/ivy/functional/backends/torch/random.py @@ -14,53 +14,6 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Extra # -# ------# - - -def random_uniform( - *, - low: Union[float, torch.Tensor] = 0.0, - high: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - device: torch.device, - seed: Optional[int] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - shape = _check_bounds_and_get_shape(low, high, shape).shape - rand_range = high - low - if seed: - torch.manual_seed(seed) - if torch.is_tensor(shape): - shape = shape.tolist() - return ( - torch.rand(shape, device=device, dtype=torch.float) * rand_range + low - ).type(dtype) - - -def random_normal( - *, - mean: Union[float, torch.Tensor] = 0.0, - std: Union[float, torch.Tensor] = 1.0, - shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, - dtype: torch.dtype, - seed: Optional[int] = None, - device: torch.device, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _check_valid_scale(std) - shape = _check_bounds_and_get_shape(mean, std, shape).shape - dtype = ivy.as_native_dtype(dtype) - if seed: - torch.manual_seed(seed) - if isinstance(mean, (int, float)) and isinstance(std, (int, float)): - return torch.normal(mean, std, shape, out=out).type(dtype).to(device) - return torch.normal(mean, std, out=out).type(dtype).to(device) - - -random_normal.support_native_out = True - @with_unsupported_dtypes({"2.0.1 and below": ("bfloat16",)}, backend_version) def multinomial( @@ -90,9 +43,6 @@ def multinomial( return torch.multinomial(probs.float(), num_samples, replace, out=out).to(device) -multinomial.support_native_out = True - - def randint( low: Union[int, torch.Tensor], high: Union[int, torch.Tensor], @@ -115,6 +65,51 @@ def randint( return (torch.rand(shape, device=device) * rand_range + low).to(dtype) +def random_normal( + *, + mean: Union[float, torch.Tensor] = 0.0, + std: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + seed: Optional[int] = None, + device: torch.device, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_valid_scale(std) + shape = _check_bounds_and_get_shape(mean, std, shape).shape + dtype = ivy.as_native_dtype(dtype) + if seed: + torch.manual_seed(seed) + if isinstance(mean, (int, float)) and isinstance(std, (int, float)): + return torch.normal(mean, std, shape, out=out).type(dtype).to(device) + return torch.normal(mean, std, out=out).type(dtype).to(device) + + +# Extra # +# ------# + + +def random_uniform( + *, + low: Union[float, torch.Tensor] = 0.0, + high: Union[float, torch.Tensor] = 1.0, + shape: Optional[Union[torch.Tensor, ivy.NativeShape, Sequence[int]]] = None, + dtype: torch.dtype, + device: torch.device, + seed: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + shape = _check_bounds_and_get_shape(low, high, shape).shape + rand_range = high - low + if seed: + torch.manual_seed(seed) + if torch.is_tensor(shape): + shape = shape.tolist() + return ( + torch.rand(shape, device=device, dtype=torch.float) * rand_range + low + ).type(dtype) + + def seed(*, seed_value: int = 0) -> None: torch.manual_seed(seed_value) torch.cuda.manual_seed(seed_value) @@ -142,4 +137,6 @@ def shuffle( return torch.index_select(x, 0, torch.randperm(batch_size), out=out) +random_normal.support_native_out = True +multinomial.support_native_out = True shuffle.support_native_out = True diff --git a/ivy/functional/backends/torch/searching.py b/ivy/functional/backends/torch/searching.py index 702100b3b5ab1..778a0a3b5d141 100644 --- a/ivy/functional/backends/torch/searching.py +++ b/ivy/functional/backends/torch/searching.py @@ -69,6 +69,19 @@ def argmin( return ret +# Extra # +# ----- # + + +def argwhere( + x: torch.Tensor, + /, + *, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return torch.argwhere(x) + + def nonzero( x: torch.Tensor, /, @@ -107,16 +120,3 @@ def where( if condition.dtype is not torch.bool: condition = condition == 1.0 return ivy.astype(torch.where(condition, x1, x2), x1.dtype, copy=False) - - -# Extra # -# ----- # - - -def argwhere( - x: torch.Tensor, - /, - *, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - return torch.argwhere(x) diff --git a/ivy/functional/backends/torch/sorting.py b/ivy/functional/backends/torch/sorting.py index 6fea3ca9794c3..835accb3e1496 100644 --- a/ivy/functional/backends/torch/sorting.py +++ b/ivy/functional/backends/torch/sorting.py @@ -26,30 +26,6 @@ def argsort( return sorted_indices -argsort.support_native_out = True - - -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def sort( - x: torch.Tensor, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if out is not None: - out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) - sorted_tensor, _ = torch.sort( - x, dim=axis, descending=descending, stable=stable, out=out - ) - return sorted_tensor - - -sort.support_native_out = True - - # msort @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def msort( @@ -58,9 +34,6 @@ def msort( return torch.msort(a, out=out) -msort.support_native_out = True - - @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) def searchsorted( x: torch.Tensor, @@ -113,4 +86,25 @@ def searchsorted( return torch.searchsorted(x, v, sorter=sorter, side=side).to(ret_dtype) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def sort( + x: torch.Tensor, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if out is not None: + out = tuple([out, torch.zeros(out.shape, dtype=torch.long)]) + sorted_tensor, _ = torch.sort( + x, dim=axis, descending=descending, stable=stable, out=out + ) + return sorted_tensor + + +argsort.support_native_out = True +sort.support_native_out = True +msort.support_native_out = True searchsorted.support_native_out = True diff --git a/ivy/functional/backends/torch/statistical.py b/ivy/functional/backends/torch/statistical.py index 04694a0bdca3f..930858211703a 100644 --- a/ivy/functional/backends/torch/statistical.py +++ b/ivy/functional/backends/torch/statistical.py @@ -1,5 +1,3 @@ -# global -torch_scatter = None from typing import Union, Optional, Sequence import torch @@ -10,30 +8,130 @@ from ivy.func_wrapper import with_unsupported_dtypes from . import backend_version -# Array API Standard # -# -------------------# +# global +torch_scatter = None -@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) -def min( +# --- Helpers --- # +# --------------- # + + +def _infer_dtype(dtype: torch.dtype) -> torch.dtype: + default_dtype = ivy.infer_default_dtype(dtype) + if default_dtype in ivy.valid_dtypes: + if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): + return ivy.as_native_dtype(default_dtype) + return ivy.as_native_dtype(dtype) + + +# --- Main --- # +# ------------ # + + +# Extra # +# ----- # + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "2.0.1 and below": ("uint8", "float16", "bfloat16"), + }, + backend_version, +) +def cumprod( x: torch.Tensor, /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if axis == (): - if ivy.exists(out): - return ivy.inplace_update(out, x) + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + dtype = _infer_dtype(x.dtype) + + if not (exclusive or reverse): + return torch.cumprod(x, axis, dtype=dtype, out=out) + elif exclusive and reverse: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + ret = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumprod(x, -1, dtype=dtype) + ret = torch.transpose(x, axis, -1) + else: + x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + ret = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret + + +# Function does support uint8, but allowing support for unsigned will cause +# the function to break the upcasting rule defined in the Array API Standard +# TODO: bfloat16 support is added in PyTorch 1.12.1 +@with_unsupported_dtypes( + { + "1.12.1 and below": ("uint8", "float16", "bfloat16"), + "1.12.1 and above": ("uint8", "float16"), + }, + backend_version, +) +def cumsum( + x: torch.Tensor, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + *, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = ivy.as_native_dtype(dtype) + if dtype is None: + if ivy.is_int_dtype(x.dtype): + dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) + dtype = _infer_dtype(x.dtype) + if exclusive or reverse: + if exclusive and reverse: + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + x = torch.transpose(x, axis, -1) + x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.transpose(x, axis, -1) + res = torch.flip(x, dims=(axis,)) + elif exclusive: + x = torch.transpose(x, axis, -1) + x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) + x = torch.cumsum(x, -1, dtype=dtype) + res = torch.transpose(x, axis, -1) else: - return x - if not keepdims and not axis and axis != 0: - return torch.amin(input=x, out=out) - return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) + x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) + res = torch.flip(x, dims=(axis,)) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res + return torch.cumsum(x, axis, dtype=dtype, out=out) -min.support_native_out = True +@with_unsupported_dtypes( + {"2.0.1 and below": ("float16",)}, + backend_version, +) +def einsum( + equation: str, + *operands: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dtype = _get_promoted_type_of_operands(operands) + return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) @with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) @@ -55,9 +153,6 @@ def max( return torch.amax(input=x, dim=axis, keepdim=keepdims, out=out) -max.support_native_out = True - - def mean( x: torch.Tensor, /, @@ -81,15 +176,27 @@ def mean( return torch.mean(x, dim=axis, keepdim=keepdims, out=out) -mean.support_native_out = True +# Array API Standard # +# -------------------# -def _infer_dtype(dtype: torch.dtype) -> torch.dtype: - default_dtype = ivy.infer_default_dtype(dtype) - if default_dtype in ivy.valid_dtypes: - if ivy.dtype_bits(dtype) < ivy.dtype_bits(default_dtype): - return ivy.as_native_dtype(default_dtype) - return ivy.as_native_dtype(dtype) +@with_unsupported_dtypes({"2.0.1 and below": ("complex",)}, backend_version) +def min( + x: torch.Tensor, + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + keepdims: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if axis == (): + if ivy.exists(out): + return ivy.inplace_update(out, x) + else: + return x + if not keepdims and not axis and axis != 0: + return torch.amin(input=x, out=out) + return torch.amin(input=x, dim=axis, keepdim=keepdims, out=out) # Function does support uint8, but allowing support for unsigned will cause @@ -214,113 +321,8 @@ def var( ).to(x.dtype) -# Extra # -# ----- # - - -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "2.0.1 and below": ("uint8", "float16", "bfloat16"), - }, - backend_version, -) -def cumprod( - x: torch.Tensor, - /, - *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - dtype = _infer_dtype(x.dtype) - - if not (exclusive or reverse): - return torch.cumprod(x, axis, dtype=dtype, out=out) - elif exclusive and reverse: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - ret = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumprod(x, -1, dtype=dtype) - ret = torch.transpose(x, axis, -1) - else: - x = torch.cumprod(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - ret = torch.flip(x, dims=(axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret - - +min.support_native_out = True +max.support_native_out = True +mean.support_native_out = True cumprod.support_native_out = True - - -# Function does support uint8, but allowing support for unsigned will cause -# the function to break the upcasting rule defined in the Array API Standard -# TODO: bfloat16 support is added in PyTorch 1.12.1 -@with_unsupported_dtypes( - { - "1.12.1 and below": ("uint8", "float16", "bfloat16"), - "1.12.1 and above": ("uint8", "float16"), - }, - backend_version, -) -def cumsum( - x: torch.Tensor, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - *, - dtype: Optional[torch.dtype] = None, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = ivy.as_native_dtype(dtype) - if dtype is None: - if ivy.is_int_dtype(x.dtype): - dtype = ivy.promote_types(x.dtype, ivy.default_int_dtype(as_native=True)) - dtype = _infer_dtype(x.dtype) - if exclusive or reverse: - if exclusive and reverse: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - x = torch.transpose(x, axis, -1) - x = torch.concat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.transpose(x, axis, -1) - res = torch.flip(x, dims=(axis,)) - elif exclusive: - x = torch.transpose(x, axis, -1) - x = torch.cat((torch.zeros_like(x[..., -1:]), x[..., :-1]), -1) - x = torch.cumsum(x, -1, dtype=dtype) - res = torch.transpose(x, axis, -1) - else: - x = torch.cumsum(torch.flip(x, dims=(axis,)), axis, dtype=dtype) - res = torch.flip(x, dims=(axis,)) - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res - return torch.cumsum(x, axis, dtype=dtype, out=out) - - cumsum.support_native_out = True - - -@with_unsupported_dtypes( - {"2.0.1 and below": ("float16",)}, - backend_version, -) -def einsum( - equation: str, - *operands: torch.Tensor, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - dtype = _get_promoted_type_of_operands(operands) - return ivy.astype(torch.einsum(equation, *operands), dtype, copy=False) diff --git a/ivy/functional/backends/torch/utility.py b/ivy/functional/backends/torch/utility.py index ca643ce0d4095..44e6f70bc2baa 100644 --- a/ivy/functional/backends/torch/utility.py +++ b/ivy/functional/backends/torch/utility.py @@ -25,9 +25,6 @@ def all( return x -all.support_native_out = True - - def any( x: torch.Tensor, /, @@ -50,4 +47,5 @@ def any( return x +all.support_native_out = True any.support_native_out = True diff --git a/ivy/functional/frontends/torch/comparison_ops.py b/ivy/functional/frontends/torch/comparison_ops.py index b743b38d135f0..eeaadadfbf864 100644 --- a/ivy/functional/frontends/torch/comparison_ops.py +++ b/ivy/functional/frontends/torch/comparison_ops.py @@ -290,7 +290,7 @@ def topk(input, k, dim=None, largest=True, sorted=True, *, out=None): gt = greater +ne = not_equal ge = greater_equal le = less_equal lt = less -ne = not_equal diff --git a/ivy/functional/ivy/activations.py b/ivy/functional/ivy/activations.py index f19180ab7a529..d55d9dfcde068 100644 --- a/ivy/functional/ivy/activations.py +++ b/ivy/functional/ivy/activations.py @@ -18,6 +18,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _gelu_jax_like( x: Union[ivy.Array, ivy.NativeArray], /, @@ -31,6 +35,87 @@ def _gelu_jax_like( return fn_original(x, approximate=True, out=out) +def _leaky_relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original: Optional[Callable] = None, + alpha: float = 0.2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.astype(x * alpha, x.dtype), + x, + ) + + +def _relu_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + return ivy.where( + ( + ivy.logical_or( + ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) + ) + ), + ivy.array(0.0, dtype=x.dtype), + x, + ) + + +def _softplus_jax_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + fn_original=None, + beta: Optional[Union[int, float]] = None, + threshold: Optional[Union[int, float]] = None, + out: Optional[ivy.Array] = None, +): + if beta is not None: + x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) + else: + x_beta = x + amax = ivy.relu(x_beta) + res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) + res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) + res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( + x.dtype + ) * ivy.astype(1j, x.dtype) + if beta is not None: + res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) + if threshold is not None: + res = ivy.where( + ivy.real(x_beta) < threshold, + res, + x, + ).astype(x.dtype) + return res + + +def _wrap_between(y, a): + """Wrap y between [-a, a]""" + a = ivy.array(a, dtype=y.dtype) + a2 = ivy.array(2 * a, dtype=y.dtype) + zero = ivy.array(0, dtype=y.dtype) + rem = ivy.remainder(ivy.add(y, a), a2) + rem = ivy.where(rem < zero, rem + a2, rem) - a + return rem + + +# --- Main --- # +# ------------ # + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -99,26 +184,52 @@ def gelu( return current_backend(x).gelu(x, approximate=approximate, out=out) -gelu.jax_like = _gelu_jax_like +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def hardswish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply the hardswish activation function element-wise. + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. -def _leaky_relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original: Optional[Callable] = None, - alpha: float = 0.2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.astype(x * alpha, x.dtype), - x, - ) + Returns + ------- + ret + an array containing the hardswish activation of each element in ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 0., 4.]) + >>> y = ivy.hardswish(x) + >>> y + ivy.array([0., 0., 4.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) + >>> x = ivy.hardswish(x, out=x) + >>> x + { + a: ivy.array([-0., 4., 5.]), + b: ivy.array([0., 5.]) + } + """ + return current_backend(x).hardswish(x, out=out) @handle_exceptions @@ -199,9 +310,6 @@ def leaky_relu( return current_backend(x).leaky_relu(x, alpha=alpha, out=out) -leaky_relu.jax_like = _leaky_relu_jax_like - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -276,22 +384,60 @@ def log_softmax( return current_backend(x).log_softmax(x, axis=axis, out=out) -def _relu_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - out: Optional[ivy.Array] = None, +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def mish( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None ) -> ivy.Array: - return ivy.where( - ( - ivy.logical_or( - ivy.real(x) < 0, ivy.logical_and(ivy.real(x) == 0, ivy.imag(x) < 0) - ) - ), - ivy.array(0.0, dtype=x.dtype), - x, - ) + """ + Apply the mish activation function element-wise. + + Parameters + ---------- + x + input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the mish activation of each element in + ``x``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.mish(x) + >>> print(y) + ivy.array([-0.30340147, 0. , 0.86509842]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.mish(x, out = y) + >>> print(y) + ivy.array([ 1.40337825, 0.56114835, -0.20788449]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.mish(x) + >>> print(x) + { + a: ivy.array([0.86509842, -0.30883577]), + b: ivy.array([0.28903052, -0.10714479]) + } + """ + return current_backend(x).mish(x, out=out) @handle_exceptions @@ -363,9 +509,6 @@ def relu( return current_backend(x).relu(x, out=out) -relu.jax_like = _relu_jax_like - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -496,46 +639,6 @@ def softmax( return current_backend(x).softmax(x, axis=axis, out=out) -def _wrap_between(y, a): - """Wrap y between [-a, a]""" - a = ivy.array(a, dtype=y.dtype) - a2 = ivy.array(2 * a, dtype=y.dtype) - zero = ivy.array(0, dtype=y.dtype) - rem = ivy.remainder(ivy.add(y, a), a2) - rem = ivy.where(rem < zero, rem + a2, rem) - a - return rem - - -def _softplus_jax_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - fn_original=None, - beta: Optional[Union[int, float]] = None, - threshold: Optional[Union[int, float]] = None, - out: Optional[ivy.Array] = None, -): - if beta is not None: - x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype)) - else: - x_beta = x - amax = ivy.relu(x_beta) - res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype))) - res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res)))) - res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype( - x.dtype - ) * ivy.astype(1j, x.dtype) - if beta is not None: - res = ivy.divide(res, ivy.array(beta, dtype=x.dtype)) - if threshold is not None: - res = ivy.where( - ivy.real(x_beta) < threshold, - res, - x, - ).astype(x.dtype) - return res - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -607,108 +710,7 @@ def softplus( return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out) +gelu.jax_like = _gelu_jax_like +leaky_relu.jax_like = _leaky_relu_jax_like +relu.jax_like = _relu_jax_like softplus.jax_like = _softplus_jax_like - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def mish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply the mish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the mish activation of each element in - ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.mish(x) - >>> print(y) - ivy.array([-0.30340147, 0. , 0.86509842]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.mish(x, out = y) - >>> print(y) - ivy.array([ 1.40337825, 0.56114835, -0.20788449]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.mish(x) - >>> print(x) - { - a: ivy.array([0.86509842, -0.30883577]), - b: ivy.array([0.28903052, -0.10714479]) - } - """ - return current_backend(x).mish(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -def hardswish( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply the hardswish activation function element-wise. - - Parameters - ---------- - x - input array - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the hardswish activation of each element in ``x``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 0., 4.]) - >>> y = ivy.hardswish(x) - >>> y - ivy.array([0., 0., 4.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-3., 4., 5.]), b=ivy.array([0., 5.])) - >>> x = ivy.hardswish(x, out=x) - >>> x - { - a: ivy.array([-0., 4., 5.]), - b: ivy.array([0., 5.]) - } - """ - return current_backend(x).hardswish(x, out=out) diff --git a/ivy/functional/ivy/constants.py b/ivy/functional/ivy/constants.py index a830e3df30d66..461877e75b3b2 100644 --- a/ivy/functional/ivy/constants.py +++ b/ivy/functional/ivy/constants.py @@ -1,64 +1,55 @@ # global import math - +atto = 1e-18 +centi = 1e-2 +deci = 1e-1 +deka = 1e1 # Array API Standard # # -------------------# e = math.e -"""IEEE 754 floating-point representation of Euler's constant.""" - -pi = math.pi -"""IEEE 754 floating-point representation of the mathematical constant π.""" - -nan = math.nan -"""IEEE 754 floating-point representation of Not a Number (NaN).""" - -inf = math.inf -"""IEEE 754 floating-point representation of (positive) infinity.""" - -newaxis = None -"""An alias for None which is useful for indexing arrays.""" - - +exa = 1e18 +exbi = 2**60 +femto = 1e-15 +gibi = 2**30 +giga = 1e9 # Mathematical constants # # ------# golden = golden_ratio = (1 + math.sqrt(5)) / 2 -quetta = 1e30 -ronna = 1e27 -yotta = 1e24 -zetta = 1e21 -exa = 1e18 -peta = 1e15 -tera = 1e12 -giga = 1e9 -mega = 1e6 -kilo = 1e3 hecto = 1e2 -deka = 1e1 -deci = 1e-1 -centi = 1e-2 -milli = 1e-3 -micro = 1e-6 -nano = 1e-9 -pico = 1e-12 -femto = 1e-15 -atto = 1e-18 -zepto = 1e-21 -yocto = 1e-24 -ronto = 1e-27 -quecto = 1e-30 - - +inf = math.inf # Binary prefixes # # ------# kibi = 2**10 +kilo = 1e3 mebi = 2**20 -gibi = 2**30 -tebi = 2**40 +mega = 1e6 +micro = 1e-6 +milli = 1e-3 +nan = math.nan +nano = 1e-9 +newaxis = None pebi = 2**50 -exbi = 2**60 -zebi = 2**70 +peta = 1e15 +pi = math.pi +pico = 1e-12 +quecto = 1e-30 +quetta = 1e30 +ronna = 1e27 +ronto = 1e-27 +tebi = 2**40 +tera = 1e12 yobi = 2**80 +yocto = 1e-24 +yotta = 1e24 +zebi = 2**70 +zepto = 1e-21 +zetta = 1e21 +"""IEEE 754 floating-point representation of Euler's constant.""" +"""IEEE 754 floating-point representation of the mathematical constant π.""" +"""IEEE 754 floating-point representation of Not a Number (NaN).""" +"""IEEE 754 floating-point representation of (positive) infinity.""" +"""An alias for None which is useful for indexing arrays.""" diff --git a/ivy/functional/ivy/control_flow_ops.py b/ivy/functional/ivy/control_flow_ops.py index 4c2d78dc13c39..e41773bbcf366 100644 --- a/ivy/functional/ivy/control_flow_ops.py +++ b/ivy/functional/ivy/control_flow_ops.py @@ -9,6 +9,100 @@ ) +# --- Helpers --- # +# --------------- # + + +def _dict_to_tuple(d): + return tuple([d[k] for k in d]) + + +def _tuple_to_dict(t): + return {k: t[k] for k in range(len(t))} + + +# --- Main --- # +# ------------ # + + +def cast_bool(x): + return bool(x) + + +# todo (nightcrab) find a better place for these cmp functions + + +def cmp_is(left, right): + return left is right + + +def cmp_isnot(left, right): + return left is not right + + +def for_loop( + iterable: Iterable[Any], + body_fn: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + """ + Loops over an iterable, passing the current iteration along with a tuple of + variables into the provided body function. + + Parameters + ---------- + iterable + The iterable to loop over. + body_fn + A function to call each iteration, first taking the iterator value + and then a tuple of extra parameters. + vars + Extra parameters to be passed to body_fn. + + Returns + ------- + ret + The loop's return value (if any). + + Example + ---- + ``` + def body_fn(k, args): + print(k+1) + return args + + lst = [5,6] + + ivy.for_loop(lst, body_fn, ()) + >>> 5 + >>> 6 + ``` + """ + iterator = iterable.__iter__() + + vars_dict = _tuple_to_dict(vars) + + def test_fn(iterator, original_body, vars_dict): + try: + val = iterator.__next__() + except StopIteration: + return False + + vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) + + for k in range(len(vars_tuple)): + vars_dict[k] = vars_tuple[k] + + return True + + def empty_function(iterator, original_body, vars_dict): + return (iterator, original_body, vars_dict) + + packed_vars = (iterator, body_fn, vars_dict) + + return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) + + def if_else( cond: Callable, body_fn: Callable, @@ -70,6 +164,17 @@ def _if_else(cond, body_fn, orelse_fn, vars): return _if_else(cond, body_fn, orelse_fn, vars) +def try_except( + body1: Callable, + body2: Callable, + vars: Iterable[Union[ivy.Array, ivy.NativeArray]], +): + try: + return body1(*vars) + except Exception as e: + return body2(*vars, e) + + def while_loop( test_fn: Callable, body_fn: Callable, @@ -124,100 +229,3 @@ def _while_loop(test_fn, body_fn, vars): body_fn = to_ivy_arrays_and_back(body_fn) return _while_loop(test_fn, body_fn, vars) - - -def for_loop( - iterable: Iterable[Any], - body_fn: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - """ - Loops over an iterable, passing the current iteration along with a tuple of - variables into the provided body function. - - Parameters - ---------- - iterable - The iterable to loop over. - body_fn - A function to call each iteration, first taking the iterator value - and then a tuple of extra parameters. - vars - Extra parameters to be passed to body_fn. - - Returns - ------- - ret - The loop's return value (if any). - - Example - ---- - ``` - def body_fn(k, args): - print(k+1) - return args - - lst = [5,6] - - ivy.for_loop(lst, body_fn, ()) - >>> 5 - >>> 6 - ``` - """ - iterator = iterable.__iter__() - - vars_dict = _tuple_to_dict(vars) - - def test_fn(iterator, original_body, vars_dict): - try: - val = iterator.__next__() - except StopIteration: - return False - - vars_tuple = original_body(val, _dict_to_tuple(vars_dict)) - - for k in range(len(vars_tuple)): - vars_dict[k] = vars_tuple[k] - - return True - - def empty_function(iterator, original_body, vars_dict): - return (iterator, original_body, vars_dict) - - packed_vars = (iterator, body_fn, vars_dict) - - return _dict_to_tuple(while_loop(test_fn, empty_function, packed_vars)[2]) - - -def try_except( - body1: Callable, - body2: Callable, - vars: Iterable[Union[ivy.Array, ivy.NativeArray]], -): - try: - return body1(*vars) - except Exception as e: - return body2(*vars, e) - - -# todo (nightcrab) find a better place for these cmp functions - - -def cmp_is(left, right): - return left is right - - -def cmp_isnot(left, right): - return left is not right - - -def cast_bool(x): - return bool(x) - - -def _tuple_to_dict(t): - return {k: t[k] for k in range(len(t))} - - -def _dict_to_tuple(d): - return tuple([d[k] for k in d]) diff --git a/ivy/functional/ivy/creation.py b/ivy/functional/ivy/creation.py index 0273b8f4b0e55..f12f82585431e 100644 --- a/ivy/functional/ivy/creation.py +++ b/ivy/functional/ivy/creation.py @@ -35,44 +35,32 @@ handle_backend_invalid, ) -# Helpers # -# --------# +# Type hints # +# -----------# -def asarray_handle_nestable(fn: Callable) -> Callable: - fn_name = fn.__name__ +SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") +_T_co = TypeVar("_T_co", covariant=True) - @functools.wraps(fn) - def _asarray_handle_nestable(*args, **kwargs): - """ - Call `fn` with the *nestable* property of the function correctly handled. This - means mapping the function to the container leaves if any containers are passed - in the input. - Parameters - ---------- - args - The arguments to be passed to the function. +class NestedSequence(Protocol[_T_co]): + def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: + ... - kwargs - The keyword arguments to be passed to the function. + def __len__(self, /) -> int: + ... - Returns - ------- - The return of the function, with the nestable property handled correctly. - """ - # This decorator should only be applied to ivy.asarray, so we know where - # the container must be if there is one. - cont_fn = getattr(ivy.Container, "static_" + fn_name) - if isinstance(args[0], ivy.Container): - return cont_fn(*args, **kwargs) - # if the passed arguments does not contain a container, the function using - # the passed arguments, returning an ivy or a native array. - return fn(*args, **kwargs) +# --- Helpers --- # +# --------------- # - _asarray_handle_nestable.handle_nestable = True - return _asarray_handle_nestable + +def _flatten_nest(xs): + for x in xs: + if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): + yield from _flatten_nest(x) + else: + yield x def _ivy_to_native(x): @@ -92,6 +80,15 @@ def _ivy_to_native(x): return x +def _remove_np_bfloat16(obj): + # unlike other frameworks, torch and paddle do not support creating tensors + # from numpy arrays that have bfloat16 dtype using any extension because + # bfloat16 in not supported natively by numpy (as of version <=1.25) + if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": + return obj.tolist() + return obj + + def _shape_to_native(x): # checks the first element of the leaf list and # converts it to a native array if it is an ivy array @@ -109,168 +106,8 @@ def _shape_to_native(x): return x -def _flatten_nest(xs): - for x in xs: - if isinstance(x, Iterable) and not isinstance(x, (str, bytes)): - yield from _flatten_nest(x) - else: - yield x - - -def _remove_np_bfloat16(obj): - # unlike other frameworks, torch and paddle do not support creating tensors - # from numpy arrays that have bfloat16 dtype using any extension because - # bfloat16 in not supported natively by numpy (as of version <=1.25) - if isinstance(obj, np.ndarray) and obj.dtype.name == "bfloat16": - return obj.tolist() - return obj - - -def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): - """ - Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances - and return arrays are all converted to `ivy.Array` instances. - - This wrapper is specifically for the backend implementations of - asarray. - - It assumes either all the elements in a leaf list are ivy arrays - or none of them are. It checks the first element of all the leaf - list. If it is an ivy array, it converts all the elements in the - leaf list to native otherwise it skips that leaf list. - """ - new_arg = _ivy_to_native(args[0]) - new_args = (new_arg,) + args[1:] - if dtype is not None: - dtype = ivy.default_dtype(dtype=dtype, as_native=True) - return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) - - return _asarray_to_native_arrays_and_back - - -def asarray_infer_dtype(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_dtype(*args, dtype=None, **kwargs): - """ - Determine the correct `dtype`, and then calls the function with the `dtype` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. - - Parameters - ---------- - args - The arguments to be passed to the function. - - dtype - The dtype for the function. - - kwargs - The keyword arguments to be passed to the function. - - Returns - ------- - The return of the function, with `dtype` passed explicitly. - """ - - def _infer_dtype(obj): - if isinstance(obj, ivy.NativeShape): - obj = list(obj) - if hasattr(obj, "dtype"): - return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype - else: - return ivy.default_dtype(item=obj) - - if not ivy.exists(dtype): - arr = args[0] - # get default dtypes for all elements - dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] - # flatten the nested structure - dtype_list = _flatten_nest(dtype_list) - # keep unique dtypes - dtype_list = list(set(dtype_list)) - if len(dtype_list) != 0: # handle the case of empty input - # promote all dtypes to a single dtype - dtype = dtype_list[0] - # we disable precise mode to avoid wider than necessary casting - # that might result from the mixing of int32 and float32 - with ivy.PreciseMode(False): - for dt in dtype_list[1:]: - dtype = ivy.promote_types(dtype, dt) - else: - dtype = ivy.default_float_dtype() - dtype = ivy.as_native_dtype(dtype) - # call the function with dtype provided explicitly - return fn(*args, dtype=dtype, **kwargs) - - _asarray_infer_dtype.infer_dtype = True - return _asarray_infer_dtype - - -def asarray_infer_device(fn: Callable) -> Callable: - @functools.wraps(fn) - def _asarray_infer_device(*args, device=None, **kwargs): - """ - Determine the correct `device`, and then calls the function with the `device` - passed explicitly. This wrapper is specifically for the backend implementations - of asarray. - - Parameters - ---------- - args - The arguments to be passed to the function. - - device - The device for the function. - - kwargs - The keyword arguments to be passed to the function. - - Returns - ------- - The return of the function, with `device` passed explicitly. - """ - if isinstance(args[0], list): - return fn( - *args, device=ivy.default_device(device, as_native=True), **kwargs - ) - - # find the first array argument, if required - arr = None if ivy.exists(device) else args[0] - # infer the correct device - device = ivy.default_device(device, item=arr, as_native=True) - # call the function with device provided explicitly - return fn(*args, device=device, **kwargs) - - _asarray_infer_device.infer_device = True - return _asarray_infer_device - - -def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: - @functools.wraps(fn) - def _inputs_to_native_shapes(*args, **kwargs): - new_arg = _shape_to_native(args[0]) - new_args = (new_arg,) + args[1:] - return fn(*new_args, **kwargs) - - _inputs_to_native_shapes.inputs_to_native_shapes = True - return _inputs_to_native_shapes - - -# Type hints # -# -----------# - -SupportsBufferProtocol = TypeVar("SupportsBufferProtocol") -_T_co = TypeVar("_T_co", covariant=True) - - -class NestedSequence(Protocol[_T_co]): - def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]: - ... - - def __len__(self, /) -> int: - ... +# --- Main --- # +# ------------ # # Array API Standard # @@ -472,434 +309,199 @@ def asarray( ) -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def zeros( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with zeros. +def asarray_handle_nestable(fn: Callable) -> Callable: + fn_name = fn.__name__ - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type must - be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + @functools.wraps(fn) + def _asarray_handle_nestable(*args, **kwargs): + """ + Call `fn` with the *nestable* property of the function correctly handled. This + means mapping the function to the container leaves if any containers are passed + in the input. - Returns - ------- - ret - an array containing zeros. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.NativeShape` input: - >>> shape = (3, 5) - >>> x = ivy.zeros(shape) - >>> print(x) - ivy.array([[0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.], - [0., 0., 0., 0., 0.]]) - - >>> x = ivy.zeros(5) - >>> print(x) - ivy.array([0., 0., 0., 0., 0.]) - """ - return current_backend().zeros(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@outputs_to_ivy_arrays -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones( - shape: Union[ivy.Shape, ivy.NativeShape], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array having a specified ``shape`` and filled with ones. - - .. note:: - - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal to - zero (i.e., ``1 + 0j``). - - Parameters - ---------- - shape - output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be the default floating-point data type. Default ``None``. - device - device on which to place the created array. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing ones. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Shape` input: - - >>> shape = (2,2) - >>> x = ivy.ones(shape) - >>> print(x) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Dtype` input: - - >>> shape = (3,2) - >>> d_type = ivy.int64 - >>> y = ivy.ones(shape, dtype=d_type) - >>> print(y) - ivy.array([[1, 1], - [1, 1], - [1, 1]]) - - With :class:`ivy.Device` input: - - >>> shape = (3,2) - >>> y = ivy.ones(shape, device="cpu") - >>> print(y) - ivy.array([[1., 1.], - [1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input: - - >>> shape = (1, 5, 2) - >>> x = ivy.zeros(shape) - >>> ivy.ones(shape, out=x) - >>> print(x) - ivy.array([[[1., 1.], - [1., 1.], - [1., 1.], - [1., 1.], - [1., 1.]]]) - """ - return current_backend().ones(shape, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def full_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - fill_value: Number, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ``fill_value`` and having the same ``shape`` as an - input array ``x`` . - - Parameters - ---------- - x - input array from which to derive the output array shape. - fill_value - Scalar fill value - dtype - output array data type. If ``dtype`` is `None`, the output array data type must - be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array having the same shape as ``x`` and where every element is equal to - ``fill_value``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :code:`int` datatype: - - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> fill_value = 1 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) - - >>> fill_value = 0.000123 - >>> x = ivy.ones(5) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) - - With float datatype: + Parameters + ---------- + args + The arguments to be passed to the function. - >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + kwargs + The keyword arguments to be passed to the function. - With :class:`ivy.NativeArray` input: + Returns + ------- + The return of the function, with the nestable property handled correctly. + """ + # This decorator should only be applied to ivy.asarray, so we know where + # the container must be if there is one. + cont_fn = getattr(ivy.Container, "static_" + fn_name) + if isinstance(args[0], ivy.Container): + return cont_fn(*args, **kwargs) - >>> x = ivy.native_array([3.0, 8.0]) - >>> fill_value = 0.000123 - >>> y = ivy.full_like(x,fill_value) - >>> print(y) - ivy.array([0.000123, 0.000123]) + # if the passed arguments does not contain a container, the function using + # the passed arguments, returning an ivy or a native array. + return fn(*args, **kwargs) - >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - ivy.array([[0.000123, 0.000123, 0.000123], - [0.000123, 0.000123, 0.000123]]) + _asarray_handle_nestable.handle_nestable = True + return _asarray_handle_nestable - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), - ... b=ivy.array([4.123, 5.23, 6.23])) - >>> fill_value = 15.0 - >>> y = ivy.full_like(x, fill_value) - >>> print(y) - { - a: ivy.array([15., 15., 15.]), - b: ivy.array([15., 15., 15.]) - } - """ - return current_backend(x).full_like( - x, fill_value, dtype=dtype, device=device, out=out - ) +def asarray_infer_device(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_device(*args, device=None, **kwargs): + """ + Determine the correct `device`, and then calls the function with the `device` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. + Parameters + ---------- + args + The arguments to be passed to the function. -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@infer_dtype -@handle_device_shifting -@infer_device -def ones_like( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array filled with ones and having the same shape as an input array - ``x``. + device + The device for the function. - .. note:: + kwargs + The keyword arguments to be passed to the function. - An output array having a complex floating-point data type must contain complex - numbers having a real component equal to one and an imaginary component equal - to zero (i.e., ``1 + 0j``). + Returns + ------- + The return of the function, with `device` passed explicitly. + """ + if isinstance(args[0], list): + return fn( + *args, device=ivy.default_device(device, as_native=True), **kwargs + ) - Parameters - ---------- - x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default ``None``. - device - device on which to place the created array. If device is ``None``, the output - array device must be inferred from ``x``. Default: ``None``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + # find the first array argument, if required + arr = None if ivy.exists(device) else args[0] + # infer the correct device + device = ivy.default_device(device, item=arr, as_native=True) + # call the function with device provided explicitly + return fn(*args, device=device, **kwargs) - Returns - ------- - ret - an array having the same shape as ``x`` and filled with ``ones``. + _asarray_infer_device.infer_device = True + return _asarray_infer_device - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. +def asarray_infer_dtype(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_infer_dtype(*args, dtype=None, **kwargs): + """ + Determine the correct `dtype`, and then calls the function with the `dtype` + passed explicitly. This wrapper is specifically for the backend implementations + of asarray. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Parameters + ---------- + args + The arguments to be passed to the function. - Functional Examples - ------------------- + dtype + The dtype for the function. - With :class:`ivy.Array` input: + kwargs + The keyword arguments to be passed to the function. - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) + Returns + ------- + The return of the function, with `dtype` passed explicitly. + """ - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1., 1., 1.], - [1., 1., 1.]]) + def _infer_dtype(obj): + if isinstance(obj, ivy.NativeShape): + obj = list(obj) + if hasattr(obj, "dtype"): + return obj.dtype.name if isinstance(obj, np.ndarray) else obj.dtype + else: + return ivy.default_dtype(item=obj) - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.zeros(3) - >>> ivy.ones_like(x, out=y) - >>> print(y) - ivy.array([1., 1., 1.]) + if not ivy.exists(dtype): + arr = args[0] + # get default dtypes for all elements + dtype_list = [ivy.nested_map(arr, lambda x: _infer_dtype(x), shallow=False)] + # flatten the nested structure + dtype_list = _flatten_nest(dtype_list) + # keep unique dtypes + dtype_list = list(set(dtype_list)) + if len(dtype_list) != 0: # handle the case of empty input + # promote all dtypes to a single dtype + dtype = dtype_list[0] + # we disable precise mode to avoid wider than necessary casting + # that might result from the mixing of int32 and float32 + with ivy.PreciseMode(False): + for dt in dtype_list[1:]: + dtype = ivy.promote_types(dtype, dt) + else: + dtype = ivy.default_float_dtype() + dtype = ivy.as_native_dtype(dtype) + # call the function with dtype provided explicitly + return fn(*args, dtype=dtype, **kwargs) - With :class:`ivy.NativeArray` input: + _asarray_infer_dtype.infer_dtype = True + return _asarray_infer_dtype - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.ones_like(x) - >>> print(y) - ivy.array([[1, 1, 1], - [1, 1, 1]]) - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([1, 1, 1, 1, 1, 1]) +def asarray_inputs_to_native_shapes(fn: Callable) -> Callable: + @functools.wraps(fn) + def _inputs_to_native_shapes(*args, **kwargs): + new_arg = _shape_to_native(args[0]) + new_args = (new_arg,) + args[1:] + return fn(*new_args, **kwargs) - With :class:`ivy.Container` input: + _inputs_to_native_shapes.inputs_to_native_shapes = True + return _inputs_to_native_shapes - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.ones_like(x) - >>> print(y) - { - a: ivy.array([1, 1, 1]), - b: ivy.array([1, 1, 1]) - } - With :class:`ivy.Array` input: +def asarray_to_native_arrays_and_back(fn: Callable) -> Callable: + @functools.wraps(fn) + def _asarray_to_native_arrays_and_back(*args, dtype=None, **kwargs): + """ + Wrap `fn` so that input arrays are all converted to `ivy.NativeArray` instances + and return arrays are all converted to `ivy.Array` instances. - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.ones_like() - >>> print(y) - ivy.array([1, 1, 1, 1, 1]) + This wrapper is specifically for the backend implementations of + asarray. - With :class:'ivy.Container' input: + It assumes either all the elements in a leaf list are ivy arrays + or none of them are. It checks the first element of all the leaf + list. If it is an ivy array, it converts all the elements in the + leaf list to native otherwise it skips that leaf list. + """ + new_arg = _ivy_to_native(args[0]) + new_args = (new_arg,) + args[1:] + if dtype is not None: + dtype = ivy.default_dtype(dtype=dtype, as_native=True) + return to_ivy(fn(*new_args, dtype=dtype, **kwargs)) - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.ones_like() - >>> print(y) - { - a: ivy.array([1., 1.]), - b: ivy.array([1., 1.]) - } - """ - return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) + return _asarray_to_native_arrays_and_back @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function -@infer_dtype @handle_device_shifting -@infer_device -def zeros_like( +def copy_array( x: Union[ivy.Array, ivy.NativeArray], /, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + to_ivy_array: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array filled with zeros and having the same ``shape`` as an input array - ``x``. + Copy an array. Parameters ---------- x - input array from which to derive the output array shape. - dtype - output array data type. If ``dtype`` is ``None``, the output array data type - must be inferred from ``x``. Default: ``None``. - device - device on which to place the created array. If ``device`` is ``None``, the - output array device must be inferred from ``x``. Default: ``None``. + array, input array containing elements to copy. + to_ivy_array + boolean, if True the returned array will be an ivy.Array object otherwise + returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., + depending on the backend), defaults to True. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -907,195 +509,83 @@ def zeros_like( Returns ------- ret - an array having the same shape as ``x`` and filled with ``zeros``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: + a copy of the input array ``x``. - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> y = ivy.zeros_like(x) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) + Examples + -------- + With one :class:`ivy.Array` input: - >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) - >>> y = ivy.zeros_like(x) + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.copy_array(x) >>> print(y) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) + ivy.array([-1, 0, 1]) - >>> x = ivy.array([3., 2., 1.]) - >>> y = ivy.ones(3) - >>> ivy.zeros_like(x, out=y) + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = ivy.copy_array(x) >>> print(y) - ivy.array([0., 0., 0.]) - - With :class:`ivy.NativeArray` input: + ivy.array([1, 0, 1, 1]) - >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) - >>> y = ivy.zeros_like(x) + >>> x = ivy.array([1, 0, 1, -1]) + >>> y = ivy.zeros((1, 4)) + >>> ivy.copy_array(x, out=y) >>> print(y) - ivy.array([[0, 0, 0],[0, 0, 0]]) - + ivy.array([1, 0, 1, -1]) - >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) - >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) - >>> print(y) - ivy.array([0, 0, 0, 0, 0, 0]) + >>> x = ivy.array([1, 0, 1, 1]) + >>> ivy.copy_array(x, out=x) + >>> print(x) + ivy.array([1, 0, 1, 1]) - With :class:`ivy.Container` input: + With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) - >>> y = ivy.zeros_like(x) + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.copy_array(x) >>> print(y) { - a: ivy.array([0, 0, 0]), - b: ivy.array([0, 0, 0]) + a: ivy.array([-1, 0, 1]) } - - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3, 8, 2, 1]) - >>> y = x.zeros_like() - >>> print(y) - ivy.array([0, 0, 0, 0, 0]) - - With :class:'ivy.Container' input: - - >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) - >>> y = x.zeros_like() + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.copy_array(x) >>> print(y) { - a: ivy.array([0., 0.]), - b: ivy.array([0., 0.]) + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) } - """ - return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def tril( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the lower triangular part of a matrix (or a stack of matrices) ``x``. - - .. note:: - - The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` - on the interval ``[0, min(M, N) - 1]``. - - Parameters - ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. - k - diagonal above which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the lower triangular part(s). The returned array must have - the same shape and data type as x. All elements above the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. + With one :class:`ivy.Container` static method: - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - return current_backend(x).tril(x, k=k, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def triu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the upper triangular part of a matrix (or a stack of matrices) ``x``. - - .. note:: - - The upper triangular part of the matrix is defined as the elements - on and above the specified diagonal ``k``. + >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) + >>> y = ivy.Container.static_copy_array(x) + >>> print(y) + { + a: ivy.array([-1, 0, 1]), + b: ivy.array([-1, 0, 1, 1, 1, 0]) + } - Parameters - ---------- - x - input array having shape (..., M, N) and whose innermost two dimensions form MxN - matrices. *, - k - diagonal below which to zero elements. If k = 0, the diagonal is the main - diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the - diagonal is above the main diagonal. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + With one :class:`ivy.Array` instance method: - Returns - ------- - ret - an array containing the upper triangular part(s). The returned array must have - the same shape and data type as x. All elements below the specified diagonal k - must be zeroed. The returned array should be allocated on the same device as x. + >>> x = ivy.array([-1, 0, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([-1, 0, 1]) + >>> x = ivy.array([1, 0, 1, 1]) + >>> y = x.copy_array() + >>> print(y) + ivy.array([1, 0, 1, 1]) - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + With :class:`ivy.Container` instance method: - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) + >>> y = x.copy_array() + >>> print(y) + { + a: ivy.array([1, 0, 1]), + b: ivy.array([-1, 0, 1, 1]) + } """ - return current_backend(x).triu(x, k=k, out=out) + return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) @handle_backend_invalid @@ -1352,43 +842,18 @@ def eye( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@infer_dtype @handle_device_shifting -@infer_device -def linspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], - /, - num: int, - *, - axis: Optional[int] = None, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - out: Optional[ivy.Array] = None, +def from_dlpack( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None ) -> ivy.Array: """ - Generate a certain number of evenly-spaced values in an interval along a given axis. - - See :math:`arange` that allows to specify the step size of evenly spaced values in - an interval. + Return a new array containing the data from another (array) object with a + ``__dlpack__`` method. Parameters ---------- - start - First entry in the range. - stop - Final entry in the range. - num - Number of values to generate. - axis - Axis along which the operation is performed. - endpoint - If True, stop is the last sample. Otherwise, it is not included. - dtype - output array data type. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + x object + input (array) object. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1396,186 +861,87 @@ def linspace( Returns ------- ret - Tensor of evenly-spaced values. + an array containing the data in `x`. + + .. admonition:: Note + :class: note + + The returned array may be either a copy or a view. See + :ref:`data-interchange` for details. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.from_dlpack.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - - Functional Examples - ------------------- - - With float input: - - >>> x = ivy.linspace(1, 2, 3) - >>> print(x) - ivy.array([1. , 1.5, 2. ]) - - >>> x = ivy.linspace(1, 2, 4, endpoint=False) - >>> print(x) - ivy.array([1., 1.25, 1.5 , 1.75]) - - >>> x = ivy.linspace(1, 10, 4, dtype="int32") - >>> print(x) - ivy.array([ 1, 4, 7, 10]) - - >>> x = ivy.linspace(1, 2, 4, device= "cpu") - >>> print(x) - ivy.array([1., 1.33333337, 1.66666663, 2.]) - - >>> y = ivy.array([0,0,0,0]) - >>> ivy.linspace(1, 2, 4, out= y) - >>> print(y) - ivy.array([1, 1, 1, 2]) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2]) - >>> y = ivy.array([4,5]) - >>> z = ivy.linspace(x, y, 4, axis = 0) - >>> print(z) - ivy.array([[1, 2], - [2, 3], - [3, 4], - [4, 5]]) """ - return current_backend(start).linspace( - start, - stop, - num, - axis=axis, - endpoint=endpoint, - dtype=dtype, - device=device, - out=out, - ) + return current_backend(x).from_dlpack(x, out=out) -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def meshgrid( - *arrays: Union[ivy.Array, ivy.NativeArray], - sparse: bool = False, - indexing: str = "xy", - out: Optional[ivy.Array] = None, -) -> List[ivy.Array]: - """ - Return coordinate matrices from coordinate vectors. - - Parameters - ---------- - arrays - an arbitrary number of one-dimensional arrays representing grid coordinates. - Each array should have the same numeric data type. - sparse - if True, a sparse grid is returned in order to conserve memory. - Default: ``False``. - indexing - Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or - one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, - respectively), the ``indexing`` keyword has no effect and should be ignored. - Default: ``'xy'``. - - Returns - ------- - ret - list of N arrays, where ``N`` is the number of provided one-dimensional input - arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional - arrays having lengths ``Ni = len(xi)``, - - - if matrix indexing ``ij``, then each returned array must have the shape - ``(N1, N2, N3, ..., Nn)``. - - if Cartesian indexing ``xy``, then each returned array must have shape - ``(N2, N1, N3, ..., Nn)``. - - Accordingly, for the two-dimensional case with input one-dimensional arrays of - length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must - have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned - array must have shape ``(N, M)``. - - Similarly, for the three-dimensional case with input one-dimensional arrays of - length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned - array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then - each returned array must have shape ``(N, M, P)``. - - Each returned array should have the same data type as the input arrays. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of - the `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) - - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) +@outputs_to_ivy_arrays +def frombuffer( + buffer: bytes, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + count: Optional[int] = -1, + offset: Optional[int] = 0, +) -> ivy.Array: + r""" + Interpret a buffer as a 1-dimensional array. - >>> x = ivy.array([1, 2, 5]) - >>> y = ivy.array([4, 1]) - >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') - >>> print(xv) - ivy.array([[1, 1], - [2, 2], - [5, 5]]) + .. note:: + Note that either of the following must be true: + 1. count is a positive non-zero number, and the total number of bytes + in the buffer is equal or greater than offset plus count times the size + (in bytes) of dtype. + 2. count is negative, and the length (number of bytes) of the buffer + subtracted by the offset is a multiple of the size (in bytes) of dtype. - >>> print(yv) - ivy.array([[4, 1], - [4, 1], - [4, 1]]) + Parameters + ---------- + buffer + An object that exposes the buffer interface. + dtype + Data-type of the returned array; default: float. + count + Number of items to read. -1 means all data in the buffer. + offset + Start reading the buffer from this offset (in bytes); default: 0. - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> xv, yv = ivy.meshgrid(x, y, sparse=True) - >>> print(xv) - ivy.array([[1, 2, 3]]) + Returns + ------- + out + 1-dimensional array. - >>> print(yv) - ivy.array([[4], [5], [6]]) + Examples + -------- + With :class:`bytes` inputs: - With :class:`ivy.NativeArray` input: + >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' + >>> y = ivy.frombuffer(x, dtype=ivy.float64) + >>> print(y) + ivy.array([1., 2.]) - >>> x = ivy.native_array([1, 2]) - >>> y = ivy.native_array([3, 4]) - >>> xv, yv = ivy.meshgrid(x, y) - >>> print(xv) - ivy.array([[1, 2], - [1, 2]]) + >>> x = b'\x01\x02\x03\x04' + >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) + >>> print(y) + ivy.array([2, 3, 4]) - >>> print(yv) - ivy.array([[3, 3], - [4, 4]]) + >>> x = b'\x00<\x00@\x00B\x00D\x00E' + >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) + >>> print(y) + ivy.array([2., 3., 4., 5.]) """ - return current_backend().meshgrid( - *arrays, sparse=sparse, indexing=indexing, out=out + return current_backend().frombuffer( + buffer, + dtype=dtype, + count=count, + offset=offset, ) @@ -1686,8 +1052,230 @@ def full( [[False, False]]]) } """ - return current_backend().full( - shape, fill_value, dtype=dtype, device=device, out=out + return current_backend().full( + shape, fill_value, dtype=dtype, device=device, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def full_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + fill_value: Number, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ``fill_value`` and having the same ``shape`` as an + input array ``x`` . + + Parameters + ---------- + x + input array from which to derive the output array shape. + fill_value + Scalar fill value + dtype + output array data type. If ``dtype`` is `None`, the output array data type must + be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and where every element is equal to + ``fill_value``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :code:`int` datatype: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> fill_value = 1 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> fill_value = 0.000123 + >>> x = ivy.ones(5) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With float datatype: + + >>> x = ivy.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123, 0.000123, 0.000123, 0.000123, 0.000123]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([3.0, 8.0]) + >>> fill_value = 0.000123 + >>> y = ivy.full_like(x,fill_value) + >>> print(y) + ivy.array([0.000123, 0.000123]) + + >>> x = ivy.native_array([[3., 8., 2.], [2., 8., 3.]]) + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + ivy.array([[0.000123, 0.000123, 0.000123], + [0.000123, 0.000123, 0.000123]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.2, 2.2324, 3.234]), + ... b=ivy.array([4.123, 5.23, 6.23])) + >>> fill_value = 15.0 + >>> y = ivy.full_like(x, fill_value) + >>> print(y) + { + a: ivy.array([15., 15., 15.]), + b: ivy.array([15., 15., 15.]) + } + """ + return current_backend(x).full_like( + x, fill_value, dtype=dtype, device=device, out=out + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def linspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], + /, + num: int, + *, + axis: Optional[int] = None, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a certain number of evenly-spaced values in an interval along a given axis. + + See :math:`arange` that allows to specify the step size of evenly spaced values in + an interval. + + Parameters + ---------- + start + First entry in the range. + stop + Final entry in the range. + num + Number of values to generate. + axis + Axis along which the operation is performed. + endpoint + If True, stop is the last sample. Otherwise, it is not included. + dtype + output array data type. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Tensor of evenly-spaced values. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With float input: + + >>> x = ivy.linspace(1, 2, 3) + >>> print(x) + ivy.array([1. , 1.5, 2. ]) + + >>> x = ivy.linspace(1, 2, 4, endpoint=False) + >>> print(x) + ivy.array([1., 1.25, 1.5 , 1.75]) + + >>> x = ivy.linspace(1, 10, 4, dtype="int32") + >>> print(x) + ivy.array([ 1, 4, 7, 10]) + + >>> x = ivy.linspace(1, 2, 4, device= "cpu") + >>> print(x) + ivy.array([1., 1.33333337, 1.66666663, 2.]) + + >>> y = ivy.array([0,0,0,0]) + >>> ivy.linspace(1, 2, 4, out= y) + >>> print(y) + ivy.array([1, 1, 1, 2]) + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2]) + >>> y = ivy.array([4,5]) + >>> z = ivy.linspace(x, y, 4, axis = 0) + >>> print(z) + ivy.array([[1, 2], + [2, 3], + [3, 4], + [4, 5]]) + """ + return current_backend(start).linspace( + start, + stop, + num, + axis=axis, + endpoint=endpoint, + dtype=dtype, + device=device, + out=out, ) @@ -1697,163 +1285,237 @@ def full( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@handle_device_shifting -def from_dlpack( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +@infer_dtype +@infer_device +def logspace( + start: Union[ivy.Array, ivy.NativeArray, float], + stop: Union[ivy.Array, ivy.NativeArray, float], + /, + num: int, + *, + base: float = 10.0, + axis: int = 0, + endpoint: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array containing the data from another (array) object with a - ``__dlpack__`` method. + Generate a certain number of evenly-spaced values in log space, in an interval along + a given axis. Parameters ---------- - x object - input (array) object. + start + First value in the range in log space. base ** start is the starting value in + the sequence. Can be an array or a float. + stop + Last value in the range in log space. base ** stop is the final value in the + sequence. Can be an array or a float. + num + Number of values to generate. + base + The base of the log space. Default is 10.0 + axis + Axis along which the operation is performed. Relevant only if start or stop are + array-like. Default is 0. + endpoint + If True, stop is the last sample. Otherwise, it is not included. Default is + True. + dtype + The data type of the output tensor. If None, the dtype of on_value is used or if + that is None, the dtype of off_value is used, or if that is None, defaults to + float32. Default is None. + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is + None. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + inputs broadcast to. Default is None. Returns ------- ret - an array containing the data in `x`. + Tensor of evenly-spaced values in log space. - .. admonition:: Note - :class: note + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - The returned array may be either a copy or a view. See - :ref:`data-interchange` for details. + Functional Examples + ------------------- + With float input: + >>> print(ivy.logspace(1, 2, 4)) + ivy.array([ 10., 21.5443469, 46.41588834, 100.]) - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. + >>> print(ivy.logspace(1, 2, 4, endpoint=False)) + ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - return current_backend(x).from_dlpack(x, out=out) + >>> print(ivy.logspace(1, 2, 4, dtype= int)) + ivy.array([ 10., 10., 10., 100.]) + >>> out = ivy.array([0,0,0,0]) + >>> ivy.logspace(1, 2, 4, out = out) + >>> print(out) + ivy.array([ 10, 21, 46, 100]) -# Extra # -# ------# + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[1.e+01, 1.e+02], + [1.e+02, 1.e+03], + [1.e+03, 1.e+04], + [1.e+04, 1.e+05]) + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4, 5]) + >>> print(ivy.logspace(x, y, 4, axis = 1)) + ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], + [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) -array = asarray + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([4]) + >>> print(ivy.logspace(x, y, 4)) + ivy.array([[ 10., 100.], + [ 100., 100.], + [ 1000., 1000.], + [10000., 10000.]]) + """ + result = base ** linspace( + start, + stop, + num, + endpoint=endpoint, + axis=axis, + dtype=dtype, + device=device, + ) + if ivy.exists(out): + return ivy.inplace_update(out, result) + return result @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def copy_array( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - to_ivy_array: bool = True, +def meshgrid( + *arrays: Union[ivy.Array, ivy.NativeArray], + sparse: bool = False, + indexing: str = "xy", out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> List[ivy.Array]: """ - Copy an array. + Return coordinate matrices from coordinate vectors. Parameters ---------- - x - array, input array containing elements to copy. - to_ivy_array - boolean, if True the returned array will be an ivy.Array object otherwise - returns an ivy.NativeArray object (i.e. a torch.tensor, np.array, etc., - depending on the backend), defaults to True. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + arrays + an arbitrary number of one-dimensional arrays representing grid coordinates. + Each array should have the same numeric data type. + sparse + if True, a sparse grid is returned in order to conserve memory. + Default: ``False``. + indexing + Cartesian ``'xy'`` or matrix ``'ij'`` indexing of output. If provided zero or + one one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, + respectively), the ``indexing`` keyword has no effect and should be ignored. + Default: ``'xy'``. Returns ------- ret - a copy of the input array ``x``. + list of N arrays, where ``N`` is the number of provided one-dimensional input + arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional + arrays having lengths ``Ni = len(xi)``, - Examples - -------- - With one :class:`ivy.Array` input: + - if matrix indexing ``ij``, then each returned array must have the shape + ``(N1, N2, N3, ..., Nn)``. + - if Cartesian indexing ``xy``, then each returned array must have shape + ``(N2, N1, N3, ..., Nn)``. - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.copy_array(x) - >>> print(y) - ivy.array([-1, 0, 1]) + Accordingly, for the two-dimensional case with input one-dimensional arrays of + length ``M`` and ``N``, if matrix indexing ``ij``, then each returned array must + have shape ``(M, N)``, and, if Cartesian indexing ``xy``, then each returned + array must have shape ``(N, M)``. - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = ivy.copy_array(x) - >>> print(y) - ivy.array([1, 0, 1, 1]) + Similarly, for the three-dimensional case with input one-dimensional arrays of + length ``M``, ``N``, and ``P``, if matrix indexing ``ij``, then each returned + array must have shape ``(M, N, P)``, and, if Cartesian indexing ``xy``, then + each returned array must have shape ``(N, M, P)``. - >>> x = ivy.array([1, 0, 1, -1]) - >>> y = ivy.zeros((1, 4)) - >>> ivy.copy_array(x, out=y) - >>> print(y) - ivy.array([1, 0, 1, -1]) + Each returned array should have the same data type as the input arrays. - >>> x = ivy.array([1, 0, 1, 1]) - >>> ivy.copy_array(x, out=x) - >>> print(x) - ivy.array([1, 0, 1, 1]) - With one :class:`ivy.Container` input: + This function conforms to the `Array API Standard + `_. This docstring is an extension of + the `docstring `_ + in the standard. - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]) - } + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) - } + Functional Examples + ------------------- - With one :class:`ivy.Container` static method: + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2]) + >>> y = ivy.array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) + + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) + + >>> x = ivy.array([1, 2, 5]) + >>> y = ivy.array([4, 1]) + >>> xv, yv = ivy.meshgrid(x, y, indexing='ij') + >>> print(xv) + ivy.array([[1, 1], + [2, 2], + [5, 5]]) - >>> x = ivy.Container(a=ivy.array([-1, 0, 1]),b=ivy.array([-1, 0, 1, 1, 1, 0])) - >>> y = ivy.Container.static_copy_array(x) - >>> print(y) - { - a: ivy.array([-1, 0, 1]), - b: ivy.array([-1, 0, 1, 1, 1, 0]) - } + >>> print(yv) + ivy.array([[4, 1], + [4, 1], + [4, 1]]) - With one :class:`ivy.Array` instance method: + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> xv, yv = ivy.meshgrid(x, y, sparse=True) + >>> print(xv) + ivy.array([[1, 2, 3]]) - >>> x = ivy.array([-1, 0, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([-1, 0, 1]) + >>> print(yv) + ivy.array([[4], [5], [6]]) - >>> x = ivy.array([1, 0, 1, 1]) - >>> y = x.copy_array() - >>> print(y) - ivy.array([1, 0, 1, 1]) + With :class:`ivy.NativeArray` input: - With :class:`ivy.Container` instance method: + >>> x = ivy.native_array([1, 2]) + >>> y = ivy.native_array([3, 4]) + >>> xv, yv = ivy.meshgrid(x, y) + >>> print(xv) + ivy.array([[1, 2], + [1, 2]]) - >>> x = ivy.Container(a=ivy.array([1, 0, 1]),b=ivy.array([-1, 0, 1, 1])) - >>> y = x.copy_array() - >>> print(y) - { - a: ivy.array([1, 0, 1]), - b: ivy.array([-1, 0, 1, 1]) - } + >>> print(yv) + ivy.array([[3, 3], + [4, 4]]) """ - return current_backend(x).copy_array(x, to_ivy_array=to_ivy_array, out=out) + return current_backend().meshgrid( + *arrays, sparse=sparse, indexing=indexing, out=out + ) @handle_backend_invalid @@ -1961,69 +1623,290 @@ def one_hot( Returns ------- ret - Tensor of zeros with the same shape and type as a, unless dtype provided which - overrides. - - Examples - -------- - With :class:`ivy.Array` inputs: + Tensor of zeros with the same shape and type as a, unless dtype provided which + overrides. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([3, 1]) + >>> y = 5 + >>> z = x.one_hot(5) + >>> print(z) + ivy.array([[0., 0., 0., 1., 0.], + ... [0., 1., 0., 0., 0.]]) + + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, y) + ivy.array([[1., 0., 0., 0., 0.]]) + + >>> x = ivy.array([0]) + >>> y = 5 + >>> ivy.one_hot(x, 5, out=z) + ivy.array([[1., 0., 0., 0., 0.]]) + >>> print(z) + ivy.array([[1., 0., 0., 0., 0.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1, 2]), \ + b=ivy.array([3, 1]), c=ivy.array([2, 3])) + >>> y = 5 + >>> z = x.one_hot(y) + >>> print(z) + { + a: ivy.array([[0., 1., 0., 0., 0.], + [0., 0., 1., 0., 0.]]), + b: ivy.array([[0., 0., 0., 1., 0.], + [0., 1., 0., 0., 0.]]), + c: ivy.array([[0., 0., 1., 0., 0.], + [0., 0., 0., 1., 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([2]), \ + b=ivy.array([]), c=ivy.native_array([4])) + >>> y = 7 + >>> z = x.one_hot(y) + >>> print(z) + { + a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), + b: ivy.array([], shape=(0, 7)), + c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) + } + """ + return current_backend(indices).one_hot( + indices, + depth, + on_value=on_value, + off_value=off_value, + axis=axis, + dtype=dtype, + device=device, + out=out, + ) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with ones. + + .. note:: + + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal to + zero (i.e., ``1 + 0j``). + + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing ones. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Shape` input: + + >>> shape = (2,2) + >>> x = ivy.ones(shape) + >>> print(x) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Dtype` input: + + >>> shape = (3,2) + >>> d_type = ivy.int64 + >>> y = ivy.ones(shape, dtype=d_type) + >>> print(y) + ivy.array([[1, 1], + [1, 1], + [1, 1]]) + + With :class:`ivy.Device` input: + + >>> shape = (3,2) + >>> y = ivy.ones(shape, device="cpu") + >>> print(y) + ivy.array([[1., 1.], + [1., 1.], + [1., 1.]]) + + With :class:`ivy.Array` input: + + >>> shape = (1, 5, 2) + >>> x = ivy.zeros(shape) + >>> ivy.ones(shape, out=x) + >>> print(x) + ivy.array([[[1., 1.], + [1., 1.], + [1., 1.], + [1., 1.], + [1., 1.]]]) + """ + return current_backend().ones(shape, dtype=dtype, device=device, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def ones_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with ones and having the same shape as an input array + ``x``. + + .. note:: + + An output array having a complex floating-point data type must contain complex + numbers having a real component equal to one and an imaginary component equal + to zero (i.e., ``1 + 0j``). + + Parameters + ---------- + x + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default ``None``. + device + device on which to place the created array. If device is ``None``, the output + array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and filled with ``ones``. - >>> x = ivy.array([3, 1]) - >>> y = 5 - >>> z = x.one_hot(5) - >>> print(z) - ivy.array([[0., 0., 0., 1., 0.], - ... [0., 1., 0., 0., 0.]]) - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, y) - ivy.array([[1., 0., 0., 0., 0.]]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = ivy.array([0]) - >>> y = 5 - >>> ivy.one_hot(x, 5, out=z) - ivy.array([[1., 0., 0., 0., 0.]]) - >>> print(z) - ivy.array([[1., 0., 0., 0., 0.]]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) + + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1., 1., 1.], + [1., 1., 1.]]) + + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.zeros(3) + >>> ivy.ones_like(x, out=y) + >>> print(y) + ivy.array([1., 1., 1.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.ones_like(x) + >>> print(y) + ivy.array([[1, 1, 1], + [1, 1, 1]]) + + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.ones_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) + >>> print(y) + ivy.array([1, 1, 1, 1, 1, 1]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 2]), \ - b=ivy.array([3, 1]), c=ivy.array([2, 3])) - >>> y = 5 - >>> z = x.one_hot(y) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.ones_like(x) + >>> print(y) { - a: ivy.array([[0., 1., 0., 0., 0.], - [0., 0., 1., 0., 0.]]), - b: ivy.array([[0., 0., 0., 1., 0.], - [0., 1., 0., 0., 0.]]), - c: ivy.array([[0., 0., 1., 0., 0.], - [0., 0., 0., 1., 0.]]) + a: ivy.array([1, 1, 1]), + b: ivy.array([1, 1, 1]) } - >>> x = ivy.Container(a=ivy.array([2]), \ - b=ivy.array([]), c=ivy.native_array([4])) - >>> y = 7 - >>> z = x.one_hot(y) - >>> print(z) + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.ones_like() + >>> print(y) + ivy.array([1, 1, 1, 1, 1]) + + With :class:'ivy.Container' input: + + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.ones_like() + >>> print(y) { - a: ivy.array([[0., 0., 1., 0., 0., 0., 0.]]), - b: ivy.array([], shape=(0, 7)), - c: ivy.array([[0., 0., 0., 0., 1., 0., 0.]]) + a: ivy.array([1., 1.]), + b: ivy.array([1., 1.]) } """ - return current_backend(indices).one_hot( - indices, - depth, - on_value=on_value, - off_value=off_value, - axis=axis, - dtype=dtype, - device=device, - out=out, - ) + return current_backend(x).ones_like(x, dtype=dtype, device=device, out=out) @handle_backend_invalid @@ -2032,178 +1915,110 @@ def one_hot( @handle_out_argument @to_native_arrays_and_back @handle_array_function -@infer_dtype -@infer_device -def logspace( - start: Union[ivy.Array, ivy.NativeArray, float], - stop: Union[ivy.Array, ivy.NativeArray, float], +@handle_device_shifting +def tril( + x: Union[ivy.Array, ivy.NativeArray], /, - num: int, *, - base: float = 10.0, - axis: int = 0, - endpoint: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + k: int = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Generate a certain number of evenly-spaced values in log space, in an interval along - a given axis. + Return the lower triangular part of a matrix (or a stack of matrices) ``x``. + + .. note:: + + The main diagonal is defined as the set of indices ``{(i, i)}`` for ``i`` + on the interval ``[0, min(M, N) - 1]``. Parameters ---------- - start - First value in the range in log space. base ** start is the starting value in - the sequence. Can be an array or a float. - stop - Last value in the range in log space. base ** stop is the final value in the - sequence. Can be an array or a float. - num - Number of values to generate. - base - The base of the log space. Default is 10.0 - axis - Axis along which the operation is performed. Relevant only if start or stop are - array-like. Default is 0. - endpoint - If True, stop is the last sample. Otherwise, it is not included. Default is - True. - dtype - The data type of the output tensor. If None, the dtype of on_value is used or if - that is None, the dtype of off_value is used, or if that is None, defaults to - float32. Default is None. - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Default is - None. + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. + k + diagonal above which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. out optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. Default is None. + inputs broadcast to. Returns ------- ret - Tensor of evenly-spaced values in log space. + an array containing the lower triangular part(s). The returned array must have + the same shape and data type as x. All elements above the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - - Functional Examples - ------------------- - With float input: - - >>> print(ivy.logspace(1, 2, 4)) - ivy.array([ 10., 21.5443469, 46.41588834, 100.]) - - >>> print(ivy.logspace(1, 2, 4, endpoint=False)) - ivy.array([10., 17.7827941, 31.6227766, 56.23413252]) - - >>> print(ivy.logspace(1, 2, 4, dtype= int)) - ivy.array([ 10., 10., 10., 100.]) - - >>> out = ivy.array([0,0,0,0]) - >>> ivy.logspace(1, 2, 4, out = out) - >>> print(out) - ivy.array([ 10, 21, 46, 100]) - - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[1.e+01, 1.e+02], - [1.e+02, 1.e+03], - [1.e+03, 1.e+04], - [1.e+04, 1.e+05]) - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4, 5]) - >>> print(ivy.logspace(x, y, 4, axis = 1)) - ivy.array([[[1.e+01, 1.e+02, 1.e+03, 1.e+04], - [1.e+02, 1.e+03, 1.e+04, 1.e+05]]]) - - >>> x = ivy.array([1, 2]) - >>> y = ivy.array([4]) - >>> print(ivy.logspace(x, y, 4)) - ivy.array([[ 10., 100.], - [ 100., 100.], - [ 1000., 1000.], - [10000., 10000.]]) """ - result = base ** linspace( - start, - stop, - num, - endpoint=endpoint, - axis=axis, - dtype=dtype, - device=device, - ) - if ivy.exists(out): - return ivy.inplace_update(out, result) - return result + return current_backend(x).tril(x, k=k, out=out) +@handle_backend_invalid @handle_nestable -@outputs_to_ivy_arrays -def frombuffer( - buffer: bytes, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - count: Optional[int] = -1, - offset: Optional[int] = 0, +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def triu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Interpret a buffer as a 1-dimensional array. + """ + Return the upper triangular part of a matrix (or a stack of matrices) ``x``. .. note:: - Note that either of the following must be true: - 1. count is a positive non-zero number, and the total number of bytes - in the buffer is equal or greater than offset plus count times the size - (in bytes) of dtype. - 2. count is negative, and the length (number of bytes) of the buffer - subtracted by the offset is a multiple of the size (in bytes) of dtype. + + The upper triangular part of the matrix is defined as the elements + on and above the specified diagonal ``k``. Parameters ---------- - buffer - An object that exposes the buffer interface. - dtype - Data-type of the returned array; default: float. - count - Number of items to read. -1 means all data in the buffer. - offset - Start reading the buffer from this offset (in bytes); default: 0. + x + input array having shape (..., M, N) and whose innermost two dimensions form MxN + matrices. *, + k + diagonal below which to zero elements. If k = 0, the diagonal is the main + diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the + diagonal is above the main diagonal. Default: ``0``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - out - 1-dimensional array. - - Examples - -------- - With :class:`bytes` inputs: + ret + an array containing the upper triangular part(s). The returned array must have + the same shape and data type as x. All elements below the specified diagonal k + must be zeroed. The returned array should be allocated on the same device as x. - >>> x = b'\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x00@' - >>> y = ivy.frombuffer(x, dtype=ivy.float64) - >>> print(y) - ivy.array([1., 2.]) - >>> x = b'\x01\x02\x03\x04' - >>> y = ivy.frombuffer(x, dtype='int8', count=-2, offset=1) - >>> print(y) - ivy.array([2, 3, 4]) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x = b'\x00<\x00@\x00B\x00D\x00E' - >>> y = ivy.frombuffer(x, dtype='float16', count=4, offset=2) - >>> print(y) - ivy.array([2., 3., 4., 5.]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. """ - return current_backend().frombuffer( - buffer, - dtype=dtype, - count=count, - offset=offset, - ) + return current_backend(x).triu(x, k=k, out=out) @handle_exceptions @@ -2300,3 +2115,193 @@ def triu_indices( (ivy.array([0, 0, 0, 0, 1, 1, 1, 1]), ivy.array([0, 1, 2, 3, 0, 1, 2, 3])) """ return current_backend().triu_indices(n_rows, n_cols, k, device=device) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@outputs_to_ivy_arrays +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def zeros( + shape: Union[ivy.Shape, ivy.NativeShape], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array having a specified ``shape`` and filled with zeros. + + Parameters + ---------- + shape + output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type must + be the default floating-point data type. Default ``None``. + device + device on which to place the created array. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing zeros. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.NativeShape` input: + >>> shape = (3, 5) + >>> x = ivy.zeros(shape) + >>> print(x) + ivy.array([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + + >>> x = ivy.zeros(5) + >>> print(x) + ivy.array([0., 0., 0., 0., 0.]) + """ + return current_backend().zeros(shape, dtype=dtype, device=device, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@infer_dtype +@handle_device_shifting +@infer_device +def zeros_like( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array filled with zeros and having the same ``shape`` as an input array + ``x``. + + Parameters + ---------- + x + input array from which to derive the output array shape. + dtype + output array data type. If ``dtype`` is ``None``, the output array data type + must be inferred from ``x``. Default: ``None``. + device + device on which to place the created array. If ``device`` is ``None``, the + output array device must be inferred from ``x``. Default: ``None``. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array having the same shape as ``x`` and filled with ``zeros``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([0, 0, 0, 0, 0, 0]) + + >>> x = ivy.array([[0, 1, 2],[3, 4, 5]], dtype = ivy.float32) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) + + >>> x = ivy.array([3., 2., 1.]) + >>> y = ivy.ones(3) + >>> ivy.zeros_like(x, out=y) + >>> print(y) + ivy.array([0., 0., 0.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([[3, 8, 2],[2, 8, 3]]) + >>> y = ivy.zeros_like(x) + >>> print(y) + ivy.array([[0, 0, 0],[0, 0, 0]]) + + + >>> x = ivy.native_array([3, 8, 2, 0, 0, 2]) + >>> y = ivy.zeros_like(x, dtype=ivy.IntDtype('int32'), device=ivy.Device('cpu')) + >>> print(y) + ivy.array([0, 0, 0, 0, 0, 0]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([3, 2, 1]), b=ivy.array([8, 2, 3])) + >>> y = ivy.zeros_like(x) + >>> print(y) + { + a: ivy.array([0, 0, 0]), + b: ivy.array([0, 0, 0]) + } + + + With :class:`ivy.Array` input: + + >>> x = ivy.array([2, 3, 8, 2, 1]) + >>> y = x.zeros_like() + >>> print(y) + ivy.array([0, 0, 0, 0, 0]) + + With :class:'ivy.Container' input: + + >>> x = ivy.Container(a=ivy.array([3., 8.]), b=ivy.array([2., 2.])) + >>> y = x.zeros_like() + >>> print(y) + { + a: ivy.array([0., 0.]), + b: ivy.array([0., 0.]) + } + """ + return current_backend(x).zeros_like(x, dtype=dtype, device=device, out=out) + + +# Extra # +# ------# + + +array = asarray diff --git a/ivy/functional/ivy/data_type.py b/ivy/functional/ivy/data_type.py index 4a084d541e98c..b5114cfda5ca2 100644 --- a/ivy/functional/ivy/data_type.py +++ b/ivy/functional/ivy/data_type.py @@ -26,52 +26,188 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# --------# +# Array API Standard # +# -------------------# +Finfo = None +Iinfo = None +default_complex_dtype_stack = list() +# Extra # +# ------# -def _is_valid_dtypes_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): - fn_supported_dtypes = fn.supported_dtypes - fn_unsupported_dtypes = fn.unsupported_dtypes - if isinstance(fn_supported_dtypes, dict): - if isinstance(fn_unsupported_dtypes, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_dtypes - and backend_str in fn_unsupported_dtypes - ): - return False - else: - if isinstance(fn_unsupported_dtypes, tuple): - return False - return True +default_dtype_stack = list() +default_float_dtype_stack = list() +default_int_dtype_stack = list() +default_uint_dtype_stack = list() -def _handle_nestable_dtype_info(fn): - def _handle_nestable_dtype_info_wrapper(type): - if isinstance(type, ivy.Container): - type = type.cont_map(lambda x, kc: fn(x)) - type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) - type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) - return type - return fn(type) +class DefaultDtype: + """Ivy's DefaultDtype class.""" - return _handle_nestable_dtype_info_wrapper + def __init__(self, dtype: ivy.Dtype): + self._dtype = dtype + def __enter__(self): + set_default_dtype(self._dtype) + return self -# Unindent every line in the source such that -# class methods can be compiled as normal methods -def _lstrip_lines(source: str) -> str: - # Separate all lines - source = source.split("\n") - # Check amount of indent before first character - indent = len(source[0]) - len(source[0].lstrip()) - # Remove same spaces from all lines - for i in range(len(source)): - source[i] = source[i][indent:] - source = "\n".join(source) - return source + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultFloatDtype: + """Ivy's DefaultFloatDtype class.""" + + def __init__(self, float_dtype: ivy.Dtype): + self._float_dtype = float_dtype + + def __enter__(self): + set_default_float_dtype(self._float_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_float_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultIntDtype: + """Ivy's DefaultIntDtype class.""" + + def __init__(self, int_dtype: ivy.Dtype): + self._int_dtype = int_dtype + + def __enter__(self): + set_default_int_dtype(self._int_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_int_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultUintDtype: + """Ivy's DefaultUintDtype class.""" + + def __init__(self, uint_dtype: ivy.UintDtype): + self._uint_dtype = uint_dtype + + def __enter__(self): + set_default_uint_dtype(self._uint_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_uint_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +class DefaultComplexDtype: + """Ivy's DefaultComplexDtype class.""" + + def __init__(self, complex_dtype: ivy.Dtype): + self._complex_dtype = complex_dtype + + def __enter__(self): + set_default_complex_dtype(self._complex_dtype) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + unset_default_complex_dtype() + if self and (exc_type is not None): + print(exc_tb) + raise exc_val + return self + + +# --- Helpers --- # +# --------------- # + + +def _check_complex128(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "complex128" + elif isinstance(input, np.ndarray): + return str(input.dtype) == "complex128" + if hasattr(input, "real") and hasattr(input, "imag"): + return _check_float64(input.real) and _check_float64(input.imag) + return False + + +def _check_float64(input) -> bool: + if ivy.is_array(input): + return ivy.dtype(input) == "float64" + if math.isfinite(input): + m, e = math.frexp(input) + return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) + return False + + +# Get the list of dtypes supported by the function +# by default returns the supported dtypes +def _get_dtypes(fn, complement=True): + supported = set(ivy.valid_dtypes) + + # We only care about getting dtype info from the base function + # if we do need to at some point use dtype information from the parent function + # we can comment out the following condition + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") + if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: + if complement: + supported = set(ivy.all_dtypes).difference(supported) + return supported + + # Their values are formatted like either + # 1. fn.supported_dtypes = ("float16",) + # Could also have the "all" value for the framework + basic = [ + ("supported_dtypes", set.intersection, ivy.valid_dtypes), + ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), + ] + + # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. + typesets = { + "valid": ivy.valid_dtypes, + "numeric": ivy.valid_numeric_dtypes, + "float": ivy.valid_float_dtypes, + "integer": ivy.valid_int_dtypes, + "unsigned": ivy.valid_uint_dtypes, + "complex": ivy.valid_complex_dtypes, + } + + for key, merge_fn, base in basic: + if hasattr(fn, key): + dtypes = getattr(fn, key) + # only einops allowed to be a dictionary + if isinstance(dtypes, dict): + dtypes = dtypes.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(dtypes, tuple) + dtypes = list(dtypes) + typeset_list = [] + for i, dtype in reversed(list(enumerate(dtypes))): + if dtype in typesets: + typeset_list.extend(typesets[dtype]) + dtypes.pop(i) + dtypes = dtypes + typeset_list + supported = merge_fn(supported, set(dtypes)) + + if complement: + supported = set(ivy.all_dtypes).difference(supported) + + return tuple(supported) # Get the list of function used the function @@ -126,6 +262,50 @@ def _get_functions_from_string(func_names, module): return ret +def _handle_nestable_dtype_info(fn): + def _handle_nestable_dtype_info_wrapper(type): + if isinstance(type, ivy.Container): + type = type.cont_map(lambda x, kc: fn(x)) + type.__dict__["max"] = type.cont_map(lambda x, kc: x.max) + type.__dict__["min"] = type.cont_map(lambda x, kc: x.min) + return type + return fn(type) + + return _handle_nestable_dtype_info_wrapper + + +def _is_valid_dtypes_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_dtypes") and hasattr(fn, "unsupported_dtypes"): + fn_supported_dtypes = fn.supported_dtypes + fn_unsupported_dtypes = fn.unsupported_dtypes + if isinstance(fn_supported_dtypes, dict): + if isinstance(fn_unsupported_dtypes, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_dtypes + and backend_str in fn_unsupported_dtypes + ): + return False + else: + if isinstance(fn_unsupported_dtypes, tuple): + return False + return True + + +# Unindent every line in the source such that +# class methods can be compiled as normal methods +def _lstrip_lines(source: str) -> str: + # Separate all lines + source = source.split("\n") + # Check amount of indent before first character + indent = len(source[0]) - len(source[0].lstrip()) + # Remove same spaces from all lines + for i in range(len(source)): + source[i] = source[i][indent:] + source = "\n".join(source) + return source + + # Get dtypes/device of nested functions, used for unsupported and supported dtypes # IMPORTANT: a few caveats: # 1. The base functions must be defined in ivy or the same module @@ -181,67 +361,44 @@ def _nested_get(f, base_set, merge_fn, get_fn, wrapper=set): return out -# Get the list of dtypes supported by the function -# by default returns the supported dtypes -def _get_dtypes(fn, complement=True): - supported = set(ivy.valid_dtypes) +# --- Main --- # +# ------------ # - # We only care about getting dtype info from the base function - # if we do need to at some point use dtype information from the parent function - # we can comment out the following condition - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - has_unsupported_dtypes_attr = hasattr(fn, "unsupported_dtypes") - if not is_backend_fn and not is_frontend_fn and not has_unsupported_dtypes_attr: - if complement: - supported = set(ivy.all_dtypes).difference(supported) - return supported - # Their values are formatted like either - # 1. fn.supported_dtypes = ("float16",) - # Could also have the "all" value for the framework - basic = [ - ("supported_dtypes", set.intersection, ivy.valid_dtypes), - ("unsupported_dtypes", set.difference, ivy.invalid_dtypes), - ] - - # allow passing "integer" if all integer dtypes are supported/unsupported for e.g. - typesets = { - "valid": ivy.valid_dtypes, - "numeric": ivy.valid_numeric_dtypes, - "float": ivy.valid_float_dtypes, - "integer": ivy.valid_int_dtypes, - "unsigned": ivy.valid_uint_dtypes, - "complex": ivy.valid_complex_dtypes, - } +@handle_exceptions +def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: + """ + Convert native data type to string representation. - for key, merge_fn, base in basic: - if hasattr(fn, key): - dtypes = getattr(fn, key) - # only einops allowed to be a dictionary - if isinstance(dtypes, dict): - dtypes = dtypes.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(dtypes, tuple) - dtypes = list(dtypes) - typeset_list = [] - for i, dtype in reversed(list(enumerate(dtypes))): - if dtype in typesets: - typeset_list.extend(typesets[dtype]) - dtypes.pop(i) - dtypes = dtypes + typeset_list - supported = merge_fn(supported, set(dtypes)) + Parameters + ---------- + dtype_in + The data type to convert to string. - if complement: - supported = set(ivy.all_dtypes).difference(supported) + Returns + ------- + ret + data type string 'float32' + """ + return current_backend(None).as_ivy_dtype(dtype_in) - return tuple(supported) +@handle_exceptions +def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: + """ + Convert data type string representation to native data type. -# Array API Standard # -# -------------------# + Parameters + ---------- + dtype_in + The data type string to convert to native data type. -Finfo = None -Iinfo = None + Returns + ------- + ret + data type e.g. ivy.float32. + """ + return current_backend(None).as_native_dtype(dtype_in) @handle_exceptions @@ -560,475 +717,49 @@ def can_cast( >>> print(ivy.can_cast(x, ivy.float64)) True - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([[-1, -1, -1], - ... [1, 1, 1]], - ... dtype='int16') - >>> print(ivy.can_cast(x, 'uint8')) - False - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3, 4, 5])) - >>> print(ivy.can_cast(x, 'int64')) - { - a: False, - b: True - } - """ - if isinstance(from_, (ivy.Array, ivy.NativeArray)): - from_ = from_.dtype - try: - dtype = ivy.promote_types(from_, to) - return dtype == to - except KeyError: - return False - - -@handle_exceptions -@handle_backend_invalid -@inputs_to_native_arrays -@handle_device_shifting -def finfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Finfo: - """ - Machine limits for floating-point data types. - - Parameters - ---------- - type - the kind of floating-point data-type about which to get information. - - Returns - ------- - ret - an object having the followng attributes: - - - **bits**: *int* - - number of bits occupied by the floating-point data type. - - - **eps**: *float* - - difference between 1.0 and the next smallest representable floating-point - number larger than 1.0 according to the IEEE-754 standard. - - - **max**: *float* - - largest representable number. - - - **min**: *float* - - smallest representable number. - - - **smallest_normal**: *float* - - smallest positive floating-point number with full precision. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Dtype` input: - - >>> y = ivy.finfo(ivy.float32) - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - - With :code:`str` input: - - >>> y = ivy.finfo('float32') - >>> print(y) - finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) - >>> print(ivy.finfo(x)) - finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) - - >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) - >>> print(ivy.finfo(x)) - finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) - - With :class:`ivy.Container` input: - - >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), - ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) - >>> print(ivy.finfo(c)) - { - x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), - y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / - max=1.7976931348623157e+308, dtype=float64) - } - """ - return current_backend(None).finfo(type) - - -@handle_exceptions -@handle_backend_invalid -@inputs_to_native_arrays -@handle_device_shifting -def iinfo( - type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], - /, -) -> Iinfo: - """ - Machine limits for integer data types. - - Parameters - ---------- - type - the kind of integer data-type about which to get information. - - Returns - ------- - ret - a class with that encapsules the following attributes: - - - **bits**: *int* - - number of bits occupied by the type. - - - **max**: *int* - - largest representable number. - - - **min**: *int* - - smallest representable number. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Dtype` input: - - >>> ivy.iinfo(ivy.int32) - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :code:`str` input: - - >>> ivy.iinfo('int32') - iinfo(min=-2147483648, max=2147483647, dtype=int32) - - With :class:`ivy.Array` input: - - >>> x = ivy.array([13,21,34], dtype=ivy.int8) - >>> ivy.iinfo(x) - iinfo(min=-128, max=127, dtype=int8) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) - >>> ivy.iinfo(x) - iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) - - With :class:`ivy.Container` input: - - >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), - ... y=ivy.array([76,81,16], dtype=ivy.uint32)) - >>> ivy.iinfo(c) - { - x: iinfo(min=0, max=65535, dtype=uint16), - y: iinfo(min=0, max=4294967295, dtype=uint32) - } - """ - return current_backend(None).iinfo(type) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def result_type( - *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] -) -> ivy.Dtype: - """ - Return the dtype that results from applying the type promotion rules (see - :ref:`type-promotion`) to the arguments. - - .. note:: - If provided mixed dtypes (e.g., integer and floating-point), the returned dtype - will be implementation-specific. - - Parameters - ---------- - arrays_and_dtypes - an arbitrary number of input arrays and/or dtypes. - - Returns - ------- - ret - the dtype resulting from an operation involving the input arrays and dtypes. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.array([3., 4., 5.]) - >>> d = ivy.result_type(x, y) - >>> print(d) - float32 - - With :class:`ivy.Dtype` input: - - >>> d = ivy.result_type(ivy.uint8, ivy.uint64) - >>> print(d) - uint64 - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = x.a.dtype - >>> print(d) - int32 - - >>> x = ivy.Container(a = ivy.array([3, 4, 5])) - >>> d = ivy.result_type(x, ivy.float64) - >>> print(d) - { - a: float64 - } - """ - return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) - - -# Extra # -# ------# - -default_dtype_stack = list() -default_float_dtype_stack = list() -default_int_dtype_stack = list() -default_uint_dtype_stack = list() -default_complex_dtype_stack = list() - - -class DefaultDtype: - """Ivy's DefaultDtype class.""" - - def __init__(self, dtype: ivy.Dtype): - self._dtype = dtype - - def __enter__(self): - set_default_dtype(self._dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultFloatDtype: - """Ivy's DefaultFloatDtype class.""" - - def __init__(self, float_dtype: ivy.Dtype): - self._float_dtype = float_dtype - - def __enter__(self): - set_default_float_dtype(self._float_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_float_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultIntDtype: - """Ivy's DefaultIntDtype class.""" - - def __init__(self, int_dtype: ivy.Dtype): - self._int_dtype = int_dtype - - def __enter__(self): - set_default_int_dtype(self._int_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_int_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultUintDtype: - """Ivy's DefaultUintDtype class.""" - - def __init__(self, uint_dtype: ivy.UintDtype): - self._uint_dtype = uint_dtype - - def __enter__(self): - set_default_uint_dtype(self._uint_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_uint_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -class DefaultComplexDtype: - """Ivy's DefaultComplexDtype class.""" - - def __init__(self, complex_dtype: ivy.Dtype): - self._complex_dtype = complex_dtype - - def __enter__(self): - set_default_complex_dtype(self._complex_dtype) - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - unset_default_complex_dtype() - if self and (exc_type is not None): - print(exc_tb) - raise exc_val - return self - - -@handle_exceptions -def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: - """ - Get the number of bits used for representing the input data type. - - Parameters - ---------- - dtype_in - The data type to determine the number of bits for. - - Returns - ------- - ret - The number of bits used to represent the data type. - - Examples - -------- - With :class:`ivy.Dtype` inputs: - - >>> x = ivy.dtype_bits(ivy.float32) - >>> print(x) - 32 - - >>> x = ivy.dtype_bits('int64') - >>> print(x) - 64 - - With :class:`ivy.NativeDtype` inputs: - - >>> x = ivy.dtype_bits(ivy.native_bool) - >>> print(x) - 1 - """ - return current_backend(dtype_in).dtype_bits(dtype_in) - - -@handle_exceptions -def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: - """ - Check if the given data type is hashable or not. - - Parameters - ---------- - dtype_in - The data type to check. - - Returns - ------- - ret - True if data type is hashable else False - """ - try: - hash(dtype_in) - return True - except TypeError: - return False - + With :class:`ivy.NativeArray` input: -@handle_exceptions -def as_ivy_dtype(dtype_in: Union[ivy.Dtype, str], /) -> ivy.Dtype: - """ - Convert native data type to string representation. + >>> x = ivy.native_array([[-1, -1, -1], + ... [1, 1, 1]], + ... dtype='int16') + >>> print(ivy.can_cast(x, 'uint8')) + False - Parameters - ---------- - dtype_in - The data type to convert to string. + With :class:`ivy.Container` input: - Returns - ------- - ret - data type string 'float32' + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3, 4, 5])) + >>> print(ivy.can_cast(x, 'int64')) + { + a: False, + b: True + } """ - return current_backend(None).as_ivy_dtype(dtype_in) + if isinstance(from_, (ivy.Array, ivy.NativeArray)): + from_ = from_.dtype + try: + dtype = ivy.promote_types(from_, to) + return dtype == to + except KeyError: + return False @handle_exceptions -def as_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> ivy.NativeDtype: +def check_float(x: Any) -> bool: """ - Convert data type string representation to native data type. + Check if the input is a float or a float-like object. Parameters ---------- - dtype_in - The data type string to convert to native data type. + x + Input to check. Returns ------- ret - data type e.g. ivy.float32. + "True" if the input is a float or a float-like object, otherwise "False". """ - return current_backend(None).as_native_dtype(dtype_in) - - -def _check_float64(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "float64" - if math.isfinite(input): - m, e = math.frexp(input) - return (abs(input) > 3.4028235e38) or (e < -126) or (e > 128) - return False - - -def _check_complex128(input) -> bool: - if ivy.is_array(input): - return ivy.dtype(input) == "complex128" - elif isinstance(input, np.ndarray): - return str(input.dtype) == "complex128" - if hasattr(input, "real") and hasattr(input, "imag"): - return _check_float64(input.real) and _check_float64(input.imag) - return False + return isinstance(x, (int, float)) and x is not bool @handle_exceptions @@ -1075,52 +806,53 @@ def closest_valid_dtype(type: Union[ivy.Dtype, str, None], /) -> Union[ivy.Dtype @handle_exceptions @handle_nestable @inputs_to_ivy_arrays -def default_float_dtype( +@handle_device_shifting +def default_complex_dtype( *, input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - float_dtype: Optional[Union[ivy.FloatDtype, ivy.NativeDtype]] = None, + complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, as_native: bool = False, ) -> Union[ivy.Dtype, str, ivy.NativeDtype]: """ Parameters ---------- input - Number or array for inferring the float dtype. - float_dtype - The float dtype to be returned. + Number or array for inferring the complex dtype. + complex_dtype + The complex dtype to be returned. as_native - Whether to return the float dtype as native dtype. + Whether to return the complex dtype as native dtype. Returns ------- - Return ``float_dtype`` as native or ivy dtype if provided, else - if ``input`` is given, return its float dtype, otherwise return the - global default float dtype. + Return ``complex_dtype`` as native or ivy dtype if provided, else + if ``input`` is given, return its complex dtype, otherwise return the + global default complex dtype. Examples -------- - >>> ivy.default_float_dtype() - 'float32' + >>> ivy.default_complex_dtype() + 'complex64' - >>> ivy.set_default_float_dtype(ivy.FloatDtype("float64")) - >>> ivy.default_float_dtype() - 'float64' + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' - >>> ivy.default_float_dtype(float_dtype=ivy.FloatDtype("float16")) - 'float16' + >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) + 'complex128' - >>> ivy.default_float_dtype(input=4294.967346) - 'float32' + >>> ivy.default_complex_dtype(input=4294.967346) + 'complex64' - >>> x = ivy.array([9.8,8.9], dtype="float16") - >>> ivy.default_float_dtype(input=x) - 'float16' + >>> x = ivy.array([9.8,8.9], dtype="complex128") + >>> ivy.default_complex_dtype(input=x) + 'complex128' """ - global default_float_dtype_stack - if ivy.exists(float_dtype): + global default_complex_dtype_stack + if ivy.exists(complex_dtype): if as_native is True: - return ivy.as_native_dtype(float_dtype) - return ivy.FloatDtype(ivy.as_ivy_dtype(float_dtype)) + return ivy.as_native_dtype(complex_dtype) + return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) as_native = ivy.default(as_native, False) if ivy.exists(input): if ivy.is_array(input): @@ -1129,90 +861,42 @@ def default_float_dtype( ret = str(input.dtype) elif isinstance(input, (list, tuple, dict)): if ivy.nested_argwhere( - input, lambda x: _check_float64(x), stop_after_n_found=1 + input, lambda x: _check_complex128(x), stop_after_n_found=1 ): - ret = ivy.float64 + ret = ivy.complex128 else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] elif isinstance(input, Number): - if _check_float64(input): - ret = ivy.float64 + if _check_complex128(input): + ret = ivy.complex128 else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] else: - if not default_float_dtype_stack: + if not default_complex_dtype_stack: def_dtype = default_dtype() - if ivy.is_float_dtype(def_dtype): + if ivy.is_complex_dtype(def_dtype): ret = def_dtype else: - ret = "float32" + ret = "complex64" else: - ret = default_float_dtype_stack[-1] + ret = default_complex_dtype_stack[-1] if as_native: return ivy.as_native_dtype(ret) - return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) - - -@handle_exceptions -def infer_default_dtype( - dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: - """ - Summary. - - Parameters - ---------- - dtype - - as_native - (Default value = False) - - Returns - ------- - Return the default data type for the “kind” (integer or floating-point) of dtype - - Examples - -------- - >>> ivy.set_default_int_dtype("int32") - >>> ivy.infer_default_dtype("int8") - 'int8' - - >>> ivy.set_default_float_dtype("float64") - >>> ivy.infer_default_dtype("float32") - 'float64' - - >>> ivy.set_default_uint_dtype("uint32") - >>> x = ivy.array([0], dtype="uint64") - >>> ivy.infer_default_dtype(x.dtype) - 'uint32' - """ - if ivy.is_complex_dtype(dtype): - default_dtype = ivy.default_complex_dtype(as_native=as_native) - elif ivy.is_float_dtype(dtype): - default_dtype = ivy.default_float_dtype(as_native=as_native) - elif ivy.is_uint_dtype(dtype): - default_dtype = ivy.default_uint_dtype(as_native=as_native) - elif ivy.is_int_dtype(dtype): - default_dtype = ivy.default_int_dtype(as_native=as_native) - elif as_native: - default_dtype = ivy.as_native_dtype("bool") - else: - default_dtype = ivy.as_ivy_dtype("bool") - return default_dtype + return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) @handle_exceptions @@ -1293,19 +977,114 @@ def default_dtype( elif as_native: return ivy.as_native_dtype("bool") else: - return "bool" - global default_dtype_stack - if not default_dtype_stack: - global default_float_dtype_stack - if default_float_dtype_stack: + return "bool" + global default_dtype_stack + if not default_dtype_stack: + global default_float_dtype_stack + if default_float_dtype_stack: + ret = default_float_dtype_stack[-1] + else: + ret = "float32" + else: + ret = default_dtype_stack[-1] + if as_native: + return ivy.as_native_dtype(ret) + return ivy.as_ivy_dtype(ret) + + +@handle_exceptions +@handle_nestable +@inputs_to_ivy_arrays +def default_float_dtype( + *, + input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + float_dtype: Optional[Union[ivy.FloatDtype, ivy.NativeDtype]] = None, + as_native: bool = False, +) -> Union[ivy.Dtype, str, ivy.NativeDtype]: + """ + Parameters + ---------- + input + Number or array for inferring the float dtype. + float_dtype + The float dtype to be returned. + as_native + Whether to return the float dtype as native dtype. + + Returns + ------- + Return ``float_dtype`` as native or ivy dtype if provided, else + if ``input`` is given, return its float dtype, otherwise return the + global default float dtype. + + Examples + -------- + >>> ivy.default_float_dtype() + 'float32' + + >>> ivy.set_default_float_dtype(ivy.FloatDtype("float64")) + >>> ivy.default_float_dtype() + 'float64' + + >>> ivy.default_float_dtype(float_dtype=ivy.FloatDtype("float16")) + 'float16' + + >>> ivy.default_float_dtype(input=4294.967346) + 'float32' + + >>> x = ivy.array([9.8,8.9], dtype="float16") + >>> ivy.default_float_dtype(input=x) + 'float16' + """ + global default_float_dtype_stack + if ivy.exists(float_dtype): + if as_native is True: + return ivy.as_native_dtype(float_dtype) + return ivy.FloatDtype(ivy.as_ivy_dtype(float_dtype)) + as_native = ivy.default(as_native, False) + if ivy.exists(input): + if ivy.is_array(input): + ret = ivy.dtype(input) + elif isinstance(input, np.ndarray): + ret = str(input.dtype) + elif isinstance(input, (list, tuple, dict)): + if ivy.nested_argwhere( + input, lambda x: _check_float64(x), stop_after_n_found=1 + ): + ret = ivy.float64 + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: + ret = default_float_dtype_stack[-1] + elif isinstance(input, Number): + if _check_float64(input): + ret = ivy.float64 + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: + ret = default_float_dtype_stack[-1] + else: + if not default_float_dtype_stack: + def_dtype = default_dtype() + if ivy.is_float_dtype(def_dtype): + ret = def_dtype + else: + ret = "float32" + else: ret = default_float_dtype_stack[-1] - else: - ret = "float32" - else: - ret = default_dtype_stack[-1] if as_native: return ivy.as_native_dtype(ret) - return ivy.as_ivy_dtype(ret) + return ivy.FloatDtype(ivy.as_ivy_dtype(ret)) @handle_exceptions @@ -1522,149 +1301,177 @@ def default_uint_dtype( @handle_exceptions +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_device_shifting -def default_complex_dtype( - *, - input: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - complex_dtype: Optional[Union[ivy.ComplexDtype, ivy.NativeDtype]] = None, - as_native: bool = False, -) -> Union[ivy.Dtype, str, ivy.NativeDtype]: +def dtype( + x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: """ + Get the data type for input array x. + Parameters ---------- - input - Number or array for inferring the complex dtype. - complex_dtype - The complex dtype to be returned. + x + Tensor for which to get the data type. as_native - Whether to return the complex dtype as native dtype. + Whether or not to return the dtype in string format. Default is ``False``. Returns ------- - Return ``complex_dtype`` as native or ivy dtype if provided, else - if ``input`` is given, return its complex dtype, otherwise return the - global default complex dtype. + ret + Data type of the array. Examples -------- - >>> ivy.default_complex_dtype() - 'complex64' + With :class:`ivy.Array` inputs: - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' + >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) + >>> y = ivy.dtype(x1) + >>> print(y) + float32 - >>> ivy.default_complex_dtype(complex_dtype=ivy.ComplexDtype("complex128")) - 'complex128' + With :class:`ivy.NativeArray` inputs: - >>> ivy.default_complex_dtype(input=4294.967346) - 'complex64' + >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) + >>> y = ivy.dtype(x1) + >>> print(y) + int32 - >>> x = ivy.array([9.8,8.9], dtype="complex128") - >>> ivy.default_complex_dtype(input=x) - 'complex128' + With :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), + ... b=ivy.native_array([1, 0, 0, 0, 1])) + >>> y = ivy.dtype(x.a) + >>> print(y) + float32 """ - global default_complex_dtype_stack - if ivy.exists(complex_dtype): - if as_native is True: - return ivy.as_native_dtype(complex_dtype) - return ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - as_native = ivy.default(as_native, False) - if ivy.exists(input): - if ivy.is_array(input): - ret = ivy.dtype(input) - elif isinstance(input, np.ndarray): - ret = str(input.dtype) - elif isinstance(input, (list, tuple, dict)): - if ivy.nested_argwhere( - input, lambda x: _check_complex128(x), stop_after_n_found=1 - ): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - elif isinstance(input, Number): - if _check_complex128(input): - ret = ivy.complex128 - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - else: - if not default_complex_dtype_stack: - def_dtype = default_dtype() - if ivy.is_complex_dtype(def_dtype): - ret = def_dtype - else: - ret = "complex64" - else: - ret = default_complex_dtype_stack[-1] - if as_native: - return ivy.as_native_dtype(ret) - return ivy.ComplexDtype(ivy.as_ivy_dtype(ret)) + return current_backend(x).dtype(x, as_native=as_native) + + +@handle_exceptions +def dtype_bits(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str], /) -> int: + """ + Get the number of bits used for representing the input data type. + + Parameters + ---------- + dtype_in + The data type to determine the number of bits for. + + Returns + ------- + ret + The number of bits used to represent the data type. + + Examples + -------- + With :class:`ivy.Dtype` inputs: + + >>> x = ivy.dtype_bits(ivy.float32) + >>> print(x) + 32 + + >>> x = ivy.dtype_bits('int64') + >>> print(x) + 64 + + With :class:`ivy.NativeDtype` inputs: + + >>> x = ivy.dtype_bits(ivy.native_bool) + >>> print(x) + 1 + """ + return current_backend(dtype_in).dtype_bits(dtype_in) @handle_exceptions @handle_backend_invalid -@handle_nestable @inputs_to_native_arrays @handle_device_shifting -def dtype( - x: Union[ivy.Array, ivy.NativeArray], *, as_native: bool = False -) -> Union[ivy.Dtype, ivy.NativeDtype]: +def finfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Finfo: """ - Get the data type for input array x. + Machine limits for floating-point data types. Parameters ---------- - x - Tensor for which to get the data type. - as_native - Whether or not to return the dtype in string format. Default is ``False``. + type + the kind of floating-point data-type about which to get information. Returns ------- ret - Data type of the array. + an object having the followng attributes: + + - **bits**: *int* + + number of bits occupied by the floating-point data type. + + - **eps**: *float* + + difference between 1.0 and the next smallest representable floating-point + number larger than 1.0 according to the IEEE-754 standard. + + - **max**: *float* + + largest representable number. + + - **min**: *float* + + smallest representable number. + + - **smallest_normal**: *float* + + smallest positive floating-point number with full precision. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Dtype` input: - >>> x1 = ivy.array([1.0, 2.0, 3.5, 4.5, 5, 6]) - >>> y = ivy.dtype(x1) + >>> y = ivy.finfo(ivy.float32) + >>> print(y) + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) + + With :code:`str` input: + + >>> y = ivy.finfo('float32') >>> print(y) - float32 + finfo(resolution=1e-06, min=-3.4028235e+38, max=3.4028235e+38, dtype=float32) - With :class:`ivy.NativeArray` inputs: + With :class:`ivy.Array` input: - >>> x1 = ivy.native_array([1, 0, 1, -1, 0]) - >>> y = ivy.dtype(x1) - >>> print(y) - int32 + >>> x = ivy.array([1.3,2.1,3.4], dtype=ivy.float64) + >>> print(ivy.finfo(x)) + finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) - With :class:`ivy.Container` inputs: + >>> x = ivy.array([0.7,8.4,3.14], dtype=ivy.float16) + >>> print(ivy.finfo(x)) + finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16) - >>> x = ivy.Container(a=ivy.native_array([1.0, 2.0, -1.0, 4.0, 1.0]), - ... b=ivy.native_array([1, 0, 0, 0, 1])) - >>> y = ivy.dtype(x.a) - >>> print(y) - float32 + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([-9.5,1.8,-8.9], dtype=ivy.float16), + ... y=ivy.array([7.6,8.1,1.6], dtype=ivy.float64)) + >>> print(ivy.finfo(c)) + { + x: finfo(resolution=0.001, min=-6.55040e+04, max=6.55040e+04, dtype=float16), + y: finfo(resolution=1e-15, min=-1.7976931348623157e+308, / + max=1.7976931348623157e+308, dtype=float64) + } """ - return current_backend(x).dtype(x, as_native=as_native) + return current_backend(None).finfo(type) @handle_exceptions @@ -1768,6 +1575,131 @@ def function_unsupported_dtypes( ) +@handle_exceptions +@handle_backend_invalid +@inputs_to_native_arrays +@handle_device_shifting +def iinfo( + type: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray], + /, +) -> Iinfo: + """ + Machine limits for integer data types. + + Parameters + ---------- + type + the kind of integer data-type about which to get information. + + Returns + ------- + ret + a class with that encapsules the following attributes: + + - **bits**: *int* + + number of bits occupied by the type. + + - **max**: *int* + + largest representable number. + + - **min**: *int* + + smallest representable number. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Examples + -------- + With :class:`ivy.Dtype` input: + + >>> ivy.iinfo(ivy.int32) + iinfo(min=-2147483648, max=2147483647, dtype=int32) + + With :code:`str` input: + + >>> ivy.iinfo('int32') + iinfo(min=-2147483648, max=2147483647, dtype=int32) + + With :class:`ivy.Array` input: + + >>> x = ivy.array([13,21,34], dtype=ivy.int8) + >>> ivy.iinfo(x) + iinfo(min=-128, max=127, dtype=int8) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([7,84,314], dtype=ivy.int64) + >>> ivy.iinfo(x) + iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64) + + With :class:`ivy.Container` input: + + >>> c = ivy.Container(x=ivy.array([0,1800,89], dtype=ivy.uint16), + ... y=ivy.array([76,81,16], dtype=ivy.uint32)) + >>> ivy.iinfo(c) + { + x: iinfo(min=0, max=65535, dtype=uint16), + y: iinfo(min=0, max=4294967295, dtype=uint32) + } + """ + return current_backend(None).iinfo(type) + + +@handle_exceptions +def infer_default_dtype( + dtype: Union[ivy.Dtype, ivy.NativeDtype, str], as_native: bool = False +) -> Union[ivy.Dtype, ivy.NativeDtype]: + """ + Summary. + + Parameters + ---------- + dtype + + as_native + (Default value = False) + + Returns + ------- + Return the default data type for the “kind” (integer or floating-point) of dtype + + Examples + -------- + >>> ivy.set_default_int_dtype("int32") + >>> ivy.infer_default_dtype("int8") + 'int8' + + >>> ivy.set_default_float_dtype("float64") + >>> ivy.infer_default_dtype("float32") + 'float64' + + >>> ivy.set_default_uint_dtype("uint32") + >>> x = ivy.array([0], dtype="uint64") + >>> ivy.infer_default_dtype(x.dtype) + 'uint32' + """ + if ivy.is_complex_dtype(dtype): + default_dtype = ivy.default_complex_dtype(as_native=as_native) + elif ivy.is_float_dtype(dtype): + default_dtype = ivy.default_float_dtype(as_native=as_native) + elif ivy.is_uint_dtype(dtype): + default_dtype = ivy.default_uint_dtype(as_native=as_native) + elif ivy.is_int_dtype(dtype): + default_dtype = ivy.default_int_dtype(as_native=as_native) + elif as_native: + default_dtype = ivy.as_native_dtype("bool") + else: + default_dtype = ivy.as_ivy_dtype("bool") + return default_dtype + + @handle_exceptions def invalid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -1822,30 +1754,151 @@ def is_bool_dtype( ret "True" if the input data type is a bool, otherwise "False". - Both the description and the type hints above assumes an array input for - simplicity but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Both the description and the type hints above assumes an array input for + simplicity but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. + """ + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, np.ndarray): + return "bool" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return ( + True + if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) + else False + ) + elif isinstance(dtype_in, (list, tuple, dict)): + return ( + True + if ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (bool, np.bool)) and x is not int, + ) + else False + ) + return "bool" in ivy.as_ivy_dtype(dtype_in) + + +@handle_exceptions +@handle_nestable +@inputs_to_ivy_arrays +def is_complex_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: + """ + Determine whether the input data type is a complex dtype. + + Parameters + ---------- + dtype_in + The array or data type to check + + Returns + ------- + ret + Whether or not the array or data type is of a complex dtype + + Examples + -------- + >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) + True + + >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) + True + + >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) + False + """ + if ivy.is_array(dtype_in): + dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() + elif isinstance(dtype_in, np.ndarray): + return "complex" in dtype_in.dtype.name + elif isinstance(dtype_in, Number): + return isinstance(dtype_in, (complex, np.complexfloating)) + elif isinstance(dtype_in, (list, tuple, dict)): + return ivy.nested_argwhere( + dtype_in, + lambda x: isinstance(x, (complex, np.complexfloating)) + or (ivy.is_array(x) and "complex" in ivy.dtype(x)), + ) + return "complex" in as_ivy_dtype(dtype_in) + + +@handle_exceptions +@handle_nestable +@inputs_to_native_arrays +def is_float_dtype( + dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], + /, +) -> bool: + """ + Determine whether the input data type is a float dtype. + + Parameters + ---------- + dtype_in + The array or data type to check + + Returns + ------- + ret + Whether or not the array or data type is of a floating point dtype + + Examples + -------- + >>> x = ivy.is_float_dtype(ivy.float32) + >>> print(x) + True + + >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) + >>> print(ivy.is_float_dtype(arr)) + True """ if ivy.is_array(dtype_in): dtype_in = ivy.dtype(dtype_in) + elif isinstance(dtype_in, ivy.Shape): + dtype_in = ivy.default_int_dtype() elif isinstance(dtype_in, np.ndarray): - return "bool" in dtype_in.dtype.name + return "float" in dtype_in.dtype.name elif isinstance(dtype_in, Number): - return ( - True - if isinstance(dtype_in, (bool, np.bool)) and not isinstance(dtype_in, bool) - else False - ) + return True if isinstance(dtype_in, (float, np.floating)) else False elif isinstance(dtype_in, (list, tuple, dict)): return ( True if ivy.nested_argwhere( dtype_in, - lambda x: isinstance(x, (bool, np.bool)) and x is not int, + lambda x: isinstance(x, (float, np.floating)) + or (ivy.is_array(x) and "float" in ivy.dtype(x)), ) else False ) - return "bool" in ivy.as_ivy_dtype(dtype_in) + return "float" in as_ivy_dtype(dtype_in) + + +@handle_exceptions +def is_hashable_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: + """ + Check if the given data type is hashable or not. + + Parameters + ---------- + dtype_in + The data type to check. + + Returns + ------- + ret + True if data type is hashable else False + """ + try: + hash(dtype_in) + return True + except TypeError: + return False @handle_exceptions @@ -1937,72 +1990,34 @@ def is_int_dtype( @handle_exceptions -def check_float(x: Any) -> bool: - """ - Check if the input is a float or a float-like object. - - Parameters - ---------- - x - Input to check. - - Returns - ------- - ret - "True" if the input is a float or a float-like object, otherwise "False". - """ - return isinstance(x, (int, float)) and x is not bool - - -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -def is_float_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], - /, -) -> bool: +def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: """ - Determine whether the input data type is a float dtype. + Determine whether the input dtype is a Native dtype. Parameters ---------- dtype_in - The array or data type to check + Determine whether the input data type is a native data type object. Returns ------- ret - Whether or not the array or data type is of a floating point dtype + Boolean, whether or not dtype_in is a native data type. Examples -------- - >>> x = ivy.is_float_dtype(ivy.float32) - >>> print(x) + >>> ivy.set_backend('numpy') + >>> ivy.is_native_dtype(np.int32) True - >>> arr = ivy.array([1.2, 3.2, 4.3], dtype=ivy.float32) - >>> print(ivy.is_float_dtype(arr)) - True + >>> ivy.set_backend('numpy') + >>> ivy.is_native_array(ivy.float64) + False """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "float" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return True if isinstance(dtype_in, (float, np.floating)) else False - elif isinstance(dtype_in, (list, tuple, dict)): - return ( - True - if ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (float, np.floating)) - or (ivy.is_array(x) and "float" in ivy.dtype(x)), - ) - else False - ) - return "float" in as_ivy_dtype(dtype_in) + try: + return current_backend(None).is_native_dtype(dtype_in) + except ValueError: + return False @handle_exceptions @@ -2054,101 +2069,203 @@ def is_uint_dtype( @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -def is_complex_dtype( - dtype_in: Union[ivy.Dtype, str, ivy.Array, ivy.NativeArray, Number], +def promote_types( + type1: Union[ivy.Dtype, ivy.NativeDtype], + type2: Union[ivy.Dtype, ivy.NativeDtype], /, -) -> bool: + *, + array_api_promotion: bool = False, +) -> ivy.Dtype: """ - Determine whether the input data type is a complex dtype. + Promote the datatypes type1 and type2, returning the data type they promote to. Parameters ---------- - dtype_in - The array or data type to check + type1 + the first of the two types to promote + type2 + the second of the two types to promote + array_api_promotion + whether to only use the array api promotion rules Returns ------- ret - Whether or not the array or data type is of a complex dtype + The type that both input types promote to + """ + query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] + query.sort(key=lambda x: str(x)) + query = tuple(query) + + def _promote(query): + if array_api_promotion: + return ivy.array_api_promotion_table[query] + return ivy.promotion_table[query] + + try: + ret = _promote(query) + except KeyError: + # try again with the dtypes swapped + query = (query[1], query[0]) + try: + ret = _promote(query) + except KeyError: + raise ivy.utils.exceptions.IvyDtypePromotionError( + "these dtypes ({} and {}) are not type promotable, ".format( + type1, type2 + ) + ) + return ret + + +@handle_exceptions +def promote_types_of_inputs( + x1: Union[ivy.NativeArray, Number, Iterable[Number]], + x2: Union[ivy.NativeArray, Number, Iterable[Number]], + /, + *, + array_api_promotion: bool = False, +) -> Tuple[ivy.NativeArray, ivy.NativeArray]: + """ + Promote the dtype of the given native array inputs to a common dtype based on type + promotion rules. + + While passing float or integer values or any other non-array input + to this function, it should be noted that the return will be an + array-like object. Therefore, outputs from this function should be + used as inputs only for those functions that expect an array-like or + tensor-like objects, otherwise it might give unexpected results. + """ + + def _special_case(a1, a2): + # check for float number and integer array case + return isinstance(a1, float) and "int" in str(a2.dtype) + + def _get_target_dtype(scalar, arr): + # identify a good dtype to give the scalar value, + # based on it's own type and that of the arr value + if _special_case(scalar, arr): + return "float64" + elif arr.dtype == bool and not isinstance(scalar, bool): + return None # let ivy infer a dtype + elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): + return "complex128" + else: + return arr.dtype + + if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): + device = ivy.default_device(item=x1, as_native=True) + x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) + elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): + device = ivy.default_device(item=x2, as_native=True) + x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) + elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): + x1 = ivy.asarray(x1) + x2 = ivy.asarray(x2) + + if x1.dtype != x2.dtype: + promoted = promote_types( + x1.dtype, x2.dtype, array_api_promotion=array_api_promotion + ) + x1 = ivy.astype(x1, promoted, copy=False) + x2 = ivy.astype(x2, promoted, copy=False) + + ivy.utils.assertions._check_jax_x64_flag(x1.dtype) + return ivy.to_native(x1), ivy.to_native(x2) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def result_type( + *arrays_and_dtypes: Union[ivy.Array, ivy.NativeArray, ivy.Dtype] +) -> ivy.Dtype: + """ + Return the dtype that results from applying the type promotion rules (see + :ref:`type-promotion`) to the arguments. + + .. note:: + If provided mixed dtypes (e.g., integer and floating-point), the returned dtype + will be implementation-specific. + + Parameters + ---------- + arrays_and_dtypes + an arbitrary number of input arrays and/or dtypes. + + Returns + ------- + ret + the dtype resulting from an operation involving the input arrays and dtypes. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.array([3., 4., 5.]) + >>> d = ivy.result_type(x, y) + >>> print(d) + float32 - Examples - -------- - >>> ivy.is_complex_dtype(ivy.ComplexDtype("complex64")) - True + With :class:`ivy.Dtype` input: - >>> ivy.is_complex_dtype(ivy.Dtype("complex128")) - True + >>> d = ivy.result_type(ivy.uint8, ivy.uint64) + >>> print(d) + uint64 - >>> ivy.is_complex_dtype(ivy.IntDtype("int64")) - False + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = x.a.dtype + >>> print(d) + int32 + + >>> x = ivy.Container(a = ivy.array([3, 4, 5])) + >>> d = ivy.result_type(x, ivy.float64) + >>> print(d) + { + a: float64 + } """ - if ivy.is_array(dtype_in): - dtype_in = ivy.dtype(dtype_in) - elif isinstance(dtype_in, ivy.Shape): - dtype_in = ivy.default_int_dtype() - elif isinstance(dtype_in, np.ndarray): - return "complex" in dtype_in.dtype.name - elif isinstance(dtype_in, Number): - return isinstance(dtype_in, (complex, np.complexfloating)) - elif isinstance(dtype_in, (list, tuple, dict)): - return ivy.nested_argwhere( - dtype_in, - lambda x: isinstance(x, (complex, np.complexfloating)) - or (ivy.is_array(x) and "complex" in ivy.dtype(x)), - ) - return "complex" in as_ivy_dtype(dtype_in) + return current_backend(arrays_and_dtypes[0]).result_type(*arrays_and_dtypes) @handle_exceptions -def promote_types( - type1: Union[ivy.Dtype, ivy.NativeDtype], - type2: Union[ivy.Dtype, ivy.NativeDtype], - /, - *, - array_api_promotion: bool = False, -) -> ivy.Dtype: +def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): """ - Promote the datatypes type1 and type2, returning the data type they promote to. + Set the 'complex_dtype' as the default data type. Parameters ---------- - type1 - the first of the two types to promote - type2 - the second of the two types to promote - array_api_promotion - whether to only use the array api promotion rules + complex_dtype + The complex data type to be set as the default. - Returns - ------- - ret - The type that both input types promote to - """ - query = [ivy.as_ivy_dtype(type1), ivy.as_ivy_dtype(type2)] - query.sort(key=lambda x: str(x)) - query = tuple(query) + Examples + -------- + With :class: `ivy.Dtype` input: - def _promote(query): - if array_api_promotion: - return ivy.array_api_promotion_table[query] - return ivy.promotion_table[query] + >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) + >>> ivy.default_complex_dtype() + 'complex64' - try: - ret = _promote(query) - except KeyError: - # try again with the dtypes swapped - query = (query[1], query[0]) - try: - ret = _promote(query) - except KeyError: - raise ivy.utils.exceptions.IvyDtypePromotionError( - "these dtypes ({} and {}) are not type promotable, ".format( - type1, type2 - ) - ) - return ret + >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) + >>> ivy.default_complex_dtype() + 'complex128' + """ + complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) + ivy.utils.assertions._check_jax_x64_flag(complex_dtype) + global default_complex_dtype_stack + default_complex_dtype_stack.append(complex_dtype) @handle_exceptions @@ -2269,34 +2386,6 @@ def set_default_uint_dtype(uint_dtype: Union[ivy.Dtype, str], /): default_uint_dtype_stack.append(uint_dtype) -@handle_exceptions -def set_default_complex_dtype(complex_dtype: Union[ivy.Dtype, str], /): - """ - Set the 'complex_dtype' as the default data type. - - Parameters - ---------- - complex_dtype - The complex data type to be set as the default. - - Examples - -------- - With :class: `ivy.Dtype` input: - - >>> ivy.set_default_complex_dtype(ivy.ComplexDtype("complex64")) - >>> ivy.default_complex_dtype() - 'complex64' - - >>> ivy.set_default_float_dtype(ivy.ComplexDtype("complex128")) - >>> ivy.default_complex_dtype() - 'complex128' - """ - complex_dtype = ivy.ComplexDtype(ivy.as_ivy_dtype(complex_dtype)) - ivy.utils.assertions._check_jax_x64_flag(complex_dtype) - global default_complex_dtype_stack - default_complex_dtype_stack.append(complex_dtype) - - @handle_exceptions def type_promote_arrays( x1: Union[ivy.Array, ivy.NativeArray], @@ -2323,6 +2412,27 @@ def type_promote_arrays( return ivy.astype(x1, new_type), ivy.astype(x2, new_type) +@handle_exceptions +def unset_default_complex_dtype(): + """ + Reset the current default complex dtype to the previous state. + + Examples + -------- + >>> ivy.set_default_complex_dtype(ivy.complex64) + >>> ivy.set_default_complex_dtype(ivy.complex128) + >>> ivy.default_complex_dtype_stack + ['complex64','complex128'] + + >>> ivy.unset_default_complex_dtype() + >>> ivy.default_complex_dtype_stack + ['complex64'] + """ + global default_complex_dtype_stack + if default_complex_dtype_stack: + default_complex_dtype_stack.pop(-1) + + @handle_exceptions def unset_default_dtype(): """ @@ -2413,27 +2523,6 @@ def unset_default_uint_dtype(): default_uint_dtype_stack.pop(-1) -@handle_exceptions -def unset_default_complex_dtype(): - """ - Reset the current default complex dtype to the previous state. - - Examples - -------- - >>> ivy.set_default_complex_dtype(ivy.complex64) - >>> ivy.set_default_complex_dtype(ivy.complex128) - >>> ivy.default_complex_dtype_stack - ['complex64','complex128'] - - >>> ivy.unset_default_complex_dtype() - >>> ivy.default_complex_dtype_stack - ['complex64'] - """ - global default_complex_dtype_stack - if default_complex_dtype_stack: - default_complex_dtype_stack.pop(-1) - - @handle_exceptions def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bool: """ @@ -2466,90 +2555,3 @@ def valid_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype, str, None], /) -> bo if dtype_in is None: return True return ivy.as_ivy_dtype(dtype_in) in ivy.valid_dtypes - - -@handle_exceptions -def promote_types_of_inputs( - x1: Union[ivy.NativeArray, Number, Iterable[Number]], - x2: Union[ivy.NativeArray, Number, Iterable[Number]], - /, - *, - array_api_promotion: bool = False, -) -> Tuple[ivy.NativeArray, ivy.NativeArray]: - """ - Promote the dtype of the given native array inputs to a common dtype based on type - promotion rules. - - While passing float or integer values or any other non-array input - to this function, it should be noted that the return will be an - array-like object. Therefore, outputs from this function should be - used as inputs only for those functions that expect an array-like or - tensor-like objects, otherwise it might give unexpected results. - """ - - def _special_case(a1, a2): - # check for float number and integer array case - return isinstance(a1, float) and "int" in str(a2.dtype) - - def _get_target_dtype(scalar, arr): - # identify a good dtype to give the scalar value, - # based on it's own type and that of the arr value - if _special_case(scalar, arr): - return "float64" - elif arr.dtype == bool and not isinstance(scalar, bool): - return None # let ivy infer a dtype - elif isinstance(scalar, complex) and not ivy.is_complex_dtype(arr): - return "complex128" - else: - return arr.dtype - - if hasattr(x1, "dtype") and not hasattr(x2, "dtype"): - device = ivy.default_device(item=x1, as_native=True) - x2 = ivy.asarray(x2, dtype=_get_target_dtype(x2, x1), device=device) - elif hasattr(x2, "dtype") and not hasattr(x1, "dtype"): - device = ivy.default_device(item=x2, as_native=True) - x1 = ivy.asarray(x1, dtype=_get_target_dtype(x1, x2), device=device) - elif not (hasattr(x1, "dtype") or hasattr(x2, "dtype")): - x1 = ivy.asarray(x1) - x2 = ivy.asarray(x2) - - if x1.dtype != x2.dtype: - promoted = promote_types( - x1.dtype, x2.dtype, array_api_promotion=array_api_promotion - ) - x1 = ivy.astype(x1, promoted, copy=False) - x2 = ivy.astype(x2, promoted, copy=False) - - ivy.utils.assertions._check_jax_x64_flag(x1.dtype) - return ivy.to_native(x1), ivy.to_native(x2) - - -@handle_exceptions -def is_native_dtype(dtype_in: Union[ivy.Dtype, ivy.NativeDtype], /) -> bool: - """ - Determine whether the input dtype is a Native dtype. - - Parameters - ---------- - dtype_in - Determine whether the input data type is a native data type object. - - Returns - ------- - ret - Boolean, whether or not dtype_in is a native data type. - - Examples - -------- - >>> ivy.set_backend('numpy') - >>> ivy.is_native_dtype(np.int32) - True - - >>> ivy.set_backend('numpy') - >>> ivy.is_native_array(ivy.float64) - False - """ - try: - return current_backend(None).is_native_dtype(dtype_in) - except ValueError: - return False diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py index 1fee8c85ac25b..058b9dd513163 100644 --- a/ivy/functional/ivy/device.py +++ b/ivy/functional/ivy/device.py @@ -10,6 +10,21 @@ import types from typing import Type, Optional, Tuple +# nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. + +from typing import Union, Callable, Iterable, Any + +# local +import ivy +from ivy.func_wrapper import ( + handle_out_argument, + to_native_arrays_and_back, + handle_nestable, + handle_array_like_without_promotion, + handle_backend_invalid, +) +from ivy.utils.exceptions import handle_exceptions + # noinspection PyUnresolvedReferences try: import pynvml @@ -24,26 +39,11 @@ " of the Ivy's device module will be limited. Please install pynvml if" " you wish to use GPUs with Ivy." ) - # nvidia-ml-py (pynvml) is not installed in CPU Dockerfile. - -from typing import Union, Callable, Iterable, Any - -# local -import ivy -from ivy.func_wrapper import ( - handle_out_argument, - to_native_arrays_and_back, - handle_nestable, - handle_array_like_without_promotion, - handle_backend_invalid, -) -from ivy.utils.exceptions import handle_exceptions default_device_stack = list() +max_chunk_sizes = dict() soft_device_mode_stack = list() -dev_handles = dict() split_factors = dict() -max_chunk_sizes = dict() # Extra # @@ -136,10 +136,87 @@ def __exit__( return self -def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): - return ivy.current_backend().handle_soft_device_variable( - *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs - ) +# Profiler # + + +class Profiler(abc.ABC): + """ + The profiler class is used to profile the execution of some code. + + Parameters + ---------- + save_dir + The directory to save the profile data to. + """ + + def __init__(self, save_dir: str): + self._save_dir = save_dir + + @abc.abstractmethod + def start(self): + """ + Start the profiler. + + This should be called before the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def stop(self): + """ + Stop the profiler. + + This should be called after the code to be profiled. + """ + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def __enter__(self): + raise ivy.utils.exceptions.IvyNotImplementedException + + @abc.abstractmethod + def __exit__(self, exc_type, exc_val, exc_tb): + raise ivy.utils.exceptions.IvyNotImplementedException + + +# --- Helpers --- # +# --------------- # + + +def _get_devices(fn: Callable, complement: bool = True) -> Tuple: + valid_devices = ivy.valid_devices + invalid_devices = ivy.invalid_devices + all_devices = ivy.all_devices + + supported = set(ivy.valid_devices) + + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + supported = set(all_devices).difference(supported) + return supported + + # Their values are formated like either + # 1. fn.supported_devices = ("cpu",) + # Could also have the "all" value for the framework + basic = [ + ("supported_devices", set.intersection, valid_devices), + ("unsupported_devices", set.difference, invalid_devices), + ] + for key, merge_fn, base in basic: + if hasattr(fn, key): + v = getattr(fn, key) + if "einops" in fn.__name__ and isinstance(v, dict): + v = v.get(ivy.current_backend_str(), base) + ivy.utils.assertions.check_isinstance(v, tuple) + supported = merge_fn(supported, set(v)) + + if complement: + supported = set(all_devices).difference(supported) + + return tuple(supported) # Helpers # @@ -155,6 +232,24 @@ def _get_nvml_gpu_handle(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: return handle +def _is_valid_devices_attributes(fn: Callable) -> bool: + if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): + fn_supported_devices = fn.supported_devices + fn_unsupported_devices = fn.unsupported_devices + if isinstance(fn_supported_devices, dict): + if isinstance(fn_unsupported_devices, dict): + backend_str = ivy.current_backend_str() + if ( + backend_str in fn_supported_devices + and backend_str in fn_unsupported_devices + ): + return False + else: + if isinstance(fn_unsupported_devices, tuple): + return False + return True + + def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kwargs): with ivy.ArrayMode(False): default_device = ivy.default_device(device_shifting_dev, as_native=True) @@ -169,176 +264,177 @@ def _shift_native_arrays_on_default_device(*args, device_shifting_dev=None, **kw return args, kwargs, default_device -# Device Queries # +# --- Main --- # +# ------------ # -# Array Printing + +# Conversions @handle_exceptions -def get_all_ivy_arrays_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, -) -> ivy.Container: +def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: """ - Get all ivy arrays which are currently alive on the specified device. + Convert device to string representation. Parameters ---------- device - The device handle from which to get the arrays + The device handle to convert to string. Returns ------- ret - Container with the arrays found for the specified device [identity, array] + Device string e.g. 'cuda:0'. Examples -------- - >>> x = ivy.array([1,0,2]) - >>> y = ivy.dev(x) - >>> z = ivy.get_all_ivy_arrays_on_dev(y) - >>> print(z) - {139740789224448:ivy.array([1,0,2])}, + >>> y = ivy.as_ivy_dev('cpu') + >>> print(y) + cpu """ - device = ivy.as_ivy_dev(device) - all_arrays = list() - for obj in gc.get_objects(): - if ( - obj is ivy.data_classes.array.array.Array - and ivy.is_ivy_array(obj) - and ivy.dev(obj) == device - ): - all_arrays.append(obj) - - return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) + return ivy.current_backend().as_ivy_dev(device) @handle_exceptions -def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: +def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: """ - Return the number of arrays which are currently alive on the specified device. + Convert device string representation to native device type. Parameters ---------- device - The device handle from which to count the arrays + The device string to convert to native device handle. + A native device handle can be passed in instead - in this case + the unmodified parameter is returned. Returns ------- ret - Number of arrays on the specified device + Native device handle. Examples -------- - >>> x1 = ivy.array([-1, 0, 5.2]) - >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 2 + With :class:`ivy.Device` input: - >>> x1 = ivy.native_array([-1, 0, 5.2]) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 0 + >>> ivy.set_backend("numpy") + >>> ivy.as_native_dev("cpu") + 'cpu' - >>> x = ivy.Container(x1=ivy.array([-1]), - ... x2=ivy.native_array([-1])) - >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) - >>> print(y) - 1 + >>> ivy.set_backend("tensorflow") + >>> ivy.as_native_dev("tpu:3") + '/TPU:3' + + With :class:`ivy.NativeDevice` input: + + >>> import torch + >>> device = torch.device("cuda") + >>> device + device(type='cuda') + + >>> ivy.as_native_dev(device) + device(type='cuda') """ - return len(ivy.get_all_ivy_arrays_on_dev(device)) + return ivy.current_backend().as_native_dev(device) + + +# Memory @handle_exceptions -@handle_nestable -def print_all_ivy_arrays_on_dev( - *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - attr_only: bool = True, -) -> None: +def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: """ - Print the shape and dtype for all ivy arrays which are currently alive on the - specified device. + Clear memory cache on target device. Parameters ---------- device - The device on which to print the arrays - - attr_only - Whether or not to only print the `shape` and `dtype` attributes of the array + The device string to convert to native device handle or native device handle. Examples -------- - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y) - ((3,), 'int32') - ((3,), 'int32') - - - >>> x = ivy.array([[1,0,2], [3,2,1]]) - >>> y = ivy.dev(x) - >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) - [1,0,2] - [3,2,1] + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.clear_cached_mem_on_dev(device) """ - arrs = ivy.get_all_ivy_arrays_on_dev(device).values() - if attr_only: - [print((arr.shape, arr.dtype)) for arr in arrs] - else: - [print(arr) for arr in arrs] + ivy.current_backend().clear_cached_mem_on_dev(device) -ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False +# Default Device # +# noinspection PyShadowingNames @handle_exceptions -def set_soft_device_mode(mode: bool) -> None: - """ - Set the mode of whether to move input arrays to `ivy.default_device()` before - performing an operation. - - Parameter - --------- - mode - boolean whether to move input arrays - Examples - -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode - False - >>> ivy.set_soft_device_mode(True) - >>> ivy.soft_device_mode - True +def default_device( + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + /, + *, + item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, + as_native: bool = None, +) -> Union[ivy.Device, ivy.NativeDevice]: """ - global soft_device_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - soft_device_mode_stack.append(mode) - ivy.__setattr__("soft_device_mode", mode, True) + Return the input device or the default device. If the as_native flag is set, the + device will be converted to a native device. If the item is provided, the item's + device is returned. If the device is not provided, the last default device is + returned. If a default device has not been set, the first gpu is returned if + available, otherwise the cpu is returned. + Parameters + ---------- + device + The device to be returned or converted. + item + The item to get the device from. + as_native + Whether to convert the device to a native device. -@handle_exceptions -def unset_soft_device_mode() -> None: - """ - Reset the mode of moving input arrays to `ivy.default_device()` before performing an - operation. + Returns + ------- + ret + Device handle or string. Examples -------- - >>> ivy.set_soft_device_mode(False) - >>> ivy.soft_device_mode - False - >>> ivy.unset_soft_device_mode() - >>> ivy.soft_device_mode - True + >>> ivy.default_device() + device(type='cpu') + + >>> ivy.default_device("gpu:0") + 'gpu:0' + + >>> ivy.default_device(item=[], as_native=False) + 'cpu' + + >>> ivy.default_device(item=(), as_native=True) + device(type='cpu') + + >>> ivy.default_device(item={"a": 1}, as_native=True) + device(type='cpu') + + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'gpu:0') + >>> ivy.default_device(item=x, as_native=True) + device(type='gpu', id=0) """ - global soft_device_mode_stack - if soft_device_mode_stack: - soft_device_mode_stack.pop(-1) - mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False - ivy.__setattr__("soft_device_mode", mode, True) + if ivy.exists(device): + if as_native is True: + return ivy.as_native_dev(device) + elif as_native is False: + return ivy.as_ivy_dev(device) + return device + as_native = ivy.default(as_native, False) + if ivy.exists(item): + if isinstance(item, (list, tuple, dict)) and len(item) == 0: + pass + elif ivy.is_array(item): + return ivy.dev(item, as_native=as_native) + global default_device_stack + if not default_device_stack: + ret = "gpu:0" if ivy.gpu_is_available() else "cpu" + else: + ret = default_device_stack[-1] + if as_native: + return ivy.as_native_dev(ret) + return ivy.as_ivy_dev(ret) # Retrieval @@ -385,195 +481,301 @@ def dev( return ivy.current_backend(x).dev(x, as_native=as_native) -# Conversions +# Utilization @handle_exceptions -def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: +def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Convert device to string representation. + Get the current utilization (%) for a given device. Parameters ---------- device - The device handle to convert to string. + The device string of the device to query utilization for. Returns ------- ret - Device string e.g. 'cuda:0'. + The device utilization (%) - Examples - -------- - >>> y = ivy.as_ivy_dev('cpu') - >>> print(y) - cpu + Example + ------- + >>> ivy.dev_util('cpu') + 13.4 + >>> ivy.dev_util('gpu:0') + 7.8 + >>> ivy.dev_util('cpu') + 93.4 + >>> ivy.dev_util('gpu:2') + 57.4 + >>> ivy.dev_util('cpu') + 84.2 """ - return ivy.current_backend().as_ivy_dev(device) + if device == "cpu": + return psutil.cpu_percent() + elif "gpu" in device: + handle = _get_nvml_gpu_handle(device) + return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions -def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: +@handle_nestable +def function_supported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: """ - Convert device string representation to native device type. + Return the supported devices of the current backend's function. The function returns + a dict containing the supported devices for the compositional and primary + implementations in case of partial mixed functions. Parameters ---------- - device - The device string to convert to native device handle. - A native device handle can be passed in instead - in this case - the unmodified parameter is returned. + fn + The function to check for the supported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. Returns ------- ret - Native device handle. + Tuple or dict containing the supported devices of the function Examples -------- - With :class:`ivy.Device` input: + >>> import ivy + >>> print(ivy.function_supported_devices(ivy.ones)) + ('cpu', 'gpu') + """ + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=False), + } + else: + supported_devices = set(_get_devices(fn, complement=False)) + if recurse: + supported_devices = ivy.functional.data_type._nested_get( + fn, supported_devices, set.intersection, function_supported_devices + ) - >>> ivy.set_backend("numpy") - >>> ivy.as_native_dev("cpu") - 'cpu' + return ( + supported_devices + if isinstance(supported_devices, dict) + else tuple(supported_devices) + ) - >>> ivy.set_backend("tensorflow") - >>> ivy.as_native_dev("tpu:3") - '/TPU:3' - With :class:`ivy.NativeDevice` input: +@handle_exceptions +@handle_nestable +def function_unsupported_devices( + fn: Callable, recurse: bool = True +) -> Union[Tuple, dict]: + """ + Return the unsupported devices of the current backend's function. The function + returns a dict containing the unsupported devices for the compositional and primary + implementations in case of partial mixed functions. - >>> import torch - >>> device = torch.device("cuda") - >>> device - device(type='cuda') + Parameters + ---------- + fn + The function to check for the unsupported device attribute + recurse + Whether to recurse into used ivy functions. Default is ``True``. - >>> ivy.as_native_dev(device) - device(type='cuda') + Returns + ------- + ret + Tuple or dict containing the unsupported devices of the function + + Examples + -------- + >>> print(ivy.function_unsupported_devices(ivy.ones)) + ('tpu',) """ - return ivy.current_backend().as_native_dev(device) + ivy.utils.assertions.check_true( + _is_valid_devices_attributes(fn), + "supported_devices and unsupported_devices attributes cannot both " + "exist in a particular backend", + ) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices(fn.compos, recurse=recurse), + "primary": _get_devices(fn, complement=True), + } + else: + unsupported_devices = set(_get_devices(fn, complement=True)) + if recurse: + unsupported_devices = ivy.functional.data_type._nested_get( + fn, unsupported_devices, set.union, function_unsupported_devices + ) + return ( + unsupported_devices + if isinstance(unsupported_devices, dict) + else tuple(unsupported_devices) + ) -# Memory +# Device Queries # + +# Array Printing @handle_exceptions -def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +def get_all_ivy_arrays_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, +) -> ivy.Container: """ - Clear memory cache on target device. + Get all ivy arrays which are currently alive on the specified device. Parameters ---------- device - The device string to convert to native device handle or native device handle. + The device handle from which to get the arrays + + Returns + ------- + ret + Container with the arrays found for the specified device [identity, array] Examples -------- - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.clear_cached_mem_on_dev(device) + >>> x = ivy.array([1,0,2]) + >>> y = ivy.dev(x) + >>> z = ivy.get_all_ivy_arrays_on_dev(y) + >>> print(z) + {139740789224448:ivy.array([1,0,2])}, """ - ivy.current_backend().clear_cached_mem_on_dev(device) + device = ivy.as_ivy_dev(device) + all_arrays = list() + for obj in gc.get_objects(): + if ( + obj is ivy.data_classes.array.array.Array + and ivy.is_ivy_array(obj) + and ivy.dev(obj) == device + ): + all_arrays.append(obj) + + return ivy.Container(dict(zip([str(id(a)) for a in all_arrays], all_arrays))) + + +# Availability @handle_exceptions -def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +def gpu_is_available() -> bool: """ - Get the total amount of memory (in GB) for a given device string. In case of CPU, - the total RAM is returned. + Determine whether a GPU is available to use, with the backend framework. + + Returns + ------- + ret + Boolean, as to whether a gpu is available. + + Examples + -------- + >>> print(ivy.gpu_is_available()) + False + """ + return ivy.current_backend().gpu_is_available() + + +def handle_soft_device_variable(*args, fn, device_shifting_dev=None, **kwargs): + return ivy.current_backend().handle_soft_device_variable( + *args, fn=fn, device_shifting_dev=device_shifting_dev, **kwargs + ) + + +@handle_exceptions +def num_cpu_cores(*, logical: bool = True) -> int: + """ + Determine the number of cores available in the cpu. Parameters ---------- - device - The device string to convert to native device handle. + logical + Whether request is for number of physical or logical cores available in CPU Returns ------- ret - The total memory on the device in GB. + Number of cores available in CPU Examples -------- - >>> x = ivy.total_mem_on_dev("cpu") - >>> print(x) - 53.66700032 - - >>> x = ivy.total_mem_on_dev("gpu:0") - >>> print(x) - 8.589934592 + >>> print(ivy.num_cpu_cores(logical=False)) + 2 """ - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.total / 1e9 - elif device == "cpu": - return psutil.virtual_memory().total / 1e9 + if logical: + return psutil.cpu_count(logical=logical) else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + return psutil.cpu_count(logical=False) @handle_exceptions -def used_mem_on_dev( - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - process_specific: bool = False, -) -> float: +def num_gpus() -> int: """ - Get the used memory (in GB) for a given device string. In case of CPU, the used RAM - is returned. + Determine the number of available GPUs, with the backend framework. + + Returns + ------- + ret + Number of available GPUs. + + Examples + -------- + >>> print(ivy.num_gpus()) + 1 + """ + return ivy.current_backend().num_gpus() + + +@handle_exceptions +def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: + """ + Return the number of arrays which are currently alive on the specified device. Parameters ---------- device - The device string to convert to native device handle. - process_specific - Whether to check the memory used by this python process alone. Default is - False. + The device handle from which to count the arrays Returns ------- ret - The used memory on the device in GB. + Number of arrays on the specified device Examples -------- - >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) - >>> print(x) - 6.219563008 + >>> x1 = ivy.array([-1, 0, 5.2]) + >>> x2 = ivy.array([-1, 0, 5.2, 4, 5]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 2 - >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) - >>> print(x) - 0.902400346 + >>> x1 = ivy.native_array([-1, 0, 5.2]) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) + >>> print(y) + 0 - >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) + >>> x = ivy.Container(x1=ivy.array([-1]), + ... x2=ivy.native_array([-1])) + >>> y = ivy.num_ivy_arrays_on_dev(ivy.default_device()) >>> print(y) - 0.525205504 + 1 """ - ivy.clear_cached_mem_on_dev(device) - if "gpu" in device: - handle = _get_nvml_gpu_handle(device) - if process_specific: - pid = os.getpid() - for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): - if process.pid == pid: - return process.usedGpuMemory / 1e9 - info = pynvml.nvmlDeviceGetMemoryInfo(handle) - return info.used / 1e9 - elif device == "cpu": - if process_specific: - return psutil.Process(os.getpid()).memory_info().rss / 1e9 - vm = psutil.virtual_memory() - return (vm.total - vm.available) / 1e9 - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + return len(ivy.get_all_ivy_arrays_on_dev(device)) @handle_exceptions @@ -637,308 +839,151 @@ def percent_used_mem_on_dev( ) -# Utilization - - @handle_exceptions -def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: +@handle_nestable +def print_all_ivy_arrays_on_dev( + *, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + attr_only: bool = True, +) -> None: """ - Get the current utilization (%) for a given device. + Print the shape and dtype for all ivy arrays which are currently alive on the + specified device. Parameters ---------- device - The device string of the device to query utilization for. + The device on which to print the arrays - Returns - ------- - ret - The device utilization (%) + attr_only + Whether or not to only print the `shape` and `dtype` attributes of the array - Example - ------- - >>> ivy.dev_util('cpu') - 13.4 - >>> ivy.dev_util('gpu:0') - 7.8 - >>> ivy.dev_util('cpu') - 93.4 - >>> ivy.dev_util('gpu:2') - 57.4 - >>> ivy.dev_util('cpu') - 84.2 - """ - if device == "cpu": - return psutil.cpu_percent() - elif "gpu" in device: - handle = _get_nvml_gpu_handle(device) - return pynvml.nvmlDeviceGetUtilizationRates(handle).gpu - else: - raise ivy.utils.exceptions.IvyException( - 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' - "but found {}".format(device) - ) + Examples + -------- + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y) + ((3,), 'int32') + ((3,), 'int32') -# Availability + >>> x = ivy.array([[1,0,2], [3,2,1]]) + >>> y = ivy.dev(x) + >>> ivy.print_all_ivy_arrays_on_dev(y, attr_only = False) + [1,0,2] + [3,2,1] + """ + arrs = ivy.get_all_ivy_arrays_on_dev(device).values() + if attr_only: + [print((arr.shape, arr.dtype)) for arr in arrs] + else: + [print(arr) for arr in arrs] @handle_exceptions -def gpu_is_available() -> bool: +def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: """ - Determine whether a GPU is available to use, with the backend framework. + Set the default device to given device instance. - Returns - ------- - ret - Boolean, as to whether a gpu is available. + Parameters + ---------- + device + The device to set as the default device Examples -------- - >>> print(ivy.gpu_is_available()) - False + >>> ivy.set_default_device("cpu") + >>> ivy.default_device() + 'cpu' + + >>> ivy.set_backend("torch") + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device(as_native=True) + device(type='cuda', index=0) + + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_default_device(device) + >>> ivy.default_device(as_native=True) + device(type='cuda') """ - return ivy.current_backend().gpu_is_available() + global default_device_stack + default_device_stack.append(device) @handle_exceptions -def num_cpu_cores(*, logical: bool = True) -> int: +def set_soft_device_mode(mode: bool) -> None: """ - Determine the number of cores available in the cpu. - - Parameters - ---------- - logical - Whether request is for number of physical or logical cores available in CPU - - Returns - ------- - ret - Number of cores available in CPU - - Examples - -------- - >>> print(ivy.num_cpu_cores(logical=False)) - 2 - """ - if logical: - return psutil.cpu_count(logical=logical) - else: - return psutil.cpu_count(logical=False) - - -@handle_exceptions -def num_gpus() -> int: - """ - Determine the number of available GPUs, with the backend framework. - - Returns - ------- - ret - Number of available GPUs. - - Examples - -------- - >>> print(ivy.num_gpus()) - 1 - """ - return ivy.current_backend().num_gpus() - - -@handle_exceptions -def tpu_is_available() -> bool: - """ - Determine whether a TPU is available to use, with the backend framework. - - Returns - ------- - ret - Boolean, as to whether a tpu is available. + Set the mode of whether to move input arrays to `ivy.default_device()` before + performing an operation. + Parameter + --------- + mode + boolean whether to move input arrays Examples -------- - >>> print(ivy.tpu_is_available()) + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode False + >>> ivy.set_soft_device_mode(True) + >>> ivy.soft_device_mode + True """ - return ivy.current_backend().tpu_is_available() - - -# Default Device # - - -# noinspection PyShadowingNames -@handle_exceptions -def default_device( - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - /, - *, - item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, - as_native: bool = None, -) -> Union[ivy.Device, ivy.NativeDevice]: - """ - Return the input device or the default device. If the as_native flag is set, the - device will be converted to a native device. If the item is provided, the item's - device is returned. If the device is not provided, the last default device is - returned. If a default device has not been set, the first gpu is returned if - available, otherwise the cpu is returned. - - Parameters - ---------- - device - The device to be returned or converted. - item - The item to get the device from. - as_native - Whether to convert the device to a native device. - - Returns - ------- - ret - Device handle or string. - - Examples - -------- - >>> ivy.default_device() - device(type='cpu') - - >>> ivy.default_device("gpu:0") - 'gpu:0' - - >>> ivy.default_device(item=[], as_native=False) - 'cpu' - - >>> ivy.default_device(item=(), as_native=True) - device(type='cpu') - - >>> ivy.default_device(item={"a": 1}, as_native=True) - device(type='cpu') - - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'gpu:0') - >>> ivy.default_device(item=x, as_native=True) - device(type='gpu', id=0) - """ - if ivy.exists(device): - if as_native is True: - return ivy.as_native_dev(device) - elif as_native is False: - return ivy.as_ivy_dev(device) - return device - as_native = ivy.default(as_native, False) - if ivy.exists(item): - if isinstance(item, (list, tuple, dict)) and len(item) == 0: - pass - elif ivy.is_array(item): - return ivy.dev(item, as_native=as_native) - global default_device_stack - if not default_device_stack: - ret = "gpu:0" if ivy.gpu_is_available() else "cpu" - else: - ret = default_device_stack[-1] - if as_native: - return ivy.as_native_dev(ret) - return ivy.as_ivy_dev(ret) + global soft_device_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + soft_device_mode_stack.append(mode) + ivy.__setattr__("soft_device_mode", mode, True) @handle_exceptions -def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: +def set_split_factor( + factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None +) -> None: """ - Set the default device to given device instance. + Set the global split factor for a given device, which can be used to scale batch + splitting chunk sizes for the device across the codebase. Parameters ---------- + factor + The factor to set the device-specific split factor to. device - The device to set as the default device + The device to set the split factor for. Sets the default device by default. Examples -------- - >>> ivy.set_default_device("cpu") - >>> ivy.default_device() - 'cpu' + >>> print(ivy.default_device()) + cpu - >>> ivy.set_backend("torch") - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device(as_native=True) - device(type='cuda', index=0) + >>> ivy.set_split_factor(0.5) + >>> print(ivy.split_factors) + {'cpu': 0.5} >>> import torch >>> ivy.set_backend("torch") >>> device = torch.device("cuda") - >>> ivy.set_default_device(device) - >>> ivy.default_device(as_native=True) - device(type='cuda') - """ - global default_device_stack - default_device_stack.append(device) - - -@handle_exceptions -def unset_default_device() -> None: - """ - Reset the default device to "cpu". - - Examples - -------- - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device() - "gpu:0" - >>> ivy.unset_default_device() - >>> ivy.default_device() - "cpu" - """ - global default_device_stack - if default_device_stack: - default_device_stack.pop(-1) - - -# Device Allocation # - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def to_device( - x: Union[ivy.Array, ivy.NativeArray], - device: Union[ivy.Device, ivy.NativeDevice], - /, - *, - stream: Optional[Union[int, Any]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Move the input array x to the desired device, specified by device string. - - Parameters - ---------- - x - input array to be moved to the desired device - device - device to move the input array `x` to - stream - stream object to use during copy. In addition to the types supported in - array.__dlpack__(), implementations may choose to support any library-specific - stream object with the caveat that any code using such an object would not be - portable. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> ivy.set_split_factor(0.3, device=device) + >>> print(ivy.split_factors) + {device(type='cuda'): 0.3} - Returns - ------- - ret - input array x placed on the desired device + >>> ivy.set_split_factor(0.4, device="tpu") + >>> print(ivy.split_factors) + {'tpu': 0.4} - Examples - -------- - >>> x = ivy.array([1., 2., 3.]) - >>> x = ivy.to_device(x, 'cpu') - >>> print(x.device) - cpu + >>> import torch + >>> ivy.set_backend("torch") + >>> device = torch.device("cuda") + >>> ivy.set_split_factor(0.2) + >>> ivy.set_split_factor(0.3, device='gpu') + >>> print(ivy.split_factors) + {'cpu': 0.2, 'gpu': 0.3} """ - return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) + ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) + global split_factors + device = ivy.default(device, default_device()) + split_factors[device] = factor # Function Splitting # @@ -983,55 +1028,6 @@ def split_factor( return split_factors.setdefault(device, 0.0) -@handle_exceptions -def set_split_factor( - factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None -) -> None: - """ - Set the global split factor for a given device, which can be used to scale batch - splitting chunk sizes for the device across the codebase. - - Parameters - ---------- - factor - The factor to set the device-specific split factor to. - device - The device to set the split factor for. Sets the default device by default. - - Examples - -------- - >>> print(ivy.default_device()) - cpu - - >>> ivy.set_split_factor(0.5) - >>> print(ivy.split_factors) - {'cpu': 0.5} - - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.3, device=device) - >>> print(ivy.split_factors) - {device(type='cuda'): 0.3} - - >>> ivy.set_split_factor(0.4, device="tpu") - >>> print(ivy.split_factors) - {'tpu': 0.4} - - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_split_factor(0.2) - >>> ivy.set_split_factor(0.3, device='gpu') - >>> print(ivy.split_factors) - {'cpu': 0.2, 'gpu': 0.3} - """ - ivy.utils.assertions.check_less(0, factor, allow_equal=True, as_array=False) - global split_factors - device = ivy.default(device, default_device()) - split_factors[device] = factor - - @handle_exceptions def split_func_call( func: Callable, @@ -1166,200 +1162,213 @@ def split_func_call( return ret[0] if len(ret) == 1 else ret -def _is_valid_devices_attributes(fn: Callable) -> bool: - if hasattr(fn, "supported_devices") and hasattr(fn, "unsupported_devices"): - fn_supported_devices = fn.supported_devices - fn_unsupported_devices = fn.unsupported_devices - if isinstance(fn_supported_devices, dict): - if isinstance(fn_unsupported_devices, dict): - backend_str = ivy.current_backend_str() - if ( - backend_str in fn_supported_devices - and backend_str in fn_unsupported_devices - ): - return False - else: - if isinstance(fn_unsupported_devices, tuple): - return False - return True - - -def _get_devices(fn: Callable, complement: bool = True) -> Tuple: - valid_devices = ivy.valid_devices - invalid_devices = ivy.invalid_devices - all_devices = ivy.all_devices +# Device Allocation # - supported = set(ivy.valid_devices) - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - supported = set(all_devices).difference(supported) - return supported +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def to_device( + x: Union[ivy.Array, ivy.NativeArray], + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + stream: Optional[Union[int, Any]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Move the input array x to the desired device, specified by device string. - # Their values are formated like either - # 1. fn.supported_devices = ("cpu",) - # Could also have the "all" value for the framework - basic = [ - ("supported_devices", set.intersection, valid_devices), - ("unsupported_devices", set.difference, invalid_devices), - ] - for key, merge_fn, base in basic: - if hasattr(fn, key): - v = getattr(fn, key) - if "einops" in fn.__name__ and isinstance(v, dict): - v = v.get(ivy.current_backend_str(), base) - ivy.utils.assertions.check_isinstance(v, tuple) - supported = merge_fn(supported, set(v)) + Parameters + ---------- + x + input array to be moved to the desired device + device + device to move the input array `x` to + stream + stream object to use during copy. In addition to the types supported in + array.__dlpack__(), implementations may choose to support any library-specific + stream object with the caveat that any code using such an object would not be + portable. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - if complement: - supported = set(all_devices).difference(supported) + Returns + ------- + ret + input array x placed on the desired device - return tuple(supported) + Examples + -------- + >>> x = ivy.array([1., 2., 3.]) + >>> x = ivy.to_device(x, 'cpu') + >>> print(x.device) + cpu + """ + return ivy.current_backend(x).to_device(x, device, stream=stream, out=out) @handle_exceptions -@handle_nestable -def function_supported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: +def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: """ - Return the supported devices of the current backend's function. The function returns - a dict containing the supported devices for the compositional and primary - implementations in case of partial mixed functions. + Get the total amount of memory (in GB) for a given device string. In case of CPU, + the total RAM is returned. Parameters ---------- - fn - The function to check for the supported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + device + The device string to convert to native device handle. Returns ------- ret - Tuple or dict containing the supported devices of the function + The total memory on the device in GB. Examples -------- - >>> import ivy - >>> print(ivy.function_supported_devices(ivy.ones)) - ('cpu', 'gpu') + >>> x = ivy.total_mem_on_dev("cpu") + >>> print(x) + 53.66700032 + + >>> x = ivy.total_mem_on_dev("gpu:0") + >>> print(x) + 8.589934592 """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=False), - } + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.total / 1e9 + elif device == "cpu": + return psutil.virtual_memory().total / 1e9 else: - supported_devices = set(_get_devices(fn, complement=False)) - if recurse: - supported_devices = ivy.functional.data_type._nested_get( - fn, supported_devices, set.intersection, function_supported_devices - ) - - return ( - supported_devices - if isinstance(supported_devices, dict) - else tuple(supported_devices) - ) + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) @handle_exceptions -@handle_nestable -def function_unsupported_devices( - fn: Callable, recurse: bool = True -) -> Union[Tuple, dict]: +def tpu_is_available() -> bool: """ - Return the unsupported devices of the current backend's function. The function - returns a dict containing the unsupported devices for the compositional and primary - implementations in case of partial mixed functions. - - Parameters - ---------- - fn - The function to check for the unsupported device attribute - recurse - Whether to recurse into used ivy functions. Default is ``True``. + Determine whether a TPU is available to use, with the backend framework. Returns ------- ret - Tuple or dict containing the unsupported devices of the function + Boolean, as to whether a tpu is available. Examples -------- - >>> print(ivy.function_unsupported_devices(ivy.ones)) - ('tpu',) + >>> print(ivy.tpu_is_available()) + False """ - ivy.utils.assertions.check_true( - _is_valid_devices_attributes(fn), - "supported_devices and unsupported_devices attributes cannot both " - "exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices(fn.compos, recurse=recurse), - "primary": _get_devices(fn, complement=True), - } - else: - unsupported_devices = set(_get_devices(fn, complement=True)) - if recurse: - unsupported_devices = ivy.functional.data_type._nested_get( - fn, unsupported_devices, set.union, function_unsupported_devices - ) - return ( - unsupported_devices - if isinstance(unsupported_devices, dict) - else tuple(unsupported_devices) - ) + return ivy.current_backend().tpu_is_available() -# Profiler # +@handle_exceptions +def unset_default_device() -> None: + """ + Reset the default device to "cpu". + Examples + -------- + >>> ivy.set_default_device("gpu:0") + >>> ivy.default_device() + "gpu:0" + >>> ivy.unset_default_device() + >>> ivy.default_device() + "cpu" + """ + global default_device_stack + if default_device_stack: + default_device_stack.pop(-1) -class Profiler(abc.ABC): + +@handle_exceptions +def unset_soft_device_mode() -> None: """ - The profiler class is used to profile the execution of some code. + Reset the mode of moving input arrays to `ivy.default_device()` before performing an + operation. - Parameters - ---------- - save_dir - The directory to save the profile data to. + Examples + -------- + >>> ivy.set_soft_device_mode(False) + >>> ivy.soft_device_mode + False + >>> ivy.unset_soft_device_mode() + >>> ivy.soft_device_mode + True """ + global soft_device_mode_stack + if soft_device_mode_stack: + soft_device_mode_stack.pop(-1) + mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False + ivy.__setattr__("soft_device_mode", mode, True) - def __init__(self, save_dir: str): - self._save_dir = save_dir - @abc.abstractmethod - def start(self): - """ - Start the profiler. +@handle_exceptions +def used_mem_on_dev( + device: Union[ivy.Device, ivy.NativeDevice], + /, + *, + process_specific: bool = False, +) -> float: + """ + Get the used memory (in GB) for a given device string. In case of CPU, the used RAM + is returned. - This should be called before the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException + Parameters + ---------- + device + The device string to convert to native device handle. + process_specific + Whether to check the memory used by this python process alone. Default is + False. - @abc.abstractmethod - def stop(self): - """ - Stop the profiler. + Returns + ------- + ret + The used memory on the device in GB. - This should be called after the code to be profiled. - """ - raise ivy.utils.exceptions.IvyNotImplementedException + Examples + -------- + >>> x = ivy.used_mem_on_dev("cpu", process_specific = False) + >>> print(x) + 6.219563008 - @abc.abstractmethod - def __enter__(self): - raise ivy.utils.exceptions.IvyNotImplementedException + >>> x = ivy.used_mem_on_dev("cpu", process_specific = True) + >>> print(x) + 0.902400346 - @abc.abstractmethod - def __exit__(self, exc_type, exc_val, exc_tb): - raise ivy.utils.exceptions.IvyNotImplementedException + >>> y = ivy.used_mem_on_dev("gpu:0", process_specific = False) + >>> print(y) + 0.525205504 + """ + ivy.clear_cached_mem_on_dev(device) + if "gpu" in device: + handle = _get_nvml_gpu_handle(device) + if process_specific: + pid = os.getpid() + for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle): + if process.pid == pid: + return process.usedGpuMemory / 1e9 + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return info.used / 1e9 + elif device == "cpu": + if process_specific: + return psutil.Process(os.getpid()).memory_info().rss / 1e9 + vm = psutil.virtual_memory() + return (vm.total - vm.available) / 1e9 + else: + raise ivy.utils.exceptions.IvyException( + 'Invalid device string input, must be on the form "gpu:idx" or "cpu", ' + "but found {}".format(device) + ) + + +ivy.soft_device_mode = soft_device_mode_stack[-1] if soft_device_mode_stack else False +dev_handles = dict() diff --git a/ivy/functional/ivy/elementwise.py b/ivy/functional/ivy/elementwise.py index 5fd532ccfb9c9..dca990b0da0a7 100644 --- a/ivy/functional/ivy/elementwise.py +++ b/ivy/functional/ivy/elementwise.py @@ -498,6 +498,51 @@ def add( return ivy.current_backend(x1, x2).add(x1, x2, alpha=alpha, out=out) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def angle( + z: Union[ivy.Array, ivy.NativeArray], + /, + *, + deg: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Calculate Element-wise the angle for an array of complex numbers(x+yj). + + Parameters + ---------- + z + Array-like input. + deg + optional bool. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Returns an array of angles for each complex number in the input. + If deg is False(default), angle is calculated in radian and if + deg is True, then angle is calculated in degrees. + + Examples + -------- + >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> z + ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) + >>> ivy.angle(z) + ivy.array([ 2.35619449, 2.35619449, -0.78539816]) + >>> ivy.angle(z,deg=True) + ivy.array([135., 135., -45.]) + """ + return ivy.current_backend(z).angle(z, deg=deg, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -1895,6 +1940,85 @@ def cosh( return ivy.current_backend(x).cosh(x, out=out) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def deg2rad( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Convert the input from degrees to radians. + + Parameters + ---------- + x + input array whose elements are each expressed in degrees. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array with each element in ``x`` converted from degrees to radians. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x=ivy.array([0,90,180,270,360]) + >>> y=ivy.deg2rad(x) + >>> print(y) + ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + + >>> x=ivy.array([0,-1.5,-50,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([ 0., -0.02617994, -0.87266463, nan]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.deg2rad(x, out=x) + >>> print(x) + ivy.array([[ 0.01919862, 0.03839725, 0.05759586], + [-0.07679449, -0.09599311, -0.11519173]]) + + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.deg2rad(x,out=y) + >>> print(y) + ivy.array([0., 0.35081118, nan]) + + With :class:`ivy.Container` input: + + >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), + ... b=ivy.array([0,90,180,270,360])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 0.35081118, -0.88139129, nan]), + b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + } + + >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), + ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) + >>> y=ivy.deg2rad(x) + >>> print(y) + { + a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), + b: ivy.array([0., -0.02617994, -0.87266463, nan]) + } + """ + return ivy.current_backend(x).deg2rad(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2204,6 +2328,49 @@ def equal( return ivy.current_backend(x1, x2).equal(x1, x2, out=out) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def erf( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Gauss error function of ``x`` element-wise. + + Parameters + ---------- + x + Value to compute exponential for. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The Gauss error function of x. + + Examples + -------- + >>> x = ivy.array([0, 0.3, 0.7, 1.0]) + >>> ivy.erf(x) + ivy.array([0., 0.328, 0.677, 0.842]) + """ + return ivy.current_backend(x).erf(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2351,214 +2518,74 @@ def exp( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def imag( - val: Union[ivy.Array, ivy.NativeArray], +def exp2( + x: Union[ivy.Array, float, list, tuple], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the imaginary part of a complex number for each element ``x_i`` of the input - array ``val``. + Calculate 2**p for all p in the input array. Parameters ---------- - val - input array. Should have a complex floating-point data type. + x + Array-like input. out optional output array, for writing the result to. Returns ------- ret - Returns an array with the imaginary part of complex numbers. - The returned arrau must have a floating-point data type determined by - the precision of ``val`` (e.g., if ``val`` is ``complex64``, - the returned array must be ``float32``). - - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Element-wise 2 to the power x. This is a scalar if x is a scalar. Examples -------- - >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) - >>> b - ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) - >>> ivy.imag(b) - ivy.array([2., 4., 6.]) + >>> x = ivy.array([1, 2, 3]) + >>> ivy.exp2(x) + ivy.array([2., 4., 8.]) + >>> x = [5, 6, 7] + >>> ivy.exp2(x) + ivy.array([32., 64., 128.]) """ - return ivy.current_backend(val).imag(val, out=out) + return ivy.current_backend(x).exp2(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def angle( - z: Union[ivy.Array, ivy.NativeArray], +def expm1( + x: Union[ivy.Array, ivy.NativeArray], /, *, - deg: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate Element-wise the angle for an array of complex numbers(x+yj). + Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain + ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element + ``x_i`` of the input array ``x``. - Parameters - ---------- - z - Array-like input. - deg - optional bool. - out - optional output array, for writing the result to. + .. note:: + The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when + ``x`` is close to zero. Accordingly, conforming implementations should avoid + implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other + IEEE 754-2019 compliant mathematical library, for a potential reference + implementation. - Returns - ------- - ret - Returns an array of angles for each complex number in the input. - If deg is False(default), angle is calculated in radian and if - deg is True, then angle is calculated in degrees. + .. note:: + For complex floating-point operands, ``expm1(conj(x))`` + must equal ``conj(expm1(x))``. - Examples - -------- - >>> z = ivy.array([-1 + 1j, -2 + 2j, 3 - 3j]) - >>> z - ivy.array([-1.+1.j, -2.+2.j, 3.-3.j]) - >>> ivy.angle(z) - ivy.array([ 2.35619449, 2.35619449, -0.78539816]) - >>> ivy.angle(z,deg=True) - ivy.array([135., 135., -45.]) - """ - return ivy.current_backend(z).angle(z, deg=deg, out=out) + .. note:: + The exponential function is an entire function + in the complex plane and has no branch cuts. - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def gcd( - x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the greatest common divisor of |x1| and |x2|. - - Parameters - ---------- - x1 - First array-like input. - x2 - Second array-input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Element-wise gcd of |x1| and |x2|. - - Examples - -------- - >>> x1 = ivy.array([1, 2, 3]) - >>> x2 = ivy.array([4, 5, 6]) - >>> ivy.gcd(x1, x2) - ivy.array([1., 1., 3.]) - >>> x1 = ivy.array([1, 2, 3]) - >>> ivy.gcd(x1, 10) - ivy.array([1., 2., 1.]) - """ - return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def exp2( - x: Union[ivy.Array, float, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate 2**p for all p in the input array. - - Parameters - ---------- - x - Array-like input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Element-wise 2 to the power x. This is a scalar if x is a scalar. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.exp2(x) - ivy.array([2., 4., 8.]) - >>> x = [5, 6, 7] - >>> ivy.exp2(x) - ivy.array([32., 64., 128.]) - """ - return ivy.current_backend(x).exp2(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def expm1( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate an implementation-dependent approximation to ``exp(x)-1``, having domain - ``[-infinity, +infinity]`` and codomain ``[-1, +infinity]``, for each element - ``x_i`` of the input array ``x``. - - .. note:: - The purpose of this function is to calculate ``exp(x)-1.0`` more accurately when - ``x`` is close to zero. Accordingly, conforming implementations should avoid - implementing this function as simply ``exp(x)-1.0``. See FDLIBM, or some other - IEEE 754-2019 compliant mathematical library, for a potential reference - implementation. - - .. note:: - For complex floating-point operands, ``expm1(conj(x))`` - must equal ``conj(expm1(x))``. - - .. note:: - The exponential function is an entire function - in the complex plane and has no branch cuts. - - **Special cases** + **Special cases** For floating-point operands, @@ -2975,6 +3002,92 @@ def fmin( return ivy.current_backend(x1, x2).fmin(x1, x2, out=out) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fmod( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Compute the element-wise remainder of divisions of two arrays. + + Parameters + ---------- + x1 + First input array. + x2 + Second input array + out + optional output array, for writing the result to. + + Returns + ------- + ret + Array with element-wise remainder of divisions. + + Examples + -------- + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmod(x1, x2) + ivy.array([ 0, 3, 0]) + + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmod(x1, x2) + ivy.array([ nan, nan, nan]) + """ + return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def gcd( + x1: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + x2: Union[ivy.Array, ivy.NativeArray, int, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the greatest common divisor of |x1| and |x2|. + + Parameters + ---------- + x1 + First array-like input. + x2 + Second array-input. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Element-wise gcd of |x1| and |x2|. + + Examples + -------- + >>> x1 = ivy.array([1, 2, 3]) + >>> x2 = ivy.array([4, 5, 6]) + >>> ivy.gcd(x1, x2) + ivy.array([1., 1., 3.]) + >>> x1 = ivy.array([1, 2, 3]) + >>> ivy.gcd(x1, 10) + ivy.array([1., 2., 1.]) + """ + return ivy.current_backend(x1, x2).gcd(x1, x2, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -3164,277 +3277,57 @@ def greater_equal( return ivy.current_backend(x1, x2).greater_equal(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def less_equal( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def imag( + val: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 - with the respective element x2_i of the input array x2. + Return the imaginary part of a complex number for each element ``x_i`` of the input + array ``val``. Parameters ---------- - x1 - first input array. May have any data type. - x2 - second input array. Must be compatible with x1 (with Broadcasting). May have any - data type. + val + input array. Should have a complex floating-point data type. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- - ret - an array containing the element-wise results. The returned array must have a - data type of bool. - + ret + Returns an array with the imaginary part of complex numbers. + The returned arrau must have a floating-point data type determined by + the precision of ``val`` (e.g., if ``val`` is ``complex64``, + the returned array must be ``float32``). - This function conforms to the `Array API Standard - `_. This docstring is an extension of the + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.imag.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) - >>> print(x) - ivy.array([True, True, False]) - - >>> x = ivy.array([[10.1, 2.3, -3.6]]) - >>> y = ivy.array([[4.8], [5.2], [6.1]]) - >>> shape = (3,3) - >>> fill_value = False - >>> z = ivy.full(shape, fill_value) - >>> ivy.less_equal(x, y, out=z) - >>> print(z) - ivy.array([[False, True, True], - [False, True, True], - [False, True, True]]) - - >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) - >>> y = ivy.array([[8.4], [2.5], [1.6]]) - >>> ivy.less_equal(x, y, out=x) - >>> print(x) - ivy.array([[[1.], - [0.], - [1.]]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) - >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) - >>> z = ivy.less_equal(x, y) - >>> print(z) - { - a: ivy.array([False, False, False]), - b: ivy.array([True, True, True]) - } - """ - return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def multiply( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Calculate the product for each element x1_i of the input array x1 with the - respective element x2_i of the input array x2. - - .. note:: - Floating-point multiplication is not always associative due to finite precision. - - **Special Cases** - - For real-valued floating-point operands, - - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` and ``x2_i`` have the same mathematical sign, - the result has a positive mathematical sign, unless the result is ``NaN``. - If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` and ``x2_i`` have different mathematical signs, - the result has a negative mathematical sign, - unless the result is ``NaN``. If the result is ``NaN``, - the "sign" of ``NaN`` is implementation-defined. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and - ``x2_i`` is either ``+infinity`` or ``-infinity``, - the result is a signed infinity with the mathematical sign determined by - the rule already stated above. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` - is a nonzero finite number, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - If ``x1_i`` is a nonzero finite number and ``x2_i`` - is either ``+infinity`` or ``-infinity``, the result is a signed infinity with - the mathematical sign determined by the rule already stated above. - - In the remaining cases, where neither ``infinity`` nor ``NaN`` - is involved, the product must be computed and rounded to the nearest - representable value according to IEEE 754-2019 and a supported - rounding mode. If the magnitude is too large to represent, - the result is an `infinity` of appropriate mathematical sign. - If the magnitude is too small to represent, the result is a zero of - appropriate mathematical sign. - - For complex floating-point operands, multiplication is defined according to the - following table. For real components ``a`` and ``c`` and - imaginary components ``b`` and ``d``, - - +------------+----------------+-----------------+--------------------------+ - | | c | dj | c + dj | - +============+================+=================+==========================+ - | **a** | a * c | (a*d)j | (a*c) + (a*d)j | - +------------+----------------+-----------------+--------------------------+ - | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | - +------------+----------------+-----------------+--------------------------+ - | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | - +------------+----------------+-----------------+--------------------------+ - - In general, for complex floating-point operands, real-valued floating-point - special cases must independently apply to the real and imaginary component - operations involving real numbers as described in the above table. - - When ``a``, ``b``, ``c``, or ``d`` are all finite numbers - (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), - multiplication of complex floating-point operands should be computed - as if calculated according to the textbook formula for complex number multiplication - - .. math:: - (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j - - When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, - ``+infinity``, or ``-infinity``, - - - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, the result is implementation dependent. - - .. note:: - For complex floating-point operands, the results of special cases may be - implementation dependent depending on how an implementation chooses - to model complex numbers and complex infinity - (e.g., complex plane versus Riemann sphere). - For those implementations following C99 and its one-infinity model, - when at least one component is infinite, - even if the other component is ``NaN``, - the complex value is infinite, and the usual arithmetic - rules do not apply to complex-complex multiplication. - In the interest of performance, other implementations - may want to avoid the complex branching logic necessary - to implement the one-infinity model and choose to implement - all complex-complex multiplication according to the textbook formula. - Accordingly, special case behavior is unlikely - to be consistent across implementations. - - Parameters - ---------- - x1 - first input array. Should have a numeric data type. - - x2 - second input array. Must be compatible with ``x1`` - (see :ref'`broadcasting`). Should have a numeric data type - - out - optional output array, for writing the array result to. - It must have a shape that the inputs broadcast to. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Returns - ------- - ret - an array containing the element-wise products. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. - - Examples - -------- - With :code:`ivy.Array` inputs: - - >>> x1 = ivy.array([3., 5., 7.]) - >>> x2 = ivy.array([4., 6., 8.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([12., 30., 56.]) - - With :code:`ivy.NativeArray` inputs: - - >>> x1 = ivy.native_array([1., 3., 9.]) - >>> x2 = ivy.native_array([4., 7.2, 1.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([ 4. , 21.6, 9. ]) - - With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: - - >>> x1 = ivy.array([8., 6., 7.]) - >>> x2 = ivy.native_array([1., 2., 3.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - ivy.array([ 8., 12., 21.]) - - With :code:`ivy.Container` inputs: - - >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) - >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - { - a: ivy.array([12.,12.,24.]), - b: ivy.array([9.,3.,10.]) - } - - With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: - - >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) - >>> x2 = ivy.array([1.,2.,3.]) - >>> y = ivy.multiply(x1, x2) - >>> print(y) - { - a: ivy.array([3.,8.,15.]), - b: ivy.array([2.,4.,3.]) - } - """ - return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) + >>> b = ivy.array(np.array([1+2j, 3+4j, 5+6j])) + >>> b + ivy.array([1.+2.j, 3.+4.j, 5.+6.j]) + >>> ivy.imag(b) + ivy.array([2., 4., 6.]) + """ + return ivy.current_backend(val).imag(val, out=out) @handle_exceptions @@ -3774,28 +3667,27 @@ def isnan( @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def less( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def isreal( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + Test each element ``x_i`` of the input array ``x`` to determine whether the element + is real number. Returns a bool array, where True if input element is real. If + element has complex type with zero complex part, the return value for that element + is True. Parameters ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. + x + input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -3803,16 +3695,126 @@ def less( Returns ------- ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. + an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is + real number and ``False`` otherwise. The returned array should have a data type + of ``bool``. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances in place of + :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints + and also the examples below. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.less(ivy.array([1,2,3]),ivy.array([2,2,2])) - >>> print(x) - ivy.array([ True, False, False]) + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([[[True], [True], [True]]]) + + >>> x = ivy.array([1-0j, 3j, 7+5j]) + >>> z = ivy.isreal(x) + >>> print(z) + ivy.array([ True, False, False]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ + b=ivy.array([5j, 5-6j, 3])) + >>> z = ivy.isreal(x) + >>> print(z) + { + a: ivy.array([False, True, True]), + b: ivy.array([False, False, True]) + } + """ + return ivy.current_backend(x).isreal(x, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lcm( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the element-wise least common multiple (LCM) of x1 and x2. + + Parameters + ---------- + x1 + first input array, must be integers + x2 + second input array, must be integers + out + optional output array, for writing the result to. + + Returns + ------- + ret + an array that includes the element-wise least common multiples of x1 and x2 + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x1=ivy.array([2, 3, 4]) + >>> x2=ivy.array([5, 7, 15]) + >>> x1.lcm(x1, x2) + ivy.array([10, 21, 60]) + """ + return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def less( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the truth value of ``x1_i < x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + + Parameters + ---------- + x1 + first input array. Should have a numeric data type. + x2 + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of ``bool``. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.less(ivy.array([1,2,3]),ivy.array([2,2,2])) + >>> print(x) + ivy.array([ True, False, False]) >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) @@ -3861,6 +3863,93 @@ def less( return ivy.current_backend(x1).less(x1, x2, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def less_equal( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the truth value of x1_i <= x2_i for each element x1_i of the input array x1 + with the respective element x2_i of the input array x2. + + Parameters + ---------- + x1 + first input array. May have any data type. + x2 + second input array. Must be compatible with x1 (with Broadcasting). May have any + data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the element-wise results. The returned array must have a + data type of bool. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.less_equal(ivy.array([1,2,3]),ivy.array([2,2,2])) + >>> print(x) + ivy.array([True, True, False]) + + >>> x = ivy.array([[10.1, 2.3, -3.6]]) + >>> y = ivy.array([[4.8], [5.2], [6.1]]) + >>> shape = (3,3) + >>> fill_value = False + >>> z = ivy.full(shape, fill_value) + >>> ivy.less_equal(x, y, out=z) + >>> print(z) + ivy.array([[False, True, True], + [False, True, True], + [False, True, True]]) + + >>> x = ivy.array([[[1.1], [3.2], [-6.3]]]) + >>> y = ivy.array([[8.4], [2.5], [1.6]]) + >>> ivy.less_equal(x, y, out=x) + >>> print(x) + ivy.array([[[1.], + [0.], + [1.]]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([4, 5, 6]),b=ivy.array([2, 3, 4])) + >>> y = ivy.Container(a=ivy.array([1, 2, 3]),b=ivy.array([5, 6, 7])) + >>> z = ivy.less_equal(x, y) + >>> print(z) + { + a: ivy.array([False, False, False]), + b: ivy.array([True, True, True]) + } + """ + return ivy.current_backend(x1, x2).less_equal(x1, x2, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -4779,98 +4868,125 @@ def logical_xor( return ivy.current_backend(x1, x2).logical_xor(x1, x2, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def nan_to_num( - x: Union[ivy.Array, ivy.NativeArray], +def maximum( + x1: Union[ivy.Array, ivy.NativeArray, Number], + x2: Union[ivy.Array, ivy.NativeArray, Number], /, *, - copy: bool = True, - nan: Union[float, int] = 0.0, - posinf: Optional[Union[float, int]] = None, - neginf: Optional[Union[float, int]] = None, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Replace NaN with zero and infinity with large finite numbers (default behaviour) or - with the numbers defined by the user using the nan, posinf and/or neginf keywords. + Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. Parameters ---------- - x - Array input. - copy - Whether to create a copy of x (True) or to replace values in-place (False). - The in-place operation only occurs if casting to an array does not require - a copy. Default is True. - nan - Value to be used to fill NaN values. If no value is passed then NaN values - will be replaced with 0.0. - posinf - Value to be used to fill positive infinity values. If no value is passed - then positive infinity values will be replaced with a very large number. - neginf - Value to be used to fill negative infinity values. - If no value is passed then negative infinity values - will be replaced with a very small (or negative) number. + x1 + Input array containing elements to maximum threshold. + x2 + Tensor containing maximum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum + is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Array with the non-finite values replaced. - If copy is False, this may be x itself. + An array with the elements of x1, but clipped to not be lower than the x2 + values. Examples -------- - >>> x = ivy.array([1, 2, 3, nan]) - >>> ivy.nan_to_num(x) - ivy.array([1., 1., 3., 0.0]) - >>> x = ivy.array([1, 2, 3, inf]) - >>> ivy.nan_to_num(x, posinf=5e+100) - ivy.array([1., 2., 3., 5e+100]) - """ - return ivy.current_backend(x).nan_to_num( - x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out - ) + With :class:`ivy.Array` inputs: + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.maximum(x, y) + >>> print(z) + ivy.array([9, 9, 5]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.maximum(x, y, out=z) + >>> print(z) + ivy.array([[9., 9., 9., 9., 9., 9.], + [3., 5., 9., 8., 3., 7.], + [2., 5., 9., 8., 3., 7.]]) + + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.maximum(x, y, out=x) + >>> print(x) + ivy.array([[7, 7]]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]), + ... b=ivy.array([-5, 9])) + >>> z = ivy.maximum(x, y) + >>> print(z) + { + a: ivy.array([[1, 3], + [2, 4], + [3, 7]]), + b: ivy.array([[1, 9], + [2, 9], + [3, 9]]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) + >>> z = ivy.maximum(x, y) + >>> print(z) + { + a: ivy.array([1, 5, 6]), + b: ivy.array([5, 9, 7]) + } + """ + return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def negative( - x: Union[float, ivy.Array, ivy.NativeArray], +def minimum( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, + use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the negative value of each element in ``x``. - - .. note:: - For signed integer data types, the numerical negative of - the minimum representable integer is implementation-dependent. - - .. note:: - If ``x`` has a complex floating-point data type, - both the real and imaginary components for each ``x_i`` - must be negated (a result which follows from the rules of - complex number multiplication). + Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. Parameters ---------- - x - Input array. + x1 + Input array containing elements to minimum threshold. + x2 + Tensor containing minimum values, must be broadcastable to x1. + use_where + Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum + is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -4878,53 +4994,62 @@ def negative( Returns ------- ret - A new array with the negative value of each element in ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + An array with the elements of x1, but clipped to not exceed the x2 values. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.array([0,1,1,2]) - >>> y = ivy.negative(x) - >>> print(y) - ivy.array([ 0, -1, -1, -2]) + >>> x = ivy.array([7, 9, 5]) + >>> y = ivy.array([9, 3, 2]) + >>> z = ivy.minimum(x, y) + >>> print(z) + ivy.array([7, 3, 2]) - >>> x = ivy.array([0,-1,-0.5,2,3]) - >>> y = ivy.zeros(5) - >>> ivy.negative(x, out=y) - >>> print(y) - ivy.array([-0. , 1. , 0.5, -2. , -3. ]) + >>> x = ivy.array([1, 5, 9, 8, 3, 7]) + >>> y = ivy.array([[9], [3], [2]]) + >>> z = ivy.zeros((3, 6)) + >>> ivy.minimum(x, y, out=z) + >>> print(z) + ivy.array([[1.,5.,9.,8.,3.,7.], + [1.,3.,3.,3.,3.,3.], + [1.,2.,2.,2.,2.,2.]]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.negative(x,out=x) + >>> x = ivy.array([[7, 3]]) + >>> y = ivy.array([0, 7]) + >>> ivy.minimum(x, y, out=x) >>> print(x) - ivy.array([[-1.1, -2.2, -3.3], - [4.4, 5.5, 6.6]]) + ivy.array([[0, 3]]) - With :class:`ivy.Container` input: + With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.negative(x) - >>> print(y) + >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) + >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) + >>> z = ivy.minimum(x, y) + >>> print(z) { - a: ivy.array([-0., -1., -2.]), - b: ivy.array([-3., -4., 5.]) + a: ivy.array([[1, 0], + [1, 0], + [1, 0]]), + b: ivy.array([[-5, 3], + [-5, 4], + [-5, 7]]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([1, 3, 1]), + ... b=ivy.array([2, 8, 5])) + >>> y = ivy.Container(a=ivy.array([1, 5, 6]), + ... b=ivy.array([5, 9, 7])) + >>> z = ivy.minimum(x, y) + >>> print(z) + { + a: ivy.array([1, 3, 1]), + b: ivy.array([2, 8, 5]) } """ - return ivy.current_backend(x).negative(x, out=out) + return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) @handle_exceptions @@ -4934,406 +5059,244 @@ def negative( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def not_equal( - x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], - x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], +def multiply( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input - array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. + r""" + Calculate the product for each element x1_i of the input array x1 with the + respective element x2_i of the input array x2. + + .. note:: + Floating-point multiplication is not always associative due to finite precision. **Special Cases** For real-valued floating-point operands, - - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, - the result is ``True``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, - the result is ``True``. - - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, - and ``x1_i`` does not equal ``x2_i``, the result is ``True``. - - In the remaining cases, the result is ``False``. + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+0`` or ``-0``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` and ``x2_i`` have the same mathematical sign, + the result has a positive mathematical sign, unless the result is ``NaN``. + If the result is ``NaN``, the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` and ``x2_i`` have different mathematical signs, + the result has a negative mathematical sign, + unless the result is ``NaN``. If the result is ``NaN``, + the "sign" of ``NaN`` is implementation-defined. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and + ``x2_i`` is either ``+infinity`` or ``-infinity``, + the result is a signed infinity with the mathematical sign determined by + the rule already stated above. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` + is a nonzero finite number, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - If ``x1_i`` is a nonzero finite number and ``x2_i`` + is either ``+infinity`` or ``-infinity``, the result is a signed infinity with + the mathematical sign determined by the rule already stated above. + - In the remaining cases, where neither ``infinity`` nor ``NaN`` + is involved, the product must be computed and rounded to the nearest + representable value according to IEEE 754-2019 and a supported + rounding mode. If the magnitude is too large to represent, + the result is an `infinity` of appropriate mathematical sign. + If the magnitude is too small to represent, the result is a zero of + appropriate mathematical sign. - For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, - ``c = real(x2_i)``, ``d = imag(x2_i)``, and + For complex floating-point operands, multiplication is defined according to the + following table. For real components ``a`` and ``c`` and + imaginary components ``b`` and ``d``, - - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. - - In the remaining cases, the result is the logical OR of - the equality comparison between the real values ``a`` and ``c`` - (real components) and between the real values ``b`` and ``d`` - (imaginary components), as described above for real-valued floating-point operands - (i.e., ``a != c OR b != d``). + +------------+----------------+-----------------+--------------------------+ + | | c | dj | c + dj | + +============+================+=================+==========================+ + | **a** | a * c | (a*d)j | (a*c) + (a*d)j | + +------------+----------------+-----------------+--------------------------+ + | **bj** | (b*c)j | -(b*d) | -(b*d) + (b*c)j | + +------------+----------------+-----------------+--------------------------+ + | **a + bj** | (a*c) + (b*c)j | -(b*d) + (a*d)j | special rules | + +------------+----------------+-----------------+--------------------------+ - Parameters - ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + In general, for complex floating-point operands, real-valued floating-point + special cases must independently apply to the real and imaginary component + operations involving real numbers as described in the above table. - Returns - ------- - ret - an array containing the element-wise results. The returned array must have a - data type of ``bool``. + When ``a``, ``b``, ``c``, or ``d`` are all finite numbers + (i.e., a value other than ``NaN``, ``+infinity``, or ``-infinity``), + multiplication of complex floating-point operands should be computed + as if calculated according to the textbook formula for complex number multiplication + .. math:: + (a + bj) \cdot (c + dj) = (ac - bd) + (bc + ad)j - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + When at least one of ``a``, ``b``, ``c``, or ``d`` is ``NaN``, + ``+infinity``, or ``-infinity``, - Functional Examples - ------------------ + - If ``a``, ``b``, ``c``, and ``d`` are all ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, the result is implementation dependent. - With :class:`ivy.Array` inputs: + .. note:: + For complex floating-point operands, the results of special cases may be + implementation dependent depending on how an implementation chooses + to model complex numbers and complex infinity + (e.g., complex plane versus Riemann sphere). + For those implementations following C99 and its one-infinity model, + when at least one component is infinite, + even if the other component is ``NaN``, + the complex value is infinite, and the usual arithmetic + rules do not apply to complex-complex multiplication. + In the interest of performance, other implementations + may want to avoid the complex branching logic necessary + to implement the one-infinity model and choose to implement + all complex-complex multiplication according to the textbook formula. + Accordingly, special case behavior is unlikely + to be consistent across implementations. - >>> x1 = ivy.array([1, 0, 1, 1]) - >>> x2 = ivy.array([1, 0, 0, -1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False, True, True]) + Parameters + ---------- + x1 + first input array. Should have a numeric data type. - >>> x1 = ivy.array([1, 0, 1, 0]) - >>> x2 = ivy.array([0, 1, 0, 1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([True, True, True, True]) + x2 + second input array. Must be compatible with ``x1`` + (see :ref'`broadcasting`). Should have a numeric data type - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) + out + optional output array, for writing the array result to. + It must have a shape that the inputs broadcast to. - >>> x1 = ivy.array([1, -1, 1, -1]) - >>> x2 = ivy.array([0, -1, 1, 0]) - >>> y = ivy.not_equal(x1, x2, out=x1) - >>> print(y) - ivy.array([1, 0, 0, 1]) - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> x1 = ivy.native_array([1, 2]) - >>> x2 = ivy.array([1, 2]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([False, False]) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x1 = ivy.native_array([1, -1]) - >>> x2 = ivy.array([0, 1]) - >>> y = ivy.not_equal(x1, x2) - >>> print(y) - ivy.array([True, True]) + Returns + ------- + ret + an array containing the element-wise products. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. - >>> x1 = ivy.native_array([1, -1, 1, -1]) - >>> x2 = ivy.native_array([0, -1, 1, 0]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) - >>> print(y) - ivy.array([1., 0., 0., 1.]) + Examples + -------- + With :code:`ivy.Array` inputs: - >>> x1 = ivy.native_array([1, 2, 3, 4]) - >>> x2 = ivy.native_array([0, 2, 3, 4]) - >>> y = ivy.zeros(4) - >>> ivy.not_equal(x1, x2, out=y) + >>> x1 = ivy.array([3., 5., 7.]) + >>> x2 = ivy.array([4., 6., 8.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - ivy.array([1., 0., 0., 0.]) + ivy.array([12., 30., 56.]) - With :class:`ivy.Container` input: + With :code:`ivy.NativeArray` inputs: - >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.native_array([1., 3., 9.]) + >>> x2 = ivy.native_array([4., 7.2, 1.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - { - a: ivy.array([False, True, False]), - b: ivy.array([False, False, False]), - c: ivy.array([False, False, False]) - } + ivy.array([ 4. , 21.6, 9. ]) - >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), - ... b=ivy.array([1, 2, 3]), - ... c=ivy.native_array([1.0, 2.0, 4.0])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.native_array([1.1, 2.1, 3.1]), - ... c=ivy.native_array([1, 2, 4])) - >>> y = ivy.not_equal(x1, x2) + With mixed :code:`ivy.Array` and :code:`ivy.NativeArray` inputs: + + >>> x1 = ivy.array([8., 6., 7.]) + >>> x2 = ivy.native_array([1., 2., 3.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) - { - a: ivy.array([True, True, True]), - b: ivy.array([True, True, True]), - c: ivy.array([False, False, False]) - } + ivy.array([ 8., 12., 21.]) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With :code:`ivy.Container` inputs: - >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 3, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), - ... b=ivy.array([1, 4, 5])) - >>> y = ivy.not_equal(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([12.,4.,6.]), b=ivy.array([3.,1.,5.])) + >>> x2 = ivy.Container(a=ivy.array([1.,3.,4.]), b=ivy.array([3.,3.,2.])) + >>> y = ivy.multiply(x1, x2) >>> print(y) { - a: ivy.array([False, False, False]), - b: ivy.array([False, True, False]) + a: ivy.array([12.,12.,24.]), + b: ivy.array([9.,3.,10.]) } - >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), - ... b=ivy.array([1, 4, 5])) - >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), - ... b=ivy.array([1.0, 4.0, 5.0])) - >>> y = ivy.not_equal(x1, x2) + With mixed :code:`ivy.Container` and :code:`ivy.Array` inputs: + + >>> x1 = ivy.Container(a=ivy.array([3., 4., 5.]), b=ivy.array([2., 2., 1.])) + >>> x2 = ivy.array([1.,2.,3.]) + >>> y = ivy.multiply(x1, x2) >>> print(y) { - a: ivy.array([False, False, False]), - b: ivy.array([False, False, False]) + a: ivy.array([3.,8.,15.]), + b: ivy.array([2.,4.,3.]) } """ - return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) + return ivy.current_backend(x1, x2).multiply(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def positive( - x: Union[float, ivy.Array, ivy.NativeArray], +def nan_to_num( + x: Union[ivy.Array, ivy.NativeArray], /, *, + copy: bool = True, + nan: Union[float, int] = 0.0, + posinf: Optional[Union[float, int]] = None, + neginf: Optional[Union[float, int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a new array with the positive value of each element in ``x``. + Replace NaN with zero and infinity with large finite numbers (default behaviour) or + with the numbers defined by the user using the nan, posinf and/or neginf keywords. Parameters ---------- x - Input array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - A new array with the positive value of each element in ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: - - >>> x = ivy.array([2, 3 ,5, 7]) - >>> y = ivy.positive(x) - >>> print(y) - ivy.array([2, 3, 5, 7]) - - >>> x = ivy.array([0, -1, -0.5, 2, 3]) - >>> y = ivy.zeros(5) - >>> ivy.positive(x, out=y) - >>> print(y) - ivy.array([0., -1., -0.5, 2., 3.]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.positive(x,out=x) - >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-4.4, -5.5, -6.6]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., -5.])) - >>> y = ivy.positive(x) - >>> print(y) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., -5.]) - } - """ - return ivy.current_backend(x).positive(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def pow( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Calculate an implementation-dependent approximation of exponentiation by raising - each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` - (the exponent), where ``x2_i`` is the corresponding element of the input array - ``x2``. - - .. note:: - If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when - ``x2_i`` is negative (i.e., less than zero) is unspecified and thus - implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a - floating-point data type, behavior is implementation-dependent (type promotion - between data type "kinds" (integer versus floating-point) is unspecified). - - **Special cases** - - For floating-point operands, - - - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. - - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result - is ``+infinity``. - - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result - is ``+0``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. - - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. - - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is - ``+0``. - - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is - ``+0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an - odd integer value, the result is ``-infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not - an odd integer value, the result is ``+infinity``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an - odd integer value, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is - ``+infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd - integer value, the result is ``-0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+0``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer - value, the result is ``-infinity``. - - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd - integer value, the result is ``+infinity``. - - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite - number, and ``x2_i`` is not an integer value, the result is ``NaN``. - - For complex floating-point operands, special cases should be handled - as if the operation is implemented as ``exp(x2*log(x1))``. - - .. note:: - Conforming implementations are allowed to treat special cases involving - complex floating-point operands more carefully than as described - in this specification. - - Parameters - ---------- - x1 - first input array whose elements correspond to the exponentiation base. Should - have a numeric data type. - x2 - second input array whose elements correspond to the exponentiation exponent. - Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric - data type. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the element-wise results. The returned array must have a - data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.pow(x, 3) - >>> print(y) - ivy.array([1, 8, 27]) - - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.pow(x, 2, out=y) - >>> print(y) - ivy.array([2.25, 0.64, 0.09]) - - >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) - >>> ivy.pow(x, 2.3, out=x) - >>> print(x) - ivy.array([[ 1.52095687, 4.92457771, 13.49372482], - [ 1. , 8.22738838, 156.5877228 ]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.pow(x, 3) - >>> print(y) - { - a:ivy.array([0,1]), - b:ivy.array([8,27]) - } - """ - return ivy.current_backend(x1, x2).pow(x1, x2, out=out) + Array input. + copy + Whether to create a copy of x (True) or to replace values in-place (False). + The in-place operation only occurs if casting to an array does not require + a copy. Default is True. + nan + Value to be used to fill NaN values. If no value is passed then NaN values + will be replaced with 0.0. + posinf + Value to be used to fill positive infinity values. If no value is passed + then positive infinity values will be replaced with a very large number. + neginf + Value to be used to fill negative infinity values. + If no value is passed then negative infinity values + will be replaced with a very small (or negative) number. + out + optional output array, for writing the result to. + Returns + ------- + ret + Array with the non-finite values replaced. + If copy is False, this may be x itself. -pow.unsupported_gradients = {"torch": ["float16"]} + Examples + -------- + >>> x = ivy.array([1, 2, 3, nan]) + >>> ivy.nan_to_num(x) + ivy.array([1., 1., 3., 0.0]) + >>> x = ivy.array([1, 2, 3, inf]) + >>> ivy.nan_to_num(x, posinf=5e+100) + ivy.array([1., 2., 3., 5e+100]) + """ + return ivy.current_backend(x).nan_to_num( + x, copy=copy, nan=nan, posinf=posinf, neginf=neginf, out=out + ) @handle_exceptions @@ -5342,23 +5305,31 @@ def pow( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def real( - x: Union[ivy.Array, ivy.NativeArray], +def negative( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Test each element ``x_i`` of the input array ``x`` to take only real part from it. - Returns a float array, where it only contains . If element has complex type with - zero complex part, the return value will be that element, else it only returns real - part. + Return a new array with the negative value of each element in ``x``. + + .. note:: + For signed integer data types, the numerical negative of + the minimum representable integer is implementation-dependent. + + .. note:: + If ``x`` has a complex floating-point data type, + both the real and imaginary components for each ``x_i`` + must be negated (a result which follows from the rules of + complex number multiplication). Parameters ---------- x - input array. + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5366,45 +5337,53 @@ def real( Returns ------- ret - an array containing test results. An element ``out_i`` is - ``real number`` if ``x_i`` contain real number part only - and if it is ``real number with complex part also`` then it - returns the real number part. - The returned array must have a floating-point data type with the - same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, - the returned array must have the floating-point precision of ``float32``). + A new array with the negative value of each element in ``x``. - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input: - >>> x = ivy.array([[[1.1], [2], [-6.3]]]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([[[1.1], [2.], [-6.3]]]) + >>> x = ivy.array([0,1,1,2]) + >>> y = ivy.negative(x) + >>> print(y) + ivy.array([ 0, -1, -1, -2]) - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.real(x) - >>> print(z) - ivy.array([4.2, 0., 7.]) + >>> x = ivy.array([0,-1,-0.5,2,3]) + >>> y = ivy.zeros(5) + >>> ivy.negative(x, out=y) + >>> print(y) + ivy.array([-0. , 1. , 0.5, -2. , -3. ]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.negative(x,out=x) + >>> print(x) + ivy.array([[-1.1, -2.2, -3.3], + [4.4, 5.5, 6.6]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ - b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.real(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.negative(x) + >>> print(y) { - a: ivy.array([-6.7, 0.314, 1.23]), - b: ivy.array([0., 5.32, 3.001]) + a: ivy.array([-0., -1., -2.]), + b: ivy.array([-3., -4., 5.]) } """ - return ivy.current_backend(x).real(x, out=out) + return ivy.current_backend(x).negative(x, out=out) @handle_exceptions @@ -5414,75 +5393,47 @@ def real( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def remainder( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def not_equal( + x1: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], + x2: Union[float, ivy.Array, ivy.NativeArray, ivy.Container], /, *, - modulus: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the remainder of division for each element ``x1_i`` of the input array ``x1`` - and the respective element ``x2_i`` of the input array ``x2``. - - .. note:: - This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For - input arrays which promote to an integer data type, the result of division by - zero is unspecified and thus implementation-defined. In general, similar to - Python’s ``%`` operator, this function is not recommended for floating-point - operands as semantics do not follow IEEE 754. That this function is specified - to accept floating-point operands is primarily for reasons of backward - compatibility. + Compute the truth value of ``x1_i != x2_i`` for each element ``x1_i`` of the input + array ``x1`` with the respective element ``x2_i`` of the input array ``x2``. **Special Cases** - For floating-point operands, + For real-valued floating-point operands, - - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. - - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either - ``+infinity`` or ``-infinity``, the result is ``NaN``. - - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, - the result is ``NaN``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. - - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. - - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) - finite number, the result is ``NaN``. - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x2_i``. (note: this result matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``+infinity``, the result is ``x2_i``. (note: this results matches Python - behavior.) - - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is - ``-infinity``, the result is ``x1_i``. (note: this result matches Python - behavior.) - - In the remaining cases, the result must match that of the Python ``%`` operator. + - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``True``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``-infinity``, + the result is ``True``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``+infinity``, + the result is ``True``. + - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, + and ``x1_i`` does not equal ``x2_i``, the result is ``True``. + - In the remaining cases, the result is ``False``. + + For omplex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, + ``c = real(x2_i)``, ``d = imag(x2_i)``, and + + - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``True``. + - In the remaining cases, the result is the logical OR of + the equality comparison between the real values ``a`` and ``c`` + (real components) and between the real values ``b`` and ``d`` + (imaginary components), as described above for real-valued floating-point operands + (i.e., ``a != c OR b != d``). Parameters ---------- x1 - dividend input array. Should have a numeric data type. + first input array. Should have a numeric data type. x2 - divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). Should have a numeric data type. - modulus - whether to compute the modulus instead of the remainder. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5490,51 +5441,133 @@ def remainder( Returns ------- ret - an array containing the element-wise results. Each element-wise result must have - the same sign as the respective element ``x2_i``. The returned array must have a - data type determined by :ref:`Type Promotion Rules`. + an array containing the element-wise results. The returned array must have a + data type of ``bool``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.not_equal.html>`_ in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x1 = ivy.array([1, 0, 1, 1]) + >>> x2 = ivy.array([1, 0, 0, -1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False, True, True]) + + >>> x1 = ivy.array([1, 0, 1, 0]) + >>> x2 = ivy.array([0, 1, 0, 1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([True, True, True, True]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) + + >>> x1 = ivy.array([1, -1, 1, -1]) + >>> x2 = ivy.array([0, -1, 1, 0]) + >>> y = ivy.not_equal(x1, x2, out=x1) + >>> print(y) + ivy.array([1, 0, 0, 1]) + + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + + >>> x1 = ivy.native_array([1, 2]) + >>> x2 = ivy.array([1, 2]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([False, False]) + + >>> x1 = ivy.native_array([1, -1]) + >>> x2 = ivy.array([0, 1]) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + ivy.array([True, True]) + + >>> x1 = ivy.native_array([1, -1, 1, -1]) + >>> x2 = ivy.native_array([0, -1, 1, 0]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 1.]) + + >>> x1 = ivy.native_array([1, 2, 3, 4]) + >>> x2 = ivy.native_array([0, 2, 3, 4]) + >>> y = ivy.zeros(4) + >>> ivy.not_equal(x1, x2, out=y) + >>> print(y) + ivy.array([1., 0., 0., 0.]) + + With :class:`ivy.Container` input: - Examples - -------- - With :class:`ivy.Array` inputs: + >>> x1 = ivy.Container(a=ivy.array([1, 0, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) + >>> print(y) + { + a: ivy.array([False, True, False]), + b: ivy.array([False, False, False]), + c: ivy.array([False, False, False]) + } - >>> x1 = ivy.array([2., 5., 15.]) - >>> x2 = ivy.array([3., 2., 4.]) - >>> y = ivy.remainder(x1, x2) + >>> x1 = ivy.Container(a=ivy.native_array([0, 1, 0]), + ... b=ivy.array([1, 2, 3]), + ... c=ivy.native_array([1.0, 2.0, 4.0])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.native_array([1.1, 2.1, 3.1]), + ... c=ivy.native_array([1, 2, 4])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) - ivy.array([2., 1., 3.]) + { + a: ivy.array([True, True, True]), + b: ivy.array([True, True, True]), + c: ivy.array([False, False, False]) + } - With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x1 = ivy.array([23., 1., 6.]) - >>> x2 = ivy.native_array([11., 2., 4.]) - >>> y = ivy.remainder(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 3, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3]), + ... b=ivy.array([1, 4, 5])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) - ivy.array([1., 1., 2.]) - - With :class:`ivy.Container` inputs: + { + a: ivy.array([False, False, False]), + b: ivy.array([False, True, False]) + } - >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) - >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) - >>> y = ivy.remainder(x1, x2) + >>> x1 = ivy.Container(a=ivy.array([1.0, 2.0, 3.0]), + ... b=ivy.array([1, 4, 5])) + >>> x2 = ivy.Container(a=ivy.array([1, 2, 3.0]), + ... b=ivy.array([1.0, 4.0, 5.0])) + >>> y = ivy.not_equal(x1, x2) >>> print(y) { - a: ivy.array([0., 0., 1.]), - b: ivy.array([0., 2., 1.]) + a: ivy.array([False, False, False]), + b: ivy.array([False, False, False]) } """ - return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) + return ivy.current_backend(x1, x2).not_equal(x1, x2, out=out) @handle_exceptions @@ -5545,55 +5578,19 @@ def remainder( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def round( - x: Union[ivy.Array, ivy.NativeArray], +def positive( + x: Union[float, ivy.Array, ivy.NativeArray], /, *, - decimals: Optional[int] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued - number. - - .. note:: - For complex floating-point operands, real and imaginary components - must be independently rounded to the nearest integer-valued number. - - Rounded real and imaginary components must be equal - to their equivalent rounded real-valued floating-point - counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` - must equal ``round(real(x)))`` and ``imag(round(x))`` must equal - ``round(imag(x))``). - - **Special cases** - - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - - For floating-point operands, - - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If two integers are equally close to ``x_i``, the result is - the even integer closest to ``x_i``. - - .. note:: - For complex floating-point operands, the following special - cases apply to real and imaginary components independently - (e.g., if ``real(x_i)`` is ``NaN``, the rounded - real component is ``NaN``). - - - If ``x_i`` is already integer-valued, the result is ``x_i``. + Return a new array with the positive value of each element in ``x``. Parameters ---------- x - input array containing elements to round. - decimals - number of decimal places to round to. Default is ``0``. + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5601,121 +5598,145 @@ def round( Returns ------- ret - An array of the same shape and type as x, with the elements rounded to integers. - + A new array with the positive value of each element in ``x``. - Note: PyTorch supports an additional argument :code:`decimals` for the - `round function `_. - It has been deliberately omitted here due to the imprecise - nature of the argument in :code:`torch.round`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.positive.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments - Examples - -------- - With :class:`ivy.Array` input: + Functional Examples + ------------------- - >>> x = ivy.array([1.2, 2.4, 3.6]) - >>> y = ivy.round(x) - >>> print(y) - ivy.array([1.,2.,4.]) + With :class:`ivy.Array` input: - >>> x = ivy.array([-0, 5, 4.5]) - >>> y = ivy.round(x) + >>> x = ivy.array([2, 3 ,5, 7]) + >>> y = ivy.positive(x) >>> print(y) - ivy.array([0.,5.,4.]) + ivy.array([2, 3, 5, 7]) - >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) - >>> y = ivy.zeros(4) - >>> ivy.round(x, out=y) + >>> x = ivy.array([0, -1, -0.5, 2, 3]) + >>> y = ivy.zeros(5) + >>> ivy.positive(x, out=y) >>> print(y) - ivy.array([2.,2.,15.,-5.]) + ivy.array([0., -1., -0.5, 2., 3.]) - >>> x = ivy.array([[0, 5.433, -343.3, 1.5], - ... [-5.5, 44.2, 11.5, 12.01]]) - >>> ivy.round(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.positive(x,out=x) >>> print(x) - ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) + ivy.array([[ 1.1, 2.2, 3.3], + [-4.4, -5.5, -6.6]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), - ... b=ivy.array([-300.9, -527.3, 4.5])) - >>> y = ivy.round(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., -5.])) + >>> y = ivy.positive(x) >>> print(y) { - a:ivy.array([4.,9.,7.,0.]), - b:ivy.array([-301.,-527.,4.]) + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., -5.]) } """ - return ivy.current_backend(x).round(x, decimals=decimals, out=out) + return ivy.current_backend(x).positive(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sign( - x: Union[ivy.Array, ivy.NativeArray], +def pow( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, - np_variant: Optional[bool] = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Return an indication of the sign of a number for each element ``x_i`` of the input - array ``x``. - - The sign function (also known as the **signum function**) - of a number :math:`x_{i}` is defined as - - .. math:: - \operatorname{sign}(x_i) = \begin{cases} - 0 & \textrm{if } x_i = 0 \\ - \frac{x}{|x|} & \textrm{otherwise} - \end{cases} + """ + Calculate an implementation-dependent approximation of exponentiation by raising + each element ``x1_i`` (the base) of the input array ``x1`` to the power of ``x2_i`` + (the exponent), where ``x2_i`` is the corresponding element of the input array + ``x2``. - where :math:`|x_i|` is the absolute value of :math:`x_i`. + .. note:: + If both ``x1`` and ``x2`` have integer data types, the result of ``pow`` when + ``x2_i`` is negative (i.e., less than zero) is unspecified and thus + implementation-dependent. If ``x1`` has an integer data type and ``x2`` has a + floating-point data type, behavior is implementation-dependent (type promotion + between data type "kinds" (integer versus floating-point) is unspecified). **Special cases** - - If ``x_i`` is less than ``0``, the result is ``-1``. - - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. - - If ``x_i`` is greater than ``0``, the result is ``+1``. - - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` + For floating-point operands, - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and + - If ``x1_i`` is not equal to ``1`` and ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x2_i`` is ``+0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x2_i`` is ``-0``, the result is ``1``, even if ``x1_i`` is ``NaN``. + - If ``x1_i`` is ``NaN`` and ``x2_i`` is not equal to ``0``, the result is ``NaN``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``+infinity``, the result + is ``+infinity``. + - If ``abs(x1_i)`` is greater than ``1`` and ``x2_i`` is ``-infinity``, the result + is ``+0``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``+infinity``, the result is ``1``. + - If ``abs(x1_i)`` is ``1`` and ``x2_i`` is ``-infinity``, the result is ``1``. + - If ``x1_i`` is ``1`` and ``x2_i`` is not ``NaN``, the result is ``1``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``+infinity``, the result is + ``+0``. + - If ``abs(x1_i)`` is less than ``1`` and ``x2_i`` is ``-infinity``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is greater than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is less than ``0``, the result is + ``+0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an + odd integer value, the result is ``-infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not + an odd integer value, the result is ``+infinity``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-infinity``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an + odd integer value, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is + ``+infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is an odd + integer value, the result is ``-0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is greater than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+0``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is an odd integer + value, the result is ``-infinity``. + - If ``x1_i`` is ``-0``, ``x2_i`` is less than ``0``, and ``x2_i`` is not an odd + integer value, the result is ``+infinity``. + - If ``x1_i`` is less than ``0``, ``x1_i`` is a finite number, ``x2_i`` is a finite + number, and ``x2_i`` is not an integer value, the result is ``NaN``. - - If ``a`` is either ``-0`` or ``+0`` and ``b`` is - either ``-0`` or ``+0``, the result is ``0 + 0j``. - - If ``a`` is ``NaN`` or ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - In the remaining cases, special cases must be handled - according to the rules of complex number division. + For complex floating-point operands, special cases should be handled + as if the operation is implemented as ``exp(x2*log(x1))``. + + .. note:: + Conforming implementations are allowed to treat special cases involving + complex floating-point operands more carefully than as described + in this specification. Parameters ---------- - x - input array. Should have a numeric data type. - np_variant - Handles complex numbers like numpy does If ``True``, - ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. - otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, - otherwise y = 0.`` - + x1 + first input array whose elements correspond to the exponentiation base. Should + have a numeric data type. + x2 + second input array whose elements correspond to the exponentiation exponent. + Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric + data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5723,53 +5744,54 @@ def sign( Returns ------- ret - an array containing the evaluated result for each element in ``x``. The returned - array must have the same data type as ``x``. + an array containing the element-wise results. The returned array must have a + data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.pow.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([8.3, -0, 6.8, 0.07]) - >>> y = ivy.sign(x) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.pow(x, 3) >>> print(y) - ivy.array([1., 0., 1., 1.]) + ivy.array([1, 8, 27]) - >>> x = ivy.array([[5.78, -4., -6.9, 0], - ... [-.4, 0.5, 8, -0.01]]) - >>> y = ivy.sign(x) + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.pow(x, 2, out=y) >>> print(y) - ivy.array([[ 1., -1., -1., 0.], - [-1., 1., 1., -1.]]) + ivy.array([2.25, 0.64, 0.09]) + + >>> x = ivy.array([[1.2, 2, 3.1], [1, 2.5, 9]]) + >>> ivy.pow(x, 2.3, out=x) + >>> print(x) + ivy.array([[ 1.52095687, 4.92457771, 13.49372482], + [ 1. , 8.22738838, 156.5877228 ]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., -0.]), - ... b=ivy.array([1.46, 5.9, -0.0]), - ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) - >>> y = ivy.sign(x) + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.pow(x, 3) >>> print(y) { - a: ivy.array([0., 0.]), - b: ivy.array([1., 1., 0.]), - c: ivy.array([-1., -1., -1., 1.]) + a:ivy.array([0,1]), + b:ivy.array([8,27]) } """ - return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) + return ivy.current_backend(x1, x2).pow(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -5777,37 +5799,19 @@ def sign( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sin( +def rad2deg( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the sine, having domain - ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of - the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. - - .. note:: - The sine is an entire function on the complex plane and has no branch cuts. - - **Special cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - - For complex floating-point operands, special cases - must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. + """ + Convert the input from radians to degrees. Parameters ---------- x - input array whose elements are each expressed in radians. Should have a - floating-point data type. + input array whose elements are each expressed in radians. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5815,53 +5819,56 @@ def sin( Returns ------- ret - an array containing the sine of each element in ``x``. The returned array must - have a floating-point data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + an array with each element in ``x`` converted from radians to degrees. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.sin(x) + >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) + >>> y=ivy.rad2deg(x) >>> print(y) - ivy.array([0., 0.841, 0.909]) + ivy.array([ 0., 90., 180., 270., 360.]) - >>> x = ivy.array([0., 1.2, -2.3, 3.6]) - >>> y = ivy.zeros(4) - >>> ivy.sin(x, out=y) + >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) + >>> y=ivy.zeros(4) + >>> ivy.rad2deg(x,out=y) >>> print(y) - ivy.array([0., 0.932, -0.746, -0.443]) + ivy.array([ 0. , -1.5, -50. , nan]) - >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) - >>> ivy.sin(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.rad2deg(x, out=x) >>> print(x) - ivy.array([[0.841, 0.909, 0.141], - [0.757, 0.959, 0.279]]) + ivy.array([[ 63., 126., 189.], + [-252., -315., -378.]]) + + >>> x=ivy.native_array([-0,20.1,ivy.nan]) + >>> y=ivy.zeros(3) + >>> ivy.rad2deg(x,out=y) + >>> print(y) + ivy.array([ 0., 1150., nan]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), - ... b=ivy.array([-4., -5., -6., -7.])) - >>> y = ivy.sin(x) + >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), + ... b=ivy.array([0., 1., 2., 3., 4.])) + >>> y=ivy.rad2deg(x) >>> print(y) { - a: ivy.array([0., 0.841, 0.909, 0.141]), - b: ivy.array([0.757, 0.959, 0.279, -0.657]) + a: ivy.array([0., 1150., -2890., nan]), + b: ivy.array([0., 57.3, 115., 172., 229.]) + } + + >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), + ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) + >>> y=ivy.rad2deg(x) + >>> print(y) + { + a: ivy.array([0., 573., 10300., 487., 344.]), + b: ivy.array([0., -85.9, 28.6, nan]) } """ - return ivy.current_backend(x).sin(x, out=out) + return ivy.current_backend(x).rad2deg(x, out=out) @handle_exceptions @@ -5870,76 +5877,92 @@ def sin( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def sinh( +def real( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate an implementation-dependent approximation to the hyperbolic sine, having - domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each - element ``x_i`` of the input array ``x``. + """ + Test each element ``x_i`` of the input array ``x`` to take only real part from it. + Returns a float array, where it only contains . If element has complex type with + zero complex part, the return value will be that element, else it only returns real + part. - .. math:: - \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} + Parameters + ---------- + x + input array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - .. note:: - The hyperbolic sine is an entire function in the - complex plane and has no branch cuts. - The function is periodic, with period - :math:`2\pi j`, with respect to the imaginary component. + Returns + ------- + ret + an array containing test results. An element ``out_i`` is + ``real number`` if ``x_i`` contain real number part only + and if it is ``real number with complex part also`` then it + returns the real number part. + The returned array must have a floating-point data type with the + same floating-point precision as ``x`` (e.g., if ``x`` is ``complex64``, + the returned array must have the floating-point precision of ``float32``). - **Special cases** + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. - For floating-point operands, + Examples + -------- + With :class:`ivy.Array` inputs: - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + >>> x = ivy.array([[[1.1], [2], [-6.3]]]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([[[1.1], [2.], [-6.3]]]) - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.real(x) + >>> print(z) + ivy.array([4.2, 0., 7.]) - .. note:: - For complex floating-point operands, ``sinh(conj(x))`` - must equal ``conj(sinh(x))``. + With :class:`ivy.Container` input: - - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. - - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). - - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``0 + NaN j`` (sign of the real component is unspecified). - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. - - If ``a`` is a positive (i.e., greater than ``0``) - finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is ``+0``, - the result is ``+infinity + 0j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, - the result is ``+infinity * cis(b)``. - - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``infinity + NaN j`` (sign of the real component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``infinity + NaN j`` - (sign of the real component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]),\ + b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.real(x) + >>> print(z) + { + a: ivy.array([-6.7, 0.314, 1.23]), + b: ivy.array([0., 5.32, 3.001]) + } + """ + return ivy.current_backend(x).real(x, out=out) - where ``cis(v)`` is ``cos(v) + sin(v)*1j``. + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def reciprocal( + x: Union[float, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a new array with the reciprocal of each element in ``x``. Parameters ---------- x - input array whose elements each represent a hyperbolic angle. Should have a - floating-point data type. + Input array. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -5947,126 +5970,94 @@ def sinh( Returns ------- ret - an array containing the hyperbolic sine of each element in ``x``. The returned - array must have a floating-point data type determined by :ref:`type-promotion`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments + A new array with the positive value of each element in ``x``. Examples -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.sinh(x) - >>> print(y) - ivy.array([1.18, 3.63, 10.]) - - >>> x = ivy.array([0.23, 3., -1.2]) - >>> ivy.sinh(x, out=x) - >>> print(x) - ivy.array([0.232, 10., -1.51]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) - >>> y = ivy.sinh(x) + >>> y = ivy.reciprocal(x) >>> print(y) - { - a: ivy.array([0.232, -0.253, 1.18]), - b: ivy.array([10., -27.3, 1.62]) - } + ivy.array([1. , 0.5 , 0.33333333]) """ - return ivy.current_backend(x).sinh(x, out=out) + return ivy.current_backend(x).reciprocal(x, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sqrt( - x: Union[ivy.Array, ivy.NativeArray], +def remainder( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, + modulus: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, - +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, - each result must be indistinguishable from the infinitely precise result (as - required by IEEE 754). - - .. note:: - After rounding, each result must be indistinguishable - from the infinitely precise result (as required by IEEE 754). - - .. note:: - For complex floating-point operands, ``sqrt(conj(x))`` - must equal ``conj(sqrt(x))``. + """ + Return the remainder of division for each element ``x1_i`` of the input array ``x1`` + and the respective element ``x2_i`` of the input array ``x2``. .. note:: - By convention, the branch cut of the square root is - the negative real axis :math:`(-\infty, 0)`. - - The square root is a continuous function from above - the branch cut, taking into account the sign of the imaginary component. - - Accordingly, for complex arguments, the function returns - the square root in the range of the right half-plane, - including the imaginary axis (i.e., the plane defined by - :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` - along the imaginary axis). + This function is equivalent to the Python modulus operator ``x1_i % x2_i``. For + input arrays which promote to an integer data type, the result of division by + zero is unspecified and thus implementation-defined. In general, similar to + Python’s ``%`` operator, this function is not recommended for floating-point + operands as semantics do not follow IEEE 754. That this function is specified + to accept floating-point operands is primarily for reasons of backward + compatibility. - **Special cases** + **Special Cases** For floating-point operands, - - If ``x_i`` is ``NaN``, the result is ``NaN``. - - If ``x_i`` is less than ``0``, the result is ``NaN``. - - If ``x_i`` is ``+0``, the result is ``+0``. - - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - For complex floating-point operands, let ``a = real(x_i)``, - ``b = imag(x_i)``, and - - - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, - the result is ``+0 + 0j``. - - If ``a`` is any value (including ``NaN``) and ``b`` is - ``+infinity``, the result is ``+infinity + infinity j``. - - If ``a`` is a finite number and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - - If ``a`` ``-infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. - - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, - the result is ``NaN + infinity j`` - (sign of the imaginary component is unspecified). - - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``+infinity + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is any value, - the result is ``NaN + NaN j``. - - If ``a`` is ``NaN`` and ``b`` is ``NaN``, - the result is ``NaN + NaN j``. - + - If either ``x1_i`` or ``x2_i`` is ``NaN``, the result is ``NaN``. + - If ``x1_i`` is either ``+infinity`` or ``-infinity`` and ``x2_i`` is either + ``+infinity`` or ``-infinity``, the result is ``NaN``. + - If ``x1_i`` is either ``+0`` or ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, + the result is ``NaN``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is greater than ``0``, the result is ``+0``. + - If ``x1_i`` is ``+0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is ``-0`` and ``x2_i`` is less than ``0``, the result is ``-0``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is greater than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``+0``, the result is ``NaN``. + - If ``x1_i`` is less than ``0`` and ``x2_i`` is ``-0``, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``+infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a positive (i.e., greater than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is ``-infinity`` and ``x2_i`` is a negative (i.e., less than ``0``) + finite number, the result is ``NaN``. + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a positive (i.e., greater than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x2_i``. (note: this result matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``+infinity``, the result is ``x2_i``. (note: this results matches Python + behavior.) + - If ``x1_i`` is a negative (i.e., less than ``0``) finite number and ``x2_i`` is + ``-infinity``, the result is ``x1_i``. (note: this result matches Python + behavior.) + - In the remaining cases, the result must match that of the Python ``%`` operator. Parameters ---------- - x - input array. Should have a floating-point data type. + x1 + dividend input array. Should have a numeric data type. + x2 + divisor input array. Must be compatible with ``x1`` (see ref:`Broadcasting`). + Should have a numeric data type. + modulus + whether to compute the modulus instead of the remainder. Default is ``True``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6074,14 +6065,15 @@ def sqrt( Returns ------- ret - an array containing the square root of each element in ``x``. The returned array - must have a floating-point data type determined by :ref:`type-promotion`. + an array containing the element-wise results. Each element-wise result must have + the same sign as the respective element ``x2_i``. The returned array must have a + data type determined by :ref:`Type Promotion Rules`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.remainder.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6090,57 +6082,93 @@ def sqrt( Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.array([0, 4., 8.]) - >>> y = ivy.sqrt(x) + >>> x1 = ivy.array([2., 5., 15.]) + >>> x2 = ivy.array([3., 2., 4.]) + >>> y = ivy.remainder(x1, x2) >>> print(y) - ivy.array([0., 2., 2.83]) + ivy.array([2., 1., 3.]) - >>> x = ivy.array([1, 2., 4.]) - >>> y = ivy.zeros(3) - >>> ivy.sqrt(x, out=y) - ivy.array([1., 1.41, 2.]) + With mixed :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - >>> X = ivy.array([40., 24., 100.]) - >>> ivy.sqrt(x, out=x) - >>> ivy.array([6.32455532, 4.89897949, 10.]) + >>> x1 = ivy.array([23., 1., 6.]) + >>> x2 = ivy.native_array([11., 2., 4.]) + >>> y = ivy.remainder(x1, x2) + >>> print(y) + ivy.array([1., 1., 2.]) - With :class:`ivy.Container` input: + With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa - >>> y = ivy.sqrt(x) + >>> x1 = ivy.Container(a=ivy.array([2., 3., 5.]), b=ivy.array([2., 2., 4.])) + >>> x2 = ivy.Container(a=ivy.array([1., 3., 4.]), b=ivy.array([1., 3., 3.])) + >>> y = ivy.remainder(x1, x2) >>> print(y) { - a: ivy.array([6.63, 7.48, 13.]), - b: ivy.array([[7., 1.], - [0., 4.47]]) + a: ivy.array([0., 0., 1.]), + b: ivy.array([0., 2., 1.]) } """ - return ivy.current_backend(x).sqrt(x, out=out) + return ivy.current_backend(x1, x2).remainder(x1, x2, modulus=modulus, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def round( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + decimals: Optional[int] = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Round each element ``x_i`` of the input array ``x`` to the nearest integer-valued + number. + + .. note:: + For complex floating-point operands, real and imaginary components + must be independently rounded to the nearest integer-valued number. + + Rounded real and imaginary components must be equal + to their equivalent rounded real-valued floating-point + counterparts (i.e., for complex-valued ``x``, ``real(round(x))`` + must equal ``round(real(x)))`` and ``imag(round(x))`` must equal + ``round(imag(x))``). + + **Special cases** + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + + For floating-point operands, + + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If two integers are equally close to ``x_i``, the result is + the even integer closest to ``x_i``. + + .. note:: + For complex floating-point operands, the following special + cases apply to real and imaginary components independently + (e.g., if ``real(x_i)`` is ``NaN``, the rounded + real component is ``NaN``). -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def square( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Each element ``x_i`` of the input array ``x``. + - If ``x_i`` is already integer-valued, the result is ``x_i``. Parameters ---------- x - Input array. Should have a numeric data type. + input array containing elements to round. + decimals + number of decimal places to round to. Default is ``0``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6148,13 +6176,18 @@ def square( Returns ------- ret - an array containing the evaluated result for each element in ``x``. + An array of the same shape and type as x, with the elements rounded to integers. - This method conforms to the `Array API Standard + Note: PyTorch supports an additional argument :code:`decimals` for the + `round function `_. + It has been deliberately omitted here due to the imprecise + nature of the argument in :code:`torch.round`. + + This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.round.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6165,63 +6198,99 @@ def square( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.square(x) + >>> x = ivy.array([1.2, 2.4, 3.6]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([1, 4, 9]) + ivy.array([1.,2.,4.]) - >>> x = ivy.array([1.5, -0.8, 0.3]) - >>> y = ivy.zeros(3) - >>> ivy.square(x, out=y) + >>> x = ivy.array([-0, 5, 4.5]) + >>> y = ivy.round(x) >>> print(y) - ivy.array([2.25, 0.64, 0.09]) + ivy.array([0.,5.,4.]) - >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) - >>> ivy.square(x, out=x) + >>> x = ivy.array([1.5654, 2.034, 15.1, -5.0]) + >>> y = ivy.zeros(4) + >>> ivy.round(x, out=y) + >>> print(y) + ivy.array([2.,2.,15.,-5.]) + + >>> x = ivy.array([[0, 5.433, -343.3, 1.5], + ... [-5.5, 44.2, 11.5, 12.01]]) + >>> ivy.round(x, out=x) >>> print(x) - ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) + ivy.array([[0.,5.,-343.,2.],[-6.,44.,12.,12.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) - >>> y = ivy.square(x) + >>> x = ivy.Container(a=ivy.array([4.20, 8.6, 6.90, 0.0]), + ... b=ivy.array([-300.9, -527.3, 4.5])) + >>> y = ivy.round(x) >>> print(y) { - a:ivy.array([0,1]), - b:ivy.array([4,9]) + a:ivy.array([4.,9.,7.,0.]), + b:ivy.array([-301.,-527.,4.]) } """ - return ivy.current_backend(x).square(x, out=out) + return ivy.current_backend(x).round(x, decimals=decimals, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def subtract( - x1: Union[float, ivy.Array, ivy.NativeArray], - x2: Union[float, ivy.Array, ivy.NativeArray], +def sign( + x: Union[ivy.Array, ivy.NativeArray], /, *, - alpha: Optional[Union[int, float]] = None, + np_variant: Optional[bool] = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Calculate the difference for each element ``x1_i`` of the input array ``x1`` with - the respective element ``x2_i`` of the input array ``x2``. + r""" + Return an indication of the sign of a number for each element ``x_i`` of the input + array ``x``. + + The sign function (also known as the **signum function**) + of a number :math:`x_{i}` is defined as + + .. math:: + \operatorname{sign}(x_i) = \begin{cases} + 0 & \textrm{if } x_i = 0 \\ + \frac{x}{|x|} & \textrm{otherwise} + \end{cases} + + where :math:`|x_i|` is the absolute value of :math:`x_i`. + + **Special cases** + + - If ``x_i`` is less than ``0``, the result is ``-1``. + - If ``x_i`` is either ``-0`` or ``+0``, the result is ``0``. + - If ``x_i`` is greater than ``0``, the result is ``+1``. + - For complex numbers ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j`` + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``-0`` or ``+0`` and ``b`` is + either ``-0`` or ``+0``, the result is ``0 + 0j``. + - If ``a`` is ``NaN`` or ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - In the remaining cases, special cases must be handled + according to the rules of complex number division. Parameters ---------- - x1 - first input array. Should have a numeric data type. - x2 - second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). - Should have a numeric data type. - alpha - optional scalar multiplier for ``x2``. + x + input array. Should have a numeric data type. + np_variant + Handles complex numbers like numpy does If ``True``, + ``sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j``. + otherwise, For complex numbers, ``y = sign(x) = x / |x| if x != 0, + otherwise y = 0.`` + out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6229,13 +6298,14 @@ def subtract( Returns ------- ret - an array containing the element-wise differences. + an array containing the evaluated result for each element in ``x``. The returned + array must have the same data type as ``x``. - This method conforms to the `Array API Standard + This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sign.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6244,19 +6314,34 @@ def subtract( Examples -------- - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y) - >>> print(z) - ivy.array([ 1, 5, -3]) + With :class:`ivy.Array` input: - >>> x = ivy.array([3, 6, 3]) - >>> y = ivy.array([2, 1, 6]) - >>> z = ivy.subtract(x, y, alpha=2) - >>> print(z) - ivy.array([-1, 4, -9]) + >>> x = ivy.array([8.3, -0, 6.8, 0.07]) + >>> y = ivy.sign(x) + >>> print(y) + ivy.array([1., 0., 1., 1.]) + + >>> x = ivy.array([[5.78, -4., -6.9, 0], + ... [-.4, 0.5, 8, -0.01]]) + >>> y = ivy.sign(x) + >>> print(y) + ivy.array([[ 1., -1., -1., 0.], + [-1., 1., 1., -1.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., -0.]), + ... b=ivy.array([1.46, 5.9, -0.0]), + ... c=ivy.array([-8.23, -4.9, -2.6, 7.4])) + >>> y = ivy.sign(x) + >>> print(y) + { + a: ivy.array([0., 0.]), + b: ivy.array([1., 1., 0.]), + c: ivy.array([-1., -1., -1., 1.]) + } """ - return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) + return ivy.current_backend(x).sign(x, np_variant=np_variant, out=out) @handle_exceptions @@ -6267,28 +6352,19 @@ def subtract( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def tan( +def sin( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r"""Calculate an implementation-dependent approximation to the tangent, having - domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each - element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be - expressed in radians. - .. note:: - Tangent is an analytical function on the complex plane - and has no branch cuts. The function is periodic, - with period :math:`\pi j`, with respect to the real - component and has first order poles along the real - line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. - However, IEEE 754 binary floating-point representation - cannot represent the value :math:`\pi / 2` exactly, and, - thus, no argument value is possible for - which a pole error occurs. + r""" + Calculate an implementation-dependent approximation to the sine, having domain + ``(-infinity, +infinity)`` and codomain ``[-1, +1]``, for each element ``x_i`` of + the input array ``x``. Each element ``x_i`` is assumed to be expressed in radians. - where :math:`{tanh}` is the hyperbolic tangent. + .. note:: + The sine is an entire function on the complex plane and has no branch cuts. **Special cases** @@ -6299,68 +6375,68 @@ def tan( - If ``x_i`` is ``-0``, the result is ``-0``. - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. - For complex floating-point operands, special cases must - be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. + For complex floating-point operands, special cases + must be handled as if the operation is implemented as ``-1j * sinh(x*1j)``. Parameters ---------- x - input array whose elements are expressed in radians. Should have a + input array whose elements are each expressed in radians. Should have a floating-point data type. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the tangent of each element in ``x``. The return must have a - floating-point data type determined by :ref:`type-promotion`. + an array containing the sine of each element in ``x``. The returned array must + have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sin.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tan(x) + >>> y = ivy.sin(x) >>> print(y) - ivy.array([0., 1.56, -2.19]) + ivy.array([0., 0.841, 0.909]) - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tan(x, out=y) + >>> x = ivy.array([0., 1.2, -2.3, 3.6]) + >>> y = ivy.zeros(4) + >>> ivy.sin(x, out=y) >>> print(y) - ivy.array([0.546, -0.842, -0.916]) + ivy.array([0., 0.932, -0.746, -0.443]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tan(x, out=x) + >>> x = ivy.array([[1., 2., 3.], [-4., -5., -6.]]) + >>> ivy.sin(x, out=x) >>> print(x) - ivy.array([[1.96, -1.37, 0.16], - [-3.1, 0.996, -0.328]]) + ivy.array([[0.841, 0.909, 0.141], + [0.757, 0.959, 0.279]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.tan(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2., 3.]), + ... b=ivy.array([-4., -5., -6., -7.])) + >>> y = ivy.sin(x) >>> print(y) { - a: ivy.array([0., 1.56, -2.19]), - b: ivy.array([-0.143, 1.16, -3.38]) + a: ivy.array([0., 0.841, 0.909, 0.141]), + b: ivy.array([0.757, 0.959, 0.279, -0.657]) } """ - return ivy.current_backend(x).tan(x, out=out) + return ivy.current_backend(x).sin(x, out=out) @handle_exceptions @@ -6371,18 +6447,25 @@ def tan( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -@handle_complex_input -def tanh( +def sinh( x: Union[ivy.Array, ivy.NativeArray], /, *, - complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Calculate an implementation-dependent approximation to the hyperbolic tangent, - having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element - ``x_i`` of the input array ``x``. + r""" + Calculate an implementation-dependent approximation to the hyperbolic sine, having + domain ``[-infinity, +infinity]`` and codomain ``[-infinity, +infinity]``, for each + element ``x_i`` of the input array ``x``. + + .. math:: + \operatorname{sinh}(x) = \frac{e^x - e^{-x}}{2} + + .. note:: + The hyperbolic sine is an entire function in the + complex plane and has no branch cuts. + The function is periodic, with period + :math:`2\pi j`, with respect to the imaginary component. **Special cases** @@ -6391,80 +6474,62 @@ def tanh( - If ``x_i`` is ``NaN``, the result is ``NaN``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``+infinity``, the result is ``+1``. - - If ``x_i`` is ``-infinity``, the result is ``-1``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. For complex floating-point operands, let ``a = real(x_i)``, ``b = imag(x_i)``, and .. note:: - For complex floating-point operands, ``tanh(conj(x))`` - must equal ``conj(tanh(x))``. + For complex floating-point operands, ``sinh(conj(x))`` + must equal ``conj(sinh(x))``. - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. - - If ``a`` is a nonzero finite number and ``b`` is - ``+infinity``, the result is ``NaN + NaN j``. - If ``a`` is ``+0`` and ``b`` is ``+infinity``, - the result is ``+0 + NaN j``. - - If ``a`` is a nonzero finite number and ``b`` - is ``NaN``, the result is ``NaN + NaN j``. + the result is ``0 + NaN j`` (sign of the real component is unspecified). - If ``a`` is ``+0`` and ``b`` is ``NaN``, - the result is ``+0 + NaN j``. - - If ``a`` is ``+infinity`` and ``b`` is a positive - (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. + the result is ``0 + NaN j`` (sign of the real component is unspecified). + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``+infinity``, the result is ``NaN + NaN j``. + - If ``a`` is a positive (i.e., greater than ``0``) + finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is ``+0``, + the result is ``+infinity + 0j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive finite number, + the result is ``+infinity * cis(b)``. - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). + the result is ``infinity + NaN j`` (sign of the real component is unspecified). - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, - the result is ``1 + 0j`` (sign of the imaginary - component is unspecified). - - If ``a`` is ``NaN`` and ``b`` is ``+0``, - the result is ``NaN + 0j``. - - If ``a`` is ``NaN`` and ``b`` is a nonzero number, + the result is ``infinity + NaN j`` + (sign of the real component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero finite number, the result is ``NaN + NaN j``. - If ``a`` is ``NaN`` and ``b`` is ``NaN``, the result is ``NaN + NaN j``. - .. warning:: - For historical reasons stemming from the C standard, - array libraries may not return the expected - result when ``a`` is ``+0`` and ``b`` is either - ``+infinity`` or ``NaN``. The result should be - ``+0 + NaN j`` in both cases; however, for libraries - compiled against older C versions, the result may be - ``NaN + NaN j``. - - Array libraries are not required to patch these older - C versions, and, thus, users are advised that results - may vary across array library implementations for - these special cases. - + where ``cis(v)`` is ``cos(v) + sin(v)*1j``. Parameters ---------- x input array whose elements each represent a hyperbolic angle. Should have a - real-valued floating-point data - type. - complex_mode - optional specifier for how to handle complex data types. See - ``ivy.func_wrapper.handle_complex_input`` for more detail. + floating-point data type. out - optional output, for writing the result to. It must have a shape that the inputs - broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array containing the hyperbolic tangent of each element in ``x``. - The returned array must have a real-valued floating-point data type - determined by :ref:`type-promotion`. + an array containing the hyperbolic sine of each element in ``x``. The returned + array must have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sinh.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6475,131 +6540,108 @@ def tanh( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.tanh(x) - >>> print(y) - ivy.array([0., 0.762, 0.964]) - - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.zeros(3) - >>> ivy.tanh(x, out=y) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.sinh(x) >>> print(y) - ivy.array([0.462, -0.604, 0.984]) + ivy.array([1.18, 3.63, 10.]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.tanh(x, out=x) + >>> x = ivy.array([0.23, 3., -1.2]) + >>> ivy.sinh(x, out=x) >>> print(x) - ivy.array([[0.8, 0.976, 0.997], - [-1., -1., -1.]]) + ivy.array([0.232, 10., -1.51]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.tanh(x) + >>> x = ivy.Container(a=ivy.array([0.23, -0.25, 1]), b=ivy.array([3, -4, 1.26])) + >>> y = ivy.sinh(x) >>> print(y) { - a: ivy.array([0., 0.762, 0.964]), - b: ivy.array([0.995, 0.999, 1.]) + a: ivy.array([0.232, -0.253, 1.18]), + b: ivy.array([10., -27.3, 1.62]) } """ - return ivy.current_backend(x).tanh(x, out=out) + return ivy.current_backend(x).sinh(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back +@handle_array_function @handle_device_shifting -def trapz( - y: ivy.Array, +def sqrt( + x: Union[ivy.Array, ivy.NativeArray], /, *, - x: Optional[ivy.Array] = None, - dx: float = 1.0, - axis: int = -1, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Integrate along the given axis using the composite trapezoidal rule. - - If x is provided, the integration happens in sequence along its elements - - they are not sorted.. + r""" + Calculate the square root, having domain ``[0, +infinity]`` and codomain ``[0, + +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, + each result must be indistinguishable from the infinitely precise result (as + required by IEEE 754). - Parameters - ---------- - y - The array that should be integrated. - x - The sample points corresponding to the input array values. - If x is None, the sample points are assumed to be evenly spaced - dx apart. The default is None. - dx - The spacing between sample points when x is None. The default is 1. - axis - The axis along which to integrate. - out - optional output array, for writing the result to. + .. note:: + After rounding, each result must be indistinguishable + from the infinitely precise result (as required by IEEE 754). - Returns - ------- - ret - Definite integral of n-dimensional array as approximated along - a single axis by the trapezoidal rule. If the input array is a - 1-dimensional array, then the result is a float. If n is greater - than 1, then the result is an n-1 dimensional array. + .. note:: + For complex floating-point operands, ``sqrt(conj(x))`` + must equal ``conj(sqrt(x))``. - Examples - -------- - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3]) - 4.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], x=[4, 6, 8]) - 8.0 - >>> y = ivy.array([1, 2, 3]) - >>> ivy.trapz([1,2,3], dx=2) - 8.0 - """ - return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) + .. note:: + By convention, the branch cut of the square root is + the negative real axis :math:`(-\infty, 0)`. + The square root is a continuous function from above + the branch cut, taking into account the sign of the imaginary component. -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def trunc( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Round each element x_i of the input array x to the integer-valued number that is - closest to but no greater than x_i. + Accordingly, for complex arguments, the function returns + the square root in the range of the right half-plane, + including the imaginary axis (i.e., the plane defined by + :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` + along the imaginary axis). **Special cases** - - If ``x_i`` is already an integer-valued, the result is ``x_i``. - For floating-point operands, - - If ``x_i`` is ``+infinity``, the result is ``+infinity``. - - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is less than ``0``, the result is ``NaN``. - If ``x_i`` is ``+0``, the result is ``+0``. - If ``x_i`` is ``-0``, the result is ``-0``. - - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, + the result is ``+0 + 0j``. + - If ``a`` is any value (including ``NaN``) and ``b`` is + ``+infinity``, the result is ``+infinity + infinity j``. + - If ``a`` is a finite number and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + - If ``a`` ``-infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``. + - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, + the result is ``NaN + infinity j`` + (sign of the imaginary component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``+infinity + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is any value, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + Parameters ---------- x - input array. Should have a numeric data type. + input array. Should have a floating-point data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6607,14 +6649,14 @@ def trunc( Returns ------- ret - an array containing the rounded result for each element in ``x``. - The returned array must have the same data type as ``x``. + an array containing the square root of each element in ``x``. The returned array + must have a floating-point data type determined by :ref:`type-promotion`. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.sqrt.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -6625,38 +6667,32 @@ def trunc( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) - >>> y = ivy.trunc(x) + >>> x = ivy.array([0, 4., 8.]) + >>> y = ivy.sqrt(x) >>> print(y) - ivy.array([-1., 0., 3., -0.]) - - >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) - >>> ivy.trunc(x, out=x) - >>> print(x) - ivy.array([ 0., 7., -23., -0.]) + ivy.array([0., 2., 2.83]) - >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) - >>> y = ivy.zeros([2,3]) - >>> ivy.trunc(x, out=y) - >>> print(y) - ivy.array([[ 0., -8., 0.], - [ 0., 0., 2.]]) + >>> x = ivy.array([1, 2., 4.]) + >>> y = ivy.zeros(3) + >>> ivy.sqrt(x, out=y) + ivy.array([1., 1.41, 2.]) + + >>> X = ivy.array([40., 24., 100.]) + >>> ivy.sqrt(x, out=x) + >>> ivy.array([6.32455532, 4.89897949, 10.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) - >>> y = ivy.trunc(x) + >>> x = ivy.Container(a=ivy.array([44., 56., 169.]), b=ivy.array([[49.,1.], [0,20.]])) # noqa + >>> y = ivy.sqrt(x) >>> print(y) { - a: ivy.array([-0., 4., 1.]), - b: ivy.array([12., -3., 1.]) + a: ivy.array([6.63, 7.48, 13.]), + b: ivy.array([[7., 1.], + [0., 4.47]]) } """ - return ivy.current_backend(x).trunc(x, out=out) - - -# Extra # -# ------# + return ivy.current_backend(x).sqrt(x, out=out) @handle_exceptions @@ -6667,19 +6703,19 @@ def trunc( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def erf( +def square( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Gauss error function of ``x`` element-wise. + Each element ``x_i`` of the input array ``x``. Parameters ---------- x - Value to compute exponential for. + Input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6687,15 +6723,50 @@ def erf( Returns ------- ret - The Gauss error function of x. + an array containing the evaluated result for each element in ``x``. + + + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([0, 0.3, 0.7, 1.0]) - >>> ivy.erf(x) - ivy.array([0., 0.328, 0.677, 0.842]) + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.square(x) + >>> print(y) + ivy.array([1, 4, 9]) + + >>> x = ivy.array([1.5, -0.8, 0.3]) + >>> y = ivy.zeros(3) + >>> ivy.square(x, out=y) + >>> print(y) + ivy.array([2.25, 0.64, 0.09]) + + >>> x = ivy.array([[1.2, 2, 3.1], [-1, -2.5, -9]]) + >>> ivy.square(x, out=x) + >>> print(x) + ivy.array([[1.44,4.,9.61],[1.,6.25,81.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0, 1]), b=ivy.array([2, 3])) + >>> y = ivy.square(x) + >>> print(y) + { + a:ivy.array([0,1]), + b:ivy.array([4,9]) + } """ - return ivy.current_backend(x).erf(x, out=out) + return ivy.current_backend(x).square(x, out=out) @handle_exceptions @@ -6705,26 +6776,27 @@ def erf( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def maximum( - x1: Union[ivy.Array, ivy.NativeArray, Number], - x2: Union[ivy.Array, ivy.NativeArray, Number], +def subtract( + x1: Union[float, ivy.Array, ivy.NativeArray], + x2: Union[float, ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, + alpha: Optional[Union[int, float]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the max of x1 and x2 (i.e. x1 > x2 ? x1 : x2) element-wise. + Calculate the difference for each element ``x1_i`` of the input array ``x1`` with + the respective element ``x2_i`` of the input array ``x2``. Parameters ---------- x1 - Input array containing elements to maximum threshold. + first input array. Should have a numeric data type. x2 - Tensor containing maximum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the maximum. If ``False``, the maximum - is calculated using the ``(x + y + |x - y|)/2`` formula. Default is ``True``. + second input array. Must be compatible with ``x1`` (see ref:`broadcasting`). + Should have a numeric data type. + alpha + optional scalar multiplier for ``x2``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -6732,196 +6804,141 @@ def maximum( Returns ------- ret - An array with the elements of x1, but clipped to not be lower than the x2 - values. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.maximum(x, y) - >>> print(z) - ivy.array([9, 9, 5]) + an array containing the element-wise differences. - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.maximum(x, y, out=z) - >>> print(z) - ivy.array([[9., 9., 9., 9., 9., 9.], - [3., 5., 9., 8., 3., 7.], - [2., 5., 9., 8., 3., 7.]]) - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.maximum(x, y, out=x) - >>> print(x) - ivy.array([[7, 7]]) + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - With one :class:`ivy.Container` input: + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]), - ... b=ivy.array([-5, 9])) - >>> z = ivy.maximum(x, y) + Examples + -------- + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y) >>> print(z) - { - a: ivy.array([[1, 3], - [2, 4], - [3, 7]]), - b: ivy.array([[1, 9], - [2, 9], - [3, 9]]) - } - - With multiple :class:`ivy.Container` inputs: + ivy.array([ 1, 5, -3]) - >>> x = ivy.Container(a=ivy.array([1, 3, 1]),b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]),b=ivy.array([5, 9, 7])) - >>> z = ivy.maximum(x, y) + >>> x = ivy.array([3, 6, 3]) + >>> y = ivy.array([2, 1, 6]) + >>> z = ivy.subtract(x, y, alpha=2) >>> print(z) - { - a: ivy.array([1, 5, 6]), - b: ivy.array([5, 9, 7]) - } + ivy.array([-1, 4, -9]) """ - return ivy.current_backend(x1).maximum(x1, x2, use_where=use_where, out=out) + return ivy.current_backend(x1).subtract(x1, x2, alpha=alpha, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def minimum( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def tan( + x: Union[ivy.Array, ivy.NativeArray], /, *, - use_where: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Return the min of x1 and x2 (i.e. x1 < x2 ? x1 : x2) element-wise. - - Parameters - ---------- - x1 - Input array containing elements to minimum threshold. - x2 - Tensor containing minimum values, must be broadcastable to x1. - use_where - Whether to use :func:`where` to calculate the minimum. If ``False``, the minimum - is calculated using the ``(x + y - |x - y|)/2`` formula. Default is ``True``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - An array with the elements of x1, but clipped to not exceed the x2 values. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([7, 9, 5]) - >>> y = ivy.array([9, 3, 2]) - >>> z = ivy.minimum(x, y) - >>> print(z) - ivy.array([7, 3, 2]) - - >>> x = ivy.array([1, 5, 9, 8, 3, 7]) - >>> y = ivy.array([[9], [3], [2]]) - >>> z = ivy.zeros((3, 6)) - >>> ivy.minimum(x, y, out=z) - >>> print(z) - ivy.array([[1.,5.,9.,8.,3.,7.], - [1.,3.,3.,3.,3.,3.], - [1.,2.,2.,2.,2.,2.]]) - - >>> x = ivy.array([[7, 3]]) - >>> y = ivy.array([0, 7]) - >>> ivy.minimum(x, y, out=x) - >>> print(x) - ivy.array([[0, 3]]) - - With one :class:`ivy.Container` input: - - >>> x = ivy.array([[1, 3], [2, 4], [3, 7]]) - >>> y = ivy.Container(a=ivy.array([1, 0,]),b=ivy.array([-5, 9])) - >>> z = ivy.minimum(x, y) - >>> print(z) - { - a: ivy.array([[1, 0], - [1, 0], - [1, 0]]), - b: ivy.array([[-5, 3], - [-5, 4], - [-5, 7]]) - } + r"""Calculate an implementation-dependent approximation to the tangent, having + domain ``(-infinity, +infinity)`` and codomain ``(-infinity, +infinity)``, for each + element ``x_i`` of the input array ``x``. Each element ``x_i`` is assumed to be + expressed in radians. + .. note:: + Tangent is an analytical function on the complex plane + and has no branch cuts. The function is periodic, + with period :math:`\pi j`, with respect to the real + component and has first order poles along the real + line at coordinates :math:`(\pi (\frac{1}{2} + n), 0)`. + However, IEEE 754 binary floating-point representation + cannot represent the value :math:`\pi / 2` exactly, and, + thus, no argument value is possible for + which a pole error occurs. - With multiple :class:`ivy.Container` inputs: + where :math:`{tanh}` is the hyperbolic tangent. - >>> x = ivy.Container(a=ivy.array([1, 3, 1]), - ... b=ivy.array([2, 8, 5])) - >>> y = ivy.Container(a=ivy.array([1, 5, 6]), - ... b=ivy.array([5, 9, 7])) - >>> z = ivy.minimum(x, y) - >>> print(z) - { - a: ivy.array([1, 3, 1]), - b: ivy.array([2, 8, 5]) - } - """ - return ivy.current_backend(x1).minimum(x1, x2, use_where=use_where, out=out) + **Special cases** + For floating-point operands, -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def reciprocal( - x: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a new array with the reciprocal of each element in ``x``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + + For complex floating-point operands, special cases must + be handled as if the operation is implemented as ``-1j * tanh(x*1j)``. Parameters ---------- x - Input array. + input array whose elements are expressed in radians. Should have a + floating-point data type. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - A new array with the positive value of each element in ``x``. + an array containing the tangent of each element in ``x``. The return must have a + floating-point data type determined by :ref:`type-promotion`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.reciprocal(x) + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.tan(x) >>> print(y) - ivy.array([1. , 0.5 , 0.33333333]) + ivy.array([0., 1.56, -2.19]) + + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tan(x, out=y) + >>> print(y) + ivy.array([0.546, -0.842, -0.916]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tan(x, out=x) + >>> print(x) + ivy.array([[1.96, -1.37, 0.16], + [-3.1, 0.996, -0.328]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.tan(x) + >>> print(y) + { + a: ivy.array([0., 1.56, -2.19]), + b: ivy.array([-0.143, 1.16, -3.38]) + } """ - return ivy.current_backend(x).reciprocal(x, out=out) + return ivy.current_backend(x).tan(x, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -6929,78 +6946,202 @@ def reciprocal( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def deg2rad( +@handle_complex_input +def tanh( x: Union[ivy.Array, ivy.NativeArray], /, *, + complex_mode: Literal["split", "magnitude", "jax"] = "jax", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Convert the input from degrees to radians. + Calculate an implementation-dependent approximation to the hyperbolic tangent, + having domain ``[-infinity, +infinity]`` and codomain ``[-1, +1]``, for each element + ``x_i`` of the input array ``x``. + + **Special cases** + + For floating-point operands, + + - If ``x_i`` is ``NaN``, the result is ``NaN``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``+infinity``, the result is ``+1``. + - If ``x_i`` is ``-infinity``, the result is ``-1``. + + For complex floating-point operands, let ``a = real(x_i)``, + ``b = imag(x_i)``, and + + .. note:: + For complex floating-point operands, ``tanh(conj(x))`` + must equal ``conj(tanh(x))``. + + - If ``a`` is ``+0`` and ``b`` is ``+0``, the result is ``+0 + 0j``. + - If ``a`` is a nonzero finite number and ``b`` is + ``+infinity``, the result is ``NaN + NaN j``. + - If ``a`` is ``+0`` and ``b`` is ``+infinity``, + the result is ``+0 + NaN j``. + - If ``a`` is a nonzero finite number and ``b`` + is ``NaN``, the result is ``NaN + NaN j``. + - If ``a`` is ``+0`` and ``b`` is ``NaN``, + the result is ``+0 + NaN j``. + - If ``a`` is ``+infinity`` and ``b`` is a positive + (i.e., greater than ``0``) finite number, the result is ``1 + 0j``. + - If ``a`` is ``+infinity`` and ``b`` is ``+infinity``, + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). + - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, + the result is ``1 + 0j`` (sign of the imaginary + component is unspecified). + - If ``a`` is ``NaN`` and ``b`` is ``+0``, + the result is ``NaN + 0j``. + - If ``a`` is ``NaN`` and ``b`` is a nonzero number, + the result is ``NaN + NaN j``. + - If ``a`` is ``NaN`` and ``b`` is ``NaN``, + the result is ``NaN + NaN j``. + + .. warning:: + For historical reasons stemming from the C standard, + array libraries may not return the expected + result when ``a`` is ``+0`` and ``b`` is either + ``+infinity`` or ``NaN``. The result should be + ``+0 + NaN j`` in both cases; however, for libraries + compiled against older C versions, the result may be + ``NaN + NaN j``. + + Array libraries are not required to patch these older + C versions, and, thus, users are advised that results + may vary across array library implementations for + these special cases. + Parameters ---------- x - input array whose elements are each expressed in degrees. + input array whose elements each represent a hyperbolic angle. Should have a + real-valued floating-point data + type. + complex_mode + optional specifier for how to handle complex data types. See + ``ivy.func_wrapper.handle_complex_input`` for more detail. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output, for writing the result to. It must have a shape that the inputs + broadcast to. Returns ------- ret - an array with each element in ``x`` converted from degrees to radians. + an array containing the hyperbolic tangent of each element in ``x``. + The returned array must have a real-valued floating-point data type + determined by :ref:`type-promotion`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x=ivy.array([0,90,180,270,360]) - >>> y=ivy.deg2rad(x) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.tanh(x) >>> print(y) - ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + ivy.array([0., 0.762, 0.964]) - >>> x=ivy.array([0,-1.5,-50,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.deg2rad(x,out=y) + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.zeros(3) + >>> ivy.tanh(x, out=y) >>> print(y) - ivy.array([ 0., -0.02617994, -0.87266463, nan]) + ivy.array([0.462, -0.604, 0.984]) - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.deg2rad(x, out=x) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.tanh(x, out=x) >>> print(x) - ivy.array([[ 0.01919862, 0.03839725, 0.05759586], - [-0.07679449, -0.09599311, -0.11519173]]) - - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.deg2rad(x,out=y) - >>> print(y) - ivy.array([0., 0.35081118, nan]) + ivy.array([[0.8, 0.976, 0.997], + [-1., -1., -1.]]) With :class:`ivy.Container` input: - >>> x=ivy.Container(a=ivy.array([-0,20.1,-50.5,-ivy.nan]), - ... b=ivy.array([0,90,180,270,360])) - >>> y=ivy.deg2rad(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.tanh(x) >>> print(y) { - a: ivy.array([0., 0.35081118, -0.88139129, nan]), - b: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]) + a: ivy.array([0., 0.762, 0.964]), + b: ivy.array([0.995, 0.999, 1.]) } + """ + return ivy.current_backend(x).tanh(x, out=out) - >>> x=ivy.Container(a=ivy.array([0,90,180,270,360]), - ... b=ivy.native_array([0,-1.5,-50,ivy.nan])) - >>> y=ivy.deg2rad(x) - >>> print(y) - { - a: ivy.array([0., 1.57079633, 3.14159265, 4.71238898, 6.28318531]), - b: ivy.array([0., -0.02617994, -0.87266463, nan]) - } + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def trapz( + y: ivy.Array, + /, + *, + x: Optional[ivy.Array] = None, + dx: float = 1.0, + axis: int = -1, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Integrate along the given axis using the composite trapezoidal rule. + + If x is provided, the integration happens in sequence along its elements + - they are not sorted.. + + Parameters + ---------- + y + The array that should be integrated. + x + The sample points corresponding to the input array values. + If x is None, the sample points are assumed to be evenly spaced + dx apart. The default is None. + dx + The spacing between sample points when x is None. The default is 1. + axis + The axis along which to integrate. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Definite integral of n-dimensional array as approximated along + a single axis by the trapezoidal rule. If the input array is a + 1-dimensional array, then the result is a float. If n is greater + than 1, then the result is an n-1 dimensional array. + + Examples + -------- + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3]) + 4.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], x=[4, 6, 8]) + 8.0 + >>> y = ivy.array([1, 2, 3]) + >>> ivy.trapz([1,2,3], dx=2) + 8.0 """ - return ivy.current_backend(x).deg2rad(x, out=out) + return ivy.current_backend(y).trapz(y, x=x, dx=dx, axis=axis, out=out) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -7008,19 +7149,32 @@ def deg2rad( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def rad2deg( +def trunc( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Convert the input from radians to degrees. + Round each element x_i of the input array x to the integer-valued number that is + closest to but no greater than x_i. + + **Special cases** + + - If ``x_i`` is already an integer-valued, the result is ``x_i``. + + For floating-point operands, + + - If ``x_i`` is ``+infinity``, the result is ``+infinity``. + - If ``x_i`` is ``-infinity``, the result is ``-infinity``. + - If ``x_i`` is ``+0``, the result is ``+0``. + - If ``x_i`` is ``-0``, the result is ``-0``. + - If ``x_i`` is ``NaN``, the result is ``NaN``. Parameters ---------- x - input array whose elements are each expressed in radians. + input array. Should have a numeric data type. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -7028,56 +7182,52 @@ def rad2deg( Returns ------- ret - an array with each element in ``x`` converted from radians to degrees. + an array containing the rounded result for each element in ``x``. + The returned array must have the same data type as ``x``. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments Examples -------- With :class:`ivy.Array` input: - >>> x=ivy.array([0.,1.57,3.14,4.71,6.28]) - >>> y=ivy.rad2deg(x) - >>> print(y) - ivy.array([ 0., 90., 180., 270., 360.]) - - >>> x=ivy.array([0.,-0.0262,-0.873,ivy.nan]) - >>> y=ivy.zeros(4) - >>> ivy.rad2deg(x,out=y) + >>> x = ivy.array([-1, 0.54, 3.67, -0.025]) + >>> y = ivy.trunc(x) >>> print(y) - ivy.array([ 0. , -1.5, -50. , nan]) + ivy.array([-1., 0., 3., -0.]) - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.rad2deg(x, out=x) + >>> x = ivy.array([0.56, 7, -23.4, -0.0375]) + >>> ivy.trunc(x, out=x) >>> print(x) - ivy.array([[ 63., 126., 189.], - [-252., -315., -378.]]) + ivy.array([ 0., 7., -23., -0.]) - >>> x=ivy.native_array([-0,20.1,ivy.nan]) - >>> y=ivy.zeros(3) - >>> ivy.rad2deg(x,out=y) + >>> x = ivy.array([[0.4, -8, 0.55], [0, 0.032, 2]]) + >>> y = ivy.zeros([2,3]) + >>> ivy.trunc(x, out=y) >>> print(y) - ivy.array([ 0., 1150., nan]) + ivy.array([[ 0., -8., 0.], + [ 0., 0., 2.]]) With :class:`ivy.Container` input: - >>> x=ivy.Container(a=ivy.array([-0., 20.1, -50.5, -ivy.nan]), - ... b=ivy.array([0., 1., 2., 3., 4.])) - >>> y=ivy.rad2deg(x) - >>> print(y) - { - a: ivy.array([0., 1150., -2890., nan]), - b: ivy.array([0., 57.3, 115., 172., 229.]) - } - - >>> x=ivy.Container(a=ivy.array([0,10,180,8.5,6]), - ... b=ivy.native_array([0,-1.5,0.5,ivy.nan])) - >>> y=ivy.rad2deg(x) + >>> x = ivy.Container(a=ivy.array([-0.25, 4, 1.3]), b=ivy.array([12, -3.5, 1.234])) + >>> y = ivy.trunc(x) >>> print(y) { - a: ivy.array([0., 573., 10300., 487., 344.]), - b: ivy.array([0., -85.9, 28.6, nan]) + a: ivy.array([-0., 4., 1.]), + b: ivy.array([12., -3., 1.]) } """ - return ivy.current_backend(x).rad2deg(x, out=out) + return ivy.current_backend(x).trunc(x, out=out) @handle_exceptions @@ -7126,6 +7276,7 @@ def trunc_divide( return ivy.trunc(ivy.divide(x1, x2), out=out) +pow.unsupported_gradients = {"torch": ["float16"]} trunc_divide.mixed_backend_wrappers = { "to_add": ( "handle_backend_invalid", @@ -7137,156 +7288,3 @@ def trunc_divide( ), "to_skip": ("inputs_to_ivy_arrays",), } - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def isreal( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Test each element ``x_i`` of the input array ``x`` to determine whether the element - is real number. Returns a bool array, where True if input element is real. If - element has complex type with zero complex part, the return value for that element - is True. - - Parameters - ---------- - x - input array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing test results. An element ``out_i`` is ``True`` if ``x_i`` is - real number and ``False`` otherwise. The returned array should have a data type - of ``bool``. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances in place of - :class:`ivy.Array` or :class:`ivy.NativeArray` instances, as shown in the type hints - and also the examples below. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[[1.1], [float('inf')], [-6.3]]]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([[[True], [True], [True]]]) - - >>> x = ivy.array([1-0j, 3j, 7+5j]) - >>> z = ivy.isreal(x) - >>> print(z) - ivy.array([ True, False, False]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-6.7-7j, -np.inf, 1.23]),\ - b=ivy.array([5j, 5-6j, 3])) - >>> z = ivy.isreal(x) - >>> print(z) - { - a: ivy.array([False, True, True]), - b: ivy.array([False, False, True]) - } - """ - return ivy.current_backend(x).isreal(x, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fmod( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Compute the element-wise remainder of divisions of two arrays. - - Parameters - ---------- - x1 - First input array. - x2 - Second input array - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array with element-wise remainder of divisions. - - Examples - -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmod(x1, x2) - ivy.array([ 0, 3, 0]) - - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmod(x1, x2) - ivy.array([ nan, nan, nan]) - """ - return ivy.current_backend(x1, x2).fmod(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lcm( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the element-wise least common multiple (LCM) of x1 and x2. - - Parameters - ---------- - x1 - first input array, must be integers - x2 - second input array, must be integers - out - optional output array, for writing the result to. - - Returns - ------- - ret - an array that includes the element-wise least common multiples of x1 and x2 - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x1=ivy.array([2, 3, 4]) - >>> x2=ivy.array([5, 7, 15]) - >>> x1.lcm(x1, x2) - ivy.array([10, 21, 60]) - """ - return ivy.current_backend(x1, x2).lcm(x1, x2, out=out) diff --git a/ivy/functional/ivy/experimental/activations.py b/ivy/functional/ivy/experimental/activations.py index 34203cffbaf6d..30de525bedc40 100644 --- a/ivy/functional/ivy/experimental/activations.py +++ b/ivy/functional/ivy/experimental/activations.py @@ -17,6 +17,68 @@ ) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def elu( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + alpha: float = 1.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Apply the elu unit function element-wise. + + Parameters + ---------- + x + Input array. + alpha + scaler for controlling the slope of the function for x <= 0 Default: 1.0 + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The input array with elu applied element-wise. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([0.39, -0.85]) + >>> y = ivy.elu(x) + >>> print(y) + ivy.array([ 0.38999999, -0.57258511]) + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> y = ivy.zeros(3) + >>> ivy.elu(x, out=y) + >>> print(y) + ivy.array([ 1.5, 0.69999999, -0.90928203]) + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> ivy.elu(x, out=x) + >>> print(x) + ivy.array([[ 1.10000002, 2.20000005, 3.29999995], + [-0.98772264, -0.99591321, -0.99863964]]) + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) + >>> x = ivy.elu(x, out=x) + >>> print(x) + { + a: ivy.array([0., -0.69880581]), + b: ivy.array([0.40000001, -0.18126924]) + } + """ + return current_backend(x).elu(x, alpha=alpha, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -67,6 +129,57 @@ def logit( return current_backend(x).logit(x, eps=eps, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def logsigmoid( + input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None +) -> ivy.Array: + """ + Apply element-wise Log-sigmoid of x. + + logsigmoid(x) = log(1 / (1 + exp(-x)). + + Parameters + ---------- + input + Input array. + + Returns + ------- + Array with same shape as input with Log-sigmoid applied to every element. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([-1., 0., 1.]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-1.31326175, -0.69314718, -0.31326169]) + + >>> x = ivy.array([1.5, 0.7, -2.4]) + >>> z = x.logsigmoid() + >>> print(z) + ivy.array([-0.20141329, -0.40318608, -2.48683619]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.logsigmoid(x) + >>> print(x) + { + a: ivy.array([-0.31326169, -1.46328247]), + b: ivy.array([-0.59813893, -0.43748799]) + } + """ + return ivy.current_backend(input).logsigmoid(input, out=out) + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion @@ -125,67 +238,6 @@ def prelu( raise IvyException -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def thresholded_relu( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - threshold: Union[int, float] = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Apply the rectified linear unit function with custom threshold. - - Parameters - ---------- - x - input array - threshold - threshold value above which the activation is linear. Default: ``0``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - an array containing the rectified linear unit activation of each element in - ``x``. with custom ``threshold``. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.thresholded_relu(x, threshold=0.5) - >>> print(y) - ivy.array([0., 0. , 1.]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> y = ivy.zeros(3) - >>> ivy.thresholded_relu(x, threshold=1, out = y) - >>> print(y) - ivy.array([ 1.5, 0., 0.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.thresholded_relu(x, threshold=0.5) - >>> print(x) - { - a: ivy.array([1., 0.]), - b: ivy.array([0., 0.6]) - } - """ - return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -232,57 +284,6 @@ def relu6( return current_backend(x).relu6(x, out=out) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def logsigmoid( - input: Union[ivy.NativeArray, ivy.Array], /, *, out: Optional[ivy.Array] = None -) -> ivy.Array: - """ - Apply element-wise Log-sigmoid of x. - - logsigmoid(x) = log(1 / (1 + exp(-x)). - - Parameters - ---------- - input - Input array. - - Returns - ------- - Array with same shape as input with Log-sigmoid applied to every element. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-1.31326175, -0.69314718, -0.31326169]) - - >>> x = ivy.array([1.5, 0.7, -2.4]) - >>> z = x.logsigmoid() - >>> print(z) - ivy.array([-0.20141329, -0.40318608, -2.48683619]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) - >>> x = ivy.logsigmoid(x) - >>> print(x) - { - a: ivy.array([-0.31326169, -1.46328247]), - b: ivy.array([-0.59813893, -0.43748799]) - } - """ - return ivy.current_backend(input).logsigmoid(input, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -342,6 +343,40 @@ def selu( return current_backend(x).selu(x, out=out) +def sequence_length( + x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None +) -> ivy.int64: + """ + Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy + array input. + + Parameters + ---------- + x + Can be a sequence of any tensor type: bool, complex128, + complex64, double, float, float16, int16, int32, int64, + int8, string, uint16, uint32, uint64, uint8 + + Returns + ------- + length + Length of the input sequence, as a scalar (empty shape tensor). + + Examples + -------- + >>> x = ivy.array([True, False, True]) + >>> y = ivy.sequence_length(x) + >>> print(y) + 3 + + >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] + >>> y = ivy.sequence_length(x) + >>> print(y) + 5 + """ + return current_backend(x).sequence_length(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -398,23 +433,23 @@ def silu( @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def elu( +@handle_device_shifting +def thresholded_relu( x: Union[ivy.Array, ivy.NativeArray], /, *, - alpha: float = 1.0, + threshold: Union[int, float] = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply the elu unit function element-wise. + Apply the rectified linear unit function with custom threshold. Parameters ---------- x - Input array. - alpha - scaler for controlling the slope of the function for x <= 0 Default: 1.0 + input array + threshold + threshold value above which the activation is linear. Default: ``0``. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -422,67 +457,32 @@ def elu( Returns ------- ret - The input array with elu applied element-wise. + an array containing the rectified linear unit activation of each element in + ``x``. with custom ``threshold``. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.39, -0.85]) - >>> y = ivy.elu(x) + + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.thresholded_relu(x, threshold=0.5) >>> print(y) - ivy.array([ 0.38999999, -0.57258511]) + ivy.array([0., 0. , 1.]) + >>> x = ivy.array([1.5, 0.7, -2.4]) >>> y = ivy.zeros(3) - >>> ivy.elu(x, out=y) + >>> ivy.thresholded_relu(x, threshold=1, out = y) >>> print(y) - ivy.array([ 1.5, 0.69999999, -0.90928203]) - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> ivy.elu(x, out=x) - >>> print(x) - ivy.array([[ 1.10000002, 2.20000005, 3.29999995], - [-0.98772264, -0.99591321, -0.99863964]]) + ivy.array([ 1.5, 0., 0.]) + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.0, -1.2]), b=ivy.array([0.4, -0.2])) - >>> x = ivy.elu(x, out=x) + + >>> x = ivy.Container(a=ivy.array([1.0, -1.2]), b=ivy.array([0.2, 0.6])) + >>> x = ivy.thresholded_relu(x, threshold=0.5) >>> print(x) { - a: ivy.array([0., -0.69880581]), - b: ivy.array([0.40000001, -0.18126924]) + a: ivy.array([1., 0.]), + b: ivy.array([0., 0.6]) } """ - return current_backend(x).elu(x, alpha=alpha, out=out) - - -def sequence_length( - x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None -) -> ivy.int64: - """ - Produce a scalar (tensor of empty shape) containing the number of tensors in the ivy - array input. - - Parameters - ---------- - x - Can be a sequence of any tensor type: bool, complex128, - complex64, double, float, float16, int16, int32, int64, - int8, string, uint16, uint32, uint64, uint8 - - Returns - ------- - length - Length of the input sequence, as a scalar (empty shape tensor). - - Examples - -------- - >>> x = ivy.array([True, False, True]) - >>> y = ivy.sequence_length(x) - >>> print(y) - 3 - - >>> x = [1.0, 2.5, -3.4, 5.6, -85.3] - >>> y = ivy.sequence_length(x) - >>> print(y) - 5 - """ - return current_backend(x).sequence_length(x, out=out) + return current_backend(x).thresholded_relu(x, threshold=threshold, out=out) diff --git a/ivy/functional/ivy/experimental/creation.py b/ivy/functional/ivy/experimental/creation.py index bb3fc06d68e66..cf96237fa41af 100644 --- a/ivy/functional/ivy/experimental/creation.py +++ b/ivy/functional/ivy/experimental/creation.py @@ -21,56 +21,31 @@ ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def vorbis_window( - window_length: Union[ivy.Array, ivy.NativeArray], - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array that contains a vorbis power complementary window of size - window_length. +# --- Helpers --- # +# --------------- # - Parameters - ---------- - window_length - the length of the vorbis window. - dtype - data type of the returned array. By default float32. - out - optional output array, for writing the result to. - Returns - ------- - ret - Input array with the vorbis window. +def _iter_product(*args, repeat=1): + # itertools.product + pools = [tuple(pool) for pool in args] * repeat + result = [[]] + for pool in pools: + result = [x + [y] for x in result for y in pool] + for prod in result: + yield tuple(prod) - Examples - -------- - >>> ivy.vorbis_window(3) - ivy.array([0.38268346, 1. , 0.38268352]) - >>> ivy.vorbis_window(5) - ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) - """ - return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) +# --- Main --- # +# ------------ # @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @infer_dtype @handle_device_shifting -def hann_window( +def blackman_window( size: int, *, periodic: bool = True, @@ -78,13 +53,14 @@ def hann_window( out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Generate a Hann window. The Hanning window is a taper formed by using a weighted - cosine. + Generate a Blackman window. The Blackman window is a taper formed by using the first + three terms of a summation of cosines. It was designed to have close to the minimal + leakage possible. It is close to optimal, only slightly worse than a Kaiser window. Parameters ---------- - size - the size of the returned window. + window_length + the window_length of the returned window. periodic If True, returns a window to be used as periodic function. If False, return a symmetric window. @@ -97,122 +73,111 @@ def hann_window( ------- ret The array containing the window. - Functional Examples ------------------- - >>> ivy.hann_window(4, periodic = True) - ivy.array([0. , 0.5, 1. , 0.5]) - - >>> ivy.hann_window(7, periodic = False) - ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) + >>> ivy.blackman_window(4, periodic = True) + ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) + >>> ivy.blackman_window(7, periodic = False) + ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, + 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) """ - return ivy.current_backend().hann_window( + return ivy.current_backend().blackman_window( size, periodic=periodic, dtype=dtype, out=out ) @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @infer_dtype -@handle_device_shifting -def kaiser_window( - window_length: int, - periodic: bool = True, - beta: float = 12.0, +@infer_device +def eye_like( + x: Union[ivy.Array, ivy.NativeArray], *, + k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kaiser window with window length window_length and shape beta. + Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the + same ``shape`` as the first and last dim of input array ``x``. input array ``x`` + should to be 2D. Parameters ---------- - window_length - an int defining the length of the window. - periodic - If True, returns a periodic window suitable for use in spectral analysis. - If False, returns a symmetric window suitable for use in filter design. - beta - a float used as shape parameter for the window. + x + input array from which to derive the output array shape. + k + index of the diagonal. A positive value refers to an upper diagonal, a negative + value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. dtype - data type of the returned array. + output array data type. If dtype is None, the output array data type must be the + default floating-point data type. Default: ``None``. + device + the device on which to place the created array. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The array containing the window. + an array having the same shape as ``x`` and filled with ``ones`` in + diagonal ``k`` and ``zeros`` elsewhere. - Examples - -------- - >>> ivy.kaiser_window(5) - ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) - >>> ivy.kaiser_window(5, True, 5) - ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) - >>> ivy.kaiser_window(5, False, 5) - ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) - """ - return ivy.current_backend().kaiser_window( - window_length, periodic, beta, dtype=dtype, out=out - ) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances as a replacement to any of the arguments. -@handle_exceptions -@handle_nestable -@handle_out_argument -@infer_dtype -def kaiser_bessel_derived_window( - window_length: int, - beta: float = 12.0, - *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Kaiser bessel derived window with window length window_length and shape - beta. + Functional Examples + ------------------- - Parameters - ---------- - window_length - an int defining the length of the window. - beta - a float used as shape parameter for the window. - dtype - data type of the returned array - out - optional output array, for writing the result to. + With :class:`ivy.Array` input: - Returns - ------- - ret - The array containing the window. + >>> x1 = ivy.array([[0, 1],[2, 3]]) + >>> y1 = ivy.eye_like(x1) + >>> print(y1) + ivy.array([[1., 0.], + [0., 1.]]) - Functional Examples - ------------------- - >>> ivy.kaiser_bessel_derived_window(5) - ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) + >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) + >>> y1 = ivy.eye_like(x1, k=1) + >>> print(y1) + ivy.array([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) - >>> ivy.kaiser_bessel_derived_window(5, 5) - ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) + >>> y = x.eye_like() + >>> print(y) + { + a: ivy.array([[1., 0.], + [0., 1.]]), + b: ivy.array([[1., 0.], + [0., 1.]]) + } """ - if window_length < 2: - result = ivy.array([], dtype=dtype) - if ivy.exists(out): - ivy.inplace_update(out, result) - return result - half_len = window_length // 2 - kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) - kaiser_w_csum = ivy.cumsum(kaiser_w) - half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) - window = ivy.concat((half_w, half_w[::-1]), axis=0) - result = window.astype(dtype) - return result + shape = ivy.shape(x, as_array=True) + dim = len(shape) + if dim <= 1: + cols = dim + else: + cols = int(shape[-1]) + rows = 0 if dim < 1 else int(shape[0]) + return ivy.eye( + rows, + cols, + k=k, + dtype=dtype, + device=device, + out=out, + ) @handle_exceptions @@ -272,214 +237,211 @@ def hamming_window( return result -hamming_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "handle_device_shifting", - ), - "to_skip": (), -} - - @handle_exceptions +@handle_backend_invalid @handle_nestable -@outputs_to_ivy_arrays -@infer_device -def tril_indices( - n_rows: int, - n_cols: Optional[int] = None, - k: int = 0, +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def hann_window( + size: int, *, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, -) -> Tuple[ivy.Array, ...]: + periodic: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Return the indices of the lower triangular part of a row by col matrix in a 2-by-N - shape (tuple of two N dimensional arrays), where the first row contains row - coordinates of all indices and the second row contains column coordinates. Indices - are ordered based on rows and then columns. The lower triangular part of the matrix - is defined as the elements on and below the diagonal. The argument k controls which - diagonal to consider. If k = 0, all elements on and below the main diagonal are - retained. A positive value excludes just as many diagonals below the main diagonal, - and similarly a negative value includes just as many diagonals above the main - diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, - n_cols}−1]. - - Notes - ----- - Primary purpose of this function is to slice an array of shape (n,m). See - https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html - for examples - - Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices + Generate a Hann window. The Hanning window is a taper formed by using a weighted + cosine. Parameters ---------- - n_rows - number of rows in the 2-d matrix. - n_cols - number of columns in the 2-d matrix. If None n_cols will be the same as n_rows - k - number of shifts from the main diagonal. k = 0 includes main diagonal, - k > 0 moves downward and k < 0 moves upward - device - device on which to place the created array. Default: ``None``. + size + the size of the returned window. + periodic + If True, returns a window to be used as periodic function. + If False, return a symmetric window. + dtype + The data type to produce. Must be a floating point type. + out + optional output array, for writing the result to. Returns ------- ret - an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) - contains row coordinates of all indices and the second subarray (i.e ret[1]) - contains columns indices. - - Function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - >>> x = ivy.tril_indices(4,4,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + The array containing the window. - >>> x = ivy.tril_indices(4,4,1) - >>> print(x) - (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) + Functional Examples + ------------------- + >>> ivy.hann_window(4, periodic = True) + ivy.array([0. , 0.5, 1. , 0.5]) - >>> x = ivy.tril_indices(4,4,-2) - >>> print(x) - (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) + >>> ivy.hann_window(7, periodic = False) + ivy.array([0. , 0.25, 0.75, 1. , 0.75, 0.25, 0. ]) + """ + return ivy.current_backend().hann_window( + size, periodic=periodic, dtype=dtype, out=out + ) - >>> x = ivy.tril_indices(4,2,0) - >>> print(x) - (ivy.array([0, 1, 1, 2, 2, 3, 3]), - ivy.array([0, 0, 1, 0, 1, 0, 1])) - >>> x = ivy.tril_indices(2,4,0) - >>> print(x) - (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) +@handle_exceptions +def indices( + dimensions: Sequence[int], + *, + dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, + sparse: bool = False, +) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: + """ + Return an array representing the indices of a grid. - >>> x = ivy.tril_indices(4,-4,0) - >>> print(x) - (ivy.array([]), ivy.array([])) + Parameters + ---------- + dimensions + The shape of the grid. + dtype + The data type of the result. + sparse + Return a sparse representation of the grid instead of a dense representation. - >>> x = ivy.tril_indices(4,4,100) - >>> print(x) - (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), - ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + Returns + ------- + ret + If sparse is False, returns one grid indices array of shape + (len(dimensions),) + tuple(dimensions). + If sparse is True, returns a tuple of arrays each of shape + (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. - >>> x = ivy.tril_indices(2,4,-100) - >>> print(x) - (ivy.array([]), ivy.array([])) + Examples + -------- + >>> ivy.indices((3, 2)) + ivy.array([[[0 0] + [1 1] + [2 2]] + [[0 1] + [0 1] + [0 1]]]) + >>> ivy.indices((3, 2), sparse=True) + (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) """ - return current_backend().tril_indices(n_rows, n_cols, k, device=device) + if sparse: + return tuple( + ivy.arange(dim) + .expand_dims( + axis=[j for j in range(len(dimensions)) if i != j], + ) + .astype(dtype) + for i, dim in enumerate(dimensions) + ) + else: + grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") + return ivy.stack(grid, axis=0).astype(dtype) @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays @infer_dtype -@infer_device -def eye_like( - x: Union[ivy.Array, ivy.NativeArray], +def kaiser_bessel_derived_window( + window_length: int, + beta: float = 12.0, *, - k: int = 0, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a 2D array filled with ones on the k diagonal and zeros elsewhere. having the - same ``shape`` as the first and last dim of input array ``x``. input array ``x`` - should to be 2D. + Compute the Kaiser bessel derived window with window length window_length and shape + beta. Parameters ---------- - x - input array from which to derive the output array shape. - k - index of the diagonal. A positive value refers to an upper diagonal, a negative - value to a lower diagonal, and 0 to the main diagonal. Default: ``0``. + window_length + an int defining the length of the window. + beta + a float used as shape parameter for the window. dtype - output array data type. If dtype is None, the output array data type must be the - default floating-point data type. Default: ``None``. - device - the device on which to place the created array. + data type of the returned array out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array having the same shape as ``x`` and filled with ``ones`` in - diagonal ``k`` and ``zeros`` elsewhere. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances as a replacement to any of the arguments. + The array containing the window. Functional Examples ------------------- + >>> ivy.kaiser_bessel_derived_window(5) + ivy.array([0.00726415, 0.9999736 , 0.9999736 , 0.00726415]) - With :class:`ivy.Array` input: + >>> ivy.kaiser_bessel_derived_window(5, 5) + ivy.array([0.18493208, 0.9827513 , 0.9827513 , 0.18493208]) + """ + if window_length < 2: + result = ivy.array([], dtype=dtype) + if ivy.exists(out): + ivy.inplace_update(out, result) + return result + half_len = window_length // 2 + kaiser_w = ivy.kaiser_window(half_len + 1, False, beta, dtype=dtype) + kaiser_w_csum = ivy.cumsum(kaiser_w) + half_w = ivy.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1:]) + window = ivy.concat((half_w, half_w[::-1]), axis=0) + result = window.astype(dtype) + return result - >>> x1 = ivy.array([[0, 1],[2, 3]]) - >>> y1 = ivy.eye_like(x1) - >>> print(y1) - ivy.array([[1., 0.], - [0., 1.]]) - >>> x1 = ivy.array([[0, 1, 2],[3, 4, 5],[6, 7, 8]]) - >>> y1 = ivy.eye_like(x1, k=1) - >>> print(y1) - ivy.array([[0., 1., 0.], - [0., 0., 1.], - [0., 0., 0.]]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def kaiser_window( + window_length: int, + periodic: bool = True, + beta: float = 12.0, + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the Kaiser window with window length window_length and shape beta. - With :class:`ivy.Container` input: + Parameters + ---------- + window_length + an int defining the length of the window. + periodic + If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta + a float used as shape parameter for the window. + dtype + data type of the returned array. + out + optional output array, for writing the result to. - >>> x = ivy.Container(a=ivy.array([[3, 8],[0, 2]]), b=ivy.array([[0, 2], [8, 5]])) - >>> y = x.eye_like() - >>> print(y) - { - a: ivy.array([[1., 0.], - [0., 1.]]), - b: ivy.array([[1., 0.], - [0., 1.]]) - } + Returns + ------- + ret + The array containing the window. + + Examples + -------- + >>> ivy.kaiser_window(5) + ivy.array([5.2773e-05, 1.0172e-01, 7.9294e-01, 7.9294e-01, 1.0172e-01]]) + >>> ivy.kaiser_window(5, True, 5) + ivy.array([0.0367, 0.4149, 0.9138, 0.9138, 0.4149]) + >>> ivy.kaiser_window(5, False, 5) + ivy.array([0.0367, 0.5529, 1.0000, 0.5529, 0.0367]) """ - shape = ivy.shape(x, as_array=True) - dim = len(shape) - if dim <= 1: - cols = dim - else: - cols = int(shape[-1]) - rows = 0 if dim < 1 else int(shape[0]) - return ivy.eye( - rows, - cols, - k=k, - dtype=dtype, - device=device, - out=out, + return ivy.current_backend().kaiser_window( + window_length, periodic, beta, dtype=dtype, out=out ) -def _iter_product(*args, repeat=1): - # itertools.product - pools = [tuple(pool) for pool in args] * repeat - result = [[]] - for pool in pools: - result = [x + [y] for x in result for y in pool] - for prod in result: - yield tuple(prod) - - @handle_exceptions @inputs_to_ivy_arrays def ndenumerate( @@ -553,185 +515,60 @@ def ndindex( @handle_exceptions -def indices( - dimensions: Sequence[int], - *, - dtype: Union[ivy.Dtype, ivy.NativeDtype] = ivy.int64, - sparse: bool = False, -) -> Union[ivy.Array, Tuple[ivy.Array, ...]]: - """ - Return an array representing the indices of a grid. - - Parameters - ---------- - dimensions - The shape of the grid. - dtype - The data type of the result. - sparse - Return a sparse representation of the grid instead of a dense representation. - - Returns - ------- - ret - If sparse is False, returns one grid indices array of shape - (len(dimensions),) + tuple(dimensions). - If sparse is True, returns a tuple of arrays each of shape - (1, ..., 1, dimensions[i], 1, ..., 1) with dimensions[i] in the ith place. - - Examples - -------- - >>> ivy.indices((3, 2)) - ivy.array([[[0 0] - [1 1] - [2 2]] - [[0 1] - [0 1] - [0 1]]]) - >>> ivy.indices((3, 2), sparse=True) - (ivy.array([[0], [1], [2]]), ivy.array([[0, 1]])) - """ - if sparse: - return tuple( - ivy.arange(dim) - .expand_dims( - axis=[j for j in range(len(dimensions)) if i != j], - ) - .astype(dtype) - for i, dim in enumerate(dimensions) - ) - else: - grid = ivy.meshgrid(*[ivy.arange(dim) for dim in dimensions], indexing="ij") - return ivy.stack(grid, axis=0).astype(dtype) - - -indices.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} - - -@handle_exceptions -@handle_backend_invalid @handle_nestable -@to_native_arrays_and_back -def unsorted_segment_min( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the minimum along segments of an array. Segments are defined by an integer - array of segment IDs. - - Note - ---- - If the given segment ID `i` is negative, then the corresponding - value is dropped, and will not be included in the result. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented min operation. - For each segment, it computes the min value in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_nestable -@to_native_arrays_and_back -def unsorted_segment_sum( - data: Union[ivy.Array, ivy.NativeArray], - segment_ids: Union[ivy.Array, ivy.NativeArray], - num_segments: Union[int, ivy.Array, ivy.NativeArray], -) -> ivy.Array: - """ - Compute the sum of elements along segments of an array. Segments are defined by an - integer array of segment IDs. - - Parameters - ---------- - data - The array from which to gather values. - - segment_ids - Must be in the same size with the first dimension of `data`. Has to be - of integer data type. The index-th element of `segment_ids` array is - the segment identifier for the index-th element of `data`. - - num_segments - An integer or array representing the total number of distinct segment IDs. - - Returns - ------- - ret - The output array, representing the result of a segmented sum operation. - For each segment, it computes the sum of values in `data` where `segment_ids` - equals to segment ID. - """ - return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) - - -@handle_exceptions -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back @infer_dtype -@handle_device_shifting -def blackman_window( - size: int, +def random_cp( + shape: Sequence[int], + rank: int, + /, *, - periodic: bool = True, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + full: Optional[bool] = False, + orthogonal: Optional[bool] = False, + seed: Optional[int] = None, + normalise_factors: Optional[bool] = True, +) -> Union[ivy.CPTensor, ivy.Array]: """ - Generate a Blackman window. The Blackman window is a taper formed by using the first - three terms of a summation of cosines. It was designed to have close to the minimal - leakage possible. It is close to optimal, only slightly worse than a Kaiser window. + Generate a random CP tensor. Parameters ---------- - window_length - the window_length of the returned window. - periodic - If True, returns a window to be used as periodic function. - If False, return a symmetric window. - dtype - The data type to produce. Must be a floating point type. - out - optional output array, for writing the result to. + shape + shape of the tensor to generate + rank + rank of the CP decomposition + full + if True, a full tensor is returned + otherwise, the decomposed tensor is returned + orthogonal + if True, creates a tensor with orthogonal components + seed + seed for generating random numbers Returns ------- - ret - The array containing the window. - Functional Examples - ------------------- - >>> ivy.blackman_window(4, periodic = True) - ivy.array([-1.38777878e-17, 3.40000000e-01, 1.00000000e+00, 3.40000000e-01]) - >>> ivy.blackman_window(7, periodic = False) - ivy.array([-1.38777878e-17, 1.30000000e-01, 6.30000000e-01, 1.00000000e+00, - 6.30000000e-01, 1.30000000e-01, -1.38777878e-17]) + ivy.CPTensor """ - return ivy.current_backend().blackman_window( - size, periodic=periodic, dtype=dtype, out=out - ) + rank = ivy.CPTensor.validate_cp_rank(shape, rank) + if (rank > min(shape)) and orthogonal: + warnings.warn( + "Can only construct orthogonal tensors when rank <= min(shape) but got " + f"a tensor with min(shape)={min(shape)} < rank={rank}" + ) + + factors = [ + (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape + ] + weights = ivy.ones((rank,), dtype=dtype) + if orthogonal: + factors = [ivy.qr(factor)[0] for factor in factors] + + if full: + return ivy.CPTensor.cp_to_tensor((weights, factors)) + elif normalise_factors: + return ivy.CPTensor.cp_normalize((weights, factors)) + else: + return ivy.CPTensor((weights, factors)) @handle_exceptions @@ -806,59 +643,96 @@ def random_tucker( @handle_exceptions @handle_nestable -@infer_dtype -def random_cp( - shape: Sequence[int], - rank: int, - /, +@outputs_to_ivy_arrays +@infer_device +def tril_indices( + n_rows: int, + n_cols: Optional[int] = None, + k: int = 0, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - full: Optional[bool] = False, - orthogonal: Optional[bool] = False, - seed: Optional[int] = None, - normalise_factors: Optional[bool] = True, -) -> Union[ivy.CPTensor, ivy.Array]: + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, +) -> Tuple[ivy.Array, ...]: """ - Generate a random CP tensor. + Return the indices of the lower triangular part of a row by col matrix in a 2-by-N + shape (tuple of two N dimensional arrays), where the first row contains row + coordinates of all indices and the second row contains column coordinates. Indices + are ordered based on rows and then columns. The lower triangular part of the matrix + is defined as the elements on and below the diagonal. The argument k controls which + diagonal to consider. If k = 0, all elements on and below the main diagonal are + retained. A positive value excludes just as many diagonals below the main diagonal, + and similarly a negative value includes just as many diagonals above the main + diagonal. The main diagonal are the set of indices {(i,i)} for i∈[0,min{n_rows, + n_cols}−1]. + + Notes + ----- + Primary purpose of this function is to slice an array of shape (n,m). See + https://numpy.org/doc/stable/reference/generated/numpy.tril_indices.html + for examples + + Tensorflow does not support slicing 2-D tensor with tuple of tensor of indices Parameters ---------- - shape - shape of the tensor to generate - rank - rank of the CP decomposition - full - if True, a full tensor is returned - otherwise, the decomposed tensor is returned - orthogonal - if True, creates a tensor with orthogonal components - seed - seed for generating random numbers + n_rows + number of rows in the 2-d matrix. + n_cols + number of columns in the 2-d matrix. If None n_cols will be the same as n_rows + k + number of shifts from the main diagonal. k = 0 includes main diagonal, + k > 0 moves downward and k < 0 moves upward + device + device on which to place the created array. Default: ``None``. Returns ------- - ivy.CPTensor - """ - rank = ivy.CPTensor.validate_cp_rank(shape, rank) - if (rank > min(shape)) and orthogonal: - warnings.warn( - "Can only construct orthogonal tensors when rank <= min(shape) but got " - f"a tensor with min(shape)={min(shape)} < rank={rank}" - ) + ret + an 2xN shape, tuple of two N dimensional, where first subarray (i.e. ret[0]) + contains row coordinates of all indices and the second subarray (i.e ret[1]) + contains columns indices. - factors = [ - (ivy.random_uniform(shape=(s, rank), dtype=dtype, seed=seed)) for s in shape - ] - weights = ivy.ones((rank,), dtype=dtype) - if orthogonal: - factors = [ivy.qr(factor)[0] for factor in factors] + Function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - if full: - return ivy.CPTensor.cp_to_tensor((weights, factors)) - elif normalise_factors: - return ivy.CPTensor.cp_normalize((weights, factors)) - else: - return ivy.CPTensor((weights, factors)) + Examples + -------- + >>> x = ivy.tril_indices(4,4,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,1) + >>> print(x) + (ivy.array([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(4,4,-2) + >>> print(x) + (ivy.array([2, 3, 3]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,2,0) + >>> print(x) + (ivy.array([0, 1, 1, 2, 2, 3, 3]), + ivy.array([0, 0, 1, 0, 1, 0, 1])) + + >>> x = ivy.tril_indices(2,4,0) + >>> print(x) + (ivy.array([0, 1, 1]), ivy.array([0, 0, 1])) + + >>> x = ivy.tril_indices(4,-4,0) + >>> print(x) + (ivy.array([]), ivy.array([])) + + >>> x = ivy.tril_indices(4,4,100) + >>> print(x) + (ivy.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + ivy.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])) + + >>> x = ivy.tril_indices(2,4,-100) + >>> print(x) + (ivy.array([]), ivy.array([])) + """ + return current_backend().tril_indices(n_rows, n_cols, k, device=device) @handle_nestable @@ -910,3 +784,135 @@ def trilu( instances in place of any of the arguments. """ return current_backend(x).trilu(x, k=k, upper=upper, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_min( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the minimum along segments of an array. Segments are defined by an integer + array of segment IDs. + + Note + ---- + If the given segment ID `i` is negative, then the corresponding + value is dropped, and will not be included in the result. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented min operation. + For each segment, it computes the min value in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_min(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_nestable +@to_native_arrays_and_back +def unsorted_segment_sum( + data: Union[ivy.Array, ivy.NativeArray], + segment_ids: Union[ivy.Array, ivy.NativeArray], + num_segments: Union[int, ivy.Array, ivy.NativeArray], +) -> ivy.Array: + """ + Compute the sum of elements along segments of an array. Segments are defined by an + integer array of segment IDs. + + Parameters + ---------- + data + The array from which to gather values. + + segment_ids + Must be in the same size with the first dimension of `data`. Has to be + of integer data type. The index-th element of `segment_ids` array is + the segment identifier for the index-th element of `data`. + + num_segments + An integer or array representing the total number of distinct segment IDs. + + Returns + ------- + ret + The output array, representing the result of a segmented sum operation. + For each segment, it computes the sum of values in `data` where `segment_ids` + equals to segment ID. + """ + return ivy.current_backend().unsorted_segment_sum(data, segment_ids, num_segments) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def vorbis_window( + window_length: Union[ivy.Array, ivy.NativeArray], + *, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array that contains a vorbis power complementary window of size + window_length. + + Parameters + ---------- + window_length + the length of the vorbis window. + dtype + data type of the returned array. By default float32. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Input array with the vorbis window. + + Examples + -------- + >>> ivy.vorbis_window(3) + ivy.array([0.38268346, 1. , 0.38268352]) + + >>> ivy.vorbis_window(5) + ivy.array([0.14943586, 0.8563191 , 1. , 0.8563191, 0.14943568]) + """ + return ivy.current_backend().vorbis_window(window_length, dtype=dtype, out=out) + + +hamming_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "handle_device_shifting", + ), + "to_skip": (), +} +indices.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} diff --git a/ivy/functional/ivy/experimental/elementwise.py b/ivy/functional/ivy/experimental/elementwise.py index 613fe70acc6b2..b75f71c3450a5 100644 --- a/ivy/functional/ivy/experimental/elementwise.py +++ b/ivy/functional/ivy/experimental/elementwise.py @@ -21,223 +21,199 @@ @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument @to_native_arrays_and_back -@handle_array_function -def lgamma( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def allclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> bool: """ - Compute the natural logarithm of the absolute value of the gamma function on x. + Return a True if the two arrays are element-wise equal within given tolerance; + otherwise False. + + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(x2)) and the absolute difference + atol are added together to compare against the absolute difference + between x1 and x2. + The default atol is not appropriate for comparing numbers that are + much smaller than one Parameters ---------- - x - input array. Should have a floating-point data type. + x1 + First input array. + x2 + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in x1 will be + considered equal to NaN's in x2 in the output array. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - an array containing the natural log of Gamma(x) of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. + Returns True if the two arrays are equal within the given tolerance; + False otherwise. Examples -------- - >>> x = ivy.array([1.6, 2.6, 3.5]) - >>> y = x.lgamma() + >>> x1 = ivy.array([1e10, 1e-7]) + >>> x2 = ivy.array([1.00001e10, 1e-8]) + >>> y = ivy.allclose(x1, x2) >>> print(y) - ivy.array([-0.11259177, 0.3574118 , 1.20097363]) + ivy.array(False) - >>> x = ivy.array([1., 2., 3. ]) - >>> y = x.lgamma() + >>> x1 = ivy.array([1.0, ivy.nan]) + >>> x2 = ivy.array([1.0, ivy.nan]) + >>> y = ivy.allclose(x1, x2, equal_nan=True) >>> print(y) - ivy.array([0. ,0. ,0.69314718]) + ivy.array(True) - >>> x = ivy.array([4.5, -4, -5.6]) - >>> x.lgamma(out = x) - >>> print(x) - ivy.array([2.45373654, inf, -4.6477685 ]) + >>> x1 = ivy.array([1e-10, 1e-10]) + >>> x2 = ivy.array([1.00001e-10, 1e-10]) + >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) + >>> print(y) + ivy.array(True) """ - return ivy.current_backend(x).lgamma(x, out=out) + return ivy.current_backend().allclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def sinc( +@inputs_to_ivy_arrays +def binarizer( x: Union[ivy.Array, ivy.NativeArray], /, *, + threshold: float = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate an implementation-dependent approximation of the principal value of the - normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain - ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element - ``x_i`` is assumed to be expressed in radians. - - **Special cases** - - For floating-point operands, - - - If x_i is NaN, the result is NaN. - - If ``x_i`` is ``0``, the result is ``1``. - - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. + Map the values of the input tensor to either 0 or 1, element-wise, based on the + outcome of a comparison against a threshold value. Parameters ---------- x - input array. Should have a floating-point data type. + Data to be binarized + threshold + Values greater than this are + mapped to 1, others to 0. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - an array containing the normalized sinc function of each element in x. - The returned array must have a floating-point data type determined - by :ref:`type-promotion`. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = x.sinc() - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - >>> x = ivy.array([1.5, 0.5, -1.5]) - >>> y = ivy.zeros(3) - >>> ivy.sinc(x, out=y) - >>> print(y) - ivy.array([-0.212,0.637,-0.212]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) - >>> y = ivy.sinc(x) - >>> print(y) - ivy.array([0.637,-0.212,0.127,-0.0909]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), - ... b=ivy.array([3.5, 4.5, 5.5])) - >>> y = x.sinc() - >>> print(y) - { - a: ivy.array([0.637,-0.212,0.127]), - b: ivy.array([-0.0909,0.0707,-0.0579]) - } + Binarized output data """ - return ivy.current_backend(x).sinc(x, out=out) + xc = ivy.copy_array(x, out=out) + if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": + xc = ivy.astype(xc, ivy.default_float_dtype()) + if ivy.is_complex_dtype(xc): + xc = ivy.abs(xc) + ret = ivy.where(xc > threshold, 1, 0) + return ret +@handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def fmax( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], +def conj( + x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the - case where one of the elements is NaN. ivy.maximum returns the NaN element while - ivy.fmax returns the non-NaN element. + Return the complex conjugate for each element ``x_i`` of the input array ``x``. - Parameters - ---------- - x1 - First input array. - x2 - Second input array. - out - optional output array, for writing the result to. + For complex number of the form - Returns - ------- - ret - Array with element-wise maximums. + .. math:: + a + bj - Examples - -------- - >>> x1 = ivy.array([2, 3, 4]) - >>> x2 = ivy.array([1, 5, 2]) - >>> ivy.fmax(x1, x2) - ivy.array([ 2., 5., 4.]) + the complex conjugate is defined as - >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) - >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) - >>> ivy.fmax(x1, x2) - ivy.array([ 0., 0., nan]) - """ - return ivy.current_backend().fmax(x1, x2, out=out) + .. math:: + a - bj + Hence, the returned conjugates must be computed by negating + the imaginary component of each element ``x_i`` -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def float_power( - x1: Union[ivy.Array, float, list, tuple], - x2: Union[ivy.Array, float, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must - be broadcastable to the same shape. This differs from the power function in that - integers, float16, and float32 are promoted to floats with a minimum precision of - float64 so that the result is always inexact. + This method conforms to the + `Array API Standard `_. + This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Parameters ---------- - x1 - Array-like with elements to raise in power. - x2 - Array-like of exponents. If x1.shape != x2.shape, - they must be broadcastable to a common shape - (which becomes the shape of the output). + x + input array. out optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - The bases in x1 raised to the exponents in x2. - This is a scalar if both x1 and x2 are scalars + an arrray of the same dtype as the input array with + the complex conjugates of the complex values present + in the input array. If x is a scalar then a scalar + will be returned. + + The descriptions above assume an array input for simplicity, but + the method also accepts :class:`ivy.Container` instances + in place of: class:`ivy.Array` or :class:`ivy.NativeArray` + instances, as shown in the type hints and also the examples below. + Examples -------- - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> ivy.float_power(x1, 3) - ivy.array([1., 8., 27., 64., 125.]) - >>> x1 = ivy.array([1, 2, 3, 4, 5]) - >>> x2 = ivy.array([2, 3, 3, 2, 1]) - >>> ivy.float_power(x1, x2) - ivy.array([1., 8., 27., 16., 5.]) + With :class:`ivy.Array` inputs: + >>> x = ivy.array([4.2-0j, 3j, 7+5j]) + >>> z = ivy.conj(x) + >>> print(z) + ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) + + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), + ... b=ivy.array([5j, 5.32-6.55j, 3.001])) + >>> z = ivy.conj(x) + >>> print(z) + { + a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), + b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) + } """ - return ivy.current_backend().float_power(x1, x2, out=out) + return ivy.current_backend(x).conj(x, out=out) @handle_exceptions @@ -319,230 +295,31 @@ def count_nonzero( result as dimensions with size one. With this option, the result will broadcast correctly against the input array. dtype - optional output dtype. Default is of type integer. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Number of non-zero values in the array along a given axis. Otherwise, - the total number of non-zero values in the array is returned. - - Examples - -------- - >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) - >>> ivy.count_nonzero(a) - ivy.array(7) - >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) - >>> ivy.count_nonzero(a, axis=0) - ivy.array([1, 2, 2, 2]) - >>> a = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) - >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) - ivy.array([[[3, 4]]]) - """ - return ivy.current_backend().count_nonzero( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nansum( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[Tuple[int, ...], int]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as - zero. - - Parameters - ---------- - x - Input array. - axis - Axis or axes along which the sum is computed. - The default is to compute the sum of the flattened array. - dtype - The type of the returned array and of the accumulator in - which the elements are summed. By default, the dtype of input is used. - keepdims - If this is set to True, the axes which are reduced are left - in the result as dimensions with size one. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - A new array holding the result is returned unless out is specified, - in which it is returned. - - Examples - -------- - >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) - >>> ivy.nansum(a) - 10.0 - >>> ivy.nansum(a, axis=0) - ivy.array([2.1, 5.8, 2.1]) - >>> ivy.nansum(a, axis=1) - ivy.array([5.5, 4.5]) - """ - return ivy.current_backend().nansum( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def isclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], - /, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a boolean array where two arrays are element-wise equal within a tolerance. - - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(b)) and the absolute difference - atol are added together to compare against the absolute difference - between a and b. - The default atol is not appropriate for comparing numbers that are - much smaller than one - - Parameters - ---------- - a - First input array. - b - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in a will be - considered equal to NaN's in b in the output array. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - Returns a boolean array of where a and b are equal within the given - tolerance. If both a and b are scalars, returns a single boolean value. - - Examples - -------- - >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) - ivy.array([True, False]) - >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) - ivy.array([True, True]) - >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) - ivy.array([False, False]) - >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) - ivy.array([False, True]) - """ - return ivy.current_backend().isclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def signbit( - x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return element-wise True where signbit is set (less than zero). - - Parameters - ---------- - x - Array-like input. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Output array, or reference to out if that was supplied. - This is a scalar if x is a scalar. - - Examples - -------- - >>> x = ivy.array([1, -2, 3]) - >>> ivy.signbit(x) - ivy.array([False, True, False]) - """ - return ivy.current_backend(x).signbit(x, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def hypot( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Return the hypotenuse given the two sides of a right angle triangle. - - Parameters - ---------- - x1 - The first input array - x2 - The second input array + optional output dtype. Default is of type integer. + out + optional output array, for writing the result to. Returns ------- ret - An array with the hypotenuse + Number of non-zero values in the array along a given axis. Otherwise, + the total number of non-zero values in the array is returned. Examples -------- - >>> a = ivy.array([3.0, 4.0, 5.0]) - >>> b = ivy.array([4.0, 5.0, 6.0]) - >>> ivy.hypot(a, b) - ivy.array([5.0, 6.4031, 7.8102]) + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a) + ivy.array(7) + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a, axis=0) + ivy.array([1, 2, 2, 2]) + >>> a = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) + ivy.array([[[3, 4]]]) """ - return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) + return ivy.current_backend().count_nonzero( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) @handle_backend_invalid @@ -602,45 +379,26 @@ def diff( @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def allclose( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def digamma( + x: Union[ivy.Array, ivy.NativeArray], /, *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, out: Optional[ivy.Array] = None, -) -> bool: +) -> ivy.Array: """ - Return a True if the two arrays are element-wise equal within given tolerance; - otherwise False. + Compute the logarithmic derivative of the gamma function at x. - The tolerance values are positive, typically very small numbers. - The relative difference (rtol * abs(x2)) and the absolute difference - atol are added together to compare against the absolute difference - between x1 and x2. - The default atol is not appropriate for comparing numbers that are - much smaller than one + Note + ---- + The Ivy version only accepts real-valued inputs. Parameters ---------- - x1 - First input array. - x2 - Second input array. - rtol - The relative tolerance parameter. - atol - The absolute tolerance parameter. - equal_nan - Whether to compare NaN's as equal. If True, NaN's in x1 will be - considered equal to NaN's in x2 in the output array. + x + Input array. out Alternate output array in which to place the result. The default is None. @@ -648,32 +406,16 @@ def allclose( Returns ------- ret - Returns True if the two arrays are equal within the given tolerance; - False otherwise. + Array with values computed from digamma function from + input arrays' values, element-wise. Examples -------- - >>> x1 = ivy.array([1e10, 1e-7]) - >>> x2 = ivy.array([1.00001e10, 1e-8]) - >>> y = ivy.allclose(x1, x2) - >>> print(y) - ivy.array(False) - - >>> x1 = ivy.array([1.0, ivy.nan]) - >>> x2 = ivy.array([1.0, ivy.nan]) - >>> y = ivy.allclose(x1, x2, equal_nan=True) - >>> print(y) - ivy.array(True) - - >>> x1 = ivy.array([1e-10, 1e-10]) - >>> x2 = ivy.array([1.00001e-10, 1e-10]) - >>> y = ivy.allclose(x1, x2, rtol=0.005, atol=0.0) - >>> print(y) - ivy.array(True) + >>> x = ivy.array([.9, 3, 3.2]) + >>> y = ivy.digamma(x) + ivy.array([-0.7549271 0.92278427 0.9988394]) """ - return ivy.current_backend().allclose( - a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out - ) + return ivy.current_backend(x).digamma(x, out=out) @handle_backend_invalid @@ -714,22 +456,70 @@ def fix( return ivy.current_backend(x).fix(x, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def nextafter( +def float_power( + x1: Union[ivy.Array, float, list, tuple], + x2: Union[ivy.Array, float, list, tuple], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Raise each base in x1 to the positionally-corresponding power in x2. x1 and x2 must + be broadcastable to the same shape. This differs from the power function in that + integers, float16, and float32 are promoted to floats with a minimum precision of + float64 so that the result is always inexact. + + Parameters + ---------- + x1 + Array-like with elements to raise in power. + x2 + Array-like of exponents. If x1.shape != x2.shape, + they must be broadcastable to a common shape + (which becomes the shape of the output). + out + optional output array, for writing the result to. + + Returns + ------- + ret + The bases in x1 raised to the exponents in x2. + This is a scalar if both x1 and x2 are scalars + + Examples + -------- + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> ivy.float_power(x1, 3) + ivy.array([1., 8., 27., 64., 125.]) + >>> x1 = ivy.array([1, 2, 3, 4, 5]) + >>> x2 = ivy.array([2, 3, 3, 2, 1]) + >>> ivy.float_power(x1, x2) + ivy.array([1., 8., 27., 16., 5.]) + """ + return ivy.current_backend().float_power(x1, x2, out=out) + + +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fmax( x1: Union[ivy.Array, ivy.NativeArray], x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> bool: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Return the next floating-point value after x1 towards x2, element-wise. + Compute the element-wise maximums of two arrays. Differs from ivy.maximum in the + case where one of the elements is NaN. ivy.maximum returns the NaN element while + ivy.fmax returns the non-NaN element. Parameters ---------- @@ -738,22 +528,26 @@ def nextafter( x2 Second input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. Returns ------- ret - The next representable values of x1 in the direction of x2. + Array with element-wise maximums. Examples -------- - >>> x1 = ivy.array([1.0e-50, 2.0e+50]) - >>> x2 = ivy.array([2.0, 1.0]) - >>> ivy.nextafter(x1, x2) - ivy.array([1.4013e-45., 3.4028e+38]) + >>> x1 = ivy.array([2, 3, 4]) + >>> x2 = ivy.array([1, 5, 2]) + >>> ivy.fmax(x1, x2) + ivy.array([ 2., 5., 4.]) + + >>> x1 = ivy.array([ivy.nan, 0, ivy.nan]) + >>> x2 = ivy.array([0, ivy.nan, ivy.nan]) + >>> ivy.fmax(x1, x2) + ivy.array([ 0., 0., nan]) """ - return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) + return ivy.current_backend().fmax(x1, x2, out=out) @handle_exceptions @@ -763,41 +557,35 @@ def nextafter( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def zeta( +def frexp( x: Union[ivy.Array, ivy.NativeArray], - q: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> bool: + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Compute the Hurwitz zeta function elementwisely with each pair of floats in two - arrays. + Decompose the elements of x into mantissa and twos exponent. Parameters ---------- x - First input array. - q - Second input array, must have the same shape as the first input array + Input array. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - Array with values computed from zeta function from - input arrays' values. + A tuple of two arrays, the mantissa and the twos exponent. Examples -------- - >>> x = ivy.array([5.0, 3.0]) - >>> q = ivy.array([2.0, 2.0]) - >>> ivy.zeta(x, q) - ivy.array([0.0369, 0.2021]) + >>> x = ivy.array([1, 2, 3]) + >>> ivy.frexp(x) + (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) """ - return ivy.current_backend(x, q).zeta(x, q, out=out) + return ivy.current_backend(x).frexp(x, out=out) @handle_exceptions @@ -880,91 +668,41 @@ def gradient( ) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def xlogy( - x: Union[ivy.Array, ivy.NativeArray], - y: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> bool: - """ - Compute x*log(y) element-wise so that the result is 0 if x = 0. - - Parameters - ---------- - x - First input array. - y - Second input array. - out - Alternate output array in which to place the result. - The default is None. - - Returns - ------- - ret - The next representable values of x1 in the direction of x2. - - Examples - -------- - >>> x = ivy.zeros(3) - >>> y = ivy.array([-1.0, 0.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([0.0, 0.0, 0.0]) - - >>> x = ivy.array([1.0, 2.0, 3.0]) - >>> y = ivy.array([3.0, 2.0, 1.0]) - >>> ivy.xlogy(x, y) - ivy.array([1.0986, 1.3863, 0.0000]) - """ - return ivy.current_backend(x, y).xlogy(x, y, out=out) - - -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@inputs_to_ivy_arrays -def binarizer( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - threshold: float = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def hypot( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Map the values of the input tensor to either 0 or 1, element-wise, based on the - outcome of a comparison against a threshold value. + Return the hypotenuse given the two sides of a right angle triangle. Parameters ---------- - x - Data to be binarized - threshold - Values greater than this are - mapped to 1, others to 0. - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + x1 + The first input array + x2 + The second input array Returns ------- ret - Binarized output data + An array with the hypotenuse + + Examples + -------- + >>> a = ivy.array([3.0, 4.0, 5.0]) + >>> b = ivy.array([4.0, 5.0, 6.0]) + >>> ivy.hypot(a, b) + ivy.array([5.0, 6.4031, 7.8102]) """ - xc = ivy.copy_array(x, out=out) - if ivy.is_bool_dtype(xc) and ivy.current_backend_str() == "torch": - xc = ivy.astype(xc, ivy.default_float_dtype()) - if ivy.is_complex_dtype(xc): - xc = ivy.abs(xc) - ret = ivy.where(xc > threshold, 1, 0) - return ret + return ivy.current_backend(x1, x2).hypot(x1, x2, out=out) @handle_exceptions @@ -974,80 +712,63 @@ def binarizer( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def conj( - x: Union[ivy.Array, ivy.NativeArray], +def isclose( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], /, *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the complex conjugate for each element ``x_i`` of the input array ``x``. - - For complex number of the form - - .. math:: - a + bj - - the complex conjugate is defined as - - .. math:: - a - bj - - Hence, the returned conjugates must be computed by negating - the imaginary component of each element ``x_i`` - - This method conforms to the - `Array API Standard `_. - This docstring is an extension of the - `docstring `_ - in the standard. + Return a boolean array where two arrays are element-wise equal within a tolerance. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The tolerance values are positive, typically very small numbers. + The relative difference (rtol * abs(b)) and the absolute difference + atol are added together to compare against the absolute difference + between a and b. + The default atol is not appropriate for comparing numbers that are + much smaller than one Parameters ---------- - x - input array. + a + First input array. + b + Second input array. + rtol + The relative tolerance parameter. + atol + The absolute tolerance parameter. + equal_nan + Whether to compare NaN's as equal. If True, NaN's in a will be + considered equal to NaN's in b in the output array. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - an arrray of the same dtype as the input array with - the complex conjugates of the complex values present - in the input array. If x is a scalar then a scalar - will be returned. - - The descriptions above assume an array input for simplicity, but - the method also accepts :class:`ivy.Container` instances - in place of: class:`ivy.Array` or :class:`ivy.NativeArray` - instances, as shown in the type hints and also the examples below. - + Returns a boolean array of where a and b are equal within the given + tolerance. If both a and b are scalars, returns a single boolean value. Examples -------- - With :class:`ivy.Array` inputs: - >>> x = ivy.array([4.2-0j, 3j, 7+5j]) - >>> z = ivy.conj(x) - >>> print(z) - ivy.array([4.2-0.j, 0. -3.j, 7. -5.j]) - - With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-6.7-7j, 0.314+0.355j, 1.23]), - ... b=ivy.array([5j, 5.32-6.55j, 3.001])) - >>> z = ivy.conj(x) - >>> print(z) - { - a: ivy.array([-6.7+7.j, 0.314-0.355j, 1.23-0.j]), - b: ivy.array([0.-5.j, 5.32+6.55j, 3.001-0.j]) - } + >>> ivy.isclose([1e10,1e-7], [1.00001e10,1e-8]) + ivy.array([True, False]) + >>> ivy.isclose([1.0, ivy.nan], [1.0, ivy.nan], equal_nan=True) + ivy.array([True, True]) + >>> ivy.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) + ivy.array([False, False]) + >>> ivy.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], rtol=0.005, atol=0.0) + ivy.array([False, True]) """ - return ivy.current_backend(x).conj(x, out=out) + return ivy.current_backend().isclose( + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, out=out + ) @handle_exceptions @@ -1211,15 +932,151 @@ def lerp( return ivy.add(input, ivy.multiply(weight, ivy.subtract(end, input)), out=out) -lerp.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +def lgamma( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the natural logarithm of the absolute value of the gamma function on x. + + Parameters + ---------- + x + input array. Should have a floating-point data type. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + an array containing the natural log of Gamma(x) of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. + + Examples + -------- + >>> x = ivy.array([1.6, 2.6, 3.5]) + >>> y = x.lgamma() + >>> print(y) + ivy.array([-0.11259177, 0.3574118 , 1.20097363]) + + >>> x = ivy.array([1., 2., 3. ]) + >>> y = x.lgamma() + >>> print(y) + ivy.array([0. ,0. ,0.69314718]) + + >>> x = ivy.array([4.5, -4, -5.6]) + >>> x.lgamma(out = x) + >>> print(x) + ivy.array([2.45373654, inf, -4.6477685 ]) + """ + return ivy.current_backend(x).lgamma(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def modf( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[Tuple[ivy.Array, ivy.Array]] = None, +) -> Tuple[ivy.Array, ivy.Array]: + """ + Decompose the elements of x into fractional and integral parts. + + Parameters + ---------- + x + Input array. + out + Optional output array for writing the result to. + It must have a shape that the inputs broadcast to. + + Returns + ------- + ret + A tuple of two arrays, the fractional and integral parts. + + Examples + -------- + >>> x = ivy.array([1.5, 2.7, 3.9]) + >>> ivy.modf(x) + (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) + """ + return ivy.current_backend(x).modf(x, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@infer_dtype +@handle_device_shifting +def nansum( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[Tuple[int, ...], int]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as + zero. + + Parameters + ---------- + x + Input array. + axis + Axis or axes along which the sum is computed. + The default is to compute the sum of the flattened array. + dtype + The type of the returned array and of the accumulator in + which the elements are summed. By default, the dtype of input is used. + keepdims + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + A new array holding the result is returned unless out is specified, + in which it is returned. + + Examples + -------- + >>> a = ivy.array([[ 2.1, 3.4, ivy.nan], [ivy.nan, 2.4, 2.1]]) + >>> ivy.nansum(a) + 10.0 + >>> ivy.nansum(a, axis=0) + ivy.array([2.1, 5.8, 2.1]) + >>> ivy.nansum(a, axis=1) + ivy.array([5.5, 4.5]) + """ + return ivy.current_backend().nansum( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out + ) @handle_exceptions @@ -1229,112 +1086,154 @@ def lerp( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def frexp( - x: Union[ivy.Array, ivy.NativeArray], +def nextafter( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> bool: """ - Decompose the elements of x into mantissa and twos exponent. + Return the next floating-point value after x1 towards x2, element-wise. Parameters ---------- - x - Input array. + x1 + First input array. + x2 + Second input array. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + Alternate output array in which to place the result. + The default is None. Returns ------- ret - A tuple of two arrays, the mantissa and the twos exponent. + The next representable values of x1 in the direction of x2. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.frexp(x) - (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2])) + >>> x1 = ivy.array([1.0e-50, 2.0e+50]) + >>> x2 = ivy.array([2.0, 1.0]) + >>> ivy.nextafter(x1, x2) + ivy.array([1.4013e-45., 3.4028e+38]) """ - return ivy.current_backend(x).frexp(x, out=out) + return ivy.current_backend(x1, x2).nextafter(x1, x2, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -def modf( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def signbit( + x: Union[ivy.Array, ivy.NativeArray, float, int, list, tuple], /, *, - out: Optional[Tuple[ivy.Array, ivy.Array]] = None, -) -> Tuple[ivy.Array, ivy.Array]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Decompose the elements of x into fractional and integral parts. + Return element-wise True where signbit is set (less than zero). Parameters ---------- x - Input array. + Array-like input. out - Optional output array for writing the result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - A tuple of two arrays, the fractional and integral parts. + Output array, or reference to out if that was supplied. + This is a scalar if x is a scalar. Examples -------- - >>> x = ivy.array([1.5, 2.7, 3.9]) - >>> ivy.modf(x) - (ivy.array([0.5, 0.7, 0.9]), ivy.array([1, 2, 3])) + >>> x = ivy.array([1, -2, 3]) + >>> ivy.signbit(x) + ivy.array([False, True, False]) """ - return ivy.current_backend(x).modf(x, out=out) + return ivy.current_backend(x).signbit(x, out=out) @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -def digamma( +@handle_device_shifting +def sinc( x: Union[ivy.Array, ivy.NativeArray], /, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the logarithmic derivative of the gamma function at x. + Calculate an implementation-dependent approximation of the principal value of the + normalized sinc function, having domain ``(-infinity, +infinity)`` and codomain + ``[-0.217234, 1]``, for each element ``x_i`` of the input array ``x``. Each element + ``x_i`` is assumed to be expressed in radians. - Note - ---- - The Ivy version only accepts real-valued inputs. + **Special cases** + + For floating-point operands, + + - If x_i is NaN, the result is NaN. + - If ``x_i`` is ``0``, the result is ``1``. + - If ``x_i`` is either ``+infinity`` or ``-infinity``, the result is ``NaN``. Parameters ---------- x - Input array. + input array. Should have a floating-point data type. out - Alternate output array in which to place the result. - The default is None. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Array with values computed from digamma function from - input arrays' values, element-wise. + an array containing the normalized sinc function of each element in x. + The returned array must have a floating-point data type determined + by :ref:`type-promotion`. Examples -------- - >>> x = ivy.array([.9, 3, 3.2]) - >>> y = ivy.digamma(x) - ivy.array([-0.7549271 0.92278427 0.9988394]) + With :class:`ivy.Array` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = x.sinc() + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + >>> x = ivy.array([1.5, 0.5, -1.5]) + >>> y = ivy.zeros(3) + >>> ivy.sinc(x, out=y) + >>> print(y) + ivy.array([-0.212,0.637,-0.212]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.array([0.5, 1.5, 2.5, 3.5]) + >>> y = ivy.sinc(x) + >>> print(y) + ivy.array([0.637,-0.212,0.127,-0.0909]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0.5, 1.5, 2.5]), + ... b=ivy.array([3.5, 4.5, 5.5])) + >>> y = x.sinc() + >>> print(y) + { + a: ivy.array([0.637,-0.212,0.127]), + b: ivy.array([-0.0909,0.0707,-0.0579]) + } """ - return ivy.current_backend(x).digamma(x, out=out) + return ivy.current_backend(x).sinc(x, out=out) @handle_exceptions @@ -1388,3 +1287,104 @@ def sparsify_tensor( tensor = ivy.concat([ivy.zeros(len(x) - card, dtype=x.dtype), x[-card:]], axis=0) return ivy.reshape(tensor, _shape, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def xlogy( + x: Union[ivy.Array, ivy.NativeArray], + y: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> bool: + """ + Compute x*log(y) element-wise so that the result is 0 if x = 0. + + Parameters + ---------- + x + First input array. + y + Second input array. + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + The next representable values of x1 in the direction of x2. + + Examples + -------- + >>> x = ivy.zeros(3) + >>> y = ivy.array([-1.0, 0.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([0.0, 0.0, 0.0]) + + >>> x = ivy.array([1.0, 2.0, 3.0]) + >>> y = ivy.array([3.0, 2.0, 1.0]) + >>> ivy.xlogy(x, y) + ivy.array([1.0986, 1.3863, 0.0000]) + """ + return ivy.current_backend(x, y).xlogy(x, y, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def zeta( + x: Union[ivy.Array, ivy.NativeArray], + q: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> bool: + """ + Compute the Hurwitz zeta function elementwisely with each pair of floats in two + arrays. + + Parameters + ---------- + x + First input array. + q + Second input array, must have the same shape as the first input array + out + Alternate output array in which to place the result. + The default is None. + + Returns + ------- + ret + Array with values computed from zeta function from + input arrays' values. + + Examples + -------- + >>> x = ivy.array([5.0, 3.0]) + >>> q = ivy.array([2.0, 2.0]) + >>> ivy.zeta(x, q) + ivy.array([0.0369, 0.2021]) + """ + return ivy.current_backend(x, q).zeta(x, q, out=out) + + +lerp.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} diff --git a/ivy/functional/ivy/experimental/general.py b/ivy/functional/ivy/experimental/general.py index 0cb749b476ffa..faa8df1e69b2a 100644 --- a/ivy/functional/ivy/experimental/general.py +++ b/ivy/functional/ivy/experimental/general.py @@ -13,6 +13,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _correct_ivy_callable(func): # get the current backend of the given ivy callable if ivy.nested_any( @@ -25,6 +29,10 @@ def _correct_ivy_callable(func): return func +# --- Main --- # +# ------------ # + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py index 120ebe0279784..60cb61ec63724 100644 --- a/ivy/functional/ivy/experimental/layers.py +++ b/ivy/functional/ivy/experimental/layers.py @@ -21,1301 +21,1156 @@ from ivy.functional.ivy.experimental.general import _correct_ivy_callable from ivy.utils.exceptions import handle_exceptions -_min = builtins.min -_slice = builtins.slice -_max = builtins.max +# --- Helpers --- # +# --------------- # -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool1d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NWC", - dilation: Union[int, Tuple[int]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D max pool given 3-D input x. - Parameters - ---------- - x - Input image *[batch_size, w, d_in]* if data_format is "NWC". - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) - data_format - "NWC" or "NCW". Defaults to "NWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. +def _cast_init(init, dtype): + if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): + if ivy.is_float_dtype(dtype): + info = ivy.finfo(dtype) + else: + info = ivy.iinfo(dtype) + if "float64" not in str(dtype): + init = info.max if init > 0 else info.min + return ivy.array(init, dtype=dtype) - Returns - ------- - ret - The result of the pooling operation. - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. +def _compute_idx(in_size, out_size, device): + out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) + i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) + maxlength = in_size // out_size + 1 + in_size_mod = in_size % out_size + # adaptive = True iff there are kernels with different lengths + adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) + if adaptive: + maxlength += 1 + elif in_size_mod == 0: + maxlength -= 1 + range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) + idx = ivy.expand_dims(i0, axis=-1) + range_max + if adaptive: + maxval = ivy.full_like(idx, fill_value=in_size - 1) + idx = ivy.minimum(idx, maxval) + i1 = ivy.trunc_divide( + (out_range + 1) * in_size + out_size - 1, out_size + ).astype(ivy.int64) + length = i1 - i0 + else: + length = maxlength + return idx, length, range_max, adaptive - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 4., 5., 6., 7.]], +def _compute_weight_mat( + input_size, + output_size, + scale, + align_corners, + kernel_fn, + antialias: bool, + dim_scale_factor, +): + inv_scale = 1.0 / scale + kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 + if not align_corners: + sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 + x = ( + ivy.abs( + ivy.expand_dims(sample_f) + - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) + / kernel_scale + ) + else: + sample_f = ivy.arange(output_size) * dim_scale_factor + x = ivy.abs( + ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) + ) / (kernel_scale) + weights = kernel_fn(x) + total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) + weights = ivy.where( + ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), + ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), + 0, + ) + input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 + return ivy.where( + ivy.expand_dims( + ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) + ), + weights, + 0, + ) - [[16., 17., 18., 19.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa - ivy.array([[[ 1., 3.], - [ 5., 7.], - [ 9., 11.]], - [[13., 15.], - [17., 19.], - [21., 23.]]]) - """ - return ivy.current_backend(x).max_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): + def _pad(arr, pads, pad_value): + out = ivy.astype( + ivy.pad( + arr, + ivy.maximum(0, pads).to_list(), + mode="constant", + constant_values=ivy.to_scalar(pad_value), + ), + arr.dtype, + ) + slices = tuple( + _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) + for (lo, hi), dim in zip(pads, arr.shape) + ) + return out[slices] + if ( + _min(lhs.ndim, len(rhs_shape)) < 2 + or lhs.ndim != len(rhs_shape) + or lhs.shape[1] != rhs_shape[1] + ): + raise ValueError("Dimension mismatch") + if len(window_strides) != len(rhs_shape) - 2: + raise ValueError("Wrong number of strides for spatial dimensions") + if len(pads) != len(rhs_shape) - 2: + raise ValueError("Wrong number of pads for spatial dimensions") -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_unpool1d( - x: ivy.Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D max unpooling given the 1-D pooled input x and its indices. + lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) + in_shape = lhs.shape[2:] + filter_shape = rhs_shape[2:] + dim = len(filter_shape) - Parameters - ---------- - x - Pooled input image *[batch_size, w, d_in]*. - indices - Indices obtained from the corresponding max pooling operation. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - out - optional output array, for writing the result to. + out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() + view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] - Returns - ------- - ret - The result of the unpooling operation. + out_shape = [ + (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) + ] + view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] - Both the description and the type hints above assume an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + view = ivy.as_strided(lhs, view_shape, view_strides) - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') - >>> print(pool_result) - ivy.array([[[ 4., 5., 6., 7.], - [ 8., 9., 10., 11.]], + view_axes = list(range(view.ndim)) + sum_axes = view_axes[-dim - 1 :] + rhs_axes = [view.ndim] + sum_axes + out_axes = [0, view.ndim] + list(range(1, dim + 1)) - [[16., 17., 18., 19.], - [20., 21., 22., 23.]]]) - >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') - >>> print(unpool_result) - ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], - [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], + return view, view_axes, rhs_axes, out_axes - [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], - [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) - """ - return ivy.current_backend(x).max_unpool1d( - x, indices, kernel, strides, padding, data_format=data_format, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool2d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 2-D max pool given 4-D input x. - - Parameters - ---------- - x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. - - Returns - ------- - ret - The result of the pooling operation. - - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 2., 3.], - [ 4., 5.], - [ 4., 5.]]], +def _cubic_kernel(x): + out = ((1.5 * x - 2.5) * x) * x + 1.0 + out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) + return ivy.where(x >= 2.0, 0.0, out) - [[[ 8., 9.], - [10., 11.], - [10., 11.]]]]) - - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[16., 17.]], - [[22., 23.]]], +def _dilate(operand, factors, fill_value): + outspace = list(operand.shape[:2]) + [ + shape + (factors[i] - 1) * (shape - 1) + for i, shape in enumerate(operand.shape[2:]) + ] + out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) + lhs_slices = tuple(_slice(None, None, step) for step in factors) + out[(_slice(None),) * 2 + lhs_slices] = operand + return out - [[[40., 41.]], +def _dim_scale_factor(input_size, output_size, align_corners, scales): + if align_corners: + if output_size > 1: + dim_scale_factor = (input_size - 1) / (output_size - 1) + else: + dim_scale_factor = 0.0 + else: + dim_scale_factor = ( + input_size / (input_size * scales) + if scales is not None + else input_size / output_size + ) + return dim_scale_factor - [[46., 47.]]]]) - """ - return ivy.current_backend(x).max_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _expand_to_dim(x, dim): + for _ in range(dim - len(x.shape)): + x = ivy.expand_dims(x, axis=-1) + return x -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def max_pool3d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]], - padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], - /, - *, - data_format: str = "NDHWC", - dilation: Union[int, Tuple[int, ...]] = 1, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 3-D max pool given 5-D input x. - Parameters - ---------- - x - Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - "SAME" or "VALID" indicating the algorithm; int, or list of tuple - indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". - dilaton - The stride between elements within a sliding window, must be > 0. - ceil_mode - If True, ceil is used instead of floor to compute the output shape. - This ensures that every element in 'x' is covered by a sliding window. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. +def _get_identity(func, dtype, init): + func_name = func.__name__ + if func_name in identities: + identity = identities[func_name] + return _cast_init(identity, dtype) + return init - Returns - ------- - ret - The result of the pooling operation. - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. +def _get_size(scale_factor, size, dims, x_shape): + if scale_factor is not None: + if isinstance(scale_factor, (float, int)): + scale_factor = [scale_factor] * dims + elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: + scale_factor = [scale_factor[0]] * dims - Examples - -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) - ivy.array([[[[[14., 15.]]]], + size = tuple( + [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] + ) + else: + size = (size,) * dims if isinstance(size, int) else tuple(size) + return size +def _interpolate_with_kernel( + x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode +): + spatial_dims = [2 + i for i in range(dims)] + equation = generate_einsum_equation(dims) + kernel_func = get_interpolate_kernel(mode) + output_shape = tuple(input_shape[:2]) + size + operands = [] + for i, d in enumerate(spatial_dims): + m = input_shape[d] + n = output_shape[d] + dim_scale_factor = _dim_scale_factor( + m, + n, + align_corners, + scale_factor[i] if scale_factor is not None else None, + ) + w = _compute_weight_mat( + m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor + ).astype(x.dtype) + operands.append(w) + return ivy.einsum(equation, x, *operands) - [[[[38., 39.]]]]]) - >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) - ivy.array([[[[[14., 15.]]], +def _lanczos_kernel(radius, x): + y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) + out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) + return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) - [[[22., 23.]]]], +def _mask(vals, length, range_max, dim, mask_value=0.0): + if isinstance(length, int): + return vals, length + else: + assert dim < 0 + mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) + if dim == -2: + mask = _expand_to_dim(mask, 4) + vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) + length = _expand_to_dim(length, -dim) + return vals, length - [[[[38., 39.]]], +def _mitchellcubic_kernel(x): + absx = abs(x) + if absx < 1: + return (7 * absx**3 - 12 * absx**2 + 6) / 6 + elif absx < 2: + return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 + else: + return 0 - [[[46., 47.]]]]]) - """ - return ivy.current_backend(x).max_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - dilation=dilation, - ceil_mode=ceil_mode, - out=out, - ) +def _output_ceil_shape(w, f, p, s): + return math.ceil((w - f + p) / s) + 1 -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool1d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int]], - strides: Union[int, Tuple[int]], - padding: str, - /, - *, - data_format: str = "NWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - division_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 1-D avg pool given 3-D input x. +def _padding_ceil_mode(w, f, p, s, return_added_padding=False): + remaining_pixels = (w - f + sum(p)) % s + added_padding = 0 + if s > 1 and remaining_pixels != 0 and f > 1: + input_size = w + sum(p) + # making sure that the remaining pixels are supposed + # to be covered by the window + # they won't be covered if stride is big enough to skip them + if input_size - remaining_pixels - (f - 1) + s > input_size: + return p + output_shape = _output_ceil_shape( + w, + f, + sum(p), + s, + ) + # calculating new padding with ceil_output_shape + new_pad = (output_shape - 1) * s + f - w + # updating pad_list with new padding by adding it to the end + added_padding = new_pad - sum(p) + p = ( + p[0], + p[1] + added_padding, + ) + if return_added_padding: + return p, added_padding + return p - Parameters - ---------- - x - Input image *[batch_size, w, d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimension paddings. - data_format - NWC" or "NCW". Defaults to "NWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - division_override - If specified, it will be used as the divisor, - otherwise kernel_size will be used. - out - optional output array, for writing the result to. - Returns - ------- - ret - The result of the pooling operation. +def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): + if padding.upper() == "SAME": + out_shape = [ + math.ceil(in_size / stride) + for in_size, stride in zip(in_shape, window_strides) + ] + pad_sizes = [ + _max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size in zip( + out_shape, window_strides, filter_shape, in_shape + ) + ] + return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] + else: + return [(0, 0)] * len(in_shape) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - Examples - -------- - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) - ivy.array([[[ 2., 3., 4., 5.], - [ 8., 9., 10., 11.]], +def _sum_tensors(ts): + return _reduce(ivy.add, ts) - [[14., 15., 16., 17.], - [20., 21., 22., 23.]]]) - >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) - >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) - ivy.array([[[ 2., 3., 4., 5.]], - [[14., 15., 16., 17.]]]) - """ - return ivy.current_backend(x).avg_pool1d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - division_override=division_override, - out=out, +def _tf_area_dim_scale(index, starting_index, scale, ending_index): + if index < starting_index: + dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index + else: + dim_scale = ending_index - index if index + 1 > ending_index else 1.0 + return dim_scale + + +def _tf_area_indices(dim_index, scale): + starting_index = dim_index * scale + ending_index = (dim_index + 1) * scale + rounded_indices = ( + int(starting_index), + math.ceil(ending_index), ) + return starting_index, ending_index, rounded_indices -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool2d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int]], - strides: Union[int, Tuple[int], Tuple[int, int]], - padding: str, - /, - *, - data_format: str = "NHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 2-D average pool given 4-D input x. +def _tf_area_interpolate(x, size, dims): + ret = ivy.zeros((x.shape[:2] + size)) + scale = ivy.divide(ivy.shape(x)[2:], size) + area = 1.0 / ivy.prod(scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) + sum_data = ivy.zeros( + ( + d_index[1] - d_index[0], + h_index[1] - h_index[0], + w_index[1] - w_index[0], + ) + ) + for d_ind in range(d_index[0], d_index[1]): + scale_z = _tf_area_dim_scale( + d_ind, d_in, scale[0], d_in1 + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale( + h_ind, h_in, scale[1], h_in1 + ) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[2], w_in1 + ) + sum_data[ + d_ind - d_index[0], + h_ind - h_index[0], + w_ind - w_index[0], + ] = ( + ivy.array(ch[d_ind, h_ind, w_ind]) + * scale_x + * scale_y + * scale_z + * area + ) + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) + sum_data = ivy.zeros( + (h_index[1] - h_index[0], w_index[1] - w_index[0]) + ) + for h_ind in range(h_index[0], h_index[1]): + scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale( + w_ind, w_in, scale[1], w_in1 + ) + sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( + ivy.array(ch[h_ind, w_ind]) + * scale_x + * scale_y + * area + ) + ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) + else: + for w_dim in range(size[0]): + w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) + sum_data = ivy.zeros((w_index[1] - w_index[0],)) + for w_ind in range(w_index[0], w_index[1]): + scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) + sum_data[w_ind - w_index[0]] = ( + ivy.array(ch[w_ind]) * scale_x * area + ) + ret[i, j, w_dim] = ivy.sum(sum_data) + return ret - Parameters - ---------- - x - Input image *[batch_size,h,w,d_in]*. - kernel - Size of the kernel i.e., the sliding window for each - dimension of input. *[h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list - indicating the per-dimensio paddings. - data_format - NHWC" or "NCHW". Defaults to "NHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - out - optional output array, for writing the result to. - Returns - ------- - ret - The result of the pooling operation. +def _triangle_kernel(x): + return ivy.maximum(0, 1 - ivy.abs(x)) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], +def _upsample_bicubic2d_default( + a, + output_size, + align_corners, + scale_h=None, + scale_w=None, +): + N, C, iH, iW = a.shape + oH, oW = output_size - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) - ivy.array([[[[ 8., 9.]], + def compute_scale(in_size, out_size, align_corners, scale=None): + if align_corners: + return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 + else: + return 1 / scale if scale is not None and scale > 0 else in_size / out_size - [[14., 15.]]], + def compute_source_index(scale, dst_index, align_corners): + if align_corners: + return scale * dst_index + else: + return scale * (dst_index + 0.5) - 0.5 + height_scale = compute_scale(iH, oH, align_corners, scale_h) + width_scale = compute_scale(iW, oW, align_corners, scale_w) - [[[32., 33.]], + N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) + C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) + out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) + out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) - [[38., 39.]]]]) - """ - return ivy.current_backend(x).avg_pool2d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, - ) + real_x = compute_source_index(width_scale, out_x, align_corners) + in_x = ivy.floor(real_x) + t_x = real_x - in_x + ix = ivy.astype(in_x, ivy.int64) + real_y = compute_source_index(height_scale, out_y, align_corners) + in_y = ivy.floor(real_y) + t_y = real_y - in_y + iy = ivy.astype(in_y, ivy.int64) -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def avg_pool3d( - x: Union[ivy.Array, ivy.NativeArray], - kernel: Union[int, Tuple[int], Tuple[int, int, int]], - strides: Union[int, Tuple[int], Tuple[int, int, int]], - padding: str, - /, - *, - data_format: str = "NDHWC", - count_include_pad: bool = False, - ceil_mode: bool = False, - divisor_override: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute a 3-D avg pool given 5-D input x. + iys_ofs = (iy - 1, iy, iy + 1, iy + 2) + ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) - Parameters - ---------- - x - Input volume *[batch_size,d,h,w,d_in]*. - kernel - Convolution filters *[d,h,w]*. - strides - The stride of the sliding window for each dimension of input. - padding - SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension - paddings. - data_format - NDHWC" or "NCDHW". Defaults to "NDHWC". - count_include_pad - Whether to include padding in the averaging calculation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - divisor_override - If specified, it will be used as divisor, otherwise kernel_size will be used. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + def load_bounded(ys, xs): + y_idx = ivy.clip(ys, 0, iH - 1) + x_idx = ivy.clip(xs, 0, iW - 1) + return a[N_idx, C_idx, y_idx, x_idx] - Returns - ------- - ret - The result of the pooling operation. + def get_x_interp(y): + coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) + return _upsample_cubic_interp1d(coeffs_x, t_x) - Both the description and the type hints above assumes an array input - for simplicity, but this function is *nestable*, and therefore - also accepts :class:`ivy.Container` instances in place of any of - the arguments. + coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) + result = _upsample_cubic_interp1d(coeffs_y, t_y) - Examples - -------- - >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) - >>> print(ivy.avg_pool3d(x,2,2,'VALID')) - ivy.array([[[[[ 7., 8.]]]], + return result +def _upsample_cubic_convolution1(x, A): + return ((A + 2) * x - (A + 3)) * x * x + 1 - [[[[31., 32.]]]]]) - >>> print(ivy.avg_pool3d(x,2,2,'SAME')) - ivy.array([[[[[ 7., 8.]]], +def _upsample_cubic_convolution2(x, A): + return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A - [[[19., 20.]]]], +def _upsample_cubic_interp1d(coeffs, ts): + coeffs2 = _upsample_get_cubic_coefficients(ts) + return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) - [[[[31., 32.]]], +def _upsample_get_cubic_coefficients(t): + A = -0.75 + return ( + _upsample_cubic_convolution2(t + 1.0, A), + _upsample_cubic_convolution1(t, A), + _upsample_cubic_convolution1(1.0 - t, A), + _upsample_cubic_convolution2(2.0 - t, A), + ) - [[[43., 44.]]]]]) - """ - return ivy.current_backend(x).avg_pool3d( - x, - kernel, - strides, - padding, - data_format=data_format, - count_include_pad=count_include_pad, - ceil_mode=ceil_mode, - divisor_override=divisor_override, - out=out, - ) +# --- Main --- # +# ------------ # -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def pool( - x: Union[ivy.Array, ivy.NativeArray], - window_shape: Union[int, Tuple[int], Tuple[int, int]], - pool_type: str, - /, - *, - strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - padding: str = "VALID", - data_format: Optional[str] = None, - dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, - ceil_mode: bool = False, - out: Optional[ivy.Array] = None, +@inputs_to_ivy_arrays +def adaptive_avg_pool1d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: int, ) -> ivy.Array: """ - Perform an N-D pooling operation. + Apply a 1D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - Input array to pool over. - window_shape - Shape of the pooling window. - pool_type - Type of pooling operation, either 'MAX' or 'AVG'. - strides - Strides of the pooling operation. - padding - Padding type, either 'VALID' or 'SAME'. - data_format - Data format of the input and output data, either 'NCHW' or 'NHWC'. - dilations - Dilation rate of the pooling operation. - ceil_mode - Whether to use ceil or floor for creating the output shape. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + input + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. Returns ------- - ret - The result of the pooling operation. - - Examples - -------- - >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) - >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) - ivy.array([[[[ 1., 2.], - [ 3., 4.], - [ 4., 5.]]], - [[[ 7., 8.], - [ 9., 10.], - [10., 11.]]]]) - >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) - >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) - ivy.array([[[[ 8., 9.]], - [[14., 15.]]], - [[[32., 33.]], - [[38., 39.]]]]) + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` """ - return ivy.current_backend(x).pool( - x, - window_shape, - pool_type, - strides=strides, - padding=padding, - data_format=data_format, - dilations=dilations, - ceil_mode=ceil_mode, - out=out, - ) + squeeze = False + if input.ndim == 2: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 3: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", + ) + if input.shape[-1] % output_size == 0: + stride = input.shape[-1] // output_size + kernel_size = input.shape[-1] - (output_size - 1) * stride + pooled_output = ivy.avg_pool1d( + input, kernel_size, stride, "VALID", data_format="NCW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output -@handle_exceptions -@handle_backend_invalid + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size, input.device + ) + + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., idxw]) + + if not adaptive_w: + ret = ivy.mean(vals, axis=-1) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret + + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + + ret = None + for i in range(vals.shape[-1]): + if ret is None: + ret = vals[..., i] + else: + ret = ret + vals[..., i] + pooled_output = ret / length_w.astype(ret.dtype) + + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output + + +@handle_exceptions @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def dct( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def adaptive_avg_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], +) -> ivy.Array: """ - Compute the 1D Discrete Cosine Tranformation of a given signal. + Apply a 2D adaptive average pooling over an input signal composed of several input + planes. Parameters ---------- - x - The input signal. - type - The type of the dct. Must be 1, 2, 3 or 4. - n - The lenght of the transform. If n is less than the input signal lenght, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the DCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - Array containing the transformed input. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.dct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, - 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], + if isinstance(output_size, int): + output_size = (output_size, output_size) - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.avg_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.dct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.dct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) - With one :class:`ivy.Container` input: + if not adaptive_h and not adaptive_w: + ret = ivy.mean(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.dct(x, type=3, n=None, norm='ortho') - >>> print(y) - { - a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475, 5.19664812, -1.95411837]), - b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, - -1.25980937, 0.64958102, -0.2442648]) - } + vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) - With multiple :class:`ivy.Container` inputs: + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ret + vals[..., i, :, j] + pooled_output = ret / (length_h * length_w).astype(vals.dtype) - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) - >>> print(y) - { - a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, - -26.00041008, 19.75149155, -16.97056389, 10.87819386, - -5.89381361]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } - """ - return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output -@handle_exceptions @handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def idct( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - type: Literal[1, 2, 3, 4] = 2, - n: Optional[int] = None, - axis: int = -1, - norm: Optional[Literal["ortho"]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +@inputs_to_ivy_arrays +def adaptive_max_pool2d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: Union[Sequence[int], int], +): """ - Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. + Apply a 2D adaptive maximum pooling over an input signal composed of several input + planes. Parameters ---------- - x - The input signal. - type - The type of the idct. Must be 1, 2, 3 or 4. - n - The length of the transform. If n is less than the input signal length, - then x is truncated, if n is larger then x is zero-padded. - axis - The axis to compute the IDCT along. - norm - The type of normalization to be applied. Must be either None or "ortho". - out - optional output array, for writing the result to. + input + Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is + the batch dimension, C is the feature dimension, and H_in and W_in are the 2 + spatial dimensions. + output_size + Spatial output size. Returns ------- - ret - Array containing the transformed input. + The result of the pooling operation. Will have shape (N, C, S_0, S_1) or + (C, S_0, S_1), where S = `output_size` + """ + squeeze = False + if input.ndim == 3: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 4: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", + ) - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + if isinstance(output_size, int): + output_size = (output_size, output_size) - Examples - -------- - With :class:`ivy.Array` input: + if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): + stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) + kernel_size = stride # Mathematically identical to the previous expression + pooled_output = ivy.max_pool2d( + input, kernel_size, stride, "VALID", data_format="NCHW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output - >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) - >>> y = ivy.idct(x, type=2, n=None, norm='ortho') - >>> print(y) - ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, - 13.92713165, -10.078475 , 5.19664812, -1.95411837]) - - >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], - ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) - >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) - >>> print(y) - ivy.array([[[ 9., 18., 27., 36.], - [45., 54., 63., 72.]], - - [[ 7., 14., 21., 28.], - [35., 42., 49., 56.]]]) - - >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], - ... [40.5, 48.6, 56.7, 64.8]]) - >>> y = ivy.zeros((2, 4), dtype=ivy.float32) - >>> ivy.idct(x, type=1, n=None, norm=None, out=y) - >>> print(y) - ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, - -8.10000420e+00], - [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, - -8.09999847e+00]]) + idxh, length_h, range_max_h, adaptive_h = _compute_idx( + input.shape[-2], output_size[-2], input.device + ) + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size[-1], input.device + ) - >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) - >>> ivy.idct(x, type=4, n=None, norm=None, out=x) - >>> print(x) - ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, - 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array( + input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device + ) - With one :class:`ivy.Container` input: + if not adaptive_h and not adaptive_w: + ret = ivy.max(vals, axis=(-3, -1)) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> y = ivy.idct(x, type=3, n=None, norm='ortho') - >>> print(y) - { - a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, - -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, - -8.80319249e-08, -4.05617893e-01]), - b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, - -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, - -1.10039906e-08, -5.07022366e-02]) - } + vals, length_h = _mask( + vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") + ) + vals, length_w = _mask( + vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") + ) - With multiple :class:`ivy.Container` inputs: + ret = None + for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): + if ret is None: + ret = vals[..., i, :, j] + else: + ret = ivy.maximum(ret, vals[..., i, :, j]) + pooled_output = ret.astype(vals.dtype) - >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), - ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) - >>> container_n = ivy.Container(a=9, b=4) - >>> container_type = ivy.Container(a=2, b=1) - >>> container_norm = ivy.Container(a="ortho", b=None) - >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) - >>> print(y) - { - a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, - -16.18951607, 18.06697273, -17.57439613, 11.68861485, - -4.41308832]), - b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, - -1.00000000e+00]) - } - """ - return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output -idct.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} +def area_interpolate(x, dims, size, scale): + ret = ivy.zeros((x.shape[:2] + size)) + inv_scale = ivy.divide(1.0, scale) + for i, ba in enumerate(x): + for j, ch in enumerate(ba): + if dims == 3: + for d_dim in range(size[0]): + for h_dim in range(size[1]): + for w_dim in range(size[2]): + d_index = ( + int(d_dim * inv_scale[0]), + math.ceil((d_dim + 1) * inv_scale[0]), + ) + h_index = ( + int(h_dim * inv_scale[1]), + math.ceil((h_dim + 1) * inv_scale[1]), + ) + w_index = ( + int(w_dim * scale[2]), + math.ceil((w_dim + 1) * inv_scale[2]), + ) + scale_z = d_index[1] - d_index[0] + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_z * scale_y * scale_x + ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( + ch[ + d_index[0] : d_index[1], + h_index[0] : h_index[1], + w_index[0] : w_index[1], + ] + ) * (1 / area) + elif dims == 2: + for h_dim in range(size[0]): + for w_dim in range(size[1]): + h_index = ( + int(h_dim * inv_scale[0]), + math.ceil((h_dim + 1) * inv_scale[0]), + ) + w_index = ( + int(w_dim * inv_scale[1]), + math.ceil((w_dim + 1) * inv_scale[1]), + ) + scale_y = h_index[1] - h_index[0] + scale_x = w_index[1] - w_index[0] + area = scale_y * scale_x + ret[i, j, h_dim, w_dim] = ivy.sum( + ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] + ) * (1 / area) + else: + for w_dim in range(size[0]): + w_index = ( + int(w_dim * inv_scale[0]), + math.ceil((w_dim + 1) * inv_scale[0]), + ) + scale_x = w_index[1] - w_index[0] + ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( + 1 / scale_x + ) + return ret -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def fft( +def avg_pool1d( x: Union[ivy.Array, ivy.NativeArray], - dim: int, + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, /, *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, + data_format: str = "NWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + division_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. + """ + Compute a 1-D avg pool given 3-D input x. Parameters ---------- x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT. - dim - The dimension along which to take the one dimensional FFT. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing FFT. - Should be a integer greater than 1. + Input image *[batch_size, w, d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + division_override + If specified, it will be used as the divisor, + otherwise kernel_size will be used. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - The result of the FFT operation. + The result of the pooling operation. + + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, - 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, - 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, - 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, - 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, - 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, - 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, - 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, - 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, - 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) - >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, - 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 2., 3., 4., 5.], + [ 8., 9., 10., 11.]], + + [[14., 15., 16., 17.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.avg_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 2., 3., 4., 5.]], + + [[14., 15., 16., 17.]]]) """ - return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) + return ivy.current_backend(x).avg_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + division_override=division_override, + out=out, + ) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout1d( +def avg_pool2d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int]], + strides: Union[int, Tuple[int], Tuple[int, int]], + padding: str, /, *, - training: bool = True, - data_format: str = "NWC", + data_format: str = "NHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D - feature map. + Compute a 2-D average pool given 4-D input x. Parameters ---------- x - a 2D or 3D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout1d is performed during training or ignored - during testing. + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimensio paddings. data_format - "NWC" or "NCW". Defaults to "NWC". + NHWC" or "NCHW". Defaults to "NHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. out optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + The result of the pooling operation. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.avg_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - ivy.array([[[2., 0, 2.]]]) - >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) - >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[[1, 1, 1]]]) + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.avg_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[ 8., 9.]], - With one :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), - ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) - >>> y = ivy.dropout1d(x, 0.5) - >>> print(y) - { - a: ivy.array([[[200., 400., 0.]]]), - b: ivy.array([[[0., 0., 0.]]]) - } + [[14., 15.]]], + + + [[[32., 33.]], + + [[38., 39.]]]]) """ - return ivy.current_backend(x).dropout1d( - x, prob, training=training, data_format=data_format, out=out + return ivy.current_backend(x).avg_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, ) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout2d( +def avg_pool3d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + kernel: Union[int, Tuple[int], Tuple[int, int, int]], + strides: Union[int, Tuple[int], Tuple[int, int, int]], + padding: str, /, *, - training: bool = True, - data_format: str = "NHWC", + data_format: str = "NDHWC", + count_include_pad: bool = False, + ceil_mode: bool = False, + divisor_override: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D - feature map. + Compute a 3-D avg pool given 5-D input x. Parameters ---------- x - a 3D or 4D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout2d is performed during training or ignored - during testing. + Input volume *[batch_size,d,h,w,d_in]*. + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list indicating the per-dimension + paddings. data_format - "NHWC" or "NCHW". Defaults to "NHWC". + NDHWC" or "NCDHW". Defaults to "NDHWC". + count_include_pad + Whether to include padding in the averaging calculation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + divisor_override + If specified, it will be used as divisor, otherwise kernel_size will be used. out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + The result of the pooling operation. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. Examples -------- - With :class:`ivy.Array` input: + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.avg_pool3d(x,2,2,'VALID')) + ivy.array([[[[[ 7., 8.]]]], - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 0.5) - >>> print(y) - ivy.array([[0., 2., 2.]]) - >>> x = ivy.array([[1, 1, 1]]) - >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") - >>> print(y) - ivy.array([[1, 1, 1]]) + + [[[[31., 32.]]]]]) + >>> print(ivy.avg_pool3d(x,2,2,'SAME')) + ivy.array([[[[[ 7., 8.]]], + + + [[[19., 20.]]]], + + + + [[[[31., 32.]]], + + + [[[43., 44.]]]]]) """ - return ivy.current_backend(x).dropout2d( - x, prob, training=training, data_format=data_format, out=out + return ivy.current_backend(x).avg_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + count_include_pad=count_include_pad, + ceil_mode=ceil_mode, + divisor_override=divisor_override, + out=out, ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dropout3d( +def dct( x: Union[ivy.Array, ivy.NativeArray], - prob: float, /, *, - training: bool = True, - data_format: str = "NDHWC", - out: Optional[ivy.Array] = None, -) -> ivy.Array: + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Randomly zero out entire channels with probability prob using samples from a - Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this - case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D - feature map. + Compute the 1D Discrete Cosine Tranformation of a given signal. Parameters ---------- x - a 4D or 5D input array. Should have a floating-point data type. - prob - probability of a channel to be zero-ed. - training - controls whether dropout3d is performed during training or ignored - during testing. - data_format - "NDHWC" or "NCDHW". Defaults to "NDHWC". + The input signal. + type + The type of the dct. Must be 1, 2, 3 or 4. + n + The lenght of the transform. If n is less than the input signal lenght, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the DCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". out optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. Returns ------- ret - an array with some channels zero-ed and the rest of channels are - scaled by (1/1-prob). + Array containing the transformed input. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - """ - return ivy.current_backend(x).dropout3d( - x, prob, training=training, data_format=data_format, out=out - ) + Examples + -------- + With :class:`ivy.Array` input: -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def ifft( - x: Union[ivy.Array, ivy.NativeArray], - dim: int, - *, - norm: str = "backward", - n: Optional[Union[int, Tuple[int]]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the one dimensional discrete Fourier transform given input at least 1-D - input x. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs IFFT. - dim - The dimension along which to take the one dimensional IFFT. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - n - Optional argument indicating the sequence length, if given, the input would be - padded with zero or truncated to length n before performing IFFT. - Should be a integer greater than 1. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.dct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 1.01823380e+02, -5.15385818e+01, 1.36371466e-06, -5.38763905e+00, + 0.00000000e+00, -1.60722279e+00, -8.80319249e-08, -4.05617893e-01]) - Returns - ------- - ret - The result of the IFFT operation. + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.dct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], - Examples - -------- - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) - ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, - 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, - 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, - 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) - ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, - 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, - 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, - 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, - 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, - 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, - 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) - >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") - ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, - 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, - 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, - 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) - """ - return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) + + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.dct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.dct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([ 279.4135742 , -279.6779785 , 128.3770599 , -114.8719864 , + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def embedding( - weights: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - max_norm: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Embeds a given tensor of indices using a given tensor of weights. + With one :class:`ivy.Container` input: - Parameters - ---------- - weights - The weights tensor. - indices - The indices tensor. - max_norm - The maximum norm of the embeddings. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.dct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475, 5.19664812, -1.95411837]), + b: ivy.array([9.93732834, -8.79711437, 3.75048852, -2.94867325, 1.74089146, + -1.25980937, 0.64958102, -0.2442648]) + } - Returns - ------- - ret - The result of the embedding operation. + With multiple :class:`ivy.Container` inputs: - Examples - -------- - >>> weights = ivy.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) - >>> indices = ivy.array([0, 2]) - >>> print(ivy.embedding(weights, indices, max_norm=5)) - ivy.array([[1. , 2. , 3. ], - [2.51285338, 2.87183261, 3.2308116 ]]) + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.dct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([96., -28.1580677, -31.89422607, 22.86190414, + -26.00041008, 19.75149155, -16.97056389, 10.87819386, + -5.89381361]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } """ - ivy.utils.assertions.check_equal( - len(weights.shape), 2, message="weights must be 2-d", as_array=False - ) - return ivy.current_backend(indices).embedding( - weights, - indices, - max_norm=max_norm, - out=out, - ) + return ivy.current_backend(x).dct(x, type=type, n=n, axis=axis, norm=norm, out=out) @handle_exceptions @@ -1397,406 +1252,402 @@ def dft( @handle_exceptions +@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument -@inputs_to_ivy_arrays -def interp(x, xp, fp, left=None, right=None, period=None): - x_arr = ivy.array(x) - fix_later = False - if x_arr.shape == (): - x_arr = ivy.array([x]) - fix_later = True - x = ivy.astype(x_arr, "float64") - xp = ivy.astype(ivy.array(xp), "float64") - fp = ivy.astype(ivy.array(fp), "float64") - ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) - ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) - if period is not None: - ivy.utils.assertions.check_equal(period, 0, inverse=True) - period = ivy.abs(period) - x = ivy.remainder(x, period) - xp = ivy.remainder(xp, period) - asort_xp = ivy.argsort(xp) - xp = xp[asort_xp] - fp = fp[asort_xp] - xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) - fp = ivy.concat((fp[-1:], fp, fp[0:1])) - - def interp_inner(value): - value = ivy.array(value) - if value < xp[0]: - return left if left is not None else fp[0] - elif value > xp[-1]: - return right if right is not None else fp[-1] - else: - last = None - if xp.shape[0] < 3: - for i in range(xp.shape[0] - 1, -1, -1): - if xp[i] == value: - return fp[i] - elif xp[i] < value: - last = i - else: - first = 0 - last = xp.shape[0] - while first < last: - midpoint = (first + last) // 2 - if xp[midpoint] == value: - already_exists = ivy.argwhere(xp == value) - if already_exists.shape[0] > 0: - return fp[already_exists[-1][0]] - return fp[midpoint] - else: - if value < xp[midpoint]: - last = midpoint - 1 - else: - first = midpoint + 1 - dist = (value - xp[last]) / (xp[last + 1] - xp[last]) - return (fp[last + 1] - fp[last]) * dist + fp[last] - - ret = ivy.map(interp_inner, unique={"value": x}) - if fix_later: - return ivy.astype(ivy.array(ret[0]), "float64") - else: - return ivy.astype(ivy.array(ret), "float64") +@to_native_arrays_and_back +@handle_device_shifting +def dropout1d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout1d performs a channel-wise dropout but assumes a channel is a 1D + feature map. + Parameters + ---------- + x + a 2D or 3D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout1d is performed during training or ignored + during testing. + data_format + "NWC" or "NCW". Defaults to "NWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. -def _tf_area_dim_scale(index, starting_index, scale, ending_index): - if index < starting_index: - dim_scale = scale if index + 1 > ending_index else index + 1 - starting_index - else: - dim_scale = ending_index - index if index + 1 > ending_index else 1.0 - return dim_scale + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. -def _tf_area_indices(dim_index, scale): - starting_index = dim_index * scale - ending_index = (dim_index + 1) * scale - rounded_indices = ( - int(starting_index), - math.ceil(ending_index), - ) - return starting_index, ending_index, rounded_indices + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + ivy.array([[[2., 0, 2.]]]) -def _tf_area_interpolate(x, size, dims): - ret = ivy.zeros((x.shape[:2] + size)) - scale = ivy.divide(ivy.shape(x)[2:], size) - area = 1.0 / ivy.prod(scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_in, d_in1, d_index = _tf_area_indices(d_dim, scale[0]) - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[1]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[2]) - sum_data = ivy.zeros( - ( - d_index[1] - d_index[0], - h_index[1] - h_index[0], - w_index[1] - w_index[0], - ) - ) - for d_ind in range(d_index[0], d_index[1]): - scale_z = _tf_area_dim_scale( - d_ind, d_in, scale[0], d_in1 - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale( - h_ind, h_in, scale[1], h_in1 - ) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[2], w_in1 - ) - sum_data[ - d_ind - d_index[0], - h_ind - h_index[0], - w_ind - w_index[0], - ] = ( - ivy.array(ch[d_ind, h_ind, w_ind]) - * scale_x - * scale_y - * scale_z - * area - ) - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum(sum_data) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_in, h_in1, h_index = _tf_area_indices(h_dim, scale[0]) - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[1]) - sum_data = ivy.zeros( - (h_index[1] - h_index[0], w_index[1] - w_index[0]) - ) - for h_ind in range(h_index[0], h_index[1]): - scale_y = _tf_area_dim_scale(h_ind, h_in, scale[0], h_in1) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale( - w_ind, w_in, scale[1], w_in1 - ) - sum_data[h_ind - h_index[0], w_ind - w_index[0]] = ( - ivy.array(ch[h_ind, w_ind]) - * scale_x - * scale_y - * area - ) - ret[i, j, h_dim, w_dim] = ivy.sum(sum_data) - else: - for w_dim in range(size[0]): - w_in, w_in1, w_index = _tf_area_indices(w_dim, scale[0]) - sum_data = ivy.zeros((w_index[1] - w_index[0],)) - for w_ind in range(w_index[0], w_index[1]): - scale_x = _tf_area_dim_scale(w_ind, w_in, scale[0], w_in1) - sum_data[w_ind - w_index[0]] = ( - ivy.array(ch[w_ind]) * scale_x * area - ) - ret[i, j, w_dim] = ivy.sum(sum_data) - return ret + >>> x = ivy.array([1, 1, 1]).reshape([1, 1, 3]) + >>> y = ivy.dropout1d(x, 1, training=False, data_format="NCW") + >>> print(y) + ivy.array([[[1, 1, 1]]]) + With one :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([100, 200, 300]).reshape([1, 1, 3]), + ... b=ivy.array([400, 500, 600]).reshape([1, 1, 3])) + >>> y = ivy.dropout1d(x, 0.5) + >>> print(y) + { + a: ivy.array([[[200., 400., 0.]]]), + b: ivy.array([[[0., 0., 0.]]]) + } + """ + return ivy.current_backend(x).dropout1d( + x, prob, training=training, data_format=data_format, out=out + ) -def nearest_interpolate(x, dims, size, input_shape, exact): - off = 0.5 if exact else 0 - for d in range(dims): - m = input_shape[d + 2] - n = size[d] - offsets = (ivy.arange(n, dtype="float32") + off) * m / n - offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") - x = ivy.gather(x, offsets, axis=d + 2) - return x +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dropout2d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NHWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout2d performs a channel-wise dropout but assumes a channel is a 2D + feature map. -def _triangle_kernel(x): - return ivy.maximum(0, 1 - ivy.abs(x)) + Parameters + ---------- + x + a 3D or 4D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout2d is performed during training or ignored + during testing. + data_format + "NHWC" or "NCHW". Defaults to "NHWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). -def _cubic_kernel(x): - out = ((1.5 * x - 2.5) * x) * x + 1.0 - out = ivy.where(x >= 1.0, ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0, out) - return ivy.where(x >= 2.0, 0.0, out) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + Examples + -------- + With :class:`ivy.Array` input: -def _lanczos_kernel(radius, x): - y = radius * ivy.sin(ivy.pi * x) * ivy.sin(ivy.pi * x / radius) - out = ivy.where(x != 0, ivy.divide(y, ivy.pi**2 * x**2), 1) - return ivy.where(ivy.bitwise_and(x >= radius, x < -radius), 0.0, out) + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 0.5) + >>> print(y) + ivy.array([[0., 2., 2.]]) + >>> x = ivy.array([[1, 1, 1]]) + >>> y = ivy.dropout2d(x, 1, training=False, data_format="NCW") + >>> print(y) + ivy.array([[1, 1, 1]]) + """ + return ivy.current_backend(x).dropout2d( + x, prob, training=training, data_format=data_format, out=out + ) -def _dim_scale_factor(input_size, output_size, align_corners, scales): - if align_corners: - if output_size > 1: - dim_scale_factor = (input_size - 1) / (output_size - 1) - else: - dim_scale_factor = 0.0 - else: - dim_scale_factor = ( - input_size / (input_size * scales) - if scales is not None - else input_size / output_size - ) - return dim_scale_factor +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dropout3d( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, + /, + *, + training: bool = True, + data_format: str = "NDHWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Randomly zero out entire channels with probability prob using samples from a + Bernoulli distribution and the remaining channels are scaled by (1/1-prob). In this + case, dropout3d performs a channel-wise dropout but assumes a channel is a 1D + feature map. -def _mitchellcubic_kernel(x): - absx = abs(x) - if absx < 1: - return (7 * absx**3 - 12 * absx**2 + 6) / 6 - elif absx < 2: - return (-(absx**3) + 6 * absx**2 - 11 * absx + 6) / 6 - else: - return 0 + Parameters + ---------- + x + a 4D or 5D input array. Should have a floating-point data type. + prob + probability of a channel to be zero-ed. + training + controls whether dropout3d is performed during training or ignored + during testing. + data_format + "NDHWC" or "NCDHW". Defaults to "NDHWC". + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. + Returns + ------- + ret + an array with some channels zero-ed and the rest of channels are + scaled by (1/1-prob). -def _compute_weight_mat( - input_size, - output_size, - scale, - align_corners, - kernel_fn, - antialias: bool, - dim_scale_factor, -): - inv_scale = 1.0 / scale - kernel_scale = ivy.maximum(inv_scale, 1.0) if antialias else 1.0 - if not align_corners: - sample_f = (ivy.arange(output_size) + 0.5) * dim_scale_factor - 0.5 - x = ( - ivy.abs( - ivy.expand_dims(sample_f) - - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) - / kernel_scale - ) - else: - sample_f = ivy.arange(output_size) * dim_scale_factor - x = ivy.abs( - ivy.expand_dims(sample_f) - ivy.expand_dims(ivy.arange(input_size), axis=-1) - ) / (kernel_scale) - weights = kernel_fn(x) - total_weight_sum = ivy.sum(weights, axis=0, keepdims=True) - weights = ivy.where( - ivy.abs(total_weight_sum) > 1000.0 * float(ivy.finfo("float32").eps), - ivy.divide(weights, ivy.where(total_weight_sum != 0, total_weight_sum, 1)), - 0, - ) - input_size_minus_0_5 = input_size if align_corners else input_size - 0.5 - return ivy.where( - ivy.expand_dims( - ivy.logical_and(sample_f >= -0.5, sample_f <= input_size_minus_0_5) - ), - weights, - 0, + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + return ivy.current_backend(x).dropout3d( + x, prob, training=training, data_format=data_format, out=out ) -def _upsample_cubic_convolution1(x, A): - return ((A + 2) * x - (A + 3)) * x * x + 1 - +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def embedding( + weights: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + max_norm: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Embeds a given tensor of indices using a given tensor of weights. -def _upsample_cubic_convolution2(x, A): - return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A + Parameters + ---------- + weights + The weights tensor. + indices + The indices tensor. + max_norm + The maximum norm of the embeddings. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The result of the embedding operation. -def _upsample_get_cubic_coefficients(t): - A = -0.75 - return ( - _upsample_cubic_convolution2(t + 1.0, A), - _upsample_cubic_convolution1(t, A), - _upsample_cubic_convolution1(1.0 - t, A), - _upsample_cubic_convolution2(2.0 - t, A), + Examples + -------- + >>> weights = ivy.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) + >>> indices = ivy.array([0, 2]) + >>> print(ivy.embedding(weights, indices, max_norm=5)) + ivy.array([[1. , 2. , 3. ], + [2.51285338, 2.87183261, 3.2308116 ]]) + """ + ivy.utils.assertions.check_equal( + len(weights.shape), 2, message="weights must be 2-d", as_array=False + ) + return ivy.current_backend(indices).embedding( + weights, + indices, + max_norm=max_norm, + out=out, ) -def _upsample_cubic_interp1d(coeffs, ts): - coeffs2 = _upsample_get_cubic_coefficients(ts) - return _sum_tensors(c1 * c2 for (c1, c2) in zip(coeffs, coeffs2)) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def fft( + x: Union[ivy.Array, ivy.NativeArray], + dim: int, + /, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT. + dim + The dimension along which to take the one dimensional FFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing FFT. + Should be a integer greater than 1. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. -def _sum_tensors(ts): - return _reduce(ivy.add, ts) + Returns + ------- + ret + The result of the FFT operation. + Examples + -------- + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-3.44509285e-16+1.14423775e-17j, 8.00000000e+00-8.11483250e-16j, + 2.33486982e-16+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j, + 9.95799250e-17+2.33486982e-16j, 0.00000000e+00+7.66951701e-17j, + 1.14423775e-17+1.22464680e-16j, 0.00000000e+00+1.22464680e-16j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-3.44509285e-16+1.14423775e-17j, 1.00000000e+00+5.02733949e+00j, + 8.00000000e+00-8.11483250e-16j, 1.00000000e+00-5.02733949e+00j, + 2.33486982e-16+1.22464680e-16j, 1.00000000e+00-1.49660576e+00j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00-6.68178638e-01j, + 9.95799250e-17+2.33486982e-16j, 1.00000000e+00-1.98912367e-01j, + 0.00000000e+00+7.66951701e-17j, 1.00000000e+00+1.98912367e-01j, + 1.14423775e-17+1.22464680e-16j, 1.00000000e+00+6.68178638e-01j, + 0.00000000e+00+1.22464680e-16j, 1.00000000e+00+1.49660576e+00j]) + >>> ivy.fft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 2.82842712e+00-2.86902654e-16j, + 8.25501143e-17+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+2.71158374e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+4.32978028e-17j]) + """ + return ivy.current_backend(x).fft(x, dim, norm=norm, n=n, out=out) -def _upsample_bicubic2d_default( - a, - output_size, - align_corners, - scale_h=None, - scale_w=None, -): - N, C, iH, iW = a.shape - oH, oW = output_size - def compute_scale(in_size, out_size, align_corners, scale=None): - if align_corners: - return (in_size - 1) / (out_size - 1) if out_size > 1 else 0 - else: - return 1 / scale if scale is not None and scale > 0 else in_size / out_size +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def fft2( + x: Union[ivy.Array, ivy.NativeArray], + *, + s: Sequence[int] = None, + dim: Sequence[int] = (-2, -1), + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the 2-dimensional discrete Fourier Transform. - def compute_source_index(scale, dst_index, align_corners): - if align_corners: - return scale * dst_index - else: - return scale * (dst_index + 0.5) - 0.5 + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs FFT2. + s + sequence of ints, optional + Shape (length of each transformed axis) of the output (s[0] refers to axis 0, + s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, + if the given shape is smaller than that of the input, the input is cropped. + If it is larger, the input is padded with zeros. if s is not given, the shape + of the input along the axes specified by axes is used. + dim + Axes over which to compute the FFT2. If not given, the last two axes are used. + A repeated index in axes means the transform over that axis is performed + multiple times. A one-element sequence means that a one-dimensional FFT is + performed. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. - height_scale = compute_scale(iH, oH, align_corners, scale_h) - width_scale = compute_scale(iW, oW, align_corners, scale_w) + Returns + ------- + ret + The result of the FFT2 operation. - N_idx = ivy.reshape(ivy.arange(N), (N, 1, 1, 1)) - C_idx = ivy.reshape(ivy.arange(C), (1, C, 1, 1)) - out_y = ivy.reshape(ivy.arange(oH), ((1, 1, oH, 1))) - out_x = ivy.reshape(ivy.arange(oW), ((1, 1, 1, oW))) + Examples + -------- + >>> x = ivy.array([[0, 0, 0, 0, 0], + ... [1, 1, 1, 1, 1], + ... [2, 2, 2, 2, 2], + ... [3, 3, 3, 3, 3], + ... [4, 4, 4, 4, 4]]) + >>> y = ivy.fft2(x) + >>> print(y) + ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5+17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5-17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ]]) + """ + return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) - real_x = compute_source_index(width_scale, out_x, align_corners) - in_x = ivy.floor(real_x) - t_x = real_x - in_x - ix = ivy.astype(in_x, ivy.int64) - real_y = compute_source_index(height_scale, out_y, align_corners) - in_y = ivy.floor(real_y) - t_y = real_y - in_y - iy = ivy.astype(in_y, ivy.int64) - - iys_ofs = (iy - 1, iy, iy + 1, iy + 2) - ixs_ofs = (ix - 1, ix, ix + 1, ix + 2) - - def load_bounded(ys, xs): - y_idx = ivy.clip(ys, 0, iH - 1) - x_idx = ivy.clip(xs, 0, iW - 1) - return a[N_idx, C_idx, y_idx, x_idx] - - def get_x_interp(y): - coeffs_x = tuple((load_bounded(y, x_ofs) for x_ofs in ixs_ofs)) - return _upsample_cubic_interp1d(coeffs_x, t_x) - - coeffs_y = tuple((get_x_interp(y_ofs) for y_ofs in iys_ofs)) - result = _upsample_cubic_interp1d(coeffs_y, t_y) - - return result - - -def area_interpolate(x, dims, size, scale): - ret = ivy.zeros((x.shape[:2] + size)) - inv_scale = ivy.divide(1.0, scale) - for i, ba in enumerate(x): - for j, ch in enumerate(ba): - if dims == 3: - for d_dim in range(size[0]): - for h_dim in range(size[1]): - for w_dim in range(size[2]): - d_index = ( - int(d_dim * inv_scale[0]), - math.ceil((d_dim + 1) * inv_scale[0]), - ) - h_index = ( - int(h_dim * inv_scale[1]), - math.ceil((h_dim + 1) * inv_scale[1]), - ) - w_index = ( - int(w_dim * scale[2]), - math.ceil((w_dim + 1) * inv_scale[2]), - ) - scale_z = d_index[1] - d_index[0] - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_z * scale_y * scale_x - ret[i, j, d_dim, h_dim, w_dim] = ivy.sum( - ch[ - d_index[0] : d_index[1], - h_index[0] : h_index[1], - w_index[0] : w_index[1], - ] - ) * (1 / area) - elif dims == 2: - for h_dim in range(size[0]): - for w_dim in range(size[1]): - h_index = ( - int(h_dim * inv_scale[0]), - math.ceil((h_dim + 1) * inv_scale[0]), - ) - w_index = ( - int(w_dim * inv_scale[1]), - math.ceil((w_dim + 1) * inv_scale[1]), - ) - scale_y = h_index[1] - h_index[0] - scale_x = w_index[1] - w_index[0] - area = scale_y * scale_x - ret[i, j, h_dim, w_dim] = ivy.sum( - ch[h_index[0] : h_index[1], w_index[0] : w_index[1]] - ) * (1 / area) - else: - for w_dim in range(size[0]): - w_index = ( - int(w_dim * inv_scale[0]), - math.ceil((w_dim + 1) * inv_scale[0]), - ) - scale_x = w_index[1] - w_index[0] - ret[i, j, w_dim] = ivy.sum(ch[w_index[0] : w_index[1]]) * ( - 1 / scale_x - ) - return ret +def generate_einsum_equation(dim): + alphabet = "abcdefghijklmnopqrstuvwxyz" + input_indices = alphabet[: dim + 2] + output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] + contraction_indices = ",".join([input_indices, *output_indices]) + output = input_indices[:2] + "".join([output[-1] for output in output_indices]) + einsum_string = contraction_indices + "->" + output + return einsum_string def get_interpolate_kernel(mode): @@ -1810,38 +1661,336 @@ def get_interpolate_kernel(mode): return kernel_func -def generate_einsum_equation(dim): - alphabet = "abcdefghijklmnopqrstuvwxyz" - input_indices = alphabet[: dim + 2] - output_indices = [alphabet[2 + i] + alphabet[2 + dim + i] for i in range(dim)] - contraction_indices = ",".join([input_indices, *output_indices]) - output = input_indices[:2] + "".join([output[-1] for output in output_indices]) - einsum_string = contraction_indices + "->" + output - return einsum_string +@handle_exceptions +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def idct( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + type: Literal[1, 2, 3, 4] = 2, + n: Optional[int] = None, + axis: int = -1, + norm: Optional[Literal["ortho"]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Compute the 1D Inverse Discrete Cosine Tranformation of a given signal. + Parameters + ---------- + x + The input signal. + type + The type of the idct. Must be 1, 2, 3 or 4. + n + The length of the transform. If n is less than the input signal length, + then x is truncated, if n is larger then x is zero-padded. + axis + The axis to compute the IDCT along. + norm + The type of normalization to be applied. Must be either None or "ortho". + out + optional output array, for writing the result to. -def _interpolate_with_kernel( - x, dims, size, scale, input_shape, align_corners, antialias, scale_factor, mode -): - spatial_dims = [2 + i for i in range(dims)] - equation = generate_einsum_equation(dims) - kernel_func = get_interpolate_kernel(mode) - output_shape = tuple(input_shape[:2]) + size - operands = [] - for i, d in enumerate(spatial_dims): - m = input_shape[d] - n = output_shape[d] - dim_scale_factor = _dim_scale_factor( - m, - n, - align_corners, - scale_factor[i] if scale_factor is not None else None, - ) - w = _compute_weight_mat( - m, n, scale[i], align_corners, kernel_func, antialias, dim_scale_factor - ).astype(x.dtype) - operands.append(w) - return ivy.einsum(equation, x, *operands) + Returns + ------- + ret + Array containing the transformed input. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([8, 16, 24, 32, 40, 48, 56, 64]) + >>> y = ivy.idct(x, type=2, n=None, norm='ortho') + >>> print(y) + ivy.array([ 79.49862671, -70.37691498, 30.00390816, -23.58938599, + 13.92713165, -10.078475 , 5.19664812, -1.95411837]) + + >>> x = ivy.array([[[8, 16, 24, 32], [40, 48, 56, 64]], + ... [[1, 2, 3, 4], [ 5, 6, 7, 8]]]) + >>> y = ivy.idct(x, type=1, n=None, axis=0, norm=None) + >>> print(y) + ivy.array([[[ 9., 18., 27., 36.], + [45., 54., 63., 72.]], + + [[ 7., 14., 21., 28.], + [35., 42., 49., 56.]]]) + + >>> x = ivy.array([[ 8.1, 16.2, 24.3, 32.4], + ... [40.5, 48.6, 56.7, 64.8]]) + >>> y = ivy.zeros((2, 4), dtype=ivy.float32) + >>> ivy.idct(x, type=1, n=None, norm=None, out=y) + >>> print(y) + ivy.array([[ 1.21500000e+02, -3.24000015e+01, 1.90734863e-06, + -8.10000420e+00], + [ 3.15899994e+02, -3.24000053e+01, 3.81469727e-06, + -8.09999847e+00]]) + + >>> x = ivy.array([8., 16., 24., 32., 40., 48., 56., 64.]) + >>> ivy.idct(x, type=4, n=None, norm=None, out=x) + >>> print(x) + ivy.array([279.4135742, -279.6779785, 128.3770599, -114.8719864, + 83.72109985, -79.52869415, 69.79182434, -68.72489166]) + + With one :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> y = ivy.idct(x, type=3, n=None, norm='ortho') + >>> print(y) + { + a: ivy.array([1.01823380e+02, -5.15385818e+01, 1.36371466e-06, + -5.38763905e+00, 0.00000000e+00, -1.60722279e+00, + -8.80319249e-08, -4.05617893e-01]), + b: ivy.array([1.27279224e+01, -6.44232273e+00, 1.70464332e-07, + -6.73454881e-01, 0.00000000e+00, -2.00902849e-01, + -1.10039906e-08, -5.07022366e-02]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([8, 16, 24, 32, 40, 48, 56, 64]), + ... b=ivy.array([1, 2, 3, 4, 5, 6, 7, 8])) + >>> container_n = ivy.Container(a=9, b=4) + >>> container_type = ivy.Container(a=2, b=1) + >>> container_norm = ivy.Container(a="ortho", b=None) + >>> y = ivy.idct(x, type=container_type, n=container_n, norm=container_norm) + >>> print(y) + { + a: ivy.array([86.29723358, -66.69506073, 9.93914604, 2.88008881, + -16.18951607, 18.06697273, -17.57439613, 11.68861485, + -4.41308832]), + b: ivy.array([1.50000000e+01, -4.00000000e+00, -2.22044605e-16, + -1.00000000e+00]) + } + """ + return ivy.current_backend(x).idct(x, type=type, n=n, axis=axis, norm=norm, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def ifft( + x: Union[ivy.Array, ivy.NativeArray], + dim: int, + *, + norm: str = "backward", + n: Optional[Union[int, Tuple[int]]] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the one dimensional discrete Fourier transform given input at least 1-D + input x. + + Parameters + ---------- + x + Input volume *[...,d_in,...]*, + where d_in indicates the dimension that needs IFFT. + dim + The dimension along which to take the one dimensional IFFT. + norm + Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". + "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. + "forward" indicates normalization by $\frac{1}{n}$. + n + Optional argument indicating the sequence length, if given, the input would be + padded with zero or truncated to length n before performing IFFT. + Should be a integer greater than 1. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the IFFT operation. + + Examples + -------- + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0) + ivy.array([-4.30636606e-17+1.43029718e-18j, 0.00000000e+00+1.53080850e-17j, + 1.43029718e-18+1.53080850e-17j, 0.00000000e+00+9.58689626e-18j, + 1.24474906e-17+2.91858728e-17j, 0.00000000e+00+1.53080850e-17j, + 2.91858728e-17+1.53080850e-17j, 1.00000000e+00-1.01435406e-16j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, n=16) + ivy.array([-2.15318303e-17+7.15148591e-19j, 6.25000000e-02+9.35378602e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02+4.17611649e-02j, + 7.15148591e-19+7.65404249e-18j, 6.25000000e-02+1.24320230e-02j, + 0.00000000e+00+4.79344813e-18j, 6.25000000e-02-1.24320230e-02j, + 6.22374531e-18+1.45929364e-17j, 6.25000000e-02-4.17611649e-02j, + 0.00000000e+00+7.65404249e-18j, 6.25000000e-02-9.35378602e-02j, + 1.45929364e-17+7.65404249e-18j, 6.25000000e-02-3.14208718e-01j, + 5.00000000e-01-5.07177031e-17j, 6.25000000e-02+3.14208718e-01j]) + >>> ivy.ifft(np.exp(2j * np.pi * np.arange(8) / 8), 0, norm="ortho") + ivy.array([-1.21802426e-16+4.04549134e-18j, 0.00000000e+00+4.32978028e-17j, + 4.04549134e-18+4.32978028e-17j, 0.00000000e+00+2.71158374e-17j, + 3.52068201e-17+8.25501143e-17j, 0.00000000e+00+4.32978028e-17j, + 8.25501143e-17+4.32978028e-17j, 2.82842712e+00-2.86902654e-16j]) + """ + return ivy.current_backend(x).ifft(x, dim, norm=norm, n=n, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +def ifftn( + x: Union[ivy.Array, ivy.NativeArray], + s: Optional[Union[int, Tuple[int, ...]]] = None, + axes: Optional[Union[int, Tuple[int, ...]]] = None, + *, + norm: str = "backward", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + Compute the N-dimensional inverse discrete Fourier Transform. + + Parameters + ---------- + x + Input array of complex numbers. + s + Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, + `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, + the input is cropped. If larger, input is padded with zeros. If `s` is not + given, shape of input along axes specified by axes is used. + axes + Axes over which to compute the IFFT. If not given, last `len(s)` axes are + used, or all axes if `s` is also not specified. Repeated indices in axes + means inverse transform over that axis is performed multiple times. + norm + Indicates direction of the forward/backward pair of transforms is scaled + and with what normalization factor. "backward" indicates no normalization. + "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" + indicates normalization by $\frac{1}{n}$. + out + Optional output array for writing the result to. It must have a shape that + the inputs broadcast to. + + Returns + ------- + out + The truncated or zero-padded input, transformed along the axes indicated + by axes, or by a combination of s or x, as explained in the parameters + section above. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of axes is larger than the number of axes of x. + + Examples + -------- + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> y = ivy.ifftn(x) + >>> print(y) + ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, + -0.015561 -0.04216015j], + [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, + -0.08371392+0.17252843j], + [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, + 0.05344324+0.07972424j]]) + + >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, + ... 0.98193269+0.49560517j], + ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, + ... 0.2343787 +0.83528011j], + ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, + ... 0.44719226+0.72654048j]]) + >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') + >>> print(b) + ivy.array([[ 0.8344667 +0.98222595j], + [-0.48472244+0.30233797j]]) + """ + return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) + + +@handle_exceptions +@handle_nestable +@handle_out_argument +@inputs_to_ivy_arrays +def interp(x, xp, fp, left=None, right=None, period=None): + x_arr = ivy.array(x) + fix_later = False + if x_arr.shape == (): + x_arr = ivy.array([x]) + fix_later = True + x = ivy.astype(x_arr, "float64") + xp = ivy.astype(ivy.array(xp), "float64") + fp = ivy.astype(ivy.array(fp), "float64") + ivy.utils.assertions.check_equal(xp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(fp.ndim, 1, as_array=False) + ivy.utils.assertions.check_equal(xp.shape[0], fp.shape[0], as_array=False) + if period is not None: + ivy.utils.assertions.check_equal(period, 0, inverse=True) + period = ivy.abs(period) + x = ivy.remainder(x, period) + xp = ivy.remainder(xp, period) + asort_xp = ivy.argsort(xp) + xp = xp[asort_xp] + fp = fp[asort_xp] + xp = ivy.concat((xp[-1:] - period, xp, xp[0:1] + period)) + fp = ivy.concat((fp[-1:], fp, fp[0:1])) + + def interp_inner(value): + value = ivy.array(value) + if value < xp[0]: + return left if left is not None else fp[0] + elif value > xp[-1]: + return right if right is not None else fp[-1] + else: + last = None + if xp.shape[0] < 3: + for i in range(xp.shape[0] - 1, -1, -1): + if xp[i] == value: + return fp[i] + elif xp[i] < value: + last = i + else: + first = 0 + last = xp.shape[0] + while first < last: + midpoint = (first + last) // 2 + if xp[midpoint] == value: + already_exists = ivy.argwhere(xp == value) + if already_exists.shape[0] > 0: + return fp[already_exists[-1][0]] + return fp[midpoint] + else: + if value < xp[midpoint]: + last = midpoint - 1 + else: + first = midpoint + 1 + dist = (value - xp[last]) / (xp[last + 1] - xp[last]) + return (fp[last + 1] - fp[last]) * dist + fp[last] + + ret = ivy.map(interp_inner, unique={"value": x}) + if fix_later: + return ivy.astype(ivy.array(ret[0]), "float64") + else: + return ivy.astype(ivy.array(ret), "float64") @handle_exceptions @@ -2050,497 +2199,425 @@ def interpolate( return ivy.astype(ret, ivy.dtype(x), out=out) -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_device_shifting", - ), - "to_skip": (), -} - - -def _get_size(scale_factor, size, dims, x_shape): - if scale_factor is not None: - if isinstance(scale_factor, (float, int)): - scale_factor = [scale_factor] * dims - elif isinstance(scale_factor, (tuple, list)) and len(scale_factor) != dims: - scale_factor = [scale_factor[0]] * dims - - size = tuple( - [int(math.floor(x_shape[2 + i] * scale_factor[i])) for i in range(dims)] - ) - else: - size = (size,) * dims if isinstance(size, int) else tuple(size) - return size - - -def _output_ceil_shape(w, f, p, s): - return math.ceil((w - f + p) / s) + 1 - - -def _padding_ceil_mode(w, f, p, s, return_added_padding=False): - remaining_pixels = (w - f + sum(p)) % s - added_padding = 0 - if s > 1 and remaining_pixels != 0 and f > 1: - input_size = w + sum(p) - # making sure that the remaining pixels are supposed - # to be covered by the window - # they won't be covered if stride is big enough to skip them - if input_size - remaining_pixels - (f - 1) + s > input_size: - return p - output_shape = _output_ceil_shape( - w, - f, - sum(p), - s, - ) - # calculating new padding with ceil_output_shape - new_pad = (output_shape - 1) * s + f - w - # updating pad_list with new padding by adding it to the end - added_padding = new_pad - sum(p) - p = ( - p[0], - p[1] + added_padding, - ) - if return_added_padding: - return p, added_padding - return p - - -interpolate.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -def _compute_idx(in_size, out_size, device): - out_range = ivy.arange(out_size, device=device, dtype=ivy.int64) - i0 = ivy.trunc_divide(out_range * in_size, out_size).astype(ivy.int64) - maxlength = in_size // out_size + 1 - in_size_mod = in_size % out_size - # adaptive = True iff there are kernels with different lengths - adaptive = not (in_size_mod == 0 or out_size % in_size_mod == 0) - if adaptive: - maxlength += 1 - elif in_size_mod == 0: - maxlength -= 1 - range_max = ivy.arange(maxlength, device=device, dtype=ivy.int64) - idx = ivy.expand_dims(i0, axis=-1) + range_max - if adaptive: - maxval = ivy.full_like(idx, fill_value=in_size - 1) - idx = ivy.minimum(idx, maxval) - i1 = ivy.trunc_divide( - (out_range + 1) * in_size + out_size - 1, out_size - ).astype(ivy.int64) - length = i1 - i0 - else: - length = maxlength - return idx, length, range_max, adaptive - - -def _expand_to_dim(x, dim): - for _ in range(dim - len(x.shape)): - x = ivy.expand_dims(x, axis=-1) - return x - - -def _mask(vals, length, range_max, dim, mask_value=0.0): - if isinstance(length, int): - return vals, length - else: - assert dim < 0 - mask = ivy.greater_equal(range_max, ivy.expand_dims(length, axis=-1)) - if dim == -2: - mask = _expand_to_dim(mask, 4) - vals = ivy.where(mask, ivy.array(mask_value, device=vals.device), vals) - length = _expand_to_dim(length, -dim) - return vals, length - - +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_max_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], -): +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool1d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NWC", + dilation: Union[int, Tuple[int]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Apply a 2D adaptive maximum pooling over an input signal composed of several input - planes. + Compute a 1-D max pool given 3-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + Input image *[batch_size, w, d_in]* if data_format is "NWC". + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0)]) + data_format + "NWC" or "NCW". Defaults to "NWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) - - if isinstance(output_size, int): - output_size = (output_size, output_size) - - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.max_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output - - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device - ) - - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array( - input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw], device=input.device - ) - - if not adaptive_h and not adaptive_w: - ret = ivy.max(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + ret + The result of the pooling operation. - vals, length_h = _mask( - vals, length_h, range_max_h, dim=-2, mask_value=float("-inf") - ) - vals, length_w = _mask( - vals, length_w, range_max_w, dim=-1, mask_value=float("-inf") - ) + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ivy.maximum(ret, vals[..., i, :, j]) - pooled_output = ret.astype(vals.dtype) + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'SAME')) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, 'VALID')) + ivy.array([[[ 4., 5., 6., 7.]], + [[16., 17., 18., 19.]]]) + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> print(ivy.max_pool1d(x, 2, 2, [(1,0)], data_format="NCW", dilation=2, ceil_mode=True)) # noqa + ivy.array([[[ 1., 3.], + [ 5., 7.], + [ 9., 11.]], -adaptive_max_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[13., 15.], + [17., 19.], + [21., 23.]]]) + """ + return ivy.current_backend(x).max_pool1d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) +@handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays -def adaptive_avg_pool1d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: int, +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool2d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 1D adaptive average pooling over an input signal composed of several input - planes. + Compute a 2-D max pool given 4-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, L_in) or (C, L_in) where N is - the batch dimension, C is the feature dimension, and L_in is the spatial - dimension. - output_size - Spatial output size. + x + Input image *[batch_size,h,w,d_in]*. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NHWC" or "NCHW". Defaults to "NHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. Returns ------- - The result of the pooling operation. Will have shape (N, C, L_out) or - (C, L_out), where L_out = `output_size` - """ - squeeze = False - if input.ndim == 2: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 3: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", - ) + ret + The result of the pooling operation. - if input.shape[-1] % output_size == 0: - stride = input.shape[-1] // output_size - kernel_size = input.shape[-1] - (output_size - 1) * stride - pooled_output = ivy.avg_pool1d( - input, kernel_size, stride, "VALID", data_format="NCW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size, input.device - ) + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.max_pool2d(x, (2, 2), (1, 1), 'SAME')) + ivy.array([[[[ 2., 3.], + [ 4., 5.], + [ 4., 5.]]], - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., idxw]) - if not adaptive_w: - ret = ivy.mean(vals, axis=-1) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + [[[ 8., 9.], + [10., 11.], + [10., 11.]]]]) - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.max_pool2d(x, 3, 1, 'VALID')) + ivy.array([[[[16., 17.]], - ret = None - for i in range(vals.shape[-1]): - if ret is None: - ret = vals[..., i] - else: - ret = ret + vals[..., i] - pooled_output = ret / length_w.astype(ret.dtype) + [[22., 23.]]], - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[[40., 41.]], -adaptive_avg_pool1d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[46., 47.]]]]) + """ + return ivy.current_backend(x).max_pool2d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) -@handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def adaptive_avg_pool2d( - input: Union[ivy.Array, ivy.NativeArray], - output_size: Union[Sequence[int], int], +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_pool3d( + x: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int, ...]], + strides: Union[int, Tuple[int, ...]], + padding: Union[str, int, Tuple[int], List[Tuple[int, int]]], + /, + *, + data_format: str = "NDHWC", + dilation: Union[int, Tuple[int, ...]] = 1, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply a 2D adaptive average pooling over an input signal composed of several input - planes. + Compute a 3-D max pool given 5-D input x. Parameters ---------- - input - Input array. Must have shape (N, C, H_in, W_in) or (C, H_in, W_in) where N is - the batch dimension, C is the feature dimension, and H_in and W_in are the 2 - spatial dimensions. - output_size - Spatial output size. + x + Input tensor *[batch_size,d,h,w,d_in]* if data_format is "NDHWC". + kernel + Convolution filters *[d,h,w]*. + strides + The stride of the sliding window for each dimension of input. + padding + "SAME" or "VALID" indicating the algorithm; int, or list of tuple + indicating the per-dimension paddings. (e.g. 2, [(1, 0), (0, 1), (1, 1)]) + data_format + "NDHWC" or "NCDHW". Defaults to "NDHWC". + dilaton + The stride between elements within a sliding window, must be > 0. + ceil_mode + If True, ceil is used instead of floor to compute the output shape. + This ensures that every element in 'x' is covered by a sliding window. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - The result of the pooling operation. Will have shape (N, C, S_0, S_1) or - (C, S_0, S_1), where S = `output_size` - """ - squeeze = False - if input.ndim == 3: - input = ivy.expand_dims(input, axis=0) - squeeze = True - elif input.ndim != 4: - raise ivy.utils.exceptions.IvyException( - f"Got {len(input.shape)}D input, but only 3D and 4D inputs are supported.", - ) - - if isinstance(output_size, int): - output_size = (output_size, output_size) - - if all(i_s % o_s == 0 for i_s, o_s in zip(input.shape[-2:], output_size)): - stride = tuple(i_s // o_s for i_s, o_s in zip(input.shape[-2:], output_size)) - kernel_size = stride # Mathematically identical to the previous expression - pooled_output = ivy.avg_pool2d( - input, kernel_size, stride, "VALID", data_format="NCHW" - ) - if squeeze: - return ivy.squeeze(pooled_output, axis=0) - return pooled_output - - idxh, length_h, range_max_h, adaptive_h = _compute_idx( - input.shape[-2], output_size[-2], input.device - ) - idxw, length_w, range_max_w, adaptive_w = _compute_idx( - input.shape[-1], output_size[-1], input.device - ) - - # to numpy and back in order to bypass a slicing error in tensorflow - vals = ivy.array(input.to_numpy()[..., _expand_to_dim(idxh, 4), idxw]) + ret + The result of the pooling operation. - if not adaptive_h and not adaptive_w: - ret = ivy.mean(vals, axis=(-3, -1)) - ret = ivy.squeeze(ret, axis=0) if squeeze else ret - return ret + Both the description and the type hints above assumes an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. - vals, length_h = _mask(vals, length_h, range_max_h, dim=-2) - vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + Examples + -------- + >>> x = ivy.arange(48.).reshape((2, 3, 2, 2, 2)) + >>> print(ivy.max_pool3d(x, 2, 2, 'VALID')) + ivy.array([[[[[14., 15.]]]], - ret = None - for i, j in itertools.product(range(vals.shape[-3]), range(vals.shape[-1])): - if ret is None: - ret = vals[..., i, :, j] - else: - ret = ret + vals[..., i, :, j] - pooled_output = ret / (length_h * length_w).astype(vals.dtype) - pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output - return pooled_output + [[[[38., 39.]]]]]) + >>> print(ivy.max_pool3d(x, 2, 2, 'SAME')) + ivy.array([[[[[14., 15.]]], -adaptive_avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + [[[22., 23.]]]], -def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): - def _pad(arr, pads, pad_value): - out = ivy.astype( - ivy.pad( - arr, - ivy.maximum(0, pads).to_list(), - mode="constant", - constant_values=ivy.to_scalar(pad_value), - ), - arr.dtype, - ) - slices = tuple( - _slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) - for (lo, hi), dim in zip(pads, arr.shape) - ) - return out[slices] - if ( - _min(lhs.ndim, len(rhs_shape)) < 2 - or lhs.ndim != len(rhs_shape) - or lhs.shape[1] != rhs_shape[1] - ): - raise ValueError("Dimension mismatch") - if len(window_strides) != len(rhs_shape) - 2: - raise ValueError("Wrong number of strides for spatial dimensions") - if len(pads) != len(rhs_shape) - 2: - raise ValueError("Wrong number of pads for spatial dimensions") - lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) - in_shape = lhs.shape[2:] - filter_shape = rhs_shape[2:] - dim = len(filter_shape) + [[[[38., 39.]]], - out_strides = ivy.multiply(window_strides, lhs.strides[2:]).to_list() - view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] - out_shape = [ - (in_shape[i] - filter_shape[i]) // s + 1 for i, s in enumerate(window_strides) - ] - view_shape = list(lhs.shape[:1]) + out_shape + rhs_shape[1:] + [[[46., 47.]]]]]) + """ + return ivy.current_backend(x).max_pool3d( + x, + kernel, + strides, + padding, + data_format=data_format, + dilation=dilation, + ceil_mode=ceil_mode, + out=out, + ) - view = ivy.as_strided(lhs, view_shape, view_strides) - view_axes = list(range(view.ndim)) - sum_axes = view_axes[-dim - 1 :] - rhs_axes = [view.ndim] + sum_axes - out_axes = [0, view.ndim] + list(range(1, dim + 1)) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def max_unpool1d( + x: ivy.Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + kernel: Union[int, Tuple[int]], + strides: Union[int, Tuple[int]], + padding: str, + /, + *, + data_format: str = "NWC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 1-D max unpooling given the 1-D pooled input x and its indices. - return view, view_axes, rhs_axes, out_axes + Parameters + ---------- + x + Pooled input image *[batch_size, w, d_in]*. + indices + Indices obtained from the corresponding max pooling operation. + kernel + Size of the kernel i.e., the sliding window for each + dimension of input. *[w]*. + strides + The stride of the sliding window for each dimension of input. + padding + SAME" or "VALID" indicating the algorithm, or list + indicating the per-dimension paddings. + data_format + NWC" or "NCW". Defaults to "NWC". + out + optional output array, for writing the result to. + Returns + ------- + ret + The result of the unpooling operation. -def _dilate(operand, factors, fill_value): - outspace = list(operand.shape[:2]) + [ - shape + (factors[i] - 1) * (shape - 1) - for i, shape in enumerate(operand.shape[2:]) - ] - out = ivy.full(outspace, fill_value, dtype=fill_value.dtype) - lhs_slices = tuple(_slice(None, None, step) for step in factors) - out[(_slice(None),) * 2 + lhs_slices] = operand - return out + Both the description and the type hints above assume an array input + for simplicity, but this function is *nestable*, and therefore + also accepts :class:`ivy.Container` instances in place of any of + the arguments. + Examples + -------- + >>> x = ivy.arange(0, 24.).reshape((2, 3, 4)) + >>> pool_result = ivy.max_pool1d(x, 2, 2, 'SAME') + >>> print(pool_result) + ivy.array([[[ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], -def _padtype_to_pads(in_shape, filter_shape, window_strides, padding): - if padding.upper() == "SAME": - out_shape = [ - math.ceil(in_size / stride) - for in_size, stride in zip(in_shape, window_strides) - ] - pad_sizes = [ - _max((out_size - 1) * stride + filter_size - in_size, 0) - for out_size, stride, filter_size, in_size in zip( - out_shape, window_strides, filter_shape, in_shape - ) - ] - return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] - else: - return [(0, 0)] * len(in_shape) + [[16., 17., 18., 19.], + [20., 21., 22., 23.]]]) + >>> unpool_result = ivy.max_unpool1d(pool_result, indices, 2, 2, 'SAME') + >>> print(unpool_result) + ivy.array([[[ 0., 4., 0., 5., 0., 6., 0., 7., 0., 0., 0., 0.], + [ 0., 0., 0., 0., 8., 0., 9., 0., 10., 0., 11., 0.]], + [[ 0., 0., 0., 0., 0., 0., 0., 0., 16., 0., 17., 0.], + [ 0., 18., 0., 19., 0., 0., 0., 0., 20., 0., 21., 0.]]]) + """ + return ivy.current_backend(x).max_unpool1d( + x, indices, kernel, strides, padding, data_format=data_format, out=out + ) -identities = { - "max": -float("inf"), - "min": float("inf"), - "add": 0, - "mul": 1, - "multiply": 1, - "logical_and": True, - "logical_or": False, -} +def nearest_interpolate(x, dims, size, input_shape, exact): + off = 0.5 if exact else 0 + for d in range(dims): + m = input_shape[d + 2] + n = size[d] + offsets = (ivy.arange(n, dtype="float32") + off) * m / n + offsets = ivy.astype(ivy.floor(ivy.astype(offsets, "float32")), "int32") + x = ivy.gather(x, offsets, axis=d + 2) + return x -def _cast_init(init, dtype): - if not ivy.is_bool_dtype(dtype) and ivy.isinf(init): - if ivy.is_float_dtype(dtype): - info = ivy.finfo(dtype) - else: - info = ivy.iinfo(dtype) - if "float64" not in str(dtype): - init = info.max if init > 0 else info.min - return ivy.array(init, dtype=dtype) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def pool( + x: Union[ivy.Array, ivy.NativeArray], + window_shape: Union[int, Tuple[int], Tuple[int, int]], + pool_type: str, + /, + *, + strides: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + padding: str = "VALID", + data_format: Optional[str] = None, + dilations: Optional[Union[int, Tuple[int], Tuple[int, int]]] = None, + ceil_mode: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Perform an N-D pooling operation. -def _get_identity(func, dtype, init): - func_name = func.__name__ - if func_name in identities: - identity = identities[func_name] - return _cast_init(identity, dtype) - return init + Parameters + ---------- + x + Input array to pool over. + window_shape + Shape of the pooling window. + pool_type + Type of pooling operation, either 'MAX' or 'AVG'. + strides + Strides of the pooling operation. + padding + Padding type, either 'VALID' or 'SAME'. + data_format + Data format of the input and output data, either 'NCHW' or 'NHWC'. + dilations + Dilation rate of the pooling operation. + ceil_mode + Whether to use ceil or floor for creating the output shape. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The result of the pooling operation. -avg_pool2d.mixed_backend_wrappers = { - "to_add": ( - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Examples + -------- + >>> x = ivy.arange(12.).reshape((2, 1, 3, 2)) + >>> print(ivy.pool(x, (2, 2), 'MAX', (1, 1), 'SAME')) + ivy.array([[[[ 1., 2.], + [ 3., 4.], + [ 4., 5.]]], + [[[ 7., 8.], + [ 9., 10.], + [10., 11.]]]]) + >>> x = ivy.arange(48.).reshape((2, 4, 3, 2)) + >>> print(ivy.pool(x, 3, 'AVG', 1, 'VALID')) + ivy.array([[[[ 8., 9.]], + [[14., 15.]]], + [[[32., 33.]], + [[38., 39.]]]]) + """ + return ivy.current_backend(x).pool( + x, + window_shape, + pool_type, + strides=strides, + padding=padding, + data_format=data_format, + dilations=dilations, + ceil_mode=ceil_mode, + out=out, + ) @handle_exceptions @@ -2619,178 +2696,6 @@ def reduce_window( return ret.astype(operand.dtype) -reduce_window.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def fft2( - x: Union[ivy.Array, ivy.NativeArray], - *, - s: Sequence[int] = None, - dim: Sequence[int] = (-2, -1), - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the 2-dimensional discrete Fourier Transform. - - Parameters - ---------- - x - Input volume *[...,d_in,...]*, - where d_in indicates the dimension that needs FFT2. - s - sequence of ints, optional - Shape (length of each transformed axis) of the output (s[0] refers to axis 0, - s[1] to axis 1, etc.). This corresponds to n for fft(x, n). Along each axis, - if the given shape is smaller than that of the input, the input is cropped. - If it is larger, the input is padded with zeros. if s is not given, the shape - of the input along the axes specified by axes is used. - dim - Axes over which to compute the FFT2. If not given, the last two axes are used. - A repeated index in axes means the transform over that axis is performed - multiple times. A one-element sequence means that a one-dimensional FFT is - performed. - norm - Optional argument, "backward", "ortho" or "forward". Defaults to be "backward". - "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. - "forward" indicates normalization by $\frac{1}{n}$. - out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The result of the FFT2 operation. - - Examples - -------- - >>> x = ivy.array([[0, 0, 0, 0, 0], - ... [1, 1, 1, 1, 1], - ... [2, 2, 2, 2, 2], - ... [3, 3, 3, 3, 3], - ... [4, 4, 4, 4, 4]]) - >>> y = ivy.fft2(x) - >>> print(y) - ivy.array([[ 50. +0.j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5+17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ], - [-12.5-17.20477401j, 0. +0.j , 0. +0.j , - 0. +0.j , 0. +0.j ]]) - """ - return ivy.current_backend(x).fft2(x, s=s, dim=dim, norm=norm, out=out) - - -fft2.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -def ifftn( - x: Union[ivy.Array, ivy.NativeArray], - s: Optional[Union[int, Tuple[int, ...]]] = None, - axes: Optional[Union[int, Tuple[int, ...]]] = None, - *, - norm: str = "backward", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - Compute the N-dimensional inverse discrete Fourier Transform. - - Parameters - ---------- - x - Input array of complex numbers. - s - Shape (length of transformed axis) of the output (`s[0]` refers to axis 0, - `s[1]` to axis 1, etc.). If given shape is smaller than that of the input, - the input is cropped. If larger, input is padded with zeros. If `s` is not - given, shape of input along axes specified by axes is used. - axes - Axes over which to compute the IFFT. If not given, last `len(s)` axes are - used, or all axes if `s` is also not specified. Repeated indices in axes - means inverse transform over that axis is performed multiple times. - norm - Indicates direction of the forward/backward pair of transforms is scaled - and with what normalization factor. "backward" indicates no normalization. - "ortho" indicates normalization by $\frac{1}{\sqrt{n}}$. "forward" - indicates normalization by $\frac{1}{n}$. - out - Optional output array for writing the result to. It must have a shape that - the inputs broadcast to. - - Returns - ------- - out - The truncated or zero-padded input, transformed along the axes indicated - by axes, or by a combination of s or x, as explained in the parameters - section above. - - Raises - ------ - ValueError - If `s` and `axes` have different length. - IndexError - If an element of axes is larger than the number of axes of x. - - Examples - -------- - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> y = ivy.ifftn(x) - >>> print(y) - ivy.array([[ 0.51476765+0.66160417j, -0.04319742-0.05411636j, - -0.015561 -0.04216015j], - [ 0.06310689+0.05347854j, -0.13392983+0.16052352j, - -0.08371392+0.17252843j], - [-0.0031429 +0.05421245j, -0.10446617-0.17747098j, - 0.05344324+0.07972424j]]) - - >>> x = ivy.array([[0.24730653+0.90832391j, 0.49495562+0.9039565j, - ... 0.98193269+0.49560517j], - ... [0.93280757+0.48075343j, 0.28526384+0.3351205j, - ... 0.2343787 +0.83528011j], - ... [0.18791352+0.30690572j, 0.82115787+0.96195183j, - ... 0.44719226+0.72654048j]]) - >>> b = ivy.ifftn(x, s=[2, 1], axes=[0, 1], norm='ortho') - >>> print(b) - ivy.array([[ 0.8344667 +0.98222595j], - [-0.48472244+0.30233797j]]) - """ - return ivy.current_backend(x).ifftn(x, s=s, axes=axes, norm=norm, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2874,3 +2779,87 @@ def rfftn( raise ValueError("s and axes must have the same length.") return ivy.current_backend(x).rfftn(x, s=s, axes=axes, norm=norm, out=out) + + +idct.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_device_shifting", + ), + "to_skip": (), +} +interpolate.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +adaptive_max_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +adaptive_avg_pool1d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +adaptive_avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +avg_pool2d.mixed_backend_wrappers = { + "to_add": ( + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +reduce_window.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +fft2.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} +_max = builtins.max +_min = builtins.min +_slice = builtins.slice +identities = { + "max": -float("inf"), + "min": float("inf"), + "add": 0, + "mul": 1, + "multiply": 1, + "logical_and": True, + "logical_or": False, +} diff --git a/ivy/functional/ivy/experimental/linear_algebra.py b/ivy/functional/ivy/experimental/linear_algebra.py index 249df35212e17..94a5d48422683 100644 --- a/ivy/functional/ivy/experimental/linear_algebra.py +++ b/ivy/functional/ivy/experimental/linear_algebra.py @@ -17,149 +17,96 @@ ) from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # + +# --- Helpers --- # +# --------------- # def _check_valid_dimension_size(std): ivy.utils.assertions.check_dimensions(std) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@handle_array_function -def eigh_tridiagonal( - alpha: Union[ivy.Array, ivy.NativeArray], - beta: Union[ivy.Array, ivy.NativeArray], - /, - *, - eigvals_only: bool = True, - select: str = "a", - select_range: Optional[ - Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] - ] = None, - tol: Optional[float] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: +def _svd_checks(x, n_eigenvecs=None): """ - Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. + Run common checks to all of the SVD methods. Parameters ---------- - alpha - A real or complex array of shape (n), the diagonal elements of the - matrix. If alpha is complex, the imaginary part is ignored - (assumed zero) to satisfy the requirement that the matrix be Hermitian. - beta - A real or complex array of shape (n-1), containing the elements of - the first super-diagonal of the matrix. If beta is complex, the first - sub-diagonal of the matrix is assumed to be the conjugate of beta to - satisfy the requirement that the matrix be Hermitian. - eigvals_only - If False, both eigenvalues and corresponding eigenvectors are - computed. If True, only eigenvalues are computed. Default is True. - select - Optional string with values in {'a', 'v', 'i'} (default is 'a') that - determines which eigenvalues to calculate: 'a': all eigenvalues. - 'v': eigenvalues in the interval (min, max] given by select_range. - 'i': eigenvalues with indices min <= i <= max. - select_range - Size 2 tuple or list or array specifying the range of eigenvalues to - compute together with select. If select is 'a', select_range is ignored. - tol - Optional scalar. Ignored when backend is not Tensorflow. The absolute - tolerance to which each eigenvalue is required. An eigenvalue - (or cluster) is considered to have converged if it lies in an interval - of this width. If tol is None (default), the value eps*|T|_2 is used - where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. + matrix : 2D-array + n_eigenvecs : int, optional, default is None + if specified, number of eigen[vectors-values] to return Returns ------- - eig_vals - The eigenvalues of the matrix in non-decreasing order. - eig_vectors - If eigvals_only is False the eigenvectors are returned in the second output - argument. + n_eigenvecs : int + the number of eigenvectors to solve for + min_dim : int + the minimum dimension of matrix + max_dim : int + the maximum dimension of matrix + """ + # ndims = len(x.shape) + # if ndims != 2: + # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + dim_1, dim_2 = ivy.shape(x)[-2:] + min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) - Examples - -------- - With :class:`ivy.Array` input: + if n_eigenvecs is None: + n_eigenvecs = max_dim - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - ivy.array([0., 0.38196602, 2.61803389]) + if n_eigenvecs > max_dim: + logging.warning( + f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " + f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." + ) + n_eigenvecs = max_dim - >>> alpha = ivy.array([0., 1., 2.]) - >>> beta = ivy.array([0., 1.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, select='v', - ... select_range=[0.2,3.0]) - >>> print(y) - ivy.array([0.38196602, 2.61803389]) + return n_eigenvecs, min_dim, max_dim - >>> alpha = ivy.array([0., 1., 2., 3.]) - >>> beta = ivy.array([2., 1., 2.]) - >>> y = ivy.eigh_tridiagonal(alpha, - ... beta, - ... eigvals_only=False, - ... select='i', - ... select_range=[1,2], - ... tol=1.) - >>> print(y) - (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], - [ 0.06693714, -0.74234426], - [-0.74234426, -0.06693714], - [ 0.56710052, 0.35048741]])) - With :class:`ivy.Container` input: +# TODO uncommment the code below when these svd +# methods have been added +def _svd_interface( + matrix, + method="truncated_svd", + n_eigenvecs=None, + flip_sign=True, + u_based_flip_sign=True, + non_negative=None, + mask=None, + n_iter_mask_imputation=5, + **kwargs, +): + if method == "truncated_svd": + svd_fun = truncated_svd + # elif method == "symeig_svd": + # svd_fun = symeig_svd + # elif method == "randomized_svd": + # svd_fun = randomized_svd + elif callable(method): + svd_fun = method + else: + raise ValueError("Invalid Choice") - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.array([0.,2.]) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([0., 2., 4.]) - } + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + if mask is not None and n_eigenvecs is not None: + for _ in range(n_iter_mask_imputation): + S = S * ivy.eye(U.shape[-1], V.shape[-2]) + matrix = matrix * mask + (U @ S @ V) * (1 - mask) + U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) - >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) - >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) - >>> y = ivy.eigh_tridiagonal(alpha, beta) - >>> print(y) - { - a: ivy.array([-0.56155282, 0., 3.56155276]), - b: ivy.array([-0.82842714, 2., 4.82842731]) - } - """ - x = ivy.diag(alpha) - y = ivy.diag(beta, k=1) - z = ivy.diag(beta, k=-1) - w = x + y + z + if flip_sign: + U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) - eigh_out = ivy.linalg.eigh(w) - eigenvalues = eigh_out.eigenvalues - eigenvectors = eigh_out.eigenvectors + if non_negative is not False and non_negative is not None: + U, V = make_svd_non_negative(matrix, U, S, V) - if select == "i": - eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] - eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] - elif select == "v": - condition = ivy.logical_and( - eigenvalues.greater(select_range[0]), - eigenvalues.less_equal(select_range[1]), - ) - eigenvalues = eigenvalues[condition] - eigenvectors = eigenvectors[:, condition] + return U, S, V - if eigvals_only: - return eigenvalues - return eigenvalues, eigenvectors + +# --- Main --- # +# ------------ # @handle_exceptions @@ -169,29 +116,19 @@ def eigh_tridiagonal( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def diagflat( +def adjoint( x: Union[ivy.Array, ivy.NativeArray], /, *, - offset: int = 0, - padding_value: float = 0, - align: str = "RIGHT_LEFT", - num_rows: int = -1, - num_cols: int = -1, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a two-dimensional array with the flattened input as a diagonal. + Compute the complex conjugate transpose of x. Parameters ---------- x - Input data, which is flattened and set as the k-th diagonal of the output. - k - Diagonal to set. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. + An array with more than one dimension. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -199,35 +136,18 @@ def diagflat( Returns ------- ret - The 2-D output array. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.diagflat(x) - ivy.array([[1, 0, 0, 0], - [0, 2, 0, 0], - [0, 0, 3, 0], - [0, 0, 0, 4]]) + the complex conjugate transpose of the input. - >>> x = ivy.array([1,2]) - >>> ivy.diagflat(x, k=1) - ivy.array([[0, 1, 0], - [0, 0, 2], - [0, 0, 0]]) + Examples + -------- + >>> x = np.array([[1.-1.j, 2.+2.j], + [3.+3.j, 4.-4.j]]) + >>> x = ivy.array(x) + >>> ivy.adjoint(x) + ivy.array([[1.+1.j, 3.-3.j], + [2.-2.j, 4.+4.j]]) """ - return current_backend(x).diagflat( - x, - offset=offset, - padding_value=padding_value, - align=align, - num_rows=num_rows, - num_cols=num_cols, - out=out, - ) + return current_backend(x).adjoint(x, out=out) @handle_exceptions @@ -237,23 +157,22 @@ def diagflat( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def kron( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], +def cond( + x: Union[ivy.Array, ivy.NativeArray], /, *, + p: Optional[Union[int, float, str]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Kronecker product, a composite array made of blocks of the second array - scaled by the first. + Compute the condition number of x. Parameters ---------- - a - First input array. - b - Second input array + x + An array with more than one dimension. + p + The order of the norm of the matrix (see :func:`ivy.norm` for details). out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -261,16 +180,21 @@ def kron( Returns ------- ret - Array representing the Kronecker product of the input arrays. + the condition number of the input. Examples -------- - >>> a = ivy.array([1,2]) - >>> b = ivy.array([3,4]) - >>> ivy.kron(a, b) - ivy.array([3, 4, 6, 8]) + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x) + ivy.array(14.933034) + + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> ivy.cond(x, p=ivy.inf) + ivy.array(21.0) """ - return current_backend(a, b).kron(a, b, out=out) + return current_backend(x).cond(x, p=p, out=out) @handle_exceptions @@ -280,19 +204,29 @@ def kron( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def matrix_exp( +def diagflat( x: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, + offset: int = 0, + padding_value: float = 0, + align: str = "RIGHT_LEFT", + num_rows: int = -1, + num_cols: int = -1, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Compute the matrix exponential of a square matrix. + Return a two-dimensional array with the flattened input as a diagonal. Parameters ---------- - a - Square matrix. + x + Input data, which is flattened and set as the k-th diagonal of the output. + k + Diagonal to set. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -300,21 +234,93 @@ def matrix_exp( Returns ------- ret - the matrix exponential of the input. + The 2-D output array. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.diagflat(x) + ivy.array([[1, 0, 0, 0], + [0, 2, 0, 0], + [0, 0, 3, 0], + [0, 0, 0, 4]]) + + >>> x = ivy.array([1,2]) + >>> ivy.diagflat(x, k=1) + ivy.array([[0, 1, 0], + [0, 0, 2], + [0, 0, 0]]) + """ + return current_backend(x).diagflat( + x, + offset=offset, + padding_value=padding_value, + align=align, + num_rows=num_rows, + num_cols=num_cols, + out=out, + ) + + +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_exceptions +def dot( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the dot product between two arrays `a` and `b` using the current backend's + implementation. The dot product is defined as the sum of the element-wise product of + the input arrays. + + Parameters + ---------- + a + First input array. + b + Second input array. + out + Optional output array. If provided, the output array to store the result. + + Returns + ------- + ret + The dot product of the input arrays. Examples -------- - >>> x = ivy.array([[[1., 0.], - [0., 1.]], - [[2., 0.], - [0., 2.]]]) - >>> ivy.matrix_exp(x) - ivy.array([[[2.7183, 1.0000], - [1.0000, 2.7183]], - [[7.3891, 1.0000], - [1.0000, 7.3891]]]) + With :class:`ivy.Array` inputs: + + >>> a = ivy.array([1, 2, 3]) + >>> b = ivy.array([4, 5, 6]) + >>> result = ivy.dot(a, b) + >>> print(result) + ivy.array(32) + + >>> a = ivy.array([[1, 2], [3, 4]]) + >>> b = ivy.array([[5, 6], [7, 8]]) + >>> c = ivy.empty_like(a) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[19, 22], + [43, 50]]) + + >>> a = ivy.array([[1.1, 2.3, -3.6]]) + >>> b = ivy.array([[-4.8], [5.2], [6.1]]) + >>> c = ivy.zeros((1, 1)) + >>> ivy.dot(a, b, out=c) + >>> print(c) + ivy.array([[-15.28]]) """ - return current_backend(x).matrix_exp(x, out=out) + return current_backend(a, b).dot(a, b, out=out) @handle_exceptions @@ -382,244 +388,301 @@ def eig( @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_device_shifting -def eigvals( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> ivy.Array: - """ - Compute eigenvalues of x. Returns a set of eigenvalues. - - Parameters - ---------- - x - An array of shape (..., N, N). - - Returns - ------- - w - Not necessarily ordered array(..., N) of eigenvalues in complex type. - - Functional Examples - ------------------ - With :class:`ivy.Array` inputs: - >>> x = ivy.array([[1,2], [3,4]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array([-0.37228132+0.j, 5.37228132+0.j]) - - >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) - >>> w = ivy.eigvals(x) - >>> w - ivy.array( - [ - [-0.37228132+0.j, 5.37228132+0.j], - [ 0. +0.j, 11. +0.j] - ] - ) - """ - return current_backend(x).eigvals(x) - - -@handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def adjoint( - x: Union[ivy.Array, ivy.NativeArray], +@handle_array_function +def eigh_tridiagonal( + alpha: Union[ivy.Array, ivy.NativeArray], + beta: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + eigvals_only: bool = True, + select: str = "a", + select_range: Optional[ + Union[Tuple[int, int], List[int], ivy.Array, ivy.NativeArray] + ] = None, + tol: Optional[float] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array]]: """ - Compute the complex conjugate transpose of x. + Compute the eigenvalues and eigenvectors of a Hermitian tridiagonal matrix. Parameters ---------- - x - An array with more than one dimension. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + alpha + A real or complex array of shape (n), the diagonal elements of the + matrix. If alpha is complex, the imaginary part is ignored + (assumed zero) to satisfy the requirement that the matrix be Hermitian. + beta + A real or complex array of shape (n-1), containing the elements of + the first super-diagonal of the matrix. If beta is complex, the first + sub-diagonal of the matrix is assumed to be the conjugate of beta to + satisfy the requirement that the matrix be Hermitian. + eigvals_only + If False, both eigenvalues and corresponding eigenvectors are + computed. If True, only eigenvalues are computed. Default is True. + select + Optional string with values in {'a', 'v', 'i'} (default is 'a') that + determines which eigenvalues to calculate: 'a': all eigenvalues. + 'v': eigenvalues in the interval (min, max] given by select_range. + 'i': eigenvalues with indices min <= i <= max. + select_range + Size 2 tuple or list or array specifying the range of eigenvalues to + compute together with select. If select is 'a', select_range is ignored. + tol + Optional scalar. Ignored when backend is not Tensorflow. The absolute + tolerance to which each eigenvalue is required. An eigenvalue + (or cluster) is considered to have converged if it lies in an interval + of this width. If tol is None (default), the value eps*|T|_2 is used + where eps is the machine precision, and |T|_2 is the 2-norm of the matrix T. Returns ------- - ret - the complex conjugate transpose of the input. + eig_vals + The eigenvalues of the matrix in non-decreasing order. + eig_vectors + If eigvals_only is False the eigenvectors are returned in the second output + argument. + + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - >>> x = np.array([[1.-1.j, 2.+2.j], - [3.+3.j, 4.-4.j]]) - >>> x = ivy.array(x) - >>> ivy.adjoint(x) - ivy.array([[1.+1.j, 3.-3.j], - [2.-2.j, 4.+4.j]]) - """ - return current_backend(x).adjoint(x, out=out) + With :class:`ivy.Array` input: + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + ivy.array([0., 0.38196602, 2.61803389]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -def multi_dot( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the dot product of two or more matrices in a single function call, while - selecting the fastest evaluation order. - - Parameters - ---------- - x - sequence of matrices to multiply. - out - optional output array, for writing the result to. It must have a valid - shape, i.e. the resulting shape after applying regular matrix multiplication - to the inputs. + >>> alpha = ivy.array([0., 1., 2.]) + >>> beta = ivy.array([0., 1.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, select='v', + ... select_range=[0.2,3.0]) + >>> print(y) + ivy.array([0.38196602, 2.61803389]) - Returns - ------- - ret - dot product of the arrays. + >>> alpha = ivy.array([0., 1., 2., 3.]) + >>> beta = ivy.array([2., 1., 2.]) + >>> y = ivy.eigh_tridiagonal(alpha, + ... beta, + ... eigvals_only=False, + ... select='i', + ... select_range=[1,2], + ... tol=1.) + >>> print(y) + (ivy.array([0.38196602, 2.61803389]), ivy.array([[ 0.35048741, -0.56710052], + [ 0.06693714, -0.74234426], + [-0.74234426, -0.06693714], + [ 0.56710052, 0.35048741]])) - Examples - -------- - With :class:`ivy.Array` input: + With :class:`ivy.Container` input: - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> ivy.multi_dot((A, B, C)) - ivy.array([[ 26, 49], - [ 80, 148]]) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.array([0.,2.]) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([0., 2., 4.]) + } - >>> A = ivy.arange(2 * 3).reshape((2, 3)) - >>> B = ivy.arange(3 * 2).reshape((3, 2)) - >>> C = ivy.arange(2 * 2).reshape((2, 2)) - >>> D = ivy.zeros((2, 2)) - >>> ivy.multi_dot((A, B, C), out=D) - >>> print(D) - ivy.array([[ 26, 49], - [ 80, 148]]) + >>> alpha = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([2., 2., 2.])) + >>> beta = ivy.Container(a=ivy.array([0.,2.]), b=ivy.array([2.,2.])) + >>> y = ivy.eigh_tridiagonal(alpha, beta) + >>> print(y) + { + a: ivy.array([-0.56155282, 0., 3.56155276]), + b: ivy.array([-0.82842714, 2., 4.82842731]) + } """ - return current_backend(x).multi_dot(x, out=out) + x = ivy.diag(alpha) + y = ivy.diag(beta, k=1) + z = ivy.diag(beta, k=-1) + w = x + y + z + eigh_out = ivy.linalg.eigh(w) + eigenvalues = eigh_out.eigenvalues + eigenvectors = eigh_out.eigenvectors -multi_dot.mixed_backend_wrappers = { - "to_add": ("handle_device_shifting",), - "to_skip": (), -} + if select == "i": + eigenvalues = eigenvalues[select_range[0] : select_range[1] + 1] + eigenvectors = eigenvectors[:, select_range[0] : select_range[1] + 1] + elif select == "v": + condition = ivy.logical_and( + eigenvalues.greater(select_range[0]), + eigenvalues.less_equal(select_range[1]), + ) + eigenvalues = eigenvalues[condition] + eigenvectors = eigenvectors[:, condition] + + if eigvals_only: + return eigenvalues + return eigenvalues, eigenvectors @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def cond( +def eigvals( x: Union[ivy.Array, ivy.NativeArray], /, - *, - p: Optional[Union[int, float, str]] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the condition number of x. + Compute eigenvalues of x. Returns a set of eigenvalues. Parameters ---------- x - An array with more than one dimension. - p - The order of the norm of the matrix (see :func:`ivy.norm` for details). - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + An array of shape (..., N, N). Returns ------- - ret - the condition number of the input. + w + Not necessarily ordered array(..., N) of eigenvalues in complex type. - Examples - -------- - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x) - ivy.array(14.933034) + Functional Examples + ------------------ + With :class:`ivy.Array` inputs: + >>> x = ivy.array([[1,2], [3,4]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array([-0.37228132+0.j, 5.37228132+0.j]) - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> ivy.cond(x, p=ivy.inf) - ivy.array(21.0) + >>> x = ivy.array([[[1,2], [3,4]], [[5,6], [5,6]]]) + >>> w = ivy.eigvals(x) + >>> w + ivy.array( + [ + [-0.37228132+0.j, 5.37228132+0.j], + [ 0. +0.j, 11. +0.j] + ] + ) """ - return current_backend(x).cond(x, p=p, out=out) + return current_backend(x).eigvals(x) -# This code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py -@handle_exceptions -@handle_backend_invalid +# This function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 + + +# TODO update svd type hints when other svd methods have been added +# also update the test @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def kronecker( - x: Sequence[Union[ivy.Array, ivy.NativeArray]], - skip_matrix: Optional[int] = None, - reverse: Optional[bool] = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def initialize_tucker( + x: Union[ivy.Array, ivy.NativeArray], + rank: Sequence[int], + modes: Sequence[int], + /, + *, + init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", + seed: Optional[int] = None, + svd: Optional[Literal["truncated_svd"]] = "truncated_svd", + non_negative: Optional[bool] = False, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + svd_mask_repeats: Optional[int] = 5, +) -> Tuple[ivy.Array, Sequence[ivy.Array]]: """ - Kronecker product of a list of matrices. + Initialize core and factors used in `tucker`. The type of initialization is set + using `init`. If `init == 'random'` then initialize factor matrices using + `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the + `rank` left singular vectors of the `m`th unfolding of the input tensor. Parameters ---------- x - Sequence of matrices - - skip_matrix - if not None, index of a matrix to skip - - reverse - if True, the order of the matrices is reversed + input tensor + rank + number of components + modes + modes to consider in the input tensor + seed + Used to create a random seed distribution + when init == 'random' + init + initialization scheme for tucker decomposition. + svd + function to use to compute the SVD + non_negative + if True, non-negative factors are returned + mask + array of booleans with the same shape as ``tensor`` should be 0 where + the values are missing and 1 everywhere else. Note: if tensor is + sparse, then mask should also be sparse with a fill value of 1 (or + True). + svd_mask_repeats + number of iterations for imputing the values in the SVD matrix when + mask is not None Returns ------- - kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` - where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` - and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` + core + initialized core tensor + factors + list of factors """ - if skip_matrix is not None: - x = [x[i] for i in range(len(x)) if i != skip_matrix] + try: + assert len(x.shape) >= 2 + except ValueError: + raise ValueError( + "expected x to have atleast 2 dimensions but it has only" + f" {len(x.shape)} dimension(s)" + ) + + # Initialisation + if init == "svd": + factors = [] + for index, mode in enumerate(modes): + mask_unfold = None if mask is None else ivy.unfold(mask, mode) + U, _, _ = _svd_interface( + ivy.unfold(x, mode), + n_eigenvecs=rank[index], + method=svd, + non_negative=non_negative, + mask=mask_unfold, + n_iter_mask_imputation=svd_mask_repeats, + # random_state=random_state, + ) + factors.append(U) + + # The initial core approximation is needed here for the masking step + core = multi_mode_dot(x, factors, modes=modes, transpose=True) + + elif init == "random": + core = ( + ivy.random_uniform( + shape=[rank[index] for index in range(len(modes))], + dtype=x.dtype, + seed=seed, + ) + + 0.01 + ) + factors = [ + ivy.random_uniform( + shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed + ) + for index, mode in enumerate(modes) + ] - if reverse: - order = -1 else: - order = 1 + (core, factors) = init - for i, matrix in enumerate(x[::order]): - if not i: - res = matrix - else: - res = ivy.kron(res, matrix, out=out) - return res + if non_negative is True: + factors = [ivy.abs(f) for f in factors] + core = ivy.abs(core) + + return (core, factors) # The code has been adapated from tensorly.khatri_rao @@ -670,233 +733,54 @@ def khatri_rao( where ``prod(n_i) = prod([m.shape[0] for m in input])`` i.e. the product of the number of rows of all the input in the product. """ - if skip_matrix is not None: - x = [x[i] for i in range(len(x)) if i != skip_matrix] - - # Khatri-rao of only one matrix: just return that matrix - if len(x) == 1: - if ivy.exists(out): - return ivy.inplace_update(out, x[0]) - return x[0] - - if len(x[0].shape) == 2: - n_columns = x[0].shape[1] - else: - n_columns = 1 - x = [ivy.reshape(m, (-1, 1)) for m in x] - logging.warning( - "Khatri-rao of a series of vectors instead of input. " - "Considering each as a matrix with 1 column." - ) - - # Testing whether the input have the proper size - for i, matrix in enumerate(x): - if len(matrix.shape) != 2: - raise ValueError( - "All the input must have exactly 2 dimensions!" - f"Matrix {i} has dimension {len(matrix.shape)} != 2." - ) - if matrix.shape[1] != n_columns: - raise ValueError( - "All input must have same number of columns!" - f"Matrix {i} has {matrix.shape[1]} columns != {n_columns}." - ) - - for i, e in enumerate(x[1:]): - if not i: - if weights is None: - res = x[0] - else: - res = x[0] * ivy.reshape(weights, (1, -1)) - s1, s2 = ivy.shape(res) - s3, s4 = ivy.shape(e) - - a = ivy.reshape(res, (s1, 1, s2)) - b = ivy.reshape(e, (1, s3, s4)) - res = ivy.reshape(a * b, (-1, n_columns)) - - m = ivy.reshape(mask, (1, -1)) if mask is not None else 1 - - res = res * m - - if ivy.exists(out): - return ivy.inplace_update(out, res) - - return res - - -# The following code has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def mode_dot( - x: Union[ivy.Array, ivy.NativeArray], - /, - matrix_or_vector: Union[ivy.Array, ivy.NativeArray], - mode: int, - transpose: Optional[bool] = False, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - N-mode product of a tensor and a matrix or vector at the specified mode. - - Parameters - ---------- - x - tensor of shape ``(i_1, ..., i_k, ..., i_N)`` - matrix_or_vector - 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` - matrix or vectors to which to n-mode multiply the tensor - mode - int in the range(1, N) - transpose - If True, the matrix is transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. - - Returns - ------- - ivy.Array - `mode`-mode product of `tensor` by `matrix_or_vector` - * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` - if matrix_or_vector is a matrix - * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` - if matrix_or_vector is a vector - """ - # the mode along which to fold might decrease if we take product with a vector - fold_mode = mode - new_shape = list(x.shape) - ndims = len(matrix_or_vector.shape) - - if ndims == 2: # Tensor times matrix - # Test for the validity of the operation - dim = 0 if transpose else 1 - if matrix_or_vector.shape[dim] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" - ) - - if transpose: - matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) - - new_shape[mode] = matrix_or_vector.shape[0] - vec = False - - elif ndims == 1: # Tensor times vector - if matrix_or_vector.shape[0] != x.shape[mode]: - raise ValueError( - f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" - f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" - f" {matrix_or_vector.shape[0]} (vector size)" - ) - if len(new_shape) > 1: - new_shape.pop(mode) - else: - new_shape = [] - vec = True - - else: - raise ValueError( - "Can only take n_mode_product with a vector or a matrix." - f"Provided array of dimension {ndims} not in [1, 2]." - ) - - res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) - - if vec: # We contracted with a vector, leading to a vector - return ivy.reshape(res, new_shape, out=out) - else: # tensor times vec: refold the unfolding - return ivy.fold(res, fold_mode, new_shape, out=out) - - -# The following code has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 -@handle_nestable -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def multi_mode_dot( - x: Union[ivy.Array, ivy.NativeArray], - mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], - /, - modes: Optional[Sequence[int]] = None, - skip: Optional[Sequence[int]] = None, - transpose: Optional[bool] = False, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - r""" - N-mode product of a tensor and several matrices or vectors over several modes. - - Parameters - ---------- - x - the input tensor - - mat_or_vec_list - sequence of matrices or vectors of length ``tensor.ndim`` - - skip - None or int, optional, default is None - If not None, index of a matrix to skip. - - modes - None or int list, optional, default is None - - transpose - If True, the matrices or vectors in in the list are transposed. - For complex tensors, the conjugate transpose is used. - out - optional output array, for writing the result to. It must have a shape that the - result can broadcast to. + if skip_matrix is not None: + x = [x[i] for i in range(len(x)) if i != skip_matrix] - Returns - ------- - ivy.Array - tensor times each matrix or vector in the list at mode `mode` + # Khatri-rao of only one matrix: just return that matrix + if len(x) == 1: + if ivy.exists(out): + return ivy.inplace_update(out, x[0]) + return x[0] - Notes - ----- - If no modes are specified, just assumes there is one matrix or vector per mode and returns: - :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa - """ - if modes is None: - modes = range(len(mat_or_vec_list)) + if len(x[0].shape) == 2: + n_columns = x[0].shape[1] + else: + n_columns = 1 + x = [ivy.reshape(m, (-1, 1)) for m in x] + logging.warning( + "Khatri-rao of a series of vectors instead of input. " + "Considering each as a matrix with 1 column." + ) - decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor + # Testing whether the input have the proper size + for i, matrix in enumerate(x): + if len(matrix.shape) != 2: + raise ValueError( + "All the input must have exactly 2 dimensions!" + f"Matrix {i} has dimension {len(matrix.shape)} != 2." + ) + if matrix.shape[1] != n_columns: + raise ValueError( + "All input must have same number of columns!" + f"Matrix {i} has {matrix.shape[1]} columns != {n_columns}." + ) - res = x + for i, e in enumerate(x[1:]): + if not i: + if weights is None: + res = x[0] + else: + res = x[0] * ivy.reshape(weights, (1, -1)) + s1, s2 = ivy.shape(res) + s3, s4 = ivy.shape(e) - # Order of mode dots doesn't matter for different modes - # Sorting by mode shouldn't change order for equal modes - factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) - for i, (mat_or_vec_list, mode) in enumerate(factors_modes): - ndims = len(mat_or_vec_list.shape) - if (skip is not None) and (i == skip): - continue + a = ivy.reshape(res, (s1, 1, s2)) + b = ivy.reshape(e, (1, s3, s4)) + res = ivy.reshape(a * b, (-1, n_columns)) - if transpose and ndims == 2: - res = mode_dot( - res, - ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), - mode - decrement, - ) - else: - res = mode_dot(res, mat_or_vec_list, mode - decrement) + m = ivy.reshape(mask, (1, -1)) if mask is not None else 1 - if ndims == 1: - decrement += 1 + res = res * m if ivy.exists(out): return ivy.inplace_update(out, res) @@ -904,113 +788,98 @@ def multi_mode_dot( return res -def _svd_checks(x, n_eigenvecs=None): +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def kron( + a: Union[ivy.Array, ivy.NativeArray], + b: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Run common checks to all of the SVD methods. + Compute the Kronecker product, a composite array made of blocks of the second array + scaled by the first. Parameters ---------- - matrix : 2D-array - n_eigenvecs : int, optional, default is None - if specified, number of eigen[vectors-values] to return + a + First input array. + b + Second input array + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- - n_eigenvecs : int - the number of eigenvectors to solve for - min_dim : int - the minimum dimension of matrix - max_dim : int - the maximum dimension of matrix - """ - # ndims = len(x.shape) - # if ndims != 2: - # raise ValueError(f"matrix be a matrix. matrix.ndim is {ndims} != 2") - - dim_1, dim_2 = ivy.shape(x)[-2:] - min_dim, max_dim = min(dim_1, dim_2), max(dim_1, dim_2) - - if n_eigenvecs is None: - n_eigenvecs = max_dim - - if n_eigenvecs > max_dim: - logging.warning( - f"Trying to compute SVD with n_eigenvecs={n_eigenvecs}, which is larger " - f"than max(matrix.shape)={max_dim}. Setting n_eigenvecs to {max_dim}." - ) - n_eigenvecs = max_dim + ret + Array representing the Kronecker product of the input arrays. - return n_eigenvecs, min_dim, max_dim + Examples + -------- + >>> a = ivy.array([1,2]) + >>> b = ivy.array([3,4]) + >>> ivy.kron(a, b) + ivy.array([3, 4, 6, 8]) + """ + return current_backend(a, b).kron(a, b, out=out) -# This function has been adapated from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 -@handle_nestable +# This code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/_kronecker.py @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def svd_flip( - U: Union[ivy.Array, ivy.NativeArray], - V: Union[ivy.Array, ivy.NativeArray], - /, - u_based_decision: Optional[bool] = True, -) -> Tuple[ivy.Array, ivy.Array]: +def kronecker( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], + skip_matrix: Optional[int] = None, + reverse: Optional[bool] = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Sign correction to ensure deterministic output from SVD. Adjusts the columns of u - and the rows of v such that the loadings in the columns in u that are largest in - absolute value are always positive. This function is borrowed from scikit- - learn/utils/extmath.py. + Kronecker product of a list of matrices. Parameters ---------- - U - left singular matrix output of SVD - V - right singular matrix output of SVD - u_based_decision - If True, use the columns of u as the basis for sign flipping. - Otherwise, use the rows of v. The choice of which variable to base the - decision on is generally algorithm dependent. + x + Sequence of matrices + + skip_matrix + if not None, index of a matrix to skip + + reverse + if True, the order of the matrices is reversed Returns ------- - u_adjusted, v_adjusted : arrays with the same dimensions as the input. + kronecker_product: matrix of shape ``(prod(n_rows), prod(n_columns)`` + where ``prod(n_rows) = prod([m.shape[0] for m in matrices])`` + and ``prod(n_columns) = prod([m.shape[1] for m in matrices])`` """ - if u_based_decision: - # columns of U, rows of V - max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) - signs = ivy.sign( - ivy.array( - [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], - ) - ) - U = U * signs - if ivy.shape(V)[0] > ivy.shape(U)[1]: - signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) - V = V * signs[: ivy.shape(V)[0]][:, None] + if skip_matrix is not None: + x = [x[i] for i in range(len(x)) if i != skip_matrix] + + if reverse: + order = -1 else: - # rows of V, columns of U - max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) - signs = ivy.sign( - ivy.array( - [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], - ) - ) - V = V * signs[:, None] - if ivy.shape(U)[1] > ivy.shape(V)[0]: - signs = ivy.concat( - ( - signs, - ivy.ones( - ivy.shape(U)[1] - ivy.shape(V)[0], - ), - ) - ) - U = U * signs[: ivy.shape(U)[1]] + order = 1 - return U, V + for i, matrix in enumerate(x[::order]): + if not i: + res = matrix + else: + res = ivy.kron(res, matrix, out=out) + return res # This function has been adapted from TensorLy @@ -1105,210 +974,282 @@ def make_svd_non_negative( ' "nndsvda")' ) - return W, H + return W, H + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def matrix_exp( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the matrix exponential of a square matrix. + + Parameters + ---------- + a + Square matrix. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + the matrix exponential of the input. + + Examples + -------- + >>> x = ivy.array([[[1., 0.], + [0., 1.]], + [[2., 0.], + [0., 2.]]]) + >>> ivy.matrix_exp(x) + ivy.array([[[2.7183, 1.0000], + [1.0000, 2.7183]], + [[7.3891, 1.0000], + [1.0000, 7.3891]]]) + """ + return current_backend(x).matrix_exp(x, out=out) -# The following function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 +# The following code has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L5 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def truncated_svd( +def mode_dot( x: Union[ivy.Array, ivy.NativeArray], /, - compute_uv: bool = True, - n_eigenvecs: Optional[int] = None, -) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: + matrix_or_vector: Union[ivy.Array, ivy.NativeArray], + mode: int, + transpose: Optional[bool] = False, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Compute a truncated SVD on `x` using the standard SVD. + N-mode product of a tensor and a matrix or vector at the specified mode. Parameters ---------- x - 2D-array - compute_uv - If ``True`` then left and right singular vectors will be computed and returned - in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be - computed, which can be significantly faster. - n_eigenvecs - if specified, number of eigen[vectors-values] to return - else full matrices will be returned + tensor of shape ``(i_1, ..., i_k, ..., i_N)`` + matrix_or_vector + 1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )`` + matrix or vectors to which to n-mode multiply the tensor + mode + int in the range(1, N) + transpose + If True, the matrix is transposed. + For complex tensors, the conjugate transpose is used. + out + optional output array, for writing the result to. It must have a shape that the + result can broadcast to. Returns ------- - ret - a namedtuple ``(U, S, Vh)`` - Each returned array must have the same floating-point data type as ``x``. + ivy.Array + `mode`-mode product of `tensor` by `matrix_or_vector` + * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` + if matrix_or_vector is a matrix + * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` + if matrix_or_vector is a vector """ - n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) - full_matrices = True if n_eigenvecs > min_dim else False + # the mode along which to fold might decrease if we take product with a vector + fold_mode = mode + new_shape = list(x.shape) + ndims = len(matrix_or_vector.shape) - if compute_uv: - U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) - return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] - else: - S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) - return S[:n_eigenvecs] + if ndims == 2: # Tensor times matrix + # Test for the validity of the operation + dim = 0 if transpose else 1 + if matrix_or_vector.shape[dim] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned in" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[dim]} (dim 1 of matrix)" + ) + if transpose: + matrix_or_vector = ivy.conj(ivy.permute_dims(matrix_or_vector, (1, 0))) + + new_shape[mode] = matrix_or_vector.shape[0] + vec = False + + elif ndims == 1: # Tensor times vector + if matrix_or_vector.shape[0] != x.shape[mode]: + raise ValueError( + f"shapes {x.shape} and {matrix_or_vector.shape} not aligned for" + f" mode-{mode} multiplication: {x.shape[mode]} (mode {mode}) !=" + f" {matrix_or_vector.shape[0]} (vector size)" + ) + if len(new_shape) > 1: + new_shape.pop(mode) + else: + new_shape = [] + vec = True -# TODO uncommment the code below when these svd -# methods have been added -def _svd_interface( - matrix, - method="truncated_svd", - n_eigenvecs=None, - flip_sign=True, - u_based_flip_sign=True, - non_negative=None, - mask=None, - n_iter_mask_imputation=5, - **kwargs, -): - if method == "truncated_svd": - svd_fun = truncated_svd - # elif method == "symeig_svd": - # svd_fun = symeig_svd - # elif method == "randomized_svd": - # svd_fun = randomized_svd - elif callable(method): - svd_fun = method else: - raise ValueError("Invalid Choice") + raise ValueError( + "Can only take n_mode_product with a vector or a matrix." + f"Provided array of dimension {ndims} not in [1, 2]." + ) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) - if mask is not None and n_eigenvecs is not None: - for _ in range(n_iter_mask_imputation): - S = S * ivy.eye(U.shape[-1], V.shape[-2]) - matrix = matrix * mask + (U @ S @ V) * (1 - mask) - U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs) + res = ivy.matmul(matrix_or_vector, ivy.unfold(x, mode)) - if flip_sign: - U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign) + if vec: # We contracted with a vector, leading to a vector + return ivy.reshape(res, new_shape, out=out) + else: # tensor times vec: refold the unfolding + return ivy.fold(res, fold_mode, new_shape, out=out) - if non_negative is not False and non_negative is not None: - U, V = make_svd_non_negative(matrix, U, S, V) - return U, S, V +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +def multi_dot( + x: Sequence[Union[ivy.Array, ivy.NativeArray]], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute the dot product of two or more matrices in a single function call, while + selecting the fastest evaluation order. + + Parameters + ---------- + x + sequence of matrices to multiply. + out + optional output array, for writing the result to. It must have a valid + shape, i.e. the resulting shape after applying regular matrix multiplication + to the inputs. + + Returns + ------- + ret + dot product of the arrays. + Examples + -------- + With :class:`ivy.Array` input: -# This function has been adapted from TensorLy -# https://github.com/tensorly/tensorly/blob/main/tensorly/decomposition/_tucker.py#L22 + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> ivy.multi_dot((A, B, C)) + ivy.array([[ 26, 49], + [ 80, 148]]) + >>> A = ivy.arange(2 * 3).reshape((2, 3)) + >>> B = ivy.arange(3 * 2).reshape((3, 2)) + >>> C = ivy.arange(2 * 2).reshape((2, 2)) + >>> D = ivy.zeros((2, 2)) + >>> ivy.multi_dot((A, B, C), out=D) + >>> print(D) + ivy.array([[ 26, 49], + [ 80, 148]]) + """ + return current_backend(x).multi_dot(x, out=out) -# TODO update svd type hints when other svd methods have been added -# also update the test + +# The following code has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/core_tenalg/n_mode_product.py#L81 @handle_nestable @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def initialize_tucker( +def multi_mode_dot( x: Union[ivy.Array, ivy.NativeArray], - rank: Sequence[int], - modes: Sequence[int], + mat_or_vec_list: Sequence[Union[ivy.Array, ivy.NativeArray]], /, + modes: Optional[Sequence[int]] = None, + skip: Optional[Sequence[int]] = None, + transpose: Optional[bool] = False, *, - init: Optional[Union[Literal["svd", "random"], ivy.TuckerTensor]] = "svd", - seed: Optional[int] = None, - svd: Optional[Literal["truncated_svd"]] = "truncated_svd", - non_negative: Optional[bool] = False, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - svd_mask_repeats: Optional[int] = 5, -) -> Tuple[ivy.Array, Sequence[ivy.Array]]: - """ - Initialize core and factors used in `tucker`. The type of initialization is set - using `init`. If `init == 'random'` then initialize factor matrices using - `random_state`. If `init == 'svd'` then initialize the `m`th factor matrix using the - `rank` left singular vectors of the `m`th unfolding of the input tensor. + out: Optional[ivy.Array] = None, +) -> ivy.Array: + r""" + N-mode product of a tensor and several matrices or vectors over several modes. Parameters ---------- x - input tensor - rank - number of components + the input tensor + + mat_or_vec_list + sequence of matrices or vectors of length ``tensor.ndim`` + + skip + None or int, optional, default is None + If not None, index of a matrix to skip. + modes - modes to consider in the input tensor - seed - Used to create a random seed distribution - when init == 'random' - init - initialization scheme for tucker decomposition. - svd - function to use to compute the SVD - non_negative - if True, non-negative factors are returned - mask - array of booleans with the same shape as ``tensor`` should be 0 where - the values are missing and 1 everywhere else. Note: if tensor is - sparse, then mask should also be sparse with a fill value of 1 (or - True). - svd_mask_repeats - number of iterations for imputing the values in the SVD matrix when - mask is not None + None or int list, optional, default is None + + transpose + If True, the matrices or vectors in in the list are transposed. + For complex tensors, the conjugate transpose is used. + out + optional output array, for writing the result to. It must have a shape that the + result can broadcast to. Returns ------- - core - initialized core tensor - factors - list of factors + ivy.Array + tensor times each matrix or vector in the list at mode `mode` + + Notes + ----- + If no modes are specified, just assumes there is one matrix or vector per mode and returns: + :math:`\\text{x }\\times_0 \\text{ matrix or vec list[0] }\\times_1 \\cdots \\times_n \\text{ matrix or vec list[n] }` # noqa """ - try: - assert len(x.shape) >= 2 - except ValueError: - raise ValueError( - "expected x to have atleast 2 dimensions but it has only" - f" {len(x.shape)} dimension(s)" - ) + if modes is None: + modes = range(len(mat_or_vec_list)) - # Initialisation - if init == "svd": - factors = [] - for index, mode in enumerate(modes): - mask_unfold = None if mask is None else ivy.unfold(mask, mode) - U, _, _ = _svd_interface( - ivy.unfold(x, mode), - n_eigenvecs=rank[index], - method=svd, - non_negative=non_negative, - mask=mask_unfold, - n_iter_mask_imputation=svd_mask_repeats, - # random_state=random_state, - ) - factors.append(U) + decrement = 0 # If we multiply by a vector, we diminish the dimension of the tensor - # The initial core approximation is needed here for the masking step - core = multi_mode_dot(x, factors, modes=modes, transpose=True) + res = x - elif init == "random": - core = ( - ivy.random_uniform( - shape=[rank[index] for index in range(len(modes))], - dtype=x.dtype, - seed=seed, - ) - + 0.01 - ) - factors = [ - ivy.random_uniform( - shape=(x.shape[mode], rank[index]), dtype=x.dtype, seed=seed + # Order of mode dots doesn't matter for different modes + # Sorting by mode shouldn't change order for equal modes + factors_modes = sorted(zip(mat_or_vec_list, modes), key=lambda x: x[1]) + for i, (mat_or_vec_list, mode) in enumerate(factors_modes): + ndims = len(mat_or_vec_list.shape) + if (skip is not None) and (i == skip): + continue + + if transpose and ndims == 2: + res = mode_dot( + res, + ivy.conj(ivy.permute_dims(mat_or_vec_list, (1, 0))), + mode - decrement, ) - for index, mode in enumerate(modes) - ] + else: + res = mode_dot(res, mat_or_vec_list, mode - decrement) - else: - (core, factors) = init + if ndims == 1: + decrement += 1 - if non_negative is True: - factors = [ivy.abs(f) for f in factors] - core = ivy.abs(core) + if ivy.exists(out): + return ivy.inplace_update(out, res) - return (core, factors) + return res # This function has been adpated from TensorLy @@ -1462,6 +1403,122 @@ def partial_tucker( return (core, factors) +# This function has been adapated from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L12 +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def svd_flip( + U: Union[ivy.Array, ivy.NativeArray], + V: Union[ivy.Array, ivy.NativeArray], + /, + u_based_decision: Optional[bool] = True, +) -> Tuple[ivy.Array, ivy.Array]: + """ + Sign correction to ensure deterministic output from SVD. Adjusts the columns of u + and the rows of v such that the loadings in the columns in u that are largest in + absolute value are always positive. This function is borrowed from scikit- + learn/utils/extmath.py. + + Parameters + ---------- + U + left singular matrix output of SVD + V + right singular matrix output of SVD + u_based_decision + If True, use the columns of u as the basis for sign flipping. + Otherwise, use the rows of v. The choice of which variable to base the + decision on is generally algorithm dependent. + + Returns + ------- + u_adjusted, v_adjusted : arrays with the same dimensions as the input. + """ + if u_based_decision: + # columns of U, rows of V + max_abs_cols = ivy.argmax(ivy.abs(U), axis=0) + signs = ivy.sign( + ivy.array( + [U[i, j] for (i, j) in zip(max_abs_cols, range(ivy.shape(U)[1]))], + ) + ) + U = U * signs + if ivy.shape(V)[0] > ivy.shape(U)[1]: + signs = ivy.concat((signs, ivy.ones(ivy.shape(V)[0] - ivy.shape(U)[1]))) + V = V * signs[: ivy.shape(V)[0]][:, None] + else: + # rows of V, columns of U + max_abs_rows = ivy.argmax(ivy.abs(V), axis=1) + signs = ivy.sign( + ivy.array( + [V[i, j] for (i, j) in zip(range(ivy.shape(V)[0]), max_abs_rows)], + ) + ) + V = V * signs[:, None] + if ivy.shape(U)[1] > ivy.shape(V)[0]: + signs = ivy.concat( + ( + signs, + ivy.ones( + ivy.shape(U)[1] - ivy.shape(V)[0], + ), + ) + ) + U = U * signs[: ivy.shape(U)[1]] + + return U, V + + +# The following function has been adapted from TensorLy +# https://github.com/tensorly/tensorly/blob/main/tensorly/tenalg/svd.py#L206 +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device_shifting +def truncated_svd( + x: Union[ivy.Array, ivy.NativeArray], + /, + compute_uv: bool = True, + n_eigenvecs: Optional[int] = None, +) -> Union[ivy.Array, Tuple[ivy.Array, ivy.Array, ivy.Array]]: + """ + Compute a truncated SVD on `x` using the standard SVD. + + Parameters + ---------- + x + 2D-array + compute_uv + If ``True`` then left and right singular vectors will be computed and returned + in ``U`` and ``Vh``, respectively. Otherwise, only the singular values will be + computed, which can be significantly faster. + n_eigenvecs + if specified, number of eigen[vectors-values] to return + else full matrices will be returned + + Returns + ------- + ret + a namedtuple ``(U, S, Vh)`` + Each returned array must have the same floating-point data type as ``x``. + """ + n_eigenvecs, min_dim, _ = _svd_checks(x, n_eigenvecs=n_eigenvecs) + full_matrices = True if n_eigenvecs > min_dim else False + + if compute_uv: + U, S, Vh = ivy.svd(x, full_matrices=full_matrices, compute_uv=True) + return U[:, :n_eigenvecs], S[:n_eigenvecs], Vh[:n_eigenvecs, :] + else: + S = ivy.svd(x, full_matrices=full_matrices, compute_uv=False) + return S[:n_eigenvecs] + + @handle_nestable @handle_exceptions @handle_array_like_without_promotion @@ -1610,59 +1667,7 @@ def tucker( return ivy.TuckerTensor((core, factors)) -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_exceptions -def dot( - a: Union[ivy.Array, ivy.NativeArray], - b: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the dot product between two arrays `a` and `b` using the current backend's - implementation. The dot product is defined as the sum of the element-wise product of - the input arrays. - - Parameters - ---------- - a - First input array. - b - Second input array. - out - Optional output array. If provided, the output array to store the result. - - Returns - ------- - ret - The dot product of the input arrays. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> a = ivy.array([1, 2, 3]) - >>> b = ivy.array([4, 5, 6]) - >>> result = ivy.dot(a, b) - >>> print(result) - ivy.array(32) - - >>> a = ivy.array([[1, 2], [3, 4]]) - >>> b = ivy.array([[5, 6], [7, 8]]) - >>> c = ivy.empty_like(a) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[19, 22], - [43, 50]]) - - >>> a = ivy.array([[1.1, 2.3, -3.6]]) - >>> b = ivy.array([[-4.8], [5.2], [6.1]]) - >>> c = ivy.zeros((1, 1)) - >>> ivy.dot(a, b, out=c) - >>> print(c) - ivy.array([[-15.28]]) - """ - return current_backend(a, b).dot(a, b, out=out) +multi_dot.mixed_backend_wrappers = { + "to_add": ("handle_device_shifting",), + "to_skip": (), +} diff --git a/ivy/functional/ivy/experimental/losses.py b/ivy/functional/ivy/experimental/losses.py index 7b9f55d010c5a..f873fb6e6d1dd 100644 --- a/ivy/functional/ivy/experimental/losses.py +++ b/ivy/functional/ivy/experimental/losses.py @@ -12,86 +12,67 @@ from ivy.utils.exceptions import handle_exceptions -# log_poisson_loss @handle_exceptions @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def log_poisson_loss( +def huber_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - compute_full_loss: bool = False, - axis: int = -1, - reduction: str = "none", + delta: Optional[float] = 1.0, + reduction: Optional[str] = "mean", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the log-likelihood loss between the prediction and the target under the - assumption that the target has a Poisson distribution. Caveat: By default, this is - not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect - for optimization, but does not play well with relative loss comparisons. To compute - an approximation of the log factorial term, specify ``compute_full_loss=True`` to - enable Stirling's Approximation. + Compute the Huber loss (smooth L1 loss) between true and predicted values. Parameters ---------- - true - input array containing true labels. - pred - input array containing Predicted labels. - compute_full_loss - whether to compute the full loss. If false, a constant term is dropped - in favor of more efficient optimization. Default: ``False``. - axis - the axis along which to compute the log-likelihood loss. If axis is ``-1``, - the log-likelihood loss will be computed along the last dimension. - Default: ``-1``. - reduction - ``'none'``: No reduction will be applied to the output. - ``'mean'``: The output will be averaged. - ``'sum'``: The output will be summed. Default: ``'none'``. - out - optional output array, for writing the result to. It must have a shape + true: array_like + The true (ground truth) values. + pred : array_like + The predicted values by the model. + delta : float, optional + The threshold parameter that determines the point where the loss transitions fro + -m + squared error to absolute error. Default is 1.0. + reduction : str, optional + The type of reduction to apply to the loss. Possible values are "mean" (default) + and "sum". + out : array_like, optional + Optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret - The binary log-likelihood loss between the given distributions. - + ret : array_like + The Huber loss between the true and predicted values. Examples -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.log_poisson_loss(x, z)) - ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) + >>> true = ivy.array([2, 4, 7, 1]) + >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) + >>> huber_loss(true, pred, delta=1.0) + ivy.array([0.125, 0.125, 0.5 , 0.125]) - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) - ivy.array(1.1573164) + >>> huber_loss(true, pred, delta=2.0) + ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + + >>> huber_loss(true, pred, delta=0.5) + ivy.array([0.25 , 0.25 , 0. , 0.125]) """ - try: - assert true.shape == pred.shape - except ValueError: - raise ValueError( - "`pred` and `true` must have the same shape, received " - f"({pred.shape} vs {true.shape})." - ) + abs_diff = ivy.abs(true - pred) + quadratic_loss = 0.5 * (abs_diff**2) + linear_loss = delta * (abs_diff - 0.5 * delta) + loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) - loss = ivy.exp(pred) - pred * true - if compute_full_loss: - stirling_approx = ( - (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) - ) - cond = ivy.logical_and(true >= 0.0, true <= 1.0) - loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, axis=axis, out=out) + return ivy.sum(loss, out=out) elif reduction == "mean": - return ivy.mean(loss, axis=axis, out=out) + return ivy.mean(loss, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss @@ -154,67 +135,86 @@ def l1_loss( return ivy.inplace_update(out, loss) if out is not None else loss +# log_poisson_loss @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def huber_loss( +def log_poisson_loss( true: Union[ivy.Array, ivy.NativeArray], pred: Union[ivy.Array, ivy.NativeArray], /, *, - delta: Optional[float] = 1.0, - reduction: Optional[str] = "mean", + compute_full_loss: bool = False, + axis: int = -1, + reduction: str = "none", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the Huber loss (smooth L1 loss) between true and predicted values. + Compute the log-likelihood loss between the prediction and the target under the + assumption that the target has a Poisson distribution. Caveat: By default, this is + not the exact loss, but the loss minus a constant term [log(z!)]. That has no effect + for optimization, but does not play well with relative loss comparisons. To compute + an approximation of the log factorial term, specify ``compute_full_loss=True`` to + enable Stirling's Approximation. Parameters ---------- - true: array_like - The true (ground truth) values. - pred : array_like - The predicted values by the model. - delta : float, optional - The threshold parameter that determines the point where the loss transitions fro - -m - squared error to absolute error. Default is 1.0. - reduction : str, optional - The type of reduction to apply to the loss. Possible values are "mean" (default) - and "sum". - out : array_like, optional - Optional output array, for writing the result to. It must have a shape + true + input array containing true labels. + pred + input array containing Predicted labels. + compute_full_loss + whether to compute the full loss. If false, a constant term is dropped + in favor of more efficient optimization. Default: ``False``. + axis + the axis along which to compute the log-likelihood loss. If axis is ``-1``, + the log-likelihood loss will be computed along the last dimension. + Default: ``-1``. + reduction + ``'none'``: No reduction will be applied to the output. + ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'``. + out + optional output array, for writing the result to. It must have a shape that the inputs broadcast to. Returns ------- - ret : array_like - The Huber loss between the true and predicted values. + ret + The binary log-likelihood loss between the given distributions. + Examples -------- - >>> true = ivy.array([2, 4, 7, 1]) - >>> pred = ivy.array([2.5, 3.5, 8, 0.8]) - >>> huber_loss(true, pred, delta=1.0) - ivy.array([0.125, 0.125, 0.5 , 0.125]) - - >>> huber_loss(true, pred, delta=2.0) - ivy.array([0.125, 0.125, 0.5 , 0.2 ]) + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.log_poisson_loss(x, z)) + ivy.array([1.28402555, 1.28402555, 1.03402555, 1.28402555]) - >>> huber_loss(true, pred, delta=0.5) - ivy.array([0.25 , 0.25 , 0. , 0.125]) + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.log_poisson_loss(x, z, reduction='mean')) + ivy.array(1.1573164) """ - abs_diff = ivy.abs(true - pred) - quadratic_loss = 0.5 * (abs_diff**2) - linear_loss = delta * (abs_diff - 0.5 * delta) - loss = ivy.where(abs_diff <= delta, quadratic_loss, linear_loss) + try: + assert true.shape == pred.shape + except ValueError: + raise ValueError( + "`pred` and `true` must have the same shape, received " + f"({pred.shape} vs {true.shape})." + ) + loss = ivy.exp(pred) - pred * true + if compute_full_loss: + stirling_approx = ( + (true * ivy.log(true)) - true + (0.5 * ivy.log(2 * ivy.pi * true)) + ) + cond = ivy.logical_and(true >= 0.0, true <= 1.0) + loss += ivy.where(cond, ivy.zeros_like(loss), stirling_approx) if reduction == "sum": - return ivy.sum(loss, out=out) + return ivy.sum(loss, axis=axis, out=out) elif reduction == "mean": - return ivy.mean(loss, out=out) + return ivy.mean(loss, axis=axis, out=out) else: return ivy.inplace_update(out, loss) if out is not None else loss diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index 46b3e98ff1e70..c3a21d4a63e19 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -33,667 +33,91 @@ from ivy.utils.exceptions import handle_exceptions -@handle_exceptions -@handle_nestable -@handle_partial_mixed_function -@handle_array_like_without_promotion -@handle_view -@inputs_to_ivy_arrays -@handle_array_function -def flatten( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - start_dim: Optional[int] = 0, - end_dim: Optional[int] = -1, - order: Optional[str] = "C", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Flattens input by reshaping it into a one-dimensional tensor. If start_dim or - end_dim are passed, only dimensions starting with start_dim and ending with end_dim - are flattened. The order of elements in input is unchanged. - - Parameters - ---------- - x - input array to flatten. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - start_dim - first dim to flatten. If not set, defaults to 0. - end_dim - last dim to flatten. If not set, defaults to -1. - order - Read the elements of the input container using this index order, - and place the elements into the reshaped array using this index order. - ‘C’ means to read / write the elements using C-like index order, - with the last axis index changing fastest, back to the first axis index - changing slowest. - ‘F’ means to read / write the elements using Fortran-like index order, with - the first index changing fastest, and the last index changing slowest. - Note that the ‘C’ and ‘F’ options take no account of the memory layout - of the underlying array, and only refer to the order of indexing. - Default order is 'C' - out - optional output array, for writing the result to. - - Returns - ------- - ret - the flattened array over the specified dimensions. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x) - ivy.array([1, 2, 3, 4]) - - >>> x = ivy.array([[1,2], [3,4]]) - >>> ivy.flatten(x, order='F') - ivy.array([1, 3, 2, 4]) - - >>> x = ivy.array( - [[[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12]], - - [[ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]]], - - - [[[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17]], - - [[ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]]], - - - [[[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1]], - - [[19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]]] - ) - >>> ivy.flatten(x, start_dim = 1, end_dim = 2) - ivy.array( - [[[ 5, 5, 0, 6], - [17, 15, 11, 16], - [ 6, 3, 13, 12], - [ 6, 18, 10, 4], - [ 5, 1, 17, 3], - [14, 14, 18, 6]], - - [[12, 0, 1, 13], - [ 8, 7, 0, 3], - [19, 12, 6, 17], - [ 4, 15, 6, 15], - [ 0, 5, 17, 9], - [ 9, 3, 6, 19]], - - [[17, 13, 11, 16], - [ 4, 18, 17, 4], - [10, 10, 9, 1], - [19, 17, 13, 10], - [ 4, 19, 16, 17], - [ 2, 12, 8, 14]]])) - """ - if x.shape == (): - x = ivy.reshape(x, (1, -1))[0, :] - if start_dim == end_dim: - return ivy.inplace_update(out, x) if ivy.exists(out) else x - if start_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" - ) - if end_dim not in range(-len(x.shape), len(x.shape)): - raise IndexError( - "Dimension out of range (expected to be in range of" - f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" - ) - if start_dim < 0: - start_dim = len(x.shape) + start_dim - if end_dim < 0: - end_dim = len(x.shape) + end_dim - c = 1 - for i in range(start_dim, end_dim + 1): - c *= x.shape[i] - lst = [c] - if start_dim != 0: - for i in range(0, start_dim): - lst.insert(i, x.shape[i]) - for i in range(end_dim + 1, len(x.shape)): - lst.insert(i, x.shape[i]) - return ivy.reshape(x, tuple(lst), order=order, out=out) - - -flatten.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def moveaxis( - a: Union[ivy.Array, ivy.NativeArray], - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Move axes of an array to new positions.. - - Parameters - ---------- - a - The array whose axes should be reordered. - source - Original positions of the axes to move. These must be unique. - destination - Destination positions for each of the original axes. - These must also be unique. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array with moved axes. This array is a view of the input array. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.zeros((3, 4, 5)) - >>> ivy.moveaxis(x, 0, -1).shape - (4, 5, 3) - >>> ivy.moveaxis(x, -1, 0).shape - (5, 3, 4) - """ - return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def heaviside( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Heaviside step function for each element in x1. - - Parameters - ---------- - x1 - input array. - x2 - values to use where x1 is zero. - out - optional output array, for writing the result to. - - Returns - ------- - ret - output array with element-wise Heaviside step function of x1. - This is a scalar if both x1 and x2 are scalars. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([0.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0.0000, 0.5000, 1.0000]) - - >>> x1 = ivy.array([-1.5, 0, 2.0]) - >>> x2 = ivy.array([1.2, -2.0, 3.5]) - >>> ivy.heaviside(x1, x2) - ivy.array([0., -2., 1.]) - """ - return ivy.current_backend().heaviside(x1, x2, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def flipud( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Flip array in the up/down direction. Flip the entries in each column in the up/down - direction. Rows are preserved, but appear in a different order than before. - - Parameters - ---------- - m - The array to be flipped. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array corresponding to input array with elements - order reversed along axis 0. - - Examples - -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.flipud(m) - ivy.array([[ 0., 0., 3.], - [ 0., 2., 0.], - [ 1., 0., 0.]]) - """ - return ivy.current_backend().flipud(m, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def vstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: - """ - Stack arrays in sequence vertically (row wise). - - Parameters - ---------- - arrays - Sequence of arrays to be stacked. - - Returns - ------- - ret - The array formed by stacking the given arrays. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.vstack((x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4]]) - >>> ivy.vstack((x, y, x, y)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [1, 2, 3], - [2, 3, 4]]) - - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.vstack(y)) - ivy.array([[5, 6], - [7, 8]]) - """ - return ivy.current_backend().vstack(arrays, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def hstack( - arrays: Sequence[ivy.Array], - /, - *, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> ivy.Array: - """ - Stack arrays in sequence horizotally (column wise). - - Parameters - ---------- - arrays - Sequence of arrays to be stacked. - - Returns - ------- - ret - The array formed by stacking the given arrays. - - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.hstack((x, y)) - ivy.array([1, 2, 3, 2, 3, 4]) - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0, 0, 0]) - >>> ivy.hstack((x, y, x)) - ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) - >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] - >>> print(ivy.hstack(y)) - ivy.array([[5, 6, 7, 8]]) - """ - return ivy.current_backend().hstack(arrays, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def rot90( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - k: int = 1, - axes: Tuple[int, int] = (0, 1), - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is - from the first towards the second axis. - - Parameters - ---------- - m - Input array of two or more dimensions. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - k - Number of times the array is rotated by 90 degrees. - axes - The array is rotated in the plane defined by the axes. Axes must be - different. - out - optional output container, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - A rotated view of m. - - Examples - -------- - With :code:`ivy.Array` input: - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) - With :code:`ivy.NativeArray` input: - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m) - ivy.array([[2, 4], - [1, 3]]) - >>> m = ivy.native_array([[1,2], [3,4]]) - >>> ivy.rot90(m, k=2) - ivy.array([[4, 3], - [2, 1]]) - >>> m = ivy.native_array([[[0, 1],\ - [2, 3]],\ - [[4, 5],\ - [6, 7]]]) - >>> ivy.rot90(m, k=2, axes=(1,2)) - ivy.array([[[3, 2], - [1, 0]], - - [[7, 6], - [5, 4]]]) - """ - return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def top_k( - x: Union[ivy.Array, ivy.NativeArray], - k: int, - /, - *, - axis: int = -1, - largest: bool = True, - sorted: bool = True, - out: Optional[tuple] = None, -) -> Tuple[ivy.Array, ivy.NativeArray]: - """ - Return the `k` largest elements of the given input array along a given axis. - - Parameters - ---------- - x - The array to compute top_k for. - k - Number of top elements to retun must not exceed the array size. - axis - The axis along which we must return the top elements default value is 1. - largest - If largest is set to False we return k smallest elements of the array. - sorted - If sorted is set to True we return the elements in sorted order. - out: - Optional output tuple, for writing the result to. Must have two arrays inside, - with a shape that the returned tuple broadcast to. - - Returns - ------- - ret - A named tuple with values and indices of top k elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 2) - >>> print(y) - top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) - - >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) - >>> y = ivy.top_k(x, 2, axis=1, largest=False) - >>> print(y) - top_k(values=ivy.array([[-2., 0.], - [-8., -1.]]), indices=ivy.array([[0, 3], - [0, 2]])) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) - >>> y = ivy.top_k(x, 3) - >>> print(y) - top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) - >>> y = x.top_k(2) - >>> print(y) - [{ - a: ivy.array([2, -1]), - b: ivy.array([5., 4.]) - }, { - a: ivy.array([1, 0]), - b: ivy.array([1, 0]) - }] - """ - return current_backend(x).top_k( - x, k, axis=axis, largest=largest, sorted=sorted, out=out - ) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def fliplr( - m: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Flip array in the left/right direction. Flip the entries in each column in the - left/right direction. Columns are preserved, but appear in a different order than - before. - - Parameters - ---------- - m - The array to be flipped. Must be at least 2-D. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Array corresponding to input array with elements - order reversed along axis 1. - - Examples - -------- - >>> m = ivy.diag([1, 2, 3]) - >>> ivy.fliplr(m) - ivy.array([[0, 0, 1], - [0, 2, 0], - [3, 0, 0]]) - """ - return ivy.current_backend().fliplr(m, copy=copy, out=out) - - -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def i0( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute the Bessel i0 function of x element-wise. - - Parameters - ---------- - x - Array input. - out - optional output array, for writing the result to. +# --- Helpers --- # +# --------------- # - Returns - ------- - ret - Array with the modified Bessel function - evaluated at each of the elements of x. - Examples - -------- - >>> x = ivy.array([1, 2, 3]) - >>> ivy.i0(x) - ivy.array([1.26606588, 2.2795853 , 4.88079259]) - """ - return ivy.current_backend(x).i0(x, out=out) +def _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, +): + ivy.utils.assertions.check_true( + callable(mode) + or mode + in [ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + message="the provided mode is not supported", + ) + _check_tuple_arg(pad_width, "pad_width") + if mode not in ["dilated"]: + ivy.utils.assertions.check_true( + all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), + message="the pad_widths must be greater or equal to zero", + ) + if mode in ["maximum", "mean", "median", "minimum"]: + _check_tuple_arg(stat_length, "stat_length") + ivy.utils.assertions.check_true( + all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), + message="the stat lengths must be greater than zero", + ) + elif mode == "constant": + _check_tuple_arg(constant_values, "constant_values", force_integer=False) + elif mode == "linear_ramp": + _check_tuple_arg(end_values, "end_values", force_integer=False) + ivy.utils.assertions.check_true( + reflect_type in ["even", "odd"], + message="the provided reflect_type is not supported", + ) -def _slice_at_axis(sl, axis): - return (slice(None),) * axis + (sl,) + (...,) +def _check_bounds(shape0, shape1, strides1, itemsize): + numel0 = math.prod(shape0) + ndim1 = len(shape1) + return ( + sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize + <= numel0 * itemsize + ) -def _set_pad_area(padded, axis, width_pair, value_pair): - if width_pair[0] > 0: - left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) - padded[left_slice] = value_pair[0] - if width_pair[1] > 0: - right_slice = _slice_at_axis( - slice(padded.shape[axis] - width_pair[1], None), axis - ) - padded[right_slice] = value_pair[1] - return padded +def _check_tuple_arg(arg, name, force_integer=True): + is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype + flag_assert = False + if isinstance(arg, (tuple, list)): + for nested in arg: + if isinstance(nested, (tuple, list)): + for sub_nested in nested: + if not is_scalar(sub_nested): + flag_assert = True + break + elif not is_scalar(nested): + flag_assert = True + elif not is_scalar(arg): + flag_assert = True + if flag_assert: + if not force_integer: + raise ivy.utils.exceptions.IvyException( + name + " should be scalar, tuple of scalars or tuple of scalar tuples" + ) + else: + raise ivy.utils.exceptions.IvyException( + name + " should be int, tuple of ints or tuple of int tuples" + ) def _get_edges(padded, axis, width_pair): @@ -758,6 +182,82 @@ def _get_stats(padded, axis, width_pair, length_pair, stat_func): return left_stat, right_stat +def _interior_pad(operand, padding_value, padding_config): + for axis, (_, _, interior) in enumerate(padding_config): + if interior > 0: + new_shape = list(operand.shape) + new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior + new_array = ivy.full(new_shape, padding_value, dtype=operand.dtype) + src_indices = ivy.arange(operand.shape[axis]) + dst_indices = src_indices * (interior + 1) + index_tuple = [slice(None)] * operand.ndim + index_tuple[axis] = dst_indices + new_array[tuple(index_tuple)] = operand + operand = new_array + + start_indices = [0] * operand.ndim + limit_indices = [0] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low < 0: + start_indices[axis] = abs(low) + if high < 0: + limit_indices[axis] = high + else: + limit_indices[axis] = operand.shape[axis] + 1 + padded = _slice(operand, start_indices, limit_indices) + + pad_width = [(0, 0)] * operand.ndim + for axis, (low, high, _) in enumerate(padding_config): + if low > 0 and high > 0: + pad_width[axis] = (low, high) + elif low > 0 and not high > 0: + pad_width[axis] = (low, 0) + elif high > 0 and not low > 0: + pad_width[axis] = (0, high) + padded = ivy.constant_pad(padded, pad_width, value=padding_value) + return padded + + +def _interleave(a, b, axis): + assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 + a_pad = [(0, 0, 0)] * a.ndim + b_pad = [(0, 0, 0)] * b.ndim + a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) + b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) + a = _interior_pad(a, 0.0, a_pad) + b = _interior_pad(b, 0.0, b_pad) + return ivy.add(a, b) + + +def _pad_simple(array, pad_width, fill_value=None): + new_shape = tuple( + [left + size + right for size, (left, right) in zip(array.shape, pad_width)] + ) + padded = ivy.zeros(new_shape, dtype=array.dtype) + if fill_value is not None: + padded = ivy.ones_like(padded) * fill_value + original_area_slice = tuple( + [ + slice(left, left + size) + for size, (left, right) in zip(array.shape, pad_width) + ] + ) + padded[original_area_slice] = array + return padded, original_area_slice + + +def _set_pad_area(padded, axis, width_pair, value_pair): + if width_pair[0] > 0: + left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) + padded[left_slice] = value_pair[0] + if width_pair[1] > 0: + right_slice = _slice_at_axis( + slice(padded.shape[axis] - width_pair[1], None), axis + ) + padded[right_slice] = value_pair[1] + return padded + + def _set_reflect_both(padded, axis, width_pair, method, include_edge=False): left_pad, right_pad = width_pair old_length = padded.shape[axis] - right_pad - left_pad @@ -835,35 +335,42 @@ def _set_wrap_both(padded, axis, width_pair): return new_left_pad, new_right_pad, padded -def _pad_simple(array, pad_width, fill_value=None): - new_shape = tuple( - [left + size + right for size, (left, right) in zip(array.shape, pad_width)] - ) - padded = ivy.zeros(new_shape, dtype=array.dtype) - if fill_value is not None: - padded = ivy.ones_like(padded) * fill_value - original_area_slice = tuple( - [ - slice(left, left + size) - for size, (left, right) in zip(array.shape, pad_width) - ] - ) - padded[original_area_slice] = array - return padded, original_area_slice +def _slice(operand, start_indices, limit_indices, strides=None): + strides = [1] * len(operand.shape) if strides is None else strides + + full_slice = () + for i, _ in enumerate(operand.shape): + strides_i = int(strides[i]) + start_i = int(start_indices[i]) + limit_i = int(limit_indices[i]) + full_slice += (slice(start_i, limit_i, strides_i),) + return operand[full_slice] -def _to_pairs(x, n): +def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): + if axis >= 0: + slices = [slice(None)] * axis + [slice(start, stop, stride)] + else: + slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) + return x[tuple(slices)] + + +def _slice_at_axis(sl, axis): + return (slice(None),) * axis + (sl,) + (...,) + + +def _to_dilated(x, n): if ivy.isscalar(x): - return ((x, x),) * n - elif len(x) == 2 and ivy.isscalar(x[0]): - return ((x[0], x[1]),) * n + return ((x, x, x),) * n + elif len(x) == 3 and ivy.isscalar(x[0]): + return ((x[0], x[1], x[2]),) * n elif len(x) != n: ivy.utils.assertions.check_equal( ivy.asarray(list(x)).shape, - (n, 2), + (n, 3), message=( "tuple argument should contain " - "ndim pairs where ndim is the number of " + "ndim groups where ndim is the number of " "the input's dimensions" ), as_array=False, @@ -871,18 +378,18 @@ def _to_pairs(x, n): return x -def _to_dilated(x, n): +def _to_pairs(x, n): if ivy.isscalar(x): - return ((x, x, x),) * n - elif len(x) == 3 and ivy.isscalar(x[0]): - return ((x[0], x[1], x[2]),) * n + return ((x, x),) * n + elif len(x) == 2 and ivy.isscalar(x[0]): + return ((x[0], x[1]),) * n elif len(x) != n: ivy.utils.assertions.check_equal( ivy.asarray(list(x)).shape, - (n, 3), + (n, 2), message=( "tuple argument should contain " - "ndim groups where ndim is the number of " + "ndim pairs where ndim is the number of " "the input's dimensions" ), as_array=False, @@ -890,395 +397,422 @@ def _to_dilated(x, n): return x -def _check_tuple_arg(arg, name, force_integer=True): - is_scalar = ivy.isscalar if not force_integer else ivy.is_int_dtype - flag_assert = False - if isinstance(arg, (tuple, list)): - for nested in arg: - if isinstance(nested, (tuple, list)): - for sub_nested in nested: - if not is_scalar(sub_nested): - flag_assert = True - break - elif not is_scalar(nested): - flag_assert = True - elif not is_scalar(arg): - flag_assert = True - if flag_assert: - if not force_integer: - raise ivy.utils.exceptions.IvyException( - name + " should be scalar, tuple of scalars or tuple of scalar tuples" - ) - else: - raise ivy.utils.exceptions.IvyException( - name + " should be int, tuple of ints or tuple of int tuples" - ) +# --- Main --- # +# ------------ # + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@inputs_to_native_shapes +def as_strided( + x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + strides: Sequence[int], + /, +) -> ivy.Array: + """ + Create a copy of the input array with the given shape and strides. + + Parameters + ---------- + x + Input Array. + shape + The shape of the new array. + strides + The strides of the new array (specified in bytes). + + Returns + ------- + ret + Output Array + + Examples + -------- + >>> x = ivy.array([1, 2, 3, 4, 5, 6]) + >>> ivy.as_strided(x, (4, 3), (8, 8)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [3, 4, 5], + [4, 5, 6]]) + """ + itemsize = x.itemsize + if not _check_bounds(x.shape, shape, strides, itemsize): + raise ivy.exceptions.IvyException("attempted unsafe memory access") + if any(strides[i] % itemsize != 0 for i in range(len(strides))): + raise ivy.exceptions.IvyException("strides must be multiple of itemsize") + + src = memoryview(ivy.to_numpy(x)).cast("b") + + src_ind = ivy.inner( + ivy.indices(shape).reshape((len(shape), -1)).T, + ivy.array(strides), + ) + src_ind = ivy.expand_dims(src_ind, axis=-1) + src_ind = src_ind + ivy.arange(itemsize) + src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() + temp_list = [src[i] for i in src_ind] + temp_array = ivy.asarray(temp_list, dtype=ivy.int8) + result = bytearray(temp_array.to_numpy()) -def _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, -): - ivy.utils.assertions.check_true( - callable(mode) - or mode - in [ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - message="the provided mode is not supported", - ) - _check_tuple_arg(pad_width, "pad_width") - if mode not in ["dilated"]: - ivy.utils.assertions.check_true( - all(element[1] >= 0 for element in ivy.ndenumerate(pad_width)), - message="the pad_widths must be greater or equal to zero", - ) - if mode in ["maximum", "mean", "median", "minimum"]: - _check_tuple_arg(stat_length, "stat_length") - ivy.utils.assertions.check_true( - all(element[1] > 0 for element in ivy.ndenumerate(stat_length)), - message="the stat lengths must be greater than zero", - ) - elif mode == "constant": - _check_tuple_arg(constant_values, "constant_values", force_integer=False) - elif mode == "linear_ramp": - _check_tuple_arg(end_values, "end_values", force_integer=False) - ivy.utils.assertions.check_true( - reflect_type in ["even", "odd"], - message="the provided reflect_type is not supported", + return ivy.reshape( + ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), + shape, ) @handle_exceptions @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def pad( - input: Union[ivy.Array, ivy.NativeArray], - pad_width: Union[Iterable[Tuple[int]], int], +def associative_scan( + x: Union[ivy.Array, ivy.NativeArray], + fn: Callable, /, *, - mode: Union[ - Literal[ - "constant", - "dilated", - "edge", - "linear_ramp", - "maximum", - "mean", - "median", - "minimum", - "reflect", - "symmetric", - "wrap", - "empty", - ], - Callable, - ] = "constant", - stat_length: Union[Iterable[Tuple[int]], int] = 1, - constant_values: Union[Iterable[Tuple[Number]], Number] = 0, - end_values: Union[Iterable[Tuple[Number]], Number] = 0, - reflect_type: Literal["even", "odd"] = "even", - **kwargs: Optional[Any], + reverse: bool = False, + axis: int = 0, ) -> ivy.Array: """ - Pad an array. + Perform an associative scan over the given array. Parameters ---------- - input - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths - for each axis. - - ((before, after),) yields same before and after pad for each axis. - - pad (integer) is shortcut for before = after = pad width for all axes. - mode - One of the following string values or a user-supplied function. - - "constant": Pads with a constant value. - - "edge": Pads with the input's edge values. - - "linear_ramp": Pads with the linear ramp between end_value - and the input's edge value. - - "maximum": Pads with the maximum value of all or part of the vector - along each axis. - - "mean": Pads with the mean value of all or part of the vector along - each axis. - - "median": Pads with the median value of all or part of the vector - along each axis. - - "minimum": Pads with the minimum value of all or part of the vector - along each axis. - - "reflect": Pads with the reflection mirrored on the first and last - values of the vector along each axis. - - "symmetric": Pads with the reflection of the vector mirrored along - the edge of the input. - - "wrap": Pads with the wrap of the vector along the axis. - The first values are used to pad the end and the end values are used - to pad the beginning. - - "empty": Pads with undefined values. - - : Pads with a user-defined padding function. The padding - function should modify a rank 1 array following the signature - `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: - - `vector` is a rank 1 array already padded with zeros. Padded - values are `vector[:iaxis_pad_width[0]]` and - `vector[-iaxis_pad_width[1]:]`. - - `iaxis_pad_width` is a 2-tuple of ints, where - `iaxis_pad_width[0]` represents the number of values padded at - the beginning of `vector` and `iaxis_pad_width[1]` represents - the number of values padded at the end of `vector`. - - `iaxis` is the axis currently being calculated. - - `kwargs` is a dict of keyword arguments the function requires. - stat_length - Used in "maximum", "mean", "median", and "minimum". Number of values at edge - of each axis used to calculate the statistic value. - - ((before_1, after_1), … (before_N, after_N)) yields unique statistic - lengths for each axis. - - ((before, after),) yields same before and after statistic lengths for - each axis. - - stat_length (integer) is a shortcut for before = after = stat_length - length for all axes. - - None uses the entire axis. - constant_values - Used in "constant". The values to set the padded values for each axis. - - ((before_1, after_1), ... (before_N, after_N)) yields unique pad - constants for each axis. - - ((before, after),) yields same before and after constants for each axis. - - constant (integer) is a shortcut for before = after = constant for - all axes. - end_values - Used in "linear_ramp". The values used for the ending value of the linear_ramp - and that will form the edge of the padded array. - - ((before_1, after_1), ... (before_N, after_N)) yields unique end values - for each axis. - - ((before, after),) yields same before and after end values for each axis - - end (integer) is a shortcut for before = after = end for all axes. - reflect_type - Used in "reflect", and "symmetric". The "even" style is the default with an - unaltered reflection around the edge value. For the "odd" style, the extended - part of the array is created by subtracting the reflected values from two - times the edge value. + x + The array to scan over. + fn + The associative function to apply. + reverse + Whether to scan in reverse with respect to the given axis. + axis + The axis to scan over. + + Returns + ------- + ret + The result of the scan. + """ + elems = [x] + + if reverse: + elems = [ivy.flip(elem, axis=[axis]) for elem in elems] + + def _combine(a, b): + a = a[0] + b = b[0] + if a.shape[axis] == 0: + return [a] + c = fn(a, b) + return [c] + + def _scan(elems): + num_elems = elems[0].shape[axis] + + if num_elems < 2: + return elems + + reduced_elems = _combine( + [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], + [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], + ) + + odd_elems = _scan(reduced_elems) + + if num_elems % 2 == 0: + even_elems = _combine( + [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + else: + even_elems = _combine( + odd_elems, + [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], + ) + even_elems = [ + ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) + for (elem, result) in zip(elems, even_elems) + ] + return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) + + scans = _scan(elems) + + if reverse: + scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] + + return ivy.reshape(ivy.asarray(scans), elems[0].shape) + + +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_1d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least one dimension. Scalar inputs are converted to + 1-dimensional arrays, whilst higher-dimensional inputs are preserved. + + Parameters + ---------- + arys + One or more input arrays. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - Padded array of the same rank as the input but with shape increased according - to pad_width. + An array, or list of arrays, each with atleast 1D. + Copies are made only if necessary. + + Examples + -------- + >>> ary1 = ivy.array(5) + >>> ivy.atleast_1d(ary1) + ivy.array([5]) + >>> ary2 = ivy.array([[3,4]]) + >>> ivy.atleast_1d(ary2) + ivy.array([[3, 4]]) + >>> ivy.atleast_1d(6,7,8) + [ivy.array([6]), ivy.array([7]), ivy.array([8])] + """ + return ivy.current_backend().atleast_1d(*arys, copy=copy) - Both the description and the type hints above assume an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_2d( + *arys: Union[ivy.Array, ivy.NativeArray], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least two dimension. Scalar inputs are converted to + 2-dimensional arrays, whilst higher-dimensional inputs are preserved. + + Parameters + ---------- + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have two or more + dimensions are preserved. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + + Returns + ------- + ret + An array, or list of arrays, each with atleast 2D. + Copies are made only if necessary. Examples -------- - With :class:`ivy.Array` input: + >>> ary1 = ivy.array(5) + >>> ivy.atleast_2d(ary1) + ivy.array([[5]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_2d(ary2) + ivy.array([[[3, 4]]]) + >>> ivy.atleast_2d(6,7,8) + [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] + """ + return ivy.current_backend().atleast_2d(*arys, copy=copy) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 3, 0, 0], - [0, 0, 4, 5, 6, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="reflect") - >>> print(y) - ivy.array([[6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1], - [6, 5, 4, 5, 6, 5, 4], - [3, 2, 1, 2, 3, 2, 1]]) +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_device_shifting +def atleast_3d( + *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], + copy: Optional[bool] = None, +) -> List[ivy.Array]: + """ + Convert inputs to arrays with at least three dimension. Scalar inputs are converted + to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="symmetric") - >>> print(y) - ivy.array([[2, 1, 1, 2, 3, 3, 2], - [2, 1, 1, 2, 3, 3, 2], - [5, 4, 4, 5, 6, 6, 5], - [5, 4, 4, 5, 6, 6, 5]]) + Parameters + ---------- + arys + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have three or more + dimensions are preserved. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. - With :class:`ivy.NativeArray` input: + Returns + ------- + ret + An array, or list of arrays, each with a.ndim >= 3. Copies + are avoided where possible, and views with three or more + dimensions are returned. For example, a 1-D array of shape + (N,) becomes a view of shape (1, N, 1), and a 2-D array of + shape (M, N) becomes a view of shape (M, N, 1). - >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) - >>> padding = ((1, 1), (2, 2)) - >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) - >>> print(y) - ivy.array([[7, 7, 7, 7, 7, 7, 7], - [7, 7, 1, 2, 3, 7, 7], - [7, 7, 4, 5, 6, 7, 7], - [7, 7, 7, 7, 7, 7, 7]]) + Examples + -------- + >>> ary1 = ivy.array([5,6]) + >>> ivy.atleast_3d(ary1) + ivy.array([[[5], + [6]]]) + >>> ary2 = ivy.array([[[3,4]]]) + >>> ivy.atleast_3d(ary2) + ivy.array([[[3, 4]]]) + >>> ary3 = ivy.array([[3,4],[9,10]]) + >>> ivy.atleast_3d(6,7,ary3) + [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], + [ 4]], - With :class:`ivy.Container` input: + [[ 9], + [10]]])] + """ + return ivy.current_backend().atleast_3d(*arys, copy=copy) - >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) - >>> padding = (1, 1) - >>> y = ivy.pad(x, padding, mode="constant") - >>> print(y) - { - a: ivy.array([0, 0, 1, 2, 0]), - b: ivy.array([0, 4, 5, 6, 0]) - } + +@handle_exceptions +@inputs_to_native_shapes +def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: """ - _check_arguments( - mode, - pad_width, - stat_length, - constant_values, - end_values, - reflect_type, - ) - if mode == "dilated": - pad_width = _to_dilated(pad_width, input.ndim) - if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: - constant_values = ivy.asarray(constant_values, dtype=input.dtype) - return _interior_pad(input, constant_values, pad_width) - pad_width = _to_pairs(pad_width, len(input.shape)) - if callable(mode): - func = mode - padded, _ = _pad_simple(input, pad_width, fill_value=0) - for axis in range(padded.ndim): - padded = ivy.moveaxis(padded, axis, -1) - inds = ivy.ndindex(padded.shape[:-1]) - for ind in inds: - padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) - return padded - padded, original_area_slice = _pad_simple(input, pad_width) - axes = range(padded.ndim) - stat_functions = { - "maximum": ivy.max, - "minimum": ivy.min, - "mean": ivy.mean, - "median": ivy.median, - } - if mode == "constant": - constant_values = _to_pairs(constant_values, padded.ndim) - constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) - for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): - padded = _set_pad_area(padded, axis, width_pair, value_pair) - elif mode == "empty": - pass - elif mode == "edge": - for axis, width_pair in zip(axes, pad_width): - edge_pair = _get_edges(padded, axis, width_pair) - padded = _set_pad_area(padded, axis, width_pair, edge_pair) - elif mode == "linear_ramp": - end_values = _to_pairs(end_values, padded.ndim) - for axis, width_pair, value_pair in zip(axes, pad_width, end_values): - ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) - padded = _set_pad_area(padded, axis, width_pair, ramp_pair) - elif mode in stat_functions: - func = stat_functions[mode] - stat_length = _to_pairs(stat_length, padded.ndim) - if mode == "median": - ivy.utils.assertions.check_true( - ivy.is_float_dtype(input), - message="median interpolation is only supported for floats", - ) - for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): - stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) - padded = _set_pad_area(padded, axis, width_pair, stat_pair) - elif mode in {"reflect", "symmetric"}: - include_edge = True if mode == "symmetric" else False - for axis, (left_index, right_index) in zip(axes, pad_width): - if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): - edge_pair = _get_edges(padded, axis, (left_index, right_index)) - padded = _set_pad_area( - padded, axis, (left_index, right_index), edge_pair - ) - continue - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_reflect_both( - padded, axis, (left_index, right_index), reflect_type, include_edge - ) - elif mode == "wrap": - for axis, (left_index, right_index) in zip(axes, pad_width): - while left_index > 0 or right_index > 0: - left_index, right_index, padded = _set_wrap_both( - padded, axis, (left_index, right_index) - ) - return padded + Broadcasts shapes. + Parameters + ---------- + shapes + The shapes to broadcast. -pad.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Returns + ------- + ret + The broadcasted shape. + + Examples + -------- + >>> x = [(3, 3), (3, 1)] + >>> print(ivy.broadcast_shapes(*x)) + (3, 3) + + >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) + (3, 3) + """ + return ivy.current_backend().broadcast_shapes(*shapes) @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view +@handle_out_argument @to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def vsplit( - ary: Union[ivy.Array, ivy.NativeArray], - indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], +def choose( + arr: Union[ivy.Array, ivy.NativeArray], + choices: Union[ivy.Array, ivy.NativeArray], /, *, - copy: Optional[bool] = None, -) -> List[ivy.Array]: + out: None = None, + mode: Union[str, None] = None, +) -> ivy.Array: """ - Split an array vertically into multiple sub-arrays. + Take values from the input array by matching 1d index and data slices. Parameters ---------- - ary - Array input. - indices_or_sections - If indices_or_sections is an integer n, the array is split into n - equal sections, provided that n must be a divisor of the split axis. - If indices_or_sections is a sequence of ints or 1-D array, - then input is split at each of the indices. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + arr + The source array. + choices + The indices of the values to extract. + out + The output array. + mode + One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices + will be handled. Returns ------- ret - input array split vertically. + The returned array has the same shape as `indices`. Examples -------- - >>> ary = ivy.array( - [[[0., 1.], - [2., 3.]], - [[4., 5.], - [6., 7.]]] - ) - >>> ivy.vsplit(ary, 2) - [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) + >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], + [20, 21, 22, 23], [30, 31, 32, 33]]) + >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) + ivy.array([20, 31, 12, 3]) + >>> arr = ivy.array([2, 4, 1, 0]) + >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) + ivy.array([20, 1, 12, 3]) """ - return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) + return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def concat_from_sequence( + input_sequence: Union[ + Tuple[Union[ivy.Array, ivy.NativeArray]], + List[Union[ivy.Array, ivy.NativeArray]], + ], + /, + *, + new_axis: int = 0, + axis: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Concatenate a sequence of arrays along a new or an existing axis. + + Parameters + ---------- + input_sequence + A sequence of arrays. + new_axis + Insert and concatenate on a new axis or not, + default 0 means do not insert new axis. + new_axis = 0: concatenate + new_axis = 1: stack + axis + axis along which the arrays will be concatenated. + + out + optional output array, for writing the result to. + + Returns + ------- + ret + Output Array + """ + return current_backend(input_sequence).concat_from_sequence( + input_sequence, new_axis=new_axis, axis=axis, out=out + ) @handle_exceptions @@ -1336,243 +870,451 @@ def dsplit( return ivy.current_backend(ary).dsplit(ary, indices_or_sections, copy=copy) +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Stack arrays in sequence depth wise (along third axis). + + Parameters + ---------- + arrays + Sequence of arrays to be stacked. + + Returns + ------- + ret + The array formed by stacking the given arrays. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2], + [2, 3], + [3, 4]]]) + >>> x = ivy.array([[1], [2], [3]]) + >>> y = ivy.array([[2], [3], [4]]) + >>> ivy.dstack((x, y)) + ivy.array([[[1, 2]], + [[2, 3]], + [[3, 4]]]) + """ + return ivy.current_backend().dstack(arrays, out=out) + + +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_view +@handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_device_shifting -def atleast_1d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], +def expand( + x: Union[ivy.Array, ivy.NativeArray], + shape: Union[ivy.Shape, ivy.NativeShape], + /, + *, + copy: Optional[bool] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Broadcast the input Array following the given shape and the broadcast rule. + + Parameters + ---------- + x + Array input. + shape + A 1-D Array indicates the shape you want to expand to, + following the broadcast rule. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Output Array + """ + return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back +@handle_array_function +def fill_diagonal( + a: Union[ivy.Array, ivy.NativeArray], + v: Union[int, float], + /, + *, + wrap: bool = False, +) -> Union[ivy.Array, ivy.NativeArray]: + """ + Fill the main diagonal of the given array of any dimensionality.. + + Parameters + ---------- + a + Array at least 2D. + v + The value to write on the diagonal. + wrap + The diagonal ‘wrapped’ after N columns for tall matrices. + + Returns + ------- + ret + Array with the diagonal filled. + """ + return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) + + +@handle_exceptions +@handle_nestable +@handle_partial_mixed_function +@handle_array_like_without_promotion +@handle_view +@inputs_to_ivy_arrays +@handle_array_function +def flatten( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, copy: Optional[bool] = None, -) -> List[ivy.Array]: + start_dim: Optional[int] = 0, + end_dim: Optional[int] = -1, + order: Optional[str] = "C", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert inputs to arrays with at least one dimension. Scalar inputs are converted to - 1-dimensional arrays, whilst higher-dimensional inputs are preserved. + Flattens input by reshaping it into a one-dimensional tensor. If start_dim or + end_dim are passed, only dimensions starting with start_dim and ending with end_dim + are flattened. The order of elements in input is unchanged. Parameters ---------- - arys - One or more input arrays. + x + input array to flatten. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. + start_dim + first dim to flatten. If not set, defaults to 0. + end_dim + last dim to flatten. If not set, defaults to -1. + order + Read the elements of the input container using this index order, + and place the elements into the reshaped array using this index order. + ‘C’ means to read / write the elements using C-like index order, + with the last axis index changing fastest, back to the first axis index + changing slowest. + ‘F’ means to read / write the elements using Fortran-like index order, with + the first index changing fastest, and the last index changing slowest. + Note that the ‘C’ and ‘F’ options take no account of the memory layout + of the underlying array, and only refer to the order of indexing. + Default order is 'C' + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with atleast 1D. - Copies are made only if necessary. + the flattened array over the specified dimensions. Examples -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_1d(ary1) - ivy.array([5]) - >>> ary2 = ivy.array([[3,4]]) - >>> ivy.atleast_1d(ary2) - ivy.array([[3, 4]]) - >>> ivy.atleast_1d(6,7,8) - [ivy.array([6]), ivy.array([7]), ivy.array([8])] + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x) + ivy.array([1, 2, 3, 4]) + + >>> x = ivy.array([[1,2], [3,4]]) + >>> ivy.flatten(x, order='F') + ivy.array([1, 3, 2, 4]) + + >>> x = ivy.array( + [[[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12]], + + [[ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]]], + + + [[[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17]], + + [[ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]]], + + + [[[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1]], + + [[19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]]] + ) + >>> ivy.flatten(x, start_dim = 1, end_dim = 2) + ivy.array( + [[[ 5, 5, 0, 6], + [17, 15, 11, 16], + [ 6, 3, 13, 12], + [ 6, 18, 10, 4], + [ 5, 1, 17, 3], + [14, 14, 18, 6]], + + [[12, 0, 1, 13], + [ 8, 7, 0, 3], + [19, 12, 6, 17], + [ 4, 15, 6, 15], + [ 0, 5, 17, 9], + [ 9, 3, 6, 19]], + + [[17, 13, 11, 16], + [ 4, 18, 17, 4], + [10, 10, 9, 1], + [19, 17, 13, 10], + [ 4, 19, 16, 17], + [ 2, 12, 8, 14]]])) """ - return ivy.current_backend().atleast_1d(*arys, copy=copy) + if x.shape == (): + x = ivy.reshape(x, (1, -1))[0, :] + if start_dim == end_dim: + return ivy.inplace_update(out, x) if ivy.exists(out) else x + if start_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {start_dim}" + ) + if end_dim not in range(-len(x.shape), len(x.shape)): + raise IndexError( + "Dimension out of range (expected to be in range of" + f" {[-len(x.shape), len(x.shape) - 1]}, but got {end_dim}" + ) + if start_dim < 0: + start_dim = len(x.shape) + start_dim + if end_dim < 0: + end_dim = len(x.shape) + end_dim + c = 1 + for i in range(start_dim, end_dim + 1): + c *= x.shape[i] + lst = [c] + if start_dim != 0: + for i in range(0, start_dim): + lst.insert(i, x.shape[i]) + for i in range(end_dim + 1, len(x.shape)): + lst.insert(i, x.shape[i]) + return ivy.reshape(x, tuple(lst), order=order, out=out) @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion +@handle_view @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def dstack( - arrays: Sequence[ivy.Array], +def fliplr( + m: Union[ivy.Array, ivy.NativeArray], /, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Stack arrays in sequence depth wise (along third axis). + Flip array in the left/right direction. Flip the entries in each column in the + left/right direction. Columns are preserved, but appear in a different order than + before. Parameters ---------- - arrays - Sequence of arrays to be stacked. + m + The array to be flipped. Must be at least 2-D. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - The array formed by stacking the given arrays. + Array corresponding to input array with elements + order reversed along axis 1. Examples -------- - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([2, 3, 4]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2], - [2, 3], - [3, 4]]]) - >>> x = ivy.array([[1], [2], [3]]) - >>> y = ivy.array([[2], [3], [4]]) - >>> ivy.dstack((x, y)) - ivy.array([[[1, 2]], - [[2, 3]], - [[3, 4]]]) + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.fliplr(m) + ivy.array([[0, 0, 1], + [0, 2, 0], + [3, 0, 0]]) """ - return ivy.current_backend().dstack(arrays, out=out) + return ivy.current_backend().fliplr(m, copy=copy, out=out) @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_view +@handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def atleast_2d( - *arys: Union[ivy.Array, ivy.NativeArray], +def flipud( + m: Union[ivy.Array, ivy.NativeArray], + /, + *, copy: Optional[bool] = None, -) -> List[ivy.Array]: + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Convert inputs to arrays with at least two dimension. Scalar inputs are converted to - 2-dimensional arrays, whilst higher-dimensional inputs are preserved. + Flip array in the up/down direction. Flip the entries in each column in the up/down + direction. Rows are preserved, but appear in a different order than before. Parameters ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have two or more - dimensions are preserved. + m + The array to be flipped. copy boolean indicating whether or not to copy the input array. If True, the function must always copy. If False, the function must never copy. In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with atleast 2D. - Copies are made only if necessary. + Array corresponding to input array with elements + order reversed along axis 0. Examples -------- - >>> ary1 = ivy.array(5) - >>> ivy.atleast_2d(ary1) - ivy.array([[5]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_2d(ary2) - ivy.array([[[3, 4]]]) - >>> ivy.atleast_2d(6,7,8) - [ivy.array([[6]]), ivy.array([[7]]), ivy.array([[8]])] + >>> m = ivy.diag([1, 2, 3]) + >>> ivy.flipud(m) + ivy.array([[ 0., 0., 3.], + [ 0., 2., 0.], + [ 1., 0., 0.]]) """ - return ivy.current_backend().atleast_2d(*arys, copy=copy) + return ivy.current_backend().flipud(m, copy=copy, out=out) -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back +@inputs_to_ivy_arrays +@handle_array_function @handle_device_shifting -def atleast_3d( - *arys: Union[ivy.Array, ivy.NativeArray, bool, Number], - copy: Optional[bool] = None, -) -> List[ivy.Array]: +def fold( + x: Union[ivy.Array, ivy.NativeArray], + /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert inputs to arrays with at least three dimension. Scalar inputs are converted - to 3-dimensional arrays, whilst higher-dimensional inputs are preserved. + Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, + refolds the n-mode unfolded tensor into the original tensor of the specified shape. - Parameters - ---------- - arys - One or more array-like sequences. Non-array inputs are - converted to arrays. Arrays that already have three or more - dimensions are preserved. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. + Parameters + ---------- + input + unfolded tensor of shape ``(shape[mode], -1)`` + mode + the mode of the unfolding + shape + shape of the original tensor before unfolding + out + optional output array, for writing the result to. Returns ------- ret - An array, or list of arrays, each with a.ndim >= 3. Copies - are avoided where possible, and views with three or more - dimensions are returned. For example, a 1-D array of shape - (N,) becomes a view of shape (1, N, 1), and a 2-D array of - shape (M, N) becomes a view of shape (M, N, 1). - - Examples - -------- - >>> ary1 = ivy.array([5,6]) - >>> ivy.atleast_3d(ary1) - ivy.array([[[5], - [6]]]) - >>> ary2 = ivy.array([[[3,4]]]) - >>> ivy.atleast_3d(ary2) - ivy.array([[[3, 4]]]) - >>> ary3 = ivy.array([[3,4],[9,10]]) - >>> ivy.atleast_3d(6,7,ary3) - [ivy.array([[[6]]]), ivy.array([[[7]]]), ivy.array([[[ 3], - [ 4]], - - [[ 9], - [10]]])] + folded_tensor of shape `shape` """ - return ivy.current_backend().atleast_3d(*arys, copy=copy) + full_shape = list(shape) + mode_dim = full_shape.pop(mode) + full_shape.insert(0, mode_dim) + return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def take_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - axis: int, +def heaviside( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], /, *, - mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Take values from the input array by matching 1d index and data slices. + Compute the Heaviside step function for each element in x1. Parameters ---------- - arr - The source array. - indices - The indices of the values to extract. - axis - The axis over which to select values. - If axis is None, arr is treated as a flattened 1D array. - mode - One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices - will be handled. + x1 + input array. + x2 + values to use where x1 is zero. out - The output array. + optional output array, for writing the result to. Returns ------- ret - The returned array has the same shape as `indices`. + output array with element-wise Heaviside step function of x1. + This is a scalar if both x1 and x2 are scalars. Examples -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> y = ivy.take_along_axis(arr, indices, 1) - >>> print(y) - ivy.array([[4, 3, 3], [1, 1, 1]]) + With :class:`ivy.Array` input: + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([0.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0.0000, 0.5000, 1.0000]) + + >>> x1 = ivy.array([-1.5, 0, 2.0]) + >>> x2 = ivy.array([1.2, -2.0, 3.5]) + >>> ivy.heaviside(x1, x2) + ivy.array([0., -2., 1.]) """ - return ivy.current_backend(arr).take_along_axis( - arr, indices, axis, mode=mode, out=out - ) + return ivy.current_backend().heaviside(x1, x2, out=out) @handle_exceptions @@ -1634,524 +1376,595 @@ def hsplit( return ivy.current_backend(ary).hsplit(ary, indices_or_sections, copy=copy) -@handle_exceptions -@inputs_to_native_shapes -def broadcast_shapes(*shapes: Union[List[int], List[Tuple]]) -> Tuple[int]: +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def hstack( + arrays: Sequence[ivy.Array], + /, + *, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> ivy.Array: """ - Broadcasts shapes. + Stack arrays in sequence horizotally (column wise). Parameters ---------- - shapes - The shapes to broadcast. + arrays + Sequence of arrays to be stacked. Returns ------- ret - The broadcasted shape. + The array formed by stacking the given arrays. Examples -------- - >>> x = [(3, 3), (3, 1)] - >>> print(ivy.broadcast_shapes(*x)) - (3, 3) - - >>> print(ivy.broadcast_shapes(*[(3, 3),(3, 1),(1, 3)])) - (3, 3) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.hstack((x, y)) + ivy.array([1, 2, 3, 2, 3, 4]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0, 0, 0]) + >>> ivy.hstack((x, y, x)) + ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3]) + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.hstack(y)) + ivy.array([[5, 6, 7, 8]]) """ - return ivy.current_backend().broadcast_shapes(*shapes) + return ivy.current_backend().hstack(arrays, out=out) -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_device_shifting -def expand( +def i0( x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape], /, *, - copy: Optional[bool] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Broadcast the input Array following the given shape and the broadcast rule. + Compute the Bessel i0 function of x element-wise. Parameters ---------- x Array input. - shape - A 1-D Array indicates the shape you want to expand to, - following the broadcast rule. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. out optional output array, for writing the result to. Returns ------- ret - Output Array + Array with the modified Bessel function + evaluated at each of the elements of x. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> ivy.i0(x) + ivy.array([1.26606588, 2.2795853 , 4.88079259]) """ - return ivy.current_backend(x).expand(x, shape, out=out, copy=copy) + return ivy.current_backend(x).i0(x, out=out) -@handle_exceptions @handle_nestable +@handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays -def put_along_axis( - arr: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - values: Union[ivy.Array, ivy.NativeArray], - axis: int, +@handle_array_function +@handle_device_shifting +def matricize( + x: Union[ivy.Array, ivy.NativeArray], /, + row_modes: Sequence[int], + column_modes: Optional[Sequence[int]] = None, *, - mode: str = "raise", out: Optional[ivy.Array] = None, -) -> None: +) -> ivy.Array: """ - Put values into the input array by matching 1d index and data slices along a - specified axis. + Matricizes the given tensor. Parameters ---------- - arr : array_like - The input array to modify. - indices : array_like - The indices of the values to put into `arr`. - values : array_like - The values to put into `arr`. - axis : int - The axis over which to put the `values`. - mode : {'raise', 'wrap', 'clip'}, optional - Specifies how out-of-bounds indices will be handled. - The following modes are available: - - - 'raise': a `ValueError` is raised when an index is out of bounds. - - 'wrap': the index is wrapped around to the corresponding index - at the other end of the axis. - - 'clip': the index is clipped to the closest in-bounds index. - out : ndarray, optional - Output array in which to place the result. - If not specified, a new array is created. + x + the input tensor + row_modes + modes to use as row of the matrix (in the desired order) + column_modes + modes to use as column of the matrix, in the desired order + if None, the modes not in `row_modes` will be used in ascending order + out + optional output array, for writing the result to. - Returns + ret ------- - None - - Examples - -------- - >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) - >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) - >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) - >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') - >>> print(arr) - ivy.array([[3, 7, 5], - [6, 4, 1]]) + ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) """ - if out is None: - out = ivy.zeros_like(arr) - - indices = ivy.expand_dims(indices, axis=axis) - values = ivy.expand_dims(values, axis=axis) - - stacked = ivy.concat((arr, values), axis=axis) - - sorted_indices = ivy.argsort(indices, axis=axis) - sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), - sorted_stacked, - arr, - ) - - if mode == "clip": - indices = ivy.clip(indices, 0, arr.shape[axis] - 1) - elif mode == "wrap": - indices = ivy.mod(indices, arr.shape[axis]) - - arr = ivy.where( - ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values - ) + ndims = len(x.shape) + row_indices = list(row_modes) - ivy.assign(out, arr) + if column_modes: + column_indices = list(column_modes) + else: + column_indices = [i for i in range(ndims) if i not in row_indices] + if sorted(column_indices + row_indices) != list(range(ndims)): + msg = ( + "If you provide both column and row modes for the matricization then" + " column_modes + row_modes must contain all the modes of the tensor." + f" Yet, got row_modes={row_modes} and column_modes={column_modes}." + ) + raise ValueError(msg) + row_size, column_size = 1, 1 + row_size = int(ivy.prod([x.shape[i] for i in row_indices])) + column_size = int(ivy.prod([x.shape[i] for i in column_indices])) -def _check_bounds(shape0, shape1, strides1, itemsize): - numel0 = math.prod(shape0) - ndim1 = len(shape1) - return ( - sum((shape1[i] - 1) * strides1[i] for i in range(ndim1)) + itemsize - <= numel0 * itemsize + return ivy.reshape( + ivy.permute_dims(x, row_indices + column_indices), + (row_size, column_size), + out=out, ) -@handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@inputs_to_native_shapes -def as_strided( - x: Union[ivy.Array, ivy.NativeArray], - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - strides: Sequence[int], +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def moveaxis( + a: Union[ivy.Array, ivy.NativeArray], + source: Union[int, Sequence[int]], + destination: Union[int, Sequence[int]], /, -) -> ivy.Array: + *, + copy: Optional[bool] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Create a copy of the input array with the given shape and strides. + Move axes of an array to new positions.. Parameters ---------- - x - Input Array. - shape - The shape of the new array. - strides - The strides of the new array (specified in bytes). + a + The array whose axes should be reordered. + source + Original positions of the axes to move. These must be unique. + destination + Destination positions for each of the original axes. + These must also be unique. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + out + optional output array, for writing the result to. Returns ------- ret - Output Array + Array with moved axes. This array is a view of the input array. Examples -------- - >>> x = ivy.array([1, 2, 3, 4, 5, 6]) - >>> ivy.as_strided(x, (4, 3), (8, 8)) - ivy.array([[1, 2, 3], - [2, 3, 4], - [3, 4, 5], - [4, 5, 6]]) + With :class:`ivy.Array` input: + >>> x = ivy.zeros((3, 4, 5)) + >>> ivy.moveaxis(x, 0, -1).shape + (4, 5, 3) + >>> ivy.moveaxis(x, -1, 0).shape + (5, 3, 4) """ - itemsize = x.itemsize - if not _check_bounds(x.shape, shape, strides, itemsize): - raise ivy.exceptions.IvyException("attempted unsafe memory access") - if any(strides[i] % itemsize != 0 for i in range(len(strides))): - raise ivy.exceptions.IvyException("strides must be multiple of itemsize") - - src = memoryview(ivy.to_numpy(x)).cast("b") - - src_ind = ivy.inner( - ivy.indices(shape).reshape((len(shape), -1)).T, - ivy.array(strides), - ) - src_ind = ivy.expand_dims(src_ind, axis=-1) - src_ind = src_ind + ivy.arange(itemsize) - src_ind = ivy.reshape(src_ind, (-1,)).to_numpy() - - temp_list = [src[i] for i in src_ind] - temp_array = ivy.asarray(temp_list, dtype=ivy.int8) - result = bytearray(temp_array.to_numpy()) - - return ivy.reshape( - ivy.frombuffer(result, dtype=x.dtype, count=math.prod(shape)), - shape, - ) - - -as_strided.unsupported_dtypes = ("bfloat16",) -as_strided.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + return ivy.current_backend().moveaxis(a, source, destination, copy=copy, out=out) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_out_argument -@to_native_arrays_and_back +@handle_array_like_without_promotion +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def concat_from_sequence( - input_sequence: Union[ - Tuple[Union[ivy.Array, ivy.NativeArray]], - List[Union[ivy.Array, ivy.NativeArray]], - ], +def pad( + input: Union[ivy.Array, ivy.NativeArray], + pad_width: Union[Iterable[Tuple[int]], int], /, *, - new_axis: int = 0, - axis: int = 0, - out: Optional[ivy.Array] = None, + mode: Union[ + Literal[ + "constant", + "dilated", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", + ], + Callable, + ] = "constant", + stat_length: Union[Iterable[Tuple[int]], int] = 1, + constant_values: Union[Iterable[Tuple[Number]], Number] = 0, + end_values: Union[Iterable[Tuple[Number]], Number] = 0, + reflect_type: Literal["even", "odd"] = "even", + **kwargs: Optional[Any], ) -> ivy.Array: """ - Concatenate a sequence of arrays along a new or an existing axis. + Pad an array. Parameters ---------- - input_sequence - A sequence of arrays. - new_axis - Insert and concatenate on a new axis or not, - default 0 means do not insert new axis. - new_axis = 0: concatenate - new_axis = 1: stack - axis - axis along which the arrays will be concatenated. - - out - optional output array, for writing the result to. + input + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + - ((before_1, after_1), … (before_N, after_N)) yields unique pad widths + for each axis. + - ((before, after),) yields same before and after pad for each axis. + - pad (integer) is shortcut for before = after = pad width for all axes. + mode + One of the following string values or a user-supplied function. + - "constant": Pads with a constant value. + - "edge": Pads with the input's edge values. + - "linear_ramp": Pads with the linear ramp between end_value + and the input's edge value. + - "maximum": Pads with the maximum value of all or part of the vector + along each axis. + - "mean": Pads with the mean value of all or part of the vector along + each axis. + - "median": Pads with the median value of all or part of the vector + along each axis. + - "minimum": Pads with the minimum value of all or part of the vector + along each axis. + - "reflect": Pads with the reflection mirrored on the first and last + values of the vector along each axis. + - "symmetric": Pads with the reflection of the vector mirrored along + the edge of the input. + - "wrap": Pads with the wrap of the vector along the axis. + The first values are used to pad the end and the end values are used + to pad the beginning. + - "empty": Pads with undefined values. + - : Pads with a user-defined padding function. The padding + function should modify a rank 1 array following the signature + `padding_func(vector, iaxis_pad_width, iaxis, kwargs)`, where: + - `vector` is a rank 1 array already padded with zeros. Padded + values are `vector[:iaxis_pad_width[0]]` and + `vector[-iaxis_pad_width[1]:]`. + - `iaxis_pad_width` is a 2-tuple of ints, where + `iaxis_pad_width[0]` represents the number of values padded at + the beginning of `vector` and `iaxis_pad_width[1]` represents + the number of values padded at the end of `vector`. + - `iaxis` is the axis currently being calculated. + - `kwargs` is a dict of keyword arguments the function requires. + stat_length + Used in "maximum", "mean", "median", and "minimum". Number of values at edge + of each axis used to calculate the statistic value. + - ((before_1, after_1), … (before_N, after_N)) yields unique statistic + lengths for each axis. + - ((before, after),) yields same before and after statistic lengths for + each axis. + - stat_length (integer) is a shortcut for before = after = stat_length + length for all axes. + - None uses the entire axis. + constant_values + Used in "constant". The values to set the padded values for each axis. + - ((before_1, after_1), ... (before_N, after_N)) yields unique pad + constants for each axis. + - ((before, after),) yields same before and after constants for each axis. + - constant (integer) is a shortcut for before = after = constant for + all axes. + end_values + Used in "linear_ramp". The values used for the ending value of the linear_ramp + and that will form the edge of the padded array. + - ((before_1, after_1), ... (before_N, after_N)) yields unique end values + for each axis. + - ((before, after),) yields same before and after end values for each axis + - end (integer) is a shortcut for before = after = end for all axes. + reflect_type + Used in "reflect", and "symmetric". The "even" style is the default with an + unaltered reflection around the edge value. For the "odd" style, the extended + part of the array is created by subtracting the reflected values from two + times the edge value. Returns ------- ret - Output Array - """ - return current_backend(input_sequence).concat_from_sequence( - input_sequence, new_axis=new_axis, axis=axis, out=out - ) + Padded array of the same rank as the input but with shape increased according + to pad_width. -def _slice(operand, start_indices, limit_indices, strides=None): - strides = [1] * len(operand.shape) if strides is None else strides + Both the description and the type hints above assume an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - full_slice = () - for i, _ in enumerate(operand.shape): - strides_i = int(strides[i]) - start_i = int(start_indices[i]) - limit_i = int(limit_indices[i]) - full_slice += (slice(start_i, limit_i, strides_i),) - return operand[full_slice] + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=0) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 3, 0, 0], + [0, 0, 4, 5, 6, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="reflect") + >>> print(y) + ivy.array([[6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1], + [6, 5, 4, 5, 6, 5, 4], + [3, 2, 1, 2, 3, 2, 1]]) -def _slice_along_axis(x, start=0, stop=None, stride=1, axis=0): - if axis >= 0: - slices = [slice(None)] * axis + [slice(start, stop, stride)] - else: - slices = [Ellipsis, slice(start, stop, stride)] + [slice(None)] * (-1 - axis) - return x[tuple(slices)] + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="symmetric") + >>> print(y) + ivy.array([[2, 1, 1, 2, 3, 3, 2], + [2, 1, 1, 2, 3, 3, 2], + [5, 4, 4, 5, 6, 6, 5], + [5, 4, 4, 5, 6, 6, 5]]) + With :class:`ivy.NativeArray` input: -def _interior_pad(operand, padding_value, padding_config): - for axis, (_, _, interior) in enumerate(padding_config): - if interior > 0: - new_shape = list(operand.shape) - new_shape[axis] = new_shape[axis] + (new_shape[axis] - 1) * interior - new_array = ivy.full( - new_shape, padding_value, dtype=operand.dtype - ) - src_indices = ivy.arange(operand.shape[axis]) - dst_indices = src_indices * (interior + 1) - index_tuple = [slice(None)] * operand.ndim - index_tuple[axis] = dst_indices - new_array[tuple(index_tuple)] = operand - operand = new_array + >>> x = ivy.native_array([[1, 2, 3], [4, 5, 6]]) + >>> padding = ((1, 1), (2, 2)) + >>> y = ivy.pad(x, padding, mode="constant", constant_values=7) + >>> print(y) + ivy.array([[7, 7, 7, 7, 7, 7, 7], + [7, 7, 1, 2, 3, 7, 7], + [7, 7, 4, 5, 6, 7, 7], + [7, 7, 7, 7, 7, 7, 7]]) - start_indices = [0] * operand.ndim - limit_indices = [0] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low < 0: - start_indices[axis] = abs(low) - if high < 0: - limit_indices[axis] = high - else: - limit_indices[axis] = operand.shape[axis] + 1 - padded = _slice(operand, start_indices, limit_indices) + With :class:`ivy.Container` input: - pad_width = [(0, 0)] * operand.ndim - for axis, (low, high, _) in enumerate(padding_config): - if low > 0 and high > 0: - pad_width[axis] = (low, high) - elif low > 0 and not high > 0: - pad_width[axis] = (low, 0) - elif high > 0 and not low > 0: - pad_width[axis] = (0, high) - padded = ivy.constant_pad(padded, pad_width, value=padding_value) + >>> x = ivy.Container(a=ivy.array([0, 1, 2]), b=ivy.array([4, 5, 6])) + >>> padding = (1, 1) + >>> y = ivy.pad(x, padding, mode="constant") + >>> print(y) + { + a: ivy.array([0, 0, 1, 2, 0]), + b: ivy.array([0, 4, 5, 6, 0]) + } + """ + _check_arguments( + mode, + pad_width, + stat_length, + constant_values, + end_values, + reflect_type, + ) + if mode == "dilated": + pad_width = _to_dilated(pad_width, input.ndim) + if not ivy.is_array(constant_values) or constant_values.dtype != input.dtype: + constant_values = ivy.asarray(constant_values, dtype=input.dtype) + return _interior_pad(input, constant_values, pad_width) + pad_width = _to_pairs(pad_width, len(input.shape)) + if callable(mode): + func = mode + padded, _ = _pad_simple(input, pad_width, fill_value=0) + for axis in range(padded.ndim): + padded = ivy.moveaxis(padded, axis, -1) + inds = ivy.ndindex(padded.shape[:-1]) + for ind in inds: + padded[ind] = func(padded[ind], pad_width[axis], axis, kwargs) + return padded + padded, original_area_slice = _pad_simple(input, pad_width) + axes = range(padded.ndim) + stat_functions = { + "maximum": ivy.max, + "minimum": ivy.min, + "mean": ivy.mean, + "median": ivy.median, + } + if mode == "constant": + constant_values = _to_pairs(constant_values, padded.ndim) + constant_values = tuple(tuple(map(ivy.array, pair)) for pair in constant_values) + for axis, width_pair, value_pair in zip(axes, pad_width, constant_values): + padded = _set_pad_area(padded, axis, width_pair, value_pair) + elif mode == "empty": + pass + elif mode == "edge": + for axis, width_pair in zip(axes, pad_width): + edge_pair = _get_edges(padded, axis, width_pair) + padded = _set_pad_area(padded, axis, width_pair, edge_pair) + elif mode == "linear_ramp": + end_values = _to_pairs(end_values, padded.ndim) + for axis, width_pair, value_pair in zip(axes, pad_width, end_values): + ramp_pair = _get_linear_ramps(padded, axis, width_pair, value_pair) + padded = _set_pad_area(padded, axis, width_pair, ramp_pair) + elif mode in stat_functions: + func = stat_functions[mode] + stat_length = _to_pairs(stat_length, padded.ndim) + if mode == "median": + ivy.utils.assertions.check_true( + ivy.is_float_dtype(input), + message="median interpolation is only supported for floats", + ) + for axis, width_pair, length_pair in zip(axes, pad_width, stat_length): + stat_pair = _get_stats(padded, axis, width_pair, length_pair, func) + padded = _set_pad_area(padded, axis, width_pair, stat_pair) + elif mode in {"reflect", "symmetric"}: + include_edge = True if mode == "symmetric" else False + for axis, (left_index, right_index) in zip(axes, pad_width): + if input.shape[axis] == 1 and (left_index > 0 or right_index > 0): + edge_pair = _get_edges(padded, axis, (left_index, right_index)) + padded = _set_pad_area( + padded, axis, (left_index, right_index), edge_pair + ) + continue + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_reflect_both( + padded, axis, (left_index, right_index), reflect_type, include_edge + ) + elif mode == "wrap": + for axis, (left_index, right_index) in zip(axes, pad_width): + while left_index > 0 or right_index > 0: + left_index, right_index, padded = _set_wrap_both( + padded, axis, (left_index, right_index) + ) return padded -def _interleave(a, b, axis): - assert a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1 - a_pad = [(0, 0, 0)] * a.ndim - b_pad = [(0, 0, 0)] * b.ndim - a_pad[axis] = (0, 1 if a.shape[axis] == b.shape[axis] else 0, 1) - b_pad[axis] = (1, 0 if a.shape[axis] == b.shape[axis] else 1, 1) - a = _interior_pad(a, 0.0, a_pad) - b = _interior_pad(b, 0.0, b_pad) - return ivy.add(a, b) - - -@handle_exceptions @handle_nestable +@handle_exceptions +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def associative_scan( +@handle_device_shifting +def partial_fold( x: Union[ivy.Array, ivy.NativeArray], - fn: Callable, /, + mode: int, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, - reverse: bool = False, - axis: int = 0, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Perform an associative scan over the given array. + Re-folds a partially unfolded tensor. Parameters ---------- x - The array to scan over. - fn - The associative function to apply. - reverse - Whether to scan in reverse with respect to the given axis. - axis - The axis to scan over. + a partially unfolded tensor + mode + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions left untouched at the beginning + out + optional output array, for writing the result to. Returns ------- ret - The result of the scan. + partially re-folded tensor """ - elems = [x] - - if reverse: - elems = [ivy.flip(elem, axis=[axis]) for elem in elems] - - def _combine(a, b): - a = a[0] - b = b[0] - if a.shape[axis] == 0: - return [a] - c = fn(a, b) - return [c] - - def _scan(elems): - num_elems = elems[0].shape[axis] - - if num_elems < 2: - return elems - - reduced_elems = _combine( - [_slice_along_axis(elem, 0, -1, stride=2, axis=axis) for elem in elems], - [_slice_along_axis(elem, 1, None, stride=2, axis=axis) for elem in elems], - ) - - odd_elems = _scan(reduced_elems) - - if num_elems % 2 == 0: - even_elems = _combine( - [_slice_along_axis(e, 0, -1, axis=axis) for e in odd_elems], - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], - ) - else: - even_elems = _combine( - odd_elems, - [_slice_along_axis(e, 2, None, stride=2, axis=axis) for e in elems], - ) - even_elems = [ - ivy.concat([_slice_along_axis(elem, 0, 1, axis=axis), result], axis=axis) - for (elem, result) in zip(elems, even_elems) - ] - return list(map(partial(_interleave, axis=axis), even_elems, odd_elems)) - - scans = _scan(elems) - - if reverse: - scans = [ivy.flip(scanned, axis=[axis]) for scanned in scans] - - return ivy.reshape(ivy.asarray(scans), elems[0].shape) + transposed_shape = list(shape) + mode_dim = transposed_shape.pop(skip_begin + mode) + transposed_shape.insert(skip_begin, mode_dim) + return ivy.moveaxis( + ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unique_consecutive( +def partial_tensor_to_vec( x: Union[ivy.Array, ivy.NativeArray], /, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, *, - axis: Optional[int] = None, -) -> Tuple[ - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], - Union[ivy.Array, ivy.NativeArray], -]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Eliminates all but the first element from every consecutive group of equivalent - elements in ``x``. + Partial vectorization of a tensor while ignoring the specified dimension at the + beginning and the end. Parameters ---------- x - input array. - - axis - the axis to apply unique on. If None, unique is applied on flattened ``x``. + tensor to partially vectorise + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + out + optional output array, for writing the result to. Returns ------- ret - a namedtuple ``(output, inverse_indices, counts)`` whose - - first element has the field name ``output`` and is an array - containing ``x`` with its equivalent consecutive elements eliminated. - - second element has the field name ``inverse_indices`` and is an - array containing the indices of ``output`` that reconstruct ``x``. - - third element has the field name ``counts`` and is an array - containing the number of occurrences for each unique value or array in ``x``. - - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) - >>> ivy..unique_consecutive(x) - Results(values=ivy.array([1, 2, 3, 1, 2]), - inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), - counts=ivy.array([2, 2, 1, 2, 1])) + partially vectorised tensor with the + `skip_begin` first and `skip_end` last dimensions untouched """ - return ivy.current_backend(x).unique_consecutive(x, axis=axis) + return partial_unfold( + x, + mode=0, + skip_begin=skip_begin, + skip_end=skip_end, + ravel_tensors=True, + out=out, + ) -@handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_exceptions @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -def fill_diagonal( - a: Union[ivy.Array, ivy.NativeArray], - v: Union[int, float], +@handle_device_shifting +def partial_unfold( + x: Union[ivy.Array, ivy.NativeArray], /, + mode: Optional[int] = 0, + skip_begin: Optional[int] = 1, + skip_end: Optional[int] = 0, + ravel_tensors: Optional[bool] = False, *, - wrap: bool = False, -) -> Union[ivy.Array, ivy.NativeArray]: + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Fill the main diagonal of the given array of any dimensionality.. + Partial unfolding of a tensor while ignoring the specified number of dimensions at + the beginning and the end. For instance, if the first dimension of the tensor is the + number of samples, to unfold each sample, set skip_begin=1. This would, for each i + in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. Parameters ---------- - a - Array at least 2D. - v - The value to write on the diagonal. - wrap - The diagonal ‘wrapped’ after N columns for tall matrices. + x + tensor of shape n_samples x n_1 x n_2 x ... x n_i + mode + indexing starts at 0, therefore mode is in range(0, tensor.ndim) + skip_begin + number of dimensions to leave untouched at the beginning + skip_end + number of dimensions to leave untouched at the end + ravel_tensors + if True, the unfolded tensors are also flattened + out + optional output array, for writing the result to. Returns ------- ret - Array with the diagonal filled. + partially unfolded tensor """ - return ivy.current_backend(a).fill_diag(a, v, wrap=wrap) + if ravel_tensors: + new_shape = [-1] + else: + new_shape = [x.shape[mode + skip_begin], -1] + + if skip_begin: + new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape + + if skip_end: + new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] + + return ivy.reshape( + ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out + ) @handle_nestable @@ -2160,71 +1973,205 @@ def fill_diagonal( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def unfold( +def partial_vec_to_tensor( x: Union[ivy.Array, ivy.NativeArray], /, - mode: Optional[int] = 0, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + skip_begin: Optional[int] = 1, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. + Refolds a partially vectorised tensor into a full one. Parameters ---------- x - input tensor to be unfolded - mode - indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` + a partially vectorised tensor + shape + the shape of the original full tensor (including skipped dimensions) + skip_begin + number of dimensions to leave untouched at the beginning out optional output array, for writing the result to. Returns ------- ret - unfolded_tensor of shape ``(tensor.shape[mode], -1)`` + full tensor """ - return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) + return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) -@handle_nestable @handle_exceptions +@handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays -@handle_array_function +def put_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + values: Union[ivy.Array, ivy.NativeArray], + axis: int, + /, + *, + mode: str = "raise", + out: Optional[ivy.Array] = None, +) -> None: + """ + Put values into the input array by matching 1d index and data slices along a + specified axis. + + Parameters + ---------- + arr : array_like + The input array to modify. + indices : array_like + The indices of the values to put into `arr`. + values : array_like + The values to put into `arr`. + axis : int + The axis over which to put the `values`. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will be handled. + The following modes are available: + + - 'raise': a `ValueError` is raised when an index is out of bounds. + - 'wrap': the index is wrapped around to the corresponding index + at the other end of the axis. + - 'clip': the index is clipped to the closest in-bounds index. + out : ndarray, optional + Output array in which to place the result. + If not specified, a new array is created. + + Returns + ------- + None + + Examples + -------- + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> values = ivy.array([[9, 8, 7], [6, 5, 4]]) + >>> ivy.put_along_axis(arr, indices, values, 1, mode='clip') + >>> print(arr) + ivy.array([[3, 7, 5], + [6, 4, 1]]) + """ + if out is None: + out = ivy.zeros_like(arr) + + indices = ivy.expand_dims(indices, axis=axis) + values = ivy.expand_dims(values, axis=axis) + + stacked = ivy.concat((arr, values), axis=axis) + + sorted_indices = ivy.argsort(indices, axis=axis) + sorted_stacked = ivy.take_along_axis(stacked, sorted_indices, axis=axis) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), + sorted_stacked, + arr, + ) + + if mode == "clip": + indices = ivy.clip(indices, 0, arr.shape[axis] - 1) + elif mode == "wrap": + indices = ivy.mod(indices, arr.shape[axis]) + + arr = ivy.where( + ivy.expand_dims(sorted_indices < arr.shape[axis], axis=axis), arr, values + ) + + ivy.assign(out, arr) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def fold( - x: Union[ivy.Array, ivy.NativeArray], +def rot90( + m: Union[ivy.Array, ivy.NativeArray], /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], *, + copy: Optional[bool] = None, + k: int = 1, + axes: Tuple[int, int] = (0, 1), out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Refolds the mode-`mode` unfolding into a tensor of shape `shape` In other words, - refolds the n-mode unfolded tensor into the original tensor of the specified shape. + Rotate an array by 90 degrees in the plane specified by axes. Rotation direction is + from the first towards the second axis. Parameters ---------- - input - unfolded tensor of shape ``(shape[mode], -1)`` - mode - the mode of the unfolding - shape - shape of the original tensor before unfolding + m + Input array of two or more dimensions. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + k + Number of times the array is rotated by 90 degrees. + axes + The array is rotated in the plane defined by the axes. Axes must be + different. out - optional output array, for writing the result to. + optional output container, for writing the result to. It must have a shape + that the inputs broadcast to. Returns ------- ret - folded_tensor of shape `shape` + A rotated view of m. + + Examples + -------- + With :code:`ivy.Array` input: + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) + With :code:`ivy.NativeArray` input: + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m) + ivy.array([[2, 4], + [1, 3]]) + >>> m = ivy.native_array([[1,2], [3,4]]) + >>> ivy.rot90(m, k=2) + ivy.array([[4, 3], + [2, 1]]) + >>> m = ivy.native_array([[[0, 1],\ + [2, 3]],\ + [[4, 5],\ + [6, 7]]]) + >>> ivy.rot90(m, k=2, axes=(1,2)) + ivy.array([[[3, 2], + [1, 0]], + + [[7, 6], + [5, 4]]]) """ - full_shape = list(shape) - mode_dim = full_shape.pop(mode) - full_shape.insert(0, mode_dim) - return ivy.moveaxis(ivy.reshape(x, full_shape), 0, mode, out=out) + return ivy.current_backend(m).rot90(m, copy=copy, k=k, axes=axes, out=out) @handle_nestable @@ -2233,144 +2180,190 @@ def fold( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def partial_unfold( +def soft_thresholding( x: Union[ivy.Array, ivy.NativeArray], /, - mode: Optional[int] = 0, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, - ravel_tensors: Optional[bool] = False, + threshold: Union[float, ivy.Array, ivy.NativeArray], *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Partial unfolding of a tensor while ignoring the specified number of dimensions at - the beginning and the end. For instance, if the first dimension of the tensor is the - number of samples, to unfold each sample, set skip_begin=1. This would, for each i - in ``range(tensor.shape[0])``, unfold ``tensor[i, ...]``. + Soft-thresholding operator. + + sign(tensor) * max[abs(tensor) - threshold, 0] Parameters ---------- x - tensor of shape n_samples x n_1 x n_2 x ... x n_i - mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - ravel_tensors - if True, the unfolded tensors are also flattened + input array + threshold + float or array with shape tensor.shape + * If float the threshold is applied to the whole tensor + * If array, one threshold is applied per elements, 0 values are ignored out optional output array, for writing the result to. Returns ------- - ret - partially unfolded tensor - """ - if ravel_tensors: - new_shape = [-1] - else: - new_shape = [x.shape[mode + skip_begin], -1] + ivy.Array + thresholded tensor on which the operator has been applied - if skip_begin: - new_shape = [x.shape[i] for i in range(skip_begin)] + new_shape + Examples + -------- + Basic shrinkage - if skip_end: - new_shape += [x.shape[-i] for i in range(1, 1 + skip_end)] + >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) + >>> soft_thresholding(x, 1.1) + array([[ 0. , -0.9, 0.4], + [-2.9, 1.9, 0. ]]) - return ivy.reshape( - ivy.moveaxis(x, mode + skip_begin, skip_begin), new_shape, out=out - ) + + Example with missing values + + >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) + >>> soft_thresholding(x, mask*1.1) + array([[ 1. , -2. , 0.4], + [-2.9, 3. , 0. ]]) + """ + res = ivy.abs(x) - threshold + res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) + + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def partial_fold( - x: Union[ivy.Array, ivy.NativeArray], +def take_along_axis( + arr: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + axis: int, /, - mode: int, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, *, + mode: str = "fill", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Re-folds a partially unfolded tensor. + Take values from the input array by matching 1d index and data slices. Parameters ---------- - x - a partially unfolded tensor + arr + The source array. + indices + The indices of the values to extract. + axis + The axis over which to select values. + If axis is None, arr is treated as a flattened 1D array. mode - indexing starts at 0, therefore mode is in range(0, tensor.ndim) - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions left untouched at the beginning + One of: 'clip', 'fill', 'drop'. Parameter controlling how out-of-bounds indices + will be handled. out - optional output array, for writing the result to. + The output array. Returns ------- ret - partially re-folded tensor + The returned array has the same shape as `indices`. + + Examples + -------- + >>> arr = ivy.array([[4, 3, 5], [1, 2, 1]]) + >>> indices = ivy.array([[0, 1, 1], [2, 0, 0]]) + >>> y = ivy.take_along_axis(arr, indices, 1) + >>> print(y) + ivy.array([[4, 3, 3], [1, 1, 1]]) """ - transposed_shape = list(shape) - mode_dim = transposed_shape.pop(skip_begin + mode) - transposed_shape.insert(skip_begin, mode_dim) - return ivy.moveaxis( - ivy.reshape(x, transposed_shape), skip_begin, skip_begin + mode, out=out + return ivy.current_backend(arr).take_along_axis( + arr, indices, axis, mode=mode, out=out ) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function +@handle_out_argument +@to_native_arrays_and_back @handle_device_shifting -def partial_tensor_to_vec( +def top_k( x: Union[ivy.Array, ivy.NativeArray], + k: int, /, - skip_begin: Optional[int] = 1, - skip_end: Optional[int] = 0, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + axis: int = -1, + largest: bool = True, + sorted: bool = True, + out: Optional[tuple] = None, +) -> Tuple[ivy.Array, ivy.NativeArray]: """ - Partial vectorization of a tensor while ignoring the specified dimension at the - beginning and the end. + Return the `k` largest elements of the given input array along a given axis. Parameters ---------- x - tensor to partially vectorise - skip_begin - number of dimensions to leave untouched at the beginning - skip_end - number of dimensions to leave untouched at the end - out - optional output array, for writing the result to. + The array to compute top_k for. + k + Number of top elements to retun must not exceed the array size. + axis + The axis along which we must return the top elements default value is 1. + largest + If largest is set to False we return k smallest elements of the array. + sorted + If sorted is set to True we return the elements in sorted order. + out: + Optional output tuple, for writing the result to. Must have two arrays inside, + with a shape that the returned tuple broadcast to. Returns ------- ret - partially vectorised tensor with the - `skip_begin` first and `skip_end` last dimensions untouched - """ - return partial_unfold( - x, - mode=0, - skip_begin=skip_begin, - skip_end=skip_end, - ravel_tensors=True, - out=out, + A named tuple with values and indices of top k elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 2) + >>> print(y) + top_k(values=ivy.array([9., 5.]), indices=ivy.array([4, 3])) + + >>> x = ivy.array([[-2., 3., 4., 0.], [-8., 0., -1., 2.]]) + >>> y = ivy.top_k(x, 2, axis=1, largest=False) + >>> print(y) + top_k(values=ivy.array([[-2., 0.], + [-8., -1.]]), indices=ivy.array([[0, 3], + [0, 2]])) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([2., 1., -3., 5., 9., 0., -4]) + >>> y = ivy.top_k(x, 3) + >>> print(y) + top_k(values=ivy.array([9., 5., 2.]), indices=ivy.array([4, 3, 0])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([-1, 2, -4]), b=ivy.array([4., 5., 0.])) + >>> y = x.top_k(2) + >>> print(y) + [{ + a: ivy.array([2, -1]), + b: ivy.array([5., 4.]) + }, { + a: ivy.array([1, 0]), + b: ivy.array([1, 0]) + }] + """ + return current_backend(x).top_k( + x, k, axis=axis, largest=largest, sorted=sorted, out=out ) @@ -2380,200 +2373,209 @@ def partial_tensor_to_vec( @inputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def partial_vec_to_tensor( +def unfold( x: Union[ivy.Array, ivy.NativeArray], /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], - skip_begin: Optional[int] = 1, + mode: Optional[int] = 0, *, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Refolds a partially vectorised tensor into a full one. + Return the mode-`mode` unfolding of `tensor` with modes starting at `0`. Parameters ---------- x - a partially vectorised tensor - shape - the shape of the original full tensor (including skipped dimensions) - skip_begin - number of dimensions to leave untouched at the beginning + input tensor to be unfolded + mode + indexing starts at 0, therefore mode is in ``range(0, tensor.ndim)`` out optional output array, for writing the result to. Returns ------- ret - full tensor + unfolded_tensor of shape ``(tensor.shape[mode], -1)`` """ - return partial_fold(x, mode=0, shape=shape, skip_begin=skip_begin, out=out) + return ivy.reshape(ivy.moveaxis(x, mode, 0), (x.shape[mode], -1), out=out) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def matricize( +def unique_consecutive( x: Union[ivy.Array, ivy.NativeArray], /, - row_modes: Sequence[int], - column_modes: Optional[Sequence[int]] = None, *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + axis: Optional[int] = None, +) -> Tuple[ + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], + Union[ivy.Array, ivy.NativeArray], +]: """ - Matricizes the given tensor. + Eliminates all but the first element from every consecutive group of equivalent + elements in ``x``. Parameters ---------- x - the input tensor - row_modes - modes to use as row of the matrix (in the desired order) - column_modes - modes to use as column of the matrix, in the desired order - if None, the modes not in `row_modes` will be used in ascending order - out - optional output array, for writing the result to. + input array. - ret - ------- - ivy.Array : tensor of size (ivy.prod(x.shape[i] for i in row_modes), -1) - """ - ndims = len(x.shape) - row_indices = list(row_modes) + axis + the axis to apply unique on. If None, unique is applied on flattened ``x``. - if column_modes: - column_indices = list(column_modes) - else: - column_indices = [i for i in range(ndims) if i not in row_indices] - if sorted(column_indices + row_indices) != list(range(ndims)): - msg = ( - "If you provide both column and row modes for the matricization then" - " column_modes + row_modes must contain all the modes of the tensor." - f" Yet, got row_modes={row_modes} and column_modes={column_modes}." - ) - raise ValueError(msg) + Returns + ------- + ret + a namedtuple ``(output, inverse_indices, counts)`` whose + - first element has the field name ``output`` and is an array + containing ``x`` with its equivalent consecutive elements eliminated. + - second element has the field name ``inverse_indices`` and is an + array containing the indices of ``output`` that reconstruct ``x``. + - third element has the field name ``counts`` and is an array + containing the number of occurrences for each unique value or array in ``x``. - row_size, column_size = 1, 1 - row_size = int(ivy.prod([x.shape[i] for i in row_indices])) - column_size = int(ivy.prod([x.shape[i] for i in column_indices])) - return ivy.reshape( - ivy.permute_dims(x, row_indices + column_indices), - (row_size, column_size), - out=out, - ) + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([1, 1, 2, 2, 3, 1, 1, 2]) + >>> ivy..unique_consecutive(x) + Results(values=ivy.array([1, 2, 3, 1, 2]), + inverse_indices=ivy.array([0, 0, 1, 1, 2, 3, 3, 4]), + counts=ivy.array([2, 2, 1, 2, 1])) + """ + return ivy.current_backend(x).unique_consecutive(x, axis=axis) -@handle_nestable @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_view +@to_native_arrays_and_back @handle_array_function @handle_device_shifting -def soft_thresholding( - x: Union[ivy.Array, ivy.NativeArray], +def vsplit( + ary: Union[ivy.Array, ivy.NativeArray], + indices_or_sections: Union[int, Sequence[int], ivy.Array, ivy.NativeArray], /, - threshold: Union[float, ivy.Array, ivy.NativeArray], *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + copy: Optional[bool] = None, +) -> List[ivy.Array]: """ - Soft-thresholding operator. - - sign(tensor) * max[abs(tensor) - threshold, 0] + Split an array vertically into multiple sub-arrays. Parameters ---------- - x - input array - threshold - float or array with shape tensor.shape - * If float the threshold is applied to the whole tensor - * If array, one threshold is applied per elements, 0 values are ignored - out - optional output array, for writing the result to. + ary + Array input. + indices_or_sections + If indices_or_sections is an integer n, the array is split into n + equal sections, provided that n must be a divisor of the split axis. + If indices_or_sections is a sequence of ints or 1-D array, + then input is split at each of the indices. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- - ivy.Array - thresholded tensor on which the operator has been applied + ret + input array split vertically. Examples -------- - Basic shrinkage - - >>> x = ivy.array([[1, -2, 1.5], [-4, 3, -0.5]]) - >>> soft_thresholding(x, 1.1) - array([[ 0. , -0.9, 0.4], - [-2.9, 1.9, 0. ]]) - - - Example with missing values - - >>> mask = ivy.array([[0, 0, 1], [1, 0, 1]]) - >>> soft_thresholding(x, mask*1.1) - array([[ 1. , -2. , 0.4], - [-2.9, 3. , 0. ]]) + >>> ary = ivy.array( + [[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]] + ) + >>> ivy.vsplit(ary, 2) + [ivy.array([[[0., 1.], [2., 3.]]]), ivy.array([[[4., 5.], [6., 7.]]])]) """ - res = ivy.abs(x) - threshold - res = ivy.where(res < 0.0, 0.0, res) * ivy.sign(x) - - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + return ivy.current_backend(ary).vsplit(ary, indices_or_sections, copy=copy) -@handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def choose( - arr: Union[ivy.Array, ivy.NativeArray], - choices: Union[ivy.Array, ivy.NativeArray], +def vstack( + arrays: Sequence[ivy.Array], /, *, - out: None = None, - mode: Union[str, None] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Take values from the input array by matching 1d index and data slices. + Stack arrays in sequence vertically (row wise). Parameters ---------- - arr - The source array. - choices - The indices of the values to extract. - out - The output array. - mode - One of: 'wrap', 'clip'. Parameter controlling how out-of-bounds indices - will be handled. + arrays + Sequence of arrays to be stacked. Returns ------- ret - The returned array has the same shape as `indices`. + The array formed by stacking the given arrays. Examples -------- - >>> choices = ivy.array([[0, 1, 2, 3], [10, 11, 12, 13], - [20, 21, 22, 23], [30, 31, 32, 33]]) - >>> print(choose(ivy.array([2, 3, 1, 0]), choices)) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='clip')) # 4 goes to 3 (4-1) - ivy.array([20, 31, 12, 3]) - >>> arr = ivy.array([2, 4, 1, 0]) - >>> print(choose(arr, choices, mode='wrap')) # 4 goes to (4 mod 4) - ivy.array([20, 1, 12, 3]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([2, 3, 4]) + >>> ivy.vstack((x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4]]) + >>> ivy.vstack((x, y, x, y)) + ivy.array([[1, 2, 3], + [2, 3, 4], + [1, 2, 3], + [2, 3, 4]]) + + >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])] + >>> print(ivy.vstack(y)) + ivy.array([[5, 6], + [7, 8]]) """ - return ivy.current_backend(arr).choose(arr, choices, out=out, mode=mode) + return ivy.current_backend().vstack(arrays, out=out) + + +flatten.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +pad.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +as_strided.unsupported_dtypes = ("bfloat16",) +as_strided.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} diff --git a/ivy/functional/ivy/experimental/norms.py b/ivy/functional/ivy/experimental/norms.py index 4d59a2aaa6907..824ef22962d60 100644 --- a/ivy/functional/ivy/experimental/norms.py +++ b/ivy/functional/ivy/experimental/norms.py @@ -17,91 +17,6 @@ from ivy.utils.exceptions import handle_exceptions -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l1_normalize( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[int, Tuple[int, ...]]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """Normalize the input array along the given axis to have L1 norm equal to - 1. - - Parameters - ---------- - x - Input array. - axis - Axis or axes along which to normalize. If ``None``, - the whole array is normalized. - out - Optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - The normalized array. - - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l1_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.33333334, 1.33333337], - [1.28571439, 2.28571439]]) - """ - return current_backend(x).l1_normalize(x, axis=axis, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def l2_normalize( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """Normalize the input array along the given axis to have L2 norm equal to - 1. - - Parameters - ---------- - x - Input array. - axis - Axis along which to normalize. If ``None``, the whole array is normalized. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The normalized array. - - Examples - -------- - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = ivy.l2_normalize(x, axis=1) - >>> print(y) - ivy.array([[0.44721359, 0.89442718], - [0.60000002, 0.80000001]]) - """ - return current_backend(x).l2_normalize(x, axis=axis, out=out) - - @handle_exceptions @handle_nestable @handle_partial_mixed_function @@ -212,16 +127,80 @@ def batch_norm( return xnormalized, runningmean, runningvariance -batch_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def group_norm( + x: Union[ivy.NativeArray, ivy.Array], + num_groups: int = 1, + /, + *, + offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, + eps: Optional[float] = 1e-5, + data_format: Optional[str] = "NSC", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Apply group normalization to the input array and returns the normalized input. + + Parameters + ---------- + x + Input array of default shape (N, *S, C), where N is the batch dimension, + *S corresponds to any number of spatial dimensions and + C corresponds to the channel dimension. + num_groups + number of groups to separate the channels into + offset + An offset array of size C. If present, will be added + to the normalized input. + scale + A scale array of size C. If present, the scale is + applied to the normalized input. + eps + A small float number to avoid dividing by 0. + data_format + The ordering of the dimensions in the input, one of "NSC" or "NCS", + where N is the batch dimension, S represents any number of spatial + dimensions and C is the channel dimension. Default is "NSC". + out + optional output arrays, for writing the result to. + + Returns + ------- + ret + The normalized array. + """ + xdims = ivy.get_num_dims(x) + if data_format == "NSC": + x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) + N = x.shape[0] + C = x.shape[1] + S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 + assert C % num_groups == 0 + x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) + mean = ivy.mean(x_, axis=(2, 3), keepdims=True) + var = ivy.var(x_, axis=(2, 3), keepdims=True) + x_normalized = (x_ - mean) / ivy.sqrt(var + eps) + x_normalized = ivy.reshape(x_normalized, x.shape) + + if ivy.exists(scale): + scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized * scale + + if ivy.exists(offset): + offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) + x_normalized = x_normalized + offset + + if data_format == "NSC": + x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) + + if ivy.exists(out): + x_normalized = ivy.inplace_update(out, x_normalized) + return x_normalized @handle_exceptions @@ -344,103 +323,89 @@ def instance_norm( return (xnormalized, runningmean, runningvariance) -instance_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def group_norm( - x: Union[ivy.NativeArray, ivy.Array], - num_groups: int = 1, +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l1_normalize( + x: Union[ivy.Array, ivy.NativeArray], /, *, - offset: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - scale: Optional[Union[ivy.NativeArray, ivy.Array]] = None, - eps: Optional[float] = 1e-5, - data_format: Optional[str] = "NSC", + axis: Optional[Union[int, Tuple[int, ...]]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Apply group normalization to the input array and returns the normalized input. + """Normalize the input array along the given axis to have L1 norm equal to + 1. Parameters ---------- x - Input array of default shape (N, *S, C), where N is the batch dimension, - *S corresponds to any number of spatial dimensions and - C corresponds to the channel dimension. - num_groups - number of groups to separate the channels into - offset - An offset array of size C. If present, will be added - to the normalized input. - scale - A scale array of size C. If present, the scale is - applied to the normalized input. - eps - A small float number to avoid dividing by 0. - data_format - The ordering of the dimensions in the input, one of "NSC" or "NCS", - where N is the batch dimension, S represents any number of spatial - dimensions and C is the channel dimension. Default is "NSC". + Input array. + axis + Axis or axes along which to normalize. If ``None``, + the whole array is normalized. out - optional output arrays, for writing the result to. + Optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret The normalized array. - """ - xdims = ivy.get_num_dims(x) - if data_format == "NSC": - x = ivy.permute_dims(x, axes=(0, xdims - 1, *range(1, xdims - 1))) - N = x.shape[0] - C = x.shape[1] - S = ivy.to_scalar(ivy.prod(x.shape[2:])) if xdims > 2 else 1 - assert C % num_groups == 0 - x_ = ivy.reshape(x, [N, num_groups, C // num_groups, S]) - mean = ivy.mean(x_, axis=(2, 3), keepdims=True) - var = ivy.var(x_, axis=(2, 3), keepdims=True) - x_normalized = (x_ - mean) / ivy.sqrt(var + eps) - x_normalized = ivy.reshape(x_normalized, x.shape) - if ivy.exists(scale): - scale = ivy.expand_dims(scale, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized * scale + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l1_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.33333334, 1.33333337], + [1.28571439, 2.28571439]]) + """ + return current_backend(x).l1_normalize(x, axis=axis, out=out) - if ivy.exists(offset): - offset = ivy.expand_dims(offset, axis=[0, *(range(2, xdims))]) - x_normalized = x_normalized + offset - if data_format == "NSC": - x_normalized = ivy.permute_dims(x_normalized, axes=(0, *range(2, xdims), 1)) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def l2_normalize( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """Normalize the input array along the given axis to have L2 norm equal to + 1. - if ivy.exists(out): - x_normalized = ivy.inplace_update(out, x_normalized) - return x_normalized + Parameters + ---------- + x + Input array. + axis + Axis along which to normalize. If ``None``, the whole array is normalized. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + The normalized array. -group_norm.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + Examples + -------- + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = ivy.l2_normalize(x, axis=1) + >>> print(y) + ivy.array([[0.44721359, 0.89442718], + [0.60000002, 0.80000001]]) + """ + return current_backend(x).l2_normalize(x, axis=axis, out=out) @handle_exceptions @@ -486,3 +451,34 @@ def lp_normalize( [0.42857143, 0.5714286 ]]) """ return current_backend(x).lp_normalize(x, p=p, axis=axis, out=out) + + +batch_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +instance_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +group_norm.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} diff --git a/ivy/functional/ivy/experimental/random.py b/ivy/functional/ivy/experimental/random.py index 70586db0f3ccf..ff5a5d32f56bb 100644 --- a/ivy/functional/ivy/experimental/random.py +++ b/ivy/functional/ivy/experimental/random.py @@ -14,69 +14,65 @@ from ivy.utils.exceptions import handle_exceptions -# dirichlet @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back +@infer_dtype @handle_device_shifting -def dirichlet( - alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], - /, +@infer_device +def bernoulli( + probs: Union[float, ivy.Array, ivy.NativeArray], *, - size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- - distributed random variable can be seen as a multivariate generalization of a Beta - distribution. The Dirichlet distribution is a conjugate prior of a multinomial - distribution in Bayesian inference. + Draws samples from Bernoulli distrubution paramterized by probs or logits (but not + both) Parameters ---------- - alpha - Sequence of floats of length k - size - optional int or tuple of ints, Output shape. If the given shape is, - e.g., (m, n), then m * n * k samples are drawn. Default is None, - in which case a vector of length k is returned. + logits + An N-D Array representing the log-odds of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution where the probability of an event is sigmoid + (logits). Only one of logits or probs should be passed in. + probs + An N-D Array representing the probability of a 1 event. + Each entry in the Array parameterizes an independent Bernoulli + distribution. Only one of logits or probs should be passed in + shape + If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. + (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) + device + device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. + (Default value = None). dtype output array data type. If ``dtype`` is ``None``, the output array data type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The drawn samples, of shape (size, k). - - Functional Examples - ------------------- - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha) - ivy.array([0.10598304, 0.21537054, 0.67864642]) - - >>> alpha = [1.0, 2.0, 3.0] - >>> ivy.dirichlet(alpha, size = (2,3)) - ivy.array([[[0.48006698, 0.07472073, 0.44521229], - [0.55479872, 0.05426367, 0.39093761], - [0.19531053, 0.51675832, 0.28793114]], - - [[0.12315625, 0.29823365, 0.5786101 ], - [0.15564976, 0.50542368, 0.33892656], - [0.1325352 , 0.44439589, 0.42306891]]]) + Drawn samples from the Bernoulli distribution """ - return ivy.current_backend().dirichlet( - alpha, - size=size, + return ivy.current_backend(probs).bernoulli( + probs, + logits=logits, + shape=shape, + device=device, dtype=dtype, seed=seed, out=out, @@ -138,6 +134,75 @@ def beta( ) +# dirichlet +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def dirichlet( + alpha: Union[ivy.Array, ivy.NativeArray, float, Sequence[float]], + /, + *, + size: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Draw size samples of dimension k from a Dirichlet distribution. A Dirichlet- + distributed random variable can be seen as a multivariate generalization of a Beta + distribution. The Dirichlet distribution is a conjugate prior of a multinomial + distribution in Bayesian inference. + + Parameters + ---------- + alpha + Sequence of floats of length k + size + optional int or tuple of ints, Output shape. If the given shape is, + e.g., (m, n), then m * n * k samples are drawn. Default is None, + in which case a vector of length k is returned. + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. + + Returns + ------- + ret + The drawn samples, of shape (size, k). + + Functional Examples + ------------------- + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha) + ivy.array([0.10598304, 0.21537054, 0.67864642]) + + >>> alpha = [1.0, 2.0, 3.0] + >>> ivy.dirichlet(alpha, size = (2,3)) + ivy.array([[[0.48006698, 0.07472073, 0.44521229], + [0.55479872, 0.05426367, 0.39093761], + [0.19531053, 0.51675832, 0.28793114]], + + [[0.12315625, 0.29823365, 0.5786101 ], + [0.15564976, 0.50542368, 0.33892656], + [0.1325352 , 0.44439589, 0.42306891]]]) + """ + return ivy.current_backend().dirichlet( + alpha, + size=size, + dtype=dtype, + seed=seed, + out=out, + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -260,68 +325,3 @@ def poisson( fill_value=fill_value, out=out, ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -@infer_device -def bernoulli( - probs: Union[float, ivy.Array, ivy.NativeArray], - *, - logits: Optional[Union[float, ivy.Array, ivy.NativeArray]] = None, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Draws samples from Bernoulli distrubution paramterized by probs or logits (but not - both) - - Parameters - ---------- - logits - An N-D Array representing the log-odds of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution where the probability of an event is sigmoid - (logits). Only one of logits or probs should be passed in. - probs - An N-D Array representing the probability of a 1 event. - Each entry in the Array parameterizes an independent Bernoulli - distribution. Only one of logits or probs should be passed in - shape - If the given shape is, e.g '(m, n, k)', then 'm * n * k' samples are drawn. - (Default value = 'None', where 'ivy.shape(logits)' samples are drawn) - device - device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Drawn samples from the Bernoulli distribution - """ - return ivy.current_backend(probs).bernoulli( - probs, - logits=logits, - shape=shape, - device=device, - dtype=dtype, - seed=seed, - out=out, - ) diff --git a/ivy/functional/ivy/experimental/sparse_array.py b/ivy/functional/ivy/experimental/sparse_array.py index 46e6eb64fa428..785eccaa6b3bd 100644 --- a/ivy/functional/ivy/experimental/sparse_array.py +++ b/ivy/functional/ivy/experimental/sparse_array.py @@ -4,356 +4,6 @@ from ivy.utils.exceptions import handle_exceptions -# helpers -def _verify_coo_components(indices=None, values=None, dense_shape=None): - ivy.utils.assertions.check_all_or_any_fn( - indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message="indices, values and dense_shape must all be specified", - ) - # coordinates style (COO), must be shaped (x, y) - ivy.utils.assertions.check_equal( - len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - ivy.utils.assertions.check_equal( - len(ivy.to_ivy_shape(dense_shape)), - ivy.shape(indices)[0], - message="shape and indices shape do not match", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(values)[0], - ivy.shape(indices)[1], - message="values and indices do not match", - as_array=False, - ) - for i in range(ivy.shape(indices)[0]): - ivy.utils.assertions.check_less( - indices[i], - ivy.to_ivy_shape(dense_shape)[i], - message="indices is larger than shape", - ) - - -def _verify_common_row_format_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" -): - ivy.utils.assertions.check_all_or_any_fn( - crow_indices, - col_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "crow_indices, col_indices, values and dense_shape must all be specified." - ), - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(crow_indices)), - 1, - message="crow_indices must be 1D.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(col_indices)), - 1, - message="col_indices must be 1D.", - as_array=False, - ) - - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", - as_array=False, - ) - - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - crow_indices[-1], - message="size of col_indices does not match with last element of crow_indices", - ) - - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(col_indices)[0], - ivy.shape(values)[0], - message="values and col_indices do not match", - as_array=False, - ) - - # index in crow_indices must not exceed length of col_indices - ivy.utils.assertions.check_less( - crow_indices, - ivy.shape(col_indices)[0], - allow_equal=True, - message="index in crow_indices does not match the number of col_indices", - ) - - -def _verify_csr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="csr", - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False - ) - # number of intervals must be equal to x in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False - ) - - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1], - message="index in col_indices does not match shape", - ) - - -def _verify_bsr_components( - crow_indices=None, col_indices=None, values=None, dense_shape=None -): - _verify_common_row_format_components( - crow_indices=crow_indices, - col_indices=col_indices, - values=values, - dense_shape=dense_shape, - format="bsr", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="The number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="The number of cols of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False - ) - ivy.utils.assertions.check_less( - col_indices, - dense_shape[1] // ncolblocks, - message="index in col_indices does not match shape", - ) - - -def _verify_common_column_format_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" -): - ivy.utils.assertions.check_all_or_any_fn( - ccol_indices, - row_indices, - values, - dense_shape, - fn=ivy.exists, - type="all", - message=( - "ccol_indices, row_indices, values and dense_shape must all be specified" - ), - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(ccol_indices)), - 1, - message="ccol_indices must be 1D", - as_array=False, - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False - ) - - ivy.utils.assertions.check_equal( - len(dense_shape), - 2, - message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", - as_array=False, - ) - # number of values must match number of coordinates - ivy.utils.assertions.check_equal( - ivy.shape(row_indices)[0], - ivy.shape(values)[0], - message="values and row_indices do not match", - as_array=False, - ) - # index in ccol_indices must not exceed length of row_indices - ivy.utils.assertions.check_less( - ccol_indices, - ivy.shape(row_indices)[0], - allow_equal=True, - message="index in ccol_indices does not match the number of row_indices", - ) - - -def _verify_csc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="csc", - ) - - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 1, message="values must be 1D", as_array=False - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0], - message="index in row_indices does not match shape", - ) - - -def _verify_bsc_components( - ccol_indices=None, row_indices=None, values=None, dense_shape=None -): - _verify_common_column_format_components( - ccol_indices=ccol_indices, - row_indices=row_indices, - values=values, - dense_shape=dense_shape, - format="bsc", - ) - ivy.utils.assertions.check_equal( - len(ivy.shape(values)), 3, message="values must be 3D", as_array=False - ) - nrowblocks, ncolblocks = ivy.shape(values)[-2:] - ivy.utils.assertions.check_equal( - dense_shape[0] % nrowblocks, - 0, - message="number of rows of array must be divisible by that of block.", - as_array=False, - ) - ivy.utils.assertions.check_equal( - dense_shape[1] % ncolblocks, - 0, - message="number of cols of array must be divisible by that of block.", - as_array=False, - ) - # number of intervals must be equal to y in shape (x, y) - ivy.utils.assertions.check_equal( - ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False - ) - ivy.utils.assertions.check_less( - row_indices, - dense_shape[0] // nrowblocks, - message="index in row_indices does not match shape", - ) - - -def _is_data_not_indices_values_and_shape( - data=None, - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format=None, -): - if data is not None: - ivy.utils.assertions.check_all_or_any_fn( - coo_indices, - crow_indices, - col_indices, - ccol_indices, - row_indices, - values, - dense_shape, - format, - fn=ivy.exists, - type="any", - limit=[0], - message=( - "Only specify data, coo_indices for COO format, crow_indices and" - " col_indices for CSR and BSR, ccol_indices and row_indicesfor CSC and" - " BSC." - ), - ) - return True - return False - - -def _is_valid_format( - coo_indices=None, - crow_indices=None, - col_indices=None, - ccol_indices=None, - row_indices=None, - values=None, - dense_shape=None, - format="coo", -): - valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] - - if not isinstance(format, str) or not format.lower() in valid_formats: - return False - - if format.endswith("o"): - # format is coo - return ( - ivy.exists(coo_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and crow_indices is None - and col_indices is None - and ccol_indices is None - and row_indices is None - ) - - if format.endswith("r"): - # format is either csr or bsr - return ( - ivy.exists(crow_indices) - and ivy.exists(col_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and ccol_indices is None - and row_indices is None - ) - # format is either csc or bsc - return ( - ivy.exists(ccol_indices) - and ivy.exists(row_indices) - and ivy.exists(values) - and ivy.exists(dense_shape) - and coo_indices is None - and crow_indices is None - and col_indices is None - ) - - class SparseArray: def __init__( self, @@ -550,240 +200,598 @@ def coo_indices(self): def crow_indices(self): return self._crow_indices - @property - def col_indices(self): - return self._col_indices + @property + def col_indices(self): + return self._col_indices + + @property + def ccol_indices(self): + return self._ccol_indices + + @property + def row_indices(self): + return self._row_indices + + @property + def values(self): + return self._values + + @property + def dense_shape(self): + return self._dense_shape + + @property + def format(self): + return self._format + + # Setters # + # --------# + + @data.setter + def data(self, data): + self._init_data(data) + + @coo_indices.setter + def coo_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + _verify_coo_components( + indices=indices, values=self._values, dense_shape=self._dense_shape + ) + self._coo_indices = indices + + @crow_indices.setter + def crow_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._crow_indices = indices + + @col_indices.setter + def col_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csr": + _verify_csr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsr_components( + crow_indices=indices, + col_indices=self._col_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._col_indices = indices + + @ccol_indices.setter + def ccol_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=indices, + row_indices=self._row_indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._ccol_indices = indices + + @row_indices.setter + def row_indices(self, indices): + indices = ivy.array(indices, dtype="int64") + if self._format == "csc": + _verify_csc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + else: + _verify_bsc_components( + ccol_indices=self._ccol_indices, + row_indices=indices, + values=self._values, + dense_shape=self._dense_shape, + ) + self._row_indices = indices + + @values.setter + def values(self, values): + values = ivy.array(values) + _verify_coo_components( + indices=self._coo_indices, values=values, dense_shape=self._dense_shape + ) + self._values = values + + @dense_shape.setter + def dense_shape(self, dense_shape): + dense_shape = ivy.Shape(dense_shape) + _verify_coo_components( + indices=self._coo_indices, values=self._values, dense_shape=dense_shape + ) + self._dense_shape = dense_shape + + @format.setter + def format(self, format): + self._format = format + + # Instance Methods # + # ---------------- # + + def _coo_to_dense_coordinates(self): + all_coordinates = [] + for i in range(self._values.shape[0]): + coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) + coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) + all_coordinates.append(coordinate.to_list()) + return all_coordinates + + def _csr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._col_indices.to_list() + all_cols = self._crow_indices.to_list() + for row in range(total_rows): + cols = all_rows[all_cols[row] : all_cols[row + 1]] + for col in cols: + all_coordinates.append([row, col]) + return all_coordinates + + def _csc_to_dense_coordinates(self): + # CSC sparse array + all_coordinates = [] + total_rows = self._dense_shape[1] + all_cols = self._row_indices.to_list() + all_rows = self._ccol_indices.to_list() + for col in range(total_rows): + rows = all_cols[all_rows[col] : all_rows[col + 1]] + for row in rows: + all_coordinates.append([row, col]) + return all_coordinates + + def _bsr_to_dense_coordinates(self): + all_coordinates = [] + total_rows = self._dense_shape[0] + all_rows = self._crow_indices.to_list() + all_cols = self._col_indices.to_list() + + nblockrows, nblockcols = self._values.shape[-2:] + + for row in range(total_rows // nblockrows): + cols = all_cols[all_rows[row] : all_rows[row + 1]] + for col in cols: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates + + def _bsc_to_dense_coordinates(self): + all_coordinates = [] + total_cols = self._dense_shape[1] + all_rows = self._row_indices.to_list() + all_cols = self._ccol_indices.to_list() + + nblockrows, nblockcols = self._values.shape[-2:] + + for col in range(total_cols // nblockcols): + rows = all_rows[all_cols[col] : all_cols[col + 1]] + for row in rows: + for col_index in range(nblockcols): + for row_index in range(nblockrows): + all_coordinates.append( + [ + nblockrows * row + row_index, + nblockcols * col + col_index, + ] + ) + return all_coordinates + + def to_dense_array(self, *, native=False): + if self._format == "coo": + all_coordinates = self._coo_to_dense_coordinates() + elif self._format == "csr": + all_coordinates = self._csr_to_dense_coordinates() + elif self._format == "csc": + all_coordinates = self._csc_to_dense_coordinates() + elif self._format == "bsc": + all_coordinates = self._bsc_to_dense_coordinates() + else: + all_coordinates = self._bsr_to_dense_coordinates() + + # make dense array + ret = ivy.scatter_nd( + ivy.array(all_coordinates), + ivy.flatten(self._values), + ivy.array(self._dense_shape), + ) + return ret.to_native() if native else ret + + +class NativeSparseArray: + pass + + +# --- Helpers --- # +# --------------- # + + +def _is_data_not_indices_values_and_shape( + data=None, + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format=None, +): + if data is not None: + ivy.utils.assertions.check_all_or_any_fn( + coo_indices, + crow_indices, + col_indices, + ccol_indices, + row_indices, + values, + dense_shape, + format, + fn=ivy.exists, + type="any", + limit=[0], + message=( + "Only specify data, coo_indices for COO format, crow_indices and" + " col_indices for CSR and BSR, ccol_indices and row_indicesfor CSC and" + " BSC." + ), + ) + return True + return False + - @property - def ccol_indices(self): - return self._ccol_indices +def _is_valid_format( + coo_indices=None, + crow_indices=None, + col_indices=None, + ccol_indices=None, + row_indices=None, + values=None, + dense_shape=None, + format="coo", +): + valid_formats = ["coo", "csr", "csc", "csc", "bsc", "bsr"] - @property - def row_indices(self): - return self._row_indices + if not isinstance(format, str) or not format.lower() in valid_formats: + return False - @property - def values(self): - return self._values + if format.endswith("o"): + # format is coo + return ( + ivy.exists(coo_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and crow_indices is None + and col_indices is None + and ccol_indices is None + and row_indices is None + ) - @property - def dense_shape(self): - return self._dense_shape + if format.endswith("r"): + # format is either csr or bsr + return ( + ivy.exists(crow_indices) + and ivy.exists(col_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and ccol_indices is None + and row_indices is None + ) + # format is either csc or bsc + return ( + ivy.exists(ccol_indices) + and ivy.exists(row_indices) + and ivy.exists(values) + and ivy.exists(dense_shape) + and coo_indices is None + and crow_indices is None + and col_indices is None + ) - @property - def format(self): - return self._format - # Setters # - # --------# +def _verify_bsc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="bsc", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="number of cols of array must be divisible by that of block.", + as_array=False, + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1] // ncolblocks, as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0] // nrowblocks, + message="index in row_indices does not match shape", + ) - @data.setter - def data(self, data): - self._init_data(data) - @coo_indices.setter - def coo_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - _verify_coo_components( - indices=indices, values=self._values, dense_shape=self._dense_shape - ) - self._coo_indices = indices +def _verify_bsr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="bsr", + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 3, message="values must be 3D.", as_array=False + ) + nrowblocks, ncolblocks = ivy.shape(values)[-2:] + ivy.utils.assertions.check_equal( + dense_shape[0] % nrowblocks, + 0, + message="The number of rows of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + dense_shape[1] % ncolblocks, + 0, + message="The number of cols of array must be divisible by that of block.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0] // nrowblocks, as_array=False + ) + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1] // ncolblocks, + message="index in col_indices does not match shape", + ) - @crow_indices.setter - def crow_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._crow_indices = indices - @col_indices.setter - def col_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csr": - _verify_csr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsr_components( - crow_indices=indices, - col_indices=self._col_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._col_indices = indices +def _verify_common_column_format_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None, format="csc" +): + ivy.utils.assertions.check_all_or_any_fn( + ccol_indices, + row_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "ccol_indices, row_indices, values and dense_shape must all be specified" + ), + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(ccol_indices)), + 1, + message="ccol_indices must be 1D", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(row_indices)), 1, message="row_indices must be 1D", as_array=False + ) - @ccol_indices.setter - def ccol_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=indices, - row_indices=self._row_indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._ccol_indices = indices + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"only 2D arrays can be converted to {format.upper()} sparse arrays", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(row_indices)[0], + ivy.shape(values)[0], + message="values and row_indices do not match", + as_array=False, + ) + # index in ccol_indices must not exceed length of row_indices + ivy.utils.assertions.check_less( + ccol_indices, + ivy.shape(row_indices)[0], + allow_equal=True, + message="index in ccol_indices does not match the number of row_indices", + ) - @row_indices.setter - def row_indices(self, indices): - indices = ivy.array(indices, dtype="int64") - if self._format == "csc": - _verify_csc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - else: - _verify_bsc_components( - ccol_indices=self._ccol_indices, - row_indices=indices, - values=self._values, - dense_shape=self._dense_shape, - ) - self._row_indices = indices - @values.setter - def values(self, values): - values = ivy.array(values) - _verify_coo_components( - indices=self._coo_indices, values=values, dense_shape=self._dense_shape - ) - self._values = values +def _verify_common_row_format_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None, format="csr" +): + ivy.utils.assertions.check_all_or_any_fn( + crow_indices, + col_indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message=( + "crow_indices, col_indices, values and dense_shape must all be specified." + ), + ) - @dense_shape.setter - def dense_shape(self, dense_shape): - dense_shape = ivy.Shape(dense_shape) - _verify_coo_components( - indices=self._coo_indices, values=self._values, dense_shape=dense_shape - ) - self._dense_shape = dense_shape + ivy.utils.assertions.check_equal( + len(ivy.shape(crow_indices)), + 1, + message="crow_indices must be 1D.", + as_array=False, + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(col_indices)), + 1, + message="col_indices must be 1D.", + as_array=False, + ) - @format.setter - def format(self, format): - self._format = format + ivy.utils.assertions.check_equal( + len(dense_shape), + 2, + message=f"Only 2D arrays can be converted to {format.upper()} sparse arrays.", + as_array=False, + ) - # Instance Methods # - # ---------------- # + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + crow_indices[-1], + message="size of col_indices does not match with last element of crow_indices", + ) - def _coo_to_dense_coordinates(self): - all_coordinates = [] - for i in range(self._values.shape[0]): - coordinate = ivy.gather(self._coo_indices, ivy.array([[i]])) - coordinate = ivy.reshape(coordinate, (self._coo_indices.shape[0],)) - all_coordinates.append(coordinate.to_list()) - return all_coordinates + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(col_indices)[0], + ivy.shape(values)[0], + message="values and col_indices do not match", + as_array=False, + ) - def _csr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._col_indices.to_list() - all_cols = self._crow_indices.to_list() - for row in range(total_rows): - cols = all_rows[all_cols[row] : all_cols[row + 1]] - for col in cols: - all_coordinates.append([row, col]) - return all_coordinates + # index in crow_indices must not exceed length of col_indices + ivy.utils.assertions.check_less( + crow_indices, + ivy.shape(col_indices)[0], + allow_equal=True, + message="index in crow_indices does not match the number of col_indices", + ) - def _csc_to_dense_coordinates(self): - # CSC sparse array - all_coordinates = [] - total_rows = self._dense_shape[1] - all_cols = self._row_indices.to_list() - all_rows = self._ccol_indices.to_list() - for col in range(total_rows): - rows = all_cols[all_rows[col] : all_rows[col + 1]] - for row in rows: - all_coordinates.append([row, col]) - return all_coordinates - def _bsr_to_dense_coordinates(self): - all_coordinates = [] - total_rows = self._dense_shape[0] - all_rows = self._crow_indices.to_list() - all_cols = self._col_indices.to_list() +# helpers +def _verify_coo_components(indices=None, values=None, dense_shape=None): + ivy.utils.assertions.check_all_or_any_fn( + indices, + values, + dense_shape, + fn=ivy.exists, + type="all", + message="indices, values and dense_shape must all be specified", + ) + # coordinates style (COO), must be shaped (x, y) + ivy.utils.assertions.check_equal( + len(ivy.shape(indices)), 2, message="indices must be 2D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + ivy.utils.assertions.check_equal( + len(ivy.to_ivy_shape(dense_shape)), + ivy.shape(indices)[0], + message="shape and indices shape do not match", + as_array=False, + ) + # number of values must match number of coordinates + ivy.utils.assertions.check_equal( + ivy.shape(values)[0], + ivy.shape(indices)[1], + message="values and indices do not match", + as_array=False, + ) + for i in range(ivy.shape(indices)[0]): + ivy.utils.assertions.check_less( + indices[i], + ivy.to_ivy_shape(dense_shape)[i], + message="indices is larger than shape", + ) - nblockrows, nblockcols = self._values.shape[-2:] - for row in range(total_rows // nblockrows): - cols = all_cols[all_rows[row] : all_rows[row + 1]] - for col in cols: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates +def _verify_csc_components( + ccol_indices=None, row_indices=None, values=None, dense_shape=None +): + _verify_common_column_format_components( + ccol_indices=ccol_indices, + row_indices=row_indices, + values=values, + dense_shape=dense_shape, + format="csc", + ) - def _bsc_to_dense_coordinates(self): - all_coordinates = [] - total_cols = self._dense_shape[1] - all_rows = self._row_indices.to_list() - all_cols = self._ccol_indices.to_list() + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D", as_array=False + ) + # number of intervals must be equal to y in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(ccol_indices)[0] - 1, dense_shape[1], as_array=False + ) + ivy.utils.assertions.check_less( + row_indices, + dense_shape[0], + message="index in row_indices does not match shape", + ) - nblockrows, nblockcols = self._values.shape[-2:] - for col in range(total_cols // nblockcols): - rows = all_rows[all_cols[col] : all_cols[col + 1]] - for row in rows: - for col_index in range(nblockcols): - for row_index in range(nblockrows): - all_coordinates.append( - [ - nblockrows * row + row_index, - nblockcols * col + col_index, - ] - ) - return all_coordinates +def _verify_csr_components( + crow_indices=None, col_indices=None, values=None, dense_shape=None +): + _verify_common_row_format_components( + crow_indices=crow_indices, + col_indices=col_indices, + values=values, + dense_shape=dense_shape, + format="csr", + ) - def to_dense_array(self, *, native=False): - if self._format == "coo": - all_coordinates = self._coo_to_dense_coordinates() - elif self._format == "csr": - all_coordinates = self._csr_to_dense_coordinates() - elif self._format == "csc": - all_coordinates = self._csc_to_dense_coordinates() - elif self._format == "bsc": - all_coordinates = self._bsc_to_dense_coordinates() - else: - all_coordinates = self._bsr_to_dense_coordinates() + ivy.utils.assertions.check_equal( + len(ivy.shape(values)), 1, message="values must be 1D.", as_array=False + ) + # number of intervals must be equal to x in shape (x, y) + ivy.utils.assertions.check_equal( + ivy.shape(crow_indices)[0] - 1, dense_shape[0], as_array=False + ) - # make dense array - ret = ivy.scatter_nd( - ivy.array(all_coordinates), - ivy.flatten(self._values), - ivy.array(self._dense_shape), - ) - return ret.to_native() if native else ret + ivy.utils.assertions.check_less( + col_indices, + dense_shape[1], + message="index in col_indices does not match shape", + ) -class NativeSparseArray: - pass +# --- Main --- # +# ------------ # def is_ivy_sparse_array(x): diff --git a/ivy/functional/ivy/experimental/statistical.py b/ivy/functional/ivy/experimental/statistical.py index 1349cc5752326..e9392dcff411a 100644 --- a/ivy/functional/ivy/experimental/statistical.py +++ b/ivy/functional/ivy/experimental/statistical.py @@ -13,141 +13,47 @@ from ivy.utils.exceptions import handle_exceptions -# TODO: Make bins optional by offering an automatic bins creation like numpy. -# Make density argument work in tensorflow -# Bins as str is not defined (check Numpy implementation). -# Permit multiple axis. -# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def histogram( - a: Union[ivy.Array, ivy.NativeArray], +def bincount( + x: ivy.Array, /, *, - bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, - axis: Optional[int] = None, - extend_lower_interval: Optional[bool] = False, - extend_upper_interval: Optional[bool] = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - range: Optional[Tuple[float]] = None, - weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - density: Optional[bool] = False, + weights: Optional[ivy.Array] = None, + minlength: int = 0, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the histogram of the array ``a``. - - .. note:: - Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), - ..., I_{K-1} = [c_{K-1}, cK]. + Count the number of occurrences of each value in an integer array. Parameters ---------- - a - input array. - bins - if ``bins`` is an int, it defines the number of equal-width bins in the given - range. - if ``bins`` is an array, it defines a monotonically increasing array of bin - edges, including the rightmost edge, allowing for non-uniform bin widths. - axis - dimension along which maximum values must be computed. By default, the maximum - value must be computed over the entire array. Default: ``None``. - extend_lower_interval - if True, extend the lowest interval I0 to (-inf, c1]. - extend_upper_interval - ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). - dtype - the output type. - range - the lower and upper range of the bins. The first element of the range must be - less than or equal to the second. + self + Input array. weights - each value in ``a`` only contributes its associated weight towards the bin count - (instead of 1). Must be of the same shape as a. - density - if True, the result is the value of the probability density function at the - bin, normalized such that the integral over the range of bins is 1. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + An optional input array. + minlength + A minimum number of bins for the output array. Returns ------- ret - a tuple containing the values of the histogram and the bin edges. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The bincount of the array elements. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) - >>> z = ivy.histogram(x, bins=y) - >>> print(z) - ivy.array([1., 0., 1., 1.]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [4.4, 5.5, .6]]) - >>> bins = 4 - >>> range = (0., 5.) - >>> dtype = ivy.int32 - >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) - >>> print(y) - ivy.array([2, 1, 1, 1]) - - >>> x = ivy.array([[1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> axis = 1 - >>> extend_lower_interval = True - >>> extend_upper_interval = True - >>> dtype = ivy.float32 - >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) - >>> z = ivy.histogram( - ... x, - ... bins=y, - ... axis=axis, - ... extend_lower_interval=extend_lower_interval, - ... extend_upper_interval=extend_upper_interval, - ... dtype=dtype, - ... weights=weights) - >>> print(z) - ivy.array([[0., 3.], - [1., 0.], - [1., 0.], - [1., 0.], - [0., 0.]]) - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) - >>> dtype = ivy.int32 - >>> z = ivy.histogram(x, bins=y, dtype=dtype) - >>> print(z) - { - a: ivy.array([1, 1, 1, 0, 0]), - b: ivy.array([0, 0, 0, 1, 2]) - } + >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) + >>> a.bincount(a) + 3.0 + >>> a.bincount(a, axis=0) + array([6.5, 2. , 2.5]) """ - return ivy.current_backend(a).histogram( - a, - bins=bins, - axis=axis, - extend_lower_interval=extend_lower_interval, - extend_upper_interval=extend_upper_interval, - dtype=dtype, - range=range, - weights=weights, - density=density, - out=out, + return ivy.current_backend(x).bincount( + x, weights=weights, minlength=minlength, out=out ) @@ -157,276 +63,533 @@ def histogram( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def median( - input: ivy.Array, +def corrcoef( + x: ivy.Array, /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, + y: Optional[ivy.Array] = None, + rowvar: bool = True, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute the median along the specified axis. - - Parameters - ---------- - input - Input array. - axis - Axis or axes along which the medians are computed. The default is to compute - the median along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. - out - optional output array, for writing the result to. - - Returns - ------- - ret - The median of the array elements. - - Functional Examples - ------------------- - >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> ivy.median(a) - 3.5 - >>> ivy.median(a, axis=0) - ivy.array([6.5, 4.5, 2.5]) - """ - return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) + return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_out_argument +@handle_array_like_without_promotion @to_native_arrays_and_back -@infer_dtype -@handle_device_shifting -def nanmean( - a: ivy.Array, +def cov( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray] = None, /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, + rowVar: bool = True, + bias: bool = False, + ddof: Optional[int] = None, + fweights: Optional[ivy.Array] = None, + aweights: Optional[ivy.Array] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the mean of all non-NaN elements along the specified dimensions. + Compute the covariance of matrix x1, or variables x1 and x2. Parameters ---------- - a - Input array. - axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of sub-classes - of ndarray. If the sub-classes methods does not implement keepdims any - exceptions will be raised. + x1 + a 1D or 2D input array, with a numeric data type. + x2 + optional second 1D or 2D input array, with a numeric data type. + Must have the same shape as ``self``. + rowVar + optional variable where each row of input is interpreted as a variable + (default = True). If set to False, each column is instead interpreted + as a variable. + bias + optional variable for normalizing input (default = False) by (N - 1) where + N is the number of given observations. If set to True, then normalization + is instead by N. Can be overridden by keyword ``ddof``. + ddof + optional variable to override ``bias`` (default = None). ddof=1 will return + the unbiased estimate, even with fweights and aweights given. ddof=0 will + return the simple average. + fweights + optional 1D array of integer frequency weights; the number of times each + observation vector should be repeated. + aweights + optional 1D array of observation vector weights. These relative weights are + typically large for observations considered "important" and smaller for + observations considered less "important". If ddof=0 is specified, the array + of weights can be used to assign probabilities to observation vectors. dtype - The desired data type of returned tensor. Default is None. + optional variable to set data-type of the result. By default, data-type + will have at least ``numpy.float64`` precision. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that + the inputs broadcast to. Returns ------- ret - The nanmean of the array elements. + an array containing the covariance matrix of an input matrix, or the + covariance matrix of two variables. The returned array must have a + floating-point data type determined by Type Promotion Rules and must be + a square matrix of shape (N, N), where N is the number of variables in the + input(s). - Functional Examples - ------------------- - >>> a = ivy.array([[1, ivy.nan], [3, 4]]) - >>> ivy.nanmean(a) - 2.6666666666666665 - >>> ivy.nanmean(a, axis=0) - ivy.array([2., 4.]) - """ - return ivy.current_backend(a).nanmean( - a, axis=axis, keepdims=keepdims, dtype=dtype, out=out - ) + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + >>> x = ivy.array([[1, 2, 3], + ... [4, 5, 6]]) + >>> y = x[0].cov(x[1]) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Container` inputs: + >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) + >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) + >>> z = ivy.Container.static_cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) + >>> z = ivy.cov(x, y) + >>> print(z) + { + a: ivy.array([[1., -1.], + [-1., 1.]]), + b: ivy.array([[1., -1.], + [-1., 1.]]) + } + + With :class:`ivy.Array` input and rowVar flag set to False (True by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], rowVar=False) + >>> print(y) + ivy.array([[1., 1.], + [1., 1.]]) + + With :class:`ivy.Array` input and bias flag set to True (False by default): + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> y = x[0].cov(x[1], bias=True) + >>> print(y) + ivy.array([[0.66666667, 0.66666667], + [0.66666667, 0.66666667]]) + + With :class:`ivy.Array` input with both fweights and aweights given: + >>> x = ivy.array([[1,2,3], + ... [4,5,6]]) + >>> fw = ivy.array([1,2,3]) + >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) + >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) + >>> print(y) + ivy.array([[0.48447205, 0.48447205], + [0.48447205, 0.48447205]]) + """ + return ivy.current_backend(x1).cov( + x1, + x2, + rowVar=rowVar, + bias=bias, + ddof=ddof, + fweights=fweights, + aweights=aweights, + dtype=dtype, + ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def quantile( - a: ivy.Array, - q: Union[ivy.Array, float], +@handle_array_function +def cummax( + x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Sequence[int], int]] = None, - keepdims: bool = False, - interpolation: str = "linear", + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the q-th quantile of the data along the specified axis. + Return a tuple containing the cumulative maximum of elements of input along the + given axis and index location of each maximum value found along the given axis. Parameters ---------- - a + x Input array. - q - Quantile or sequence of quantiles to compute, which must be - between 0 and 1 inclusive. axis - Axis or axes along which the quantiles are computed. The default - is to compute the quantile(s) along a flattened version of the array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original array a. - interpolation - {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. - Default value: 'linear'. - This specifies the interpolation method to use when the desired quantile lies - between two data points i < j: - - linear: i + (j - i) * fraction, where fraction is the fractional part of the - index surrounded by i and j. - - lower: i. - - higher: j. - - nearest: i or j, whichever is nearest. - - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with - integer dtypes. - - nearest_jax: provides jax-like computation for interpolation='nearest'. + Axis along which the cumulative maximum is computed. Default is ``0``. + exclusive + Whether to perform cummax exclusively. Default is ``False``. + reverse + Whether to perform the cummax from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis - is None, a rank(q) array. The first rank(q) dimensions index quantiles for - different values of q. + Array which holds the result of applying cummax at each + original array elements along the specified axis. Examples -------- - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = ivy.array(0.5) - >>> ivy.quantile(a, q) - ivy.array(3.5) + With :class:`ivy.Array` input: - >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) - >>> q = 0.5 - >>> ivy.quantile(a, q) - ivy.array(3.5) + >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) + >>> y = ivy.cummax(x, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), + ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) - >>> ivy.quantile(a, q, axis=0) - ivy.array([6.5, 4.5, 2.5]) + >>> x = ivy.array([ 14, 15, 49, -24, -39]) + >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) + >>> print(y) + (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) - >>> ivy.quantile(a, q, axis=1) - ivy.array([7., 2.]) + >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) + >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) + >>> print(x) + ivy.array([[0, 0, 0, 0], + [0, 1, 1, 1]]) - >>> ivy.quantile(a, q, axis=1, keepdims=True) - ivy.array([[7.],[2.]]) + >>> x = ivy.array([[-36, 83, -81], + ... [ 23, 29, 63], + ... [-83, 85, 2], + ... [ 31, 25, -86], + ... [-10, -52, 0], + ... [ 22, 38, 55], + ... [ 33, 54, -16]]) + >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) + >>> print(y) + (ivy.array([[ 0, 0, 83], + [ 0, 23, 29], + [ 0, 0, 85], + [ 0, 31, 31], + [ 0, 0, 0], + [ 0, 22, 38], + [ 0, 33, 54]]), ivy.array([[0, 0, 2], + [0, 1, 2], + [0, 0, 2], + [0, 1, 1], + [0, 0, 0], + [0, 1, 2], + [0, 1, 2]])) - >>> a = ivy.array([1., 2., 3., 4.]) - >>> q = ivy.array([0.3, 0.7]) - >>> ivy.quantile(a, q, interpolation='lower') - ivy.array([1., 3.]) + >>> x = ivy.array([73, 15, 47]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) + >>> print(y) + (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) + + >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) + >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) + >>> print(y) + (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) """ - return ivy.current_backend(a).quantile( - a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out + return ivy.current_backend(x).cummax( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_device_shifting -def corrcoef( - x: ivy.Array, +@handle_array_function +def cummin( + x: Union[ivy.Array, ivy.NativeArray], /, *, - y: Optional[ivy.Array] = None, - rowvar: bool = True, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - return ivy.current_backend().corrcoef(x, y=y, rowvar=rowvar, out=out) + """ + Return the cumulative minimum of the elements along a given axis. + + Parameters + ---------- + x + Input array. + axis + Axis along which the cumulative minimum is computed. Default is ``0``. + reverse + Whether to perform the cummin from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. + out + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Array which holds the result of applying cummin at each + original array elements along the specified axis. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cummin(x) + >>> print(y) + ivy.array([1, 1, 1, 0]) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cummin(x, axis=0, reverse=True, out=y) + >>> print(y) + ivy.array([[1., 3., 0.], + [1., 3., 0.]]) + + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[ 2, 4, 5], + [ 3, 5, 5], + [ 1, 3, 10]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cummin(x, axis= 0) + >>> print(y) + { + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) + } + + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cummin(x,axis=1,reverse=True, out=y) + >>> print(y) + { + a: ivy.array([[1., 3., 4.]]), + b: ivy.array([[3., 5., 8.], + [5., 5., 5.]]), + c: ivy.array([[1., 1., 1.], + [3., 6., 9.], + [0., 2., 3.]]) + } + + >>> x = ivy.Container(a=ivy.array([[0],[5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cummin(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [0]]), + b: ivy.array([[6, 8, 7], + [4, 2, 3]]), + c: ivy.array([[1, 2], + [1, 2], + [1, 2]]) + } + """ + return ivy.current_backend(x).cummin( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + ) +# TODO: Make bins optional by offering an automatic bins creation like numpy. +# Make density argument work in tensorflow +# Bins as str is not defined (check Numpy implementation). +# Permit multiple axis. +# Modify documentation to match the above modifications. @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def nanmedian( - input: ivy.Array, +def histogram( + a: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[Tuple[int], int]] = None, - keepdims: bool = False, - overwrite_input: bool = False, + bins: Optional[Union[int, ivy.Array, ivy.NativeArray]] = None, + axis: Optional[int] = None, + extend_lower_interval: Optional[bool] = False, + extend_upper_interval: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + range: Optional[Tuple[float]] = None, + weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + density: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the - function, and so the docstring for ivy.nanmedian also applies to this method with - minimal changes. + Compute the histogram of the array ``a``. + + .. note:: + Given bins = [c0, ..., cK], defining intervals I0 = [c0, c1), I1 = [c1, c2), + ..., I_{K-1} = [c_{K-1}, cK]. Parameters ---------- - self - Input array. + a + input array. + bins + if ``bins`` is an int, it defines the number of equal-width bins in the given + range. + if ``bins`` is an array, it defines a monotonically increasing array of bin + edges, including the rightmost edge, allowing for non-uniform bin widths. axis - Axis or axes along which the means are computed. - The default is to compute the mean of the flattened array. - keepdims - If this is set to True, the axes which are reduced are left in the result - as dimensions with size one. With this option, the result will broadcast - correctly against the original a. If the value is anything but the default, - then keepdims will be passed through to the mean or sum methods of - sub-classes of ndarray. If the sub-classes methods does not implement - keepdims any exceptions will be raised. - overwrite_input - If True, then allow use of memory of input array a for calculations. - The input array will be modified by the call to median. This will - save memory when you do not need to preserve the contents of the input array. - Treat the input as undefined, but it will probably be fully or partially sorted. - Default is False. If overwrite_input is True and a is not already an ndarray, - an error will be raised. + dimension along which maximum values must be computed. By default, the maximum + value must be computed over the entire array. Default: ``None``. + extend_lower_interval + if True, extend the lowest interval I0 to (-inf, c1]. + extend_upper_interval + ff True, extend the upper interval I_{K-1} to [c_{K-1}, +inf). + dtype + the output type. + range + the lower and upper range of the bins. The first element of the range must be + less than or equal to the second. + weights + each value in ``a`` only contributes its associated weight towards the bin count + (instead of 1). Must be of the same shape as a. + density + if True, the result is the value of the probability density function at the + bin, normalized such that the integral over the range of bins is 1. out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - A new array holding the result. If the input contains integers + a tuple containing the values of the histogram and the bin edges. - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) - >>> ivy.nanmedian(x) - ivy.array(23.) - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), - b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) - >>> ivy.nanmedian(x) + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.array([0., 0.5, 1., 1.5, 2.]) + >>> z = ivy.histogram(x, bins=y) + >>> print(z) + ivy.array([1., 0., 1., 1.]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [4.4, 5.5, .6]]) + >>> bins = 4 + >>> range = (0., 5.) + >>> dtype = ivy.int32 + >>> y = ivy.histogram(x, bins=bins, range=range, dtype=dtype) + >>> print(y) + ivy.array([2, 1, 1, 1]) + + >>> x = ivy.array([[1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> axis = 1 + >>> extend_lower_interval = True + >>> extend_upper_interval = True + >>> dtype = ivy.float32 + >>> weights = ivy.array([[1., 1., 1.], [1., 1., 1.]]) + >>> z = ivy.histogram( + ... x, + ... bins=y, + ... axis=axis, + ... extend_lower_interval=extend_lower_interval, + ... extend_upper_interval=extend_upper_interval, + ... dtype=dtype, + ... weights=weights) + >>> print(z) + ivy.array([[0., 3.], + [1., 0.], + [1., 0.], + [1., 0.], + [0., 0.]]) + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.array([0., 1., 2., 3., 4., 5.]) + >>> dtype = ivy.int32 + >>> z = ivy.histogram(x, bins=y, dtype=dtype) + >>> print(z) { - a: ivy.array(3.), - b: ivy.array(23.) + a: ivy.array([1, 1, 1, 0, 0]), + b: ivy.array([0, 0, 0, 1, 2]) } """ - return ivy.current_backend().nanmedian( - input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out + return ivy.current_backend(a).histogram( + a, + bins=bins, + axis=axis, + extend_lower_interval=extend_lower_interval, + extend_upper_interval=extend_upper_interval, + dtype=dtype, + range=range, + weights=weights, + density=density, + out=out, ) @@ -436,42 +599,39 @@ def nanmedian( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def bincount( - x: ivy.Array, +def igamma( + a: Union[ivy.Array, ivy.NativeArray], /, *, - weights: Optional[ivy.Array] = None, - minlength: int = 0, - out: Optional[ivy.Array] = None, + x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, ) -> ivy.Array: """ - Count the number of occurrences of each value in an integer array. + Compute the regularized lower gamma function of ``a`` and ``x``. Parameters ---------- self Input array. - weights - An optional input array. - minlength - A minimum number of bins for the output array. + x + An additional input array. + `x` has the same type as `a`. + out + optional output array, for writing the result to. Returns ------- ret - The bincount of the array elements. + The lower incomplete gamma function of the array elements. Examples -------- - >>> a = ivy.Container([[10.0, ivy.nan, 4], [3, 2, 1]]) - >>> a.bincount(a) - 3.0 - >>> a.bincount(a, axis=0) - array([6.5, 2. , 2.5]) + >>> a = ivy.array([2.5]) + >>> x = ivy.array([1.7, 1.2]) + >>> a.igamma(x) + ivy.array([0.3614, 0.2085]) """ - return ivy.current_backend(x).bincount( - x, weights=weights, minlength=minlength, out=out - ) + return ivy.current_backend().igamma(a, x=x, out=out) @handle_exceptions @@ -480,417 +640,257 @@ def bincount( @handle_out_argument @to_native_arrays_and_back @handle_device_shifting -def igamma( - a: Union[ivy.Array, ivy.NativeArray], +def median( + input: ivy.Array, /, *, - x: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the regularized lower gamma function of ``a`` and ``x``. + Compute the median along the specified axis. Parameters ---------- - self + input Input array. - x - An additional input array. - `x` has the same type as `a`. + axis + Axis or axes along which the medians are computed. The default is to compute + the median along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. out optional output array, for writing the result to. Returns ------- ret - The lower incomplete gamma function of the array elements. + The median of the array elements. - Examples - -------- - >>> a = ivy.array([2.5]) - >>> x = ivy.array([1.7, 1.2]) - >>> a.igamma(x) - ivy.array([0.3614, 0.2085]) + Functional Examples + ------------------- + >>> a = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> ivy.median(a) + 3.5 + >>> ivy.median(a, axis=0) + ivy.array([6.5, 4.5, 2.5]) """ - return ivy.current_backend().igamma(a, x=x, out=out) + return ivy.current_backend().median(input, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion +@handle_out_argument @to_native_arrays_and_back -def cov( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray] = None, +@infer_dtype +@handle_device_shifting +def nanmean( + a: ivy.Array, /, *, - rowVar: bool = True, - bias: bool = False, - ddof: Optional[int] = None, - fweights: Optional[ivy.Array] = None, - aweights: Optional[ivy.Array] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute the covariance of matrix x1, or variables x1 and x2. + Compute the mean of all non-NaN elements along the specified dimensions. Parameters ---------- - x1 - a 1D or 2D input array, with a numeric data type. - x2 - optional second 1D or 2D input array, with a numeric data type. - Must have the same shape as ``self``. - rowVar - optional variable where each row of input is interpreted as a variable - (default = True). If set to False, each column is instead interpreted - as a variable. - bias - optional variable for normalizing input (default = False) by (N - 1) where - N is the number of given observations. If set to True, then normalization - is instead by N. Can be overridden by keyword ``ddof``. - ddof - optional variable to override ``bias`` (default = None). ddof=1 will return - the unbiased estimate, even with fweights and aweights given. ddof=0 will - return the simple average. - fweights - optional 1D array of integer frequency weights; the number of times each - observation vector should be repeated. - aweights - optional 1D array of observation vector weights. These relative weights are - typically large for observations considered "important" and smaller for - observations considered less "important". If ddof=0 is specified, the array - of weights can be used to assign probabilities to observation vectors. + a + Input array. + axis + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of sub-classes + of ndarray. If the sub-classes methods does not implement keepdims any + exceptions will be raised. dtype - optional variable to set data-type of the result. By default, data-type - will have at least ``numpy.float64`` precision. + The desired data type of returned tensor. Default is None. out - optional output array, for writing the result to. It must have a shape that - the inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - an array containing the covariance matrix of an input matrix, or the - covariance matrix of two variables. The returned array must have a - floating-point data type determined by Type Promotion Rules and must be - a square matrix of shape (N, N), where N is the number of variables in the - input(s). - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.array([[1, 2, 3], - ... [4, 5, 6]]) - >>> y = x[0].cov(x[1]) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([1., 2., 3.])) - >>> y = ivy.Container(a=ivy.array([3., 2., 1.]), b=ivy.array([3., 2., 1.])) - >>> z = ivy.Container.static_cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With a combination of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.Container(a=ivy.array([3. ,2. ,1.]), b=ivy.array([-1., -2., -3.])) - >>> z = ivy.cov(x, y) - >>> print(z) - { - a: ivy.array([[1., -1.], - [-1., 1.]]), - b: ivy.array([[1., -1.], - [-1., 1.]]) - } - - With :class:`ivy.Array` input and rowVar flag set to False (True by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], rowVar=False) - >>> print(y) - ivy.array([[1., 1.], - [1., 1.]]) - - With :class:`ivy.Array` input and bias flag set to True (False by default): - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> y = x[0].cov(x[1], bias=True) - >>> print(y) - ivy.array([[0.66666667, 0.66666667], - [0.66666667, 0.66666667]]) + The nanmean of the array elements. - With :class:`ivy.Array` input with both fweights and aweights given: - >>> x = ivy.array([[1,2,3], - ... [4,5,6]]) - >>> fw = ivy.array([1,2,3]) - >>> aw = ivy.array([ 1.2, 2.3, 3.4 ]) - >>> y = x[0].cov(x[1], fweights=fw, aweights=aw) - >>> print(y) - ivy.array([[0.48447205, 0.48447205], - [0.48447205, 0.48447205]]) + Functional Examples + ------------------- + >>> a = ivy.array([[1, ivy.nan], [3, 4]]) + >>> ivy.nanmean(a) + 2.6666666666666665 + >>> ivy.nanmean(a, axis=0) + ivy.array([2., 4.]) """ - return ivy.current_backend(x1).cov( - x1, - x2, - rowVar=rowVar, - bias=bias, - ddof=ddof, - fweights=fweights, - aweights=aweights, - dtype=dtype, + return ivy.current_backend(a).nanmean( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummax( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def nanmedian( + input: ivy.Array, /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[Tuple[int], int]] = None, + keepdims: bool = False, + overwrite_input: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return a tuple containing the cumulative maximum of elements of input along the - given axis and index location of each maximum value found along the given axis. + ivy.Array instance method variant of ivy.nanmedian. This method simply wraps the + function, and so the docstring for ivy.nanmedian also applies to this method with + minimal changes. Parameters ---------- - x + self Input array. axis - Axis along which the cumulative maximum is computed. Default is ``0``. - exclusive - Whether to perform cummax exclusively. Default is ``False``. - reverse - Whether to perform the cummax from last to first element in the selected - axis. Default is ``False`` (from first to last element) + Axis or axes along which the means are computed. + The default is to compute the mean of the flattened array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original a. If the value is anything but the default, + then keepdims will be passed through to the mean or sum methods of + sub-classes of ndarray. If the sub-classes methods does not implement + keepdims any exceptions will be raised. + overwrite_input + If True, then allow use of memory of input array a for calculations. + The input array will be modified by the call to median. This will + save memory when you do not need to preserve the contents of the input array. + Treat the input as undefined, but it will probably be fully or partially sorted. + Default is False. If overwrite_input is True and a is not already an ndarray, + an error will be raised. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cummax at each - original array elements along the specified axis. + A new array holding the result. If the input contains integers + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-86, -19, 41, 88, -5, 80, 32, 87, -90, -12]) - >>> y = ivy.cummax(x, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([-86, -19, 41, 88, 88, 88, 88, 88, 88, 88]), - ivy.array([0, 1, 2, 3, 3, 3, 3, 3, 3, 3])) - - >>> x = ivy.array([ 14, 15, 49, -24, -39]) - >>> y = ivy.cummax(x, axis=0, exclusive=False, reverse=False) - >>> print(y) - (ivy.array([14, 15, 49, 49, 49]), ivy.array([0, 1, 2, 2, 2])) - - >>> x = ivy.array([[ 63, 43, -16, -4],[ 21, 82, 59, 33]]) - >>> ivy.cummax(x, axis=0, reverse=False, dtype='int64', out=x) - >>> print(x) - ivy.array([[0, 0, 0, 0], - [0, 1, 1, 1]]) - - >>> x = ivy.array([[-36, 83, -81], - ... [ 23, 29, 63], - ... [-83, 85, 2], - ... [ 31, 25, -86], - ... [-10, -52, 0], - ... [ 22, 38, 55], - ... [ 33, 54, -16]]) - >>> y = ivy.cummax(x, axis=1, exclusive=True, reverse=False) - >>> print(y) - (ivy.array([[ 0, 0, 83], - [ 0, 23, 29], - [ 0, 0, 85], - [ 0, 31, 31], - [ 0, 0, 0], - [ 0, 22, 38], - [ 0, 33, 54]]), ivy.array([[0, 0, 2], - [0, 1, 2], - [0, 0, 2], - [0, 1, 1], - [0, 0, 0], - [0, 1, 2], - [0, 1, 2]])) - - >>> x = ivy.array([73, 15, 47]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=True) - >>> print(y) - (ivy.array([47, 47, 0]), ivy.array([0, 0, 0])) - - >>> x = ivy.array([-47, -14, -67, 15, -23, -45]) - >>> y = ivy.cummax(x, axis=0, reverse=True, exclusive=False) - >>> print(y) - (ivy.array([ 15, 15, 15, 15, -23, -45]), ivy.array([2, 2, 2, 2, 1, 0])) + >>> x = ivy.array([[12.0, 10.0, 34.0], [45.0, 23.0, ivy.nan]]) + >>> ivy.nanmedian(x) + ivy.array(23.) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + >>> x = ivy.Container(a=ivy.array([[10.0, ivy.nan, 4], [3, 2, 1]]), + b=ivy.array([[12, 10, 34], [45, 23, ivy.nan]])) + >>> ivy.nanmedian(x) + { + a: ivy.array(3.), + b: ivy.array(23.) + } """ - return ivy.current_backend(x).cummax( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + return ivy.current_backend().nanmedian( + input, axis=axis, keepdims=keepdims, overwrite_input=overwrite_input, out=out ) @handle_exceptions @handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument @to_native_arrays_and_back -@handle_array_function -def cummin( - x: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def quantile( + a: ivy.Array, + q: Union[ivy.Array, float], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[Sequence[int], int]] = None, + keepdims: bool = False, + interpolation: str = "linear", out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative minimum of the elements along a given axis. + Compute the q-th quantile of the data along the specified axis. Parameters ---------- - x + a Input array. + q + Quantile or sequence of quantiles to compute, which must be + between 0 and 1 inclusive. axis - Axis along which the cumulative minimum is computed. Default is ``0``. - reverse - Whether to perform the cummin from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. + Axis or axes along which the quantiles are computed. The default + is to compute the quantile(s) along a flattened version of the array. + keepdims + If this is set to True, the axes which are reduced are left in the result + as dimensions with size one. With this option, the result will broadcast + correctly against the original array a. + interpolation + {'nearest', 'linear', 'lower', 'higher', 'midpoint', 'nearest_jax'}. + Default value: 'linear'. + This specifies the interpolation method to use when the desired quantile lies + between two data points i < j: + - linear: i + (j - i) * fraction, where fraction is the fractional part of the + index surrounded by i and j. + - lower: i. + - higher: j. + - nearest: i or j, whichever is nearest. + - midpoint: (i + j) / 2. linear and midpoint interpolation do not work with + integer dtypes. + - nearest_jax: provides jax-like computation for interpolation='nearest'. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cummin at each - original array elements along the specified axis. + A (rank(q) + N - len(axis)) dimensional array of same dtype as a, or, if axis + is None, a rank(q) array. The first rank(q) dimensions index quantiles for + different values of q. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cummin(x) - >>> print(y) - ivy.array([1, 1, 1, 0]) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cummin(x, axis=0, reverse=True, out=y) - >>> print(y) - ivy.array([[1., 3., 0.], - [1., 3., 0.]]) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = ivy.array(0.5) + >>> ivy.quantile(a, q) + ivy.array(3.5) - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cummin(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[ 2, 4, 5], - [ 3, 5, 5], - [ 1, 3, 10]]) + >>> a = ivy.array([[10., 7., 4.], [3., 2., 1.]]) + >>> q = 0.5 + >>> ivy.quantile(a, q) + ivy.array(3.5) - With :class:`ivy.Container` input: + >>> ivy.quantile(a, q, axis=0) + ivy.array([6.5, 4.5, 2.5]) - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cummin(x, axis= 0) - >>> print(y) - { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) - } + >>> ivy.quantile(a, q, axis=1) + ivy.array([7., 2.]) - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cummin(x,axis=1,reverse=True, out=y) - >>> print(y) - { - a: ivy.array([[1., 3., 4.]]), - b: ivy.array([[3., 5., 8.], - [5., 5., 5.]]), - c: ivy.array([[1., 1., 1.], - [3., 6., 9.], - [0., 2., 3.]]) - } + >>> ivy.quantile(a, q, axis=1, keepdims=True) + ivy.array([[7.],[2.]]) - >>> x = ivy.Container(a=ivy.array([[0],[5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cummin(x,axis=0,out=x) - >>> print(x) - { - a: ivy.array([[0], - [0]]), - b: ivy.array([[6, 8, 7], - [4, 2, 3]]), - c: ivy.array([[1, 2], - [1, 2], - [1, 2]]) - } + >>> a = ivy.array([1., 2., 3., 4.]) + >>> q = ivy.array([0.3, 0.7]) + >>> ivy.quantile(a, q, interpolation='lower') + ivy.array([1., 3.]) """ - return ivy.current_backend(x).cummin( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + return ivy.current_backend(a).quantile( + a, q, axis=axis, keepdims=keepdims, interpolation=interpolation, out=out ) diff --git a/ivy/functional/ivy/general.py b/ivy/functional/ivy/general.py index 1267e86d08269..4d0aa36bfef8a 100644 --- a/ivy/functional/ivy/general.py +++ b/ivy/functional/ivy/general.py @@ -46,23 +46,18 @@ FN_CACHE = dict() INF = float("inf") - -precise_mode_stack = list() -queue_timeout_stack = list() array_mode_stack = list() -shape_array_mode_stack = list() -nestable_mode_stack = list() exception_trace_mode_stack = list() inplace_mode_stack = list() -trace_mode_dict = dict() -trace_mode_dict["frontend"] = "ivy/functional/frontends" -trace_mode_dict["ivy"] = "ivy/" -trace_mode_dict["full"] = "" -trace_mode_dict["none"] = "" -show_func_wrapper_trace_mode_stack = list() -min_denominator_stack = list() min_base_stack = list() +min_denominator_stack = list() +nestable_mode_stack = list() +precise_mode_stack = list() +queue_timeout_stack = list() +shape_array_mode_stack = list() +show_func_wrapper_trace_mode_stack = list() tmp_dir_stack = list() +trace_mode_dict = dict() # Extra # @@ -88,78 +83,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self -ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True - - -@handle_exceptions -def set_precise_mode(mode: bool) -> None: - """ - Set the mode of whether to use a promotion table that avoids any precision loss or a - compute effecient table that avoids most wider-than-necessary promotions. - - Parameter - --------- - mode - boolean whether to use high precision promtion table - - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False - - >>> ivy.set_precise_mode(True) - >>> ivy.precise_mode - True - """ - global precise_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - precise_mode_stack.append(mode) - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) - - -@handle_exceptions -def unset_precise_mode() -> None: - """ - Reset the mode of whether to use a promotion table that avoids any precision loss or - a compute effecient table that avoids most wider-than-necessary promotions. - - Examples - -------- - >>> ivy.set_precise_mode(False) - >>> ivy.precise_mode - False - - >>> ivy.unset_precise_mode() - >>> ivy.precise_mode - True - """ - global precise_mode_stack - if precise_mode_stack: - precise_mode_stack.pop(-1) - mode = precise_mode_stack[-1] if precise_mode_stack else True - ivy.__setattr__("precise_mode", mode, True) - _update_promotion_table(precise=mode) - - -def _update_promotion_table(precise): - """Update the current datatype promotion table.""" - if precise: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.precise_extra_promotion_table, - } - - else: - ivy.promotion_table = { - **ivy.array_api_promotion_table, - **ivy.common_extra_promotion_table, - **ivy.extra_promotion_table, - } - - class ArrayMode: """Array Mode Context Manager.""" @@ -179,426 +102,371 @@ def __exit__(self, exc_type, exc_val, exc_tb): return self -def get_referrers_recursive( - item, depth=0, max_depth=None, seen_set=None, local_set=None -): - """ - Summary. +# --- Helpers --- # +# --------------- # - Parameters - ---------- - item - depth - (Default value = 0) - max_depth - (Default value = None) - seen_set - (Default value = None) - local_set - (Default value = None`) - """ - seen_set = ivy.default(seen_set, set()) - local_set = ivy.default(local_set, set()) - ret_cont = ivy.Container( - repr=str(item).replace(" ", ""), - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - referrers = [ - ref - for ref in gc.get_referrers(item) - if not ( - isinstance(ref, dict) - and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) - ) - ] - local_set.add(str(id(referrers))) - for ref in referrers: - ref_id = str(id(ref)) - if ref_id in local_set or hasattr(ref, "cell_contents"): - continue - seen = ref_id in seen_set - seen_set.add(ref_id) - refs_rec = lambda: get_referrers_recursive( - ref, depth + 1, max_depth, seen_set, local_set - ) - this_repr = "tracked" if seen else str(ref).replace(" ", "") - if not seen and (not max_depth or depth < max_depth): - val = ivy.Container( - repr=this_repr, - alphabetical_keys=False, - keyword_color_dict={"repr": "magenta"}, - ) - refs = refs_rec() - for k, v in refs.items(): - val[k] = v - else: - val = this_repr - ret_cont[str(ref_id)] = val - return ret_cont +def _all_dnd_combinations(): + all_comb = {} + for device in ivy.all_devices: + all_comb[device] = ivy.all_dtypes + return all_comb -@handle_exceptions -@handle_backend_invalid -def is_native_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False -) -> bool: - """ - Determine whether the input x is an :class:`ivy.NativeArray` instance. +def _broadcast_to(input, target_shape): + input = ivy.squeeze(input) + if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): + return ivy.reshape(input, target_shape) + else: + input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input + new_dims = () + i_i = len(input.shape) - 1 + for i_t in range(len(target_shape) - 1, -1, -1): + if len(input.shape) + len(new_dims) >= len(target_shape): + break + if i_i < 0 or target_shape[i_t] != input.shape[i_i]: + new_dims += (i_t,) + else: + i_i -= 1 + input = ivy.expand_dims(input, axis=new_dims) + return ivy.broadcast_to(input, target_shape) - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. - Returns - ------- - ret - Boolean, whether or not x is an :class:`ivy.NativeArray`. +def _dnd_dict_difference(a, b): + res = a + for device in list(a): + if device in b: + difference = set.difference(set(a[device]), set(b[device])) + if difference: + res[device] = tuple(difference) + else: + del res[device] + return res - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_native_array(x) - False - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_native_array(x, exclusive=True) - True - """ - try: - return current_backend(x).is_native_array(x, exclusive=exclusive) - except ValueError: - return False +def _dnd_dict_intersection(a, b): + res = {} + for device in a: + if device in b: + intersection = set.intersection(set(a[device]), set(b[device])) + if intersection: + res[device] = tuple(intersection) + return res -@handle_exceptions -@handle_backend_invalid -def is_ivy_array( - x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False -) -> bool: - """ - Determine whether the input x is a valid Ivy Array. +def _dnd_dict_union(a, b): + res = {} + for device in set(list(a) + list(b)): + u1 = set(a.get(device, ())) + u2 = set(b.get(device, ())) + res[device] = tuple(set.union(u1, u2)) - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + return res - Returns - ------- - ret - Boolean, whether or not x is a valid Ivy Array. - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> ivy.is_ivy_array(x) - True +def _get_devices_and_dtypes(fn, recurse=False, complement=True): + supported_devices = ivy.function_supported_devices(fn, recurse=recurse) + supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> ivy.is_ivy_array(x, exclusive=True) - False - """ - return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) + if hasattr(fn, "partial_mixed_handler"): + supported_devices = supported_devices["primary"] + supported_dtypes = supported_dtypes["primary"] + supported = {} + # Generate a base supported set from other attributes + for device in supported_devices: + supported[device] = supported_dtypes -@handle_exceptions -@handle_backend_invalid -def is_array(x: Any, /, *, exclusive: bool = False) -> bool: - """ - Determine whether the input x is either an Ivy Array or a Native Array. + is_backend_fn = "backend" in fn.__module__ + is_frontend_fn = "frontend" in fn.__module__ + is_einops_fn = "einops" in fn.__name__ + if not is_backend_fn and not is_frontend_fn and not is_einops_fn: + if complement: + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported - Parameters - ---------- - x - The input to check - exclusive - Whether to check if the data type is exclusively an array, rather than a - variable or traced array. + backend = ivy.current_backend_str() - Returns - ------- - ret - Boolean, whether or not x is an array. + # Their values are formatted like either + # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype.__get__() - Examples - -------- - >>> x = ivy.array([0, 1, 2]) - >>> print(ivy.is_array(x)) - True + if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): + fn_supported_dnd = fn_supported_dnd.get(backend, supported) - >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) - >>> print(ivy.is_array(x, exclusive=True)) - True + ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) + # dict intersection + supported = _dnd_dict_intersection(supported, fn_supported_dnd) - >>> x = [2, 3] - >>> print(ivy.is_array(x)) - False - """ - return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( - x, exclusive=exclusive - ) + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() + if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) -@handle_exceptions -def is_ivy_container(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Container. + ivy.utils.assertions.check_isinstance( + list(fn_unsupported_dnd.values())[0], tuple + ) + # dict difference + supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - Parameters - ---------- - x - The input to check + if complement: + # dict difference + all_comb = _all_dnd_combinations() + supported = _dnd_dict_difference(all_comb, supported) + return supported - Returns - ------- - ret - Boolean, whether or not x is an ivy container. - Examples - -------- - >>> x = ivy.Container() - >>> print(ivy.is_ivy_container(x)) - True +def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: + fn_unsupported_dnd = {} + fn_supported_dnd = {} + backend = ivy.current_backend_str() + if hasattr(fn, "unsupported_device_and_dtype"): + fn_unsupported_dnd = fn.unsupported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): + fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) + if hasattr(fn, "supported_device_and_dtype"): + fn_supported_dnd = fn.supported_device_and_dtype + # if it's a nested dict, unwrap for the current backend + if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): + fn_supported_dnd = fn_supported_dnd.get(backend, {}) - >>> x = [2, 3] - >>> print(ivy.is_ivy_container(x)) - False - """ - return isinstance(x, ivy.Container) + ivy.utils.assertions.check_false( + fn_unsupported_dnd and fn_supported_dnd, + "unsupported_device_and_dtype and supported_device_and_dtype cannot" + " both be defined for the same function", + ) + us = "unsupported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") + _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") -ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True + ss = "supported_device_and_dtype" + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") + _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") + return True -@handle_exceptions -def set_array_mode(mode: bool) -> None: - """ - Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs - back to ivy.Array. - It Stops the conversion of ivy.NativeArray to ivy.Array in the - case when it is set to False. +def _numel(shape): + shape = tuple(shape) + return ivy.prod(shape).to_scalar() if shape != () else 1 - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions - Examples - -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False +def _parse_query(query, x_shape): + query = (query,) if not isinstance(query, tuple) else query + query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) - >>> ivy.set_array_mode(True) - >>> ivy.array_mode - True - """ - global array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - array_mode_stack.append(mode) - ivy.__setattr__("array_mode", mode, True) + # array containing all of x's flat indices + x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) + # use numpy's __getitem__ to get the queried indices + x_idxs = ivy.array(x_.to_numpy()[query_]) + target_shape = x_idxs.shape -@handle_exceptions -def unset_array_mode() -> None: - """ - Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back - to ivy.Array to the previous state. + if 0 in x_idxs.shape or 0 in x_shape: + return None, target_shape - Examples - -------- - >>> ivy.set_array_mode(False) - >>> ivy.array_mode - False + # convert the flat indices to multi-D indices + x_idxs = ivy.unravel_index(x_idxs, x_shape) - >>> ivy.unset_shape_array_mode() - >>> ivy.array_mode - True - """ - global array_mode_stack - if array_mode_stack: - array_mode_stack.pop(-1) - mode = array_mode_stack[-1] if array_mode_stack else True - ivy.__setattr__("array_mode", mode, True) + # stack the multi-D indices to bring them to gather_nd/scatter_nd format + x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) + return x_idxs, target_shape -ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True +def _update_promotion_table(precise): + """Update the current datatype promotion table.""" + if precise: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.precise_extra_promotion_table, + } -@handle_exceptions -def set_nestable_mode(mode: bool) -> None: - """ - Set the mode of whether to check if function inputs are ivy.Container. + else: + ivy.promotion_table = { + **ivy.array_api_promotion_table, + **ivy.common_extra_promotion_table, + **ivy.extra_promotion_table, + } - Parameter - --------- - mode - boolean whether to check if function inputs are ivy.Container - Examples - -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False +def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): + attr_list = () + if hasattr(fn, other_attr_name): + attr_list = getattr(fn, other_attr_name) + if isinstance(attr_list, dict): + attr_list = attr_list.get(backend, ()) + ivy.utils.assertions.check_false( + dnd_dict and attr_list, + f"Cannot specify both {first_attr_name} and {other_attr_name} " + "cannot both be defined for the same function", + ) - >>> ivy.set_nestable_mode(True) - >>> ivy.nestable_mode - True - """ - global nestable_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - nestable_mode_stack.append(mode) - ivy.__setattr__("nestable_mode", mode, True) + +# --- Main --- # +# ------------ # @handle_exceptions -def unset_nestable_mode() -> None: +@handle_nestable +@inputs_to_ivy_arrays +@handle_array_function +def all_equal( + *xs: Iterable[Any], equality_matrix: bool = False +) -> Union[bool, ivy.Array, ivy.NativeArray]: """ - Reset the mode of whether to check if function inputs are ivy.Container to the - previous state. + Determine whether the inputs are all equal. + + Parameters + ---------- + xs + inputs to compare. + equality_matrix + Whether to return a matrix of equalities comparing each input with every other. + Default is ``False``. + + Returns + ------- + ret + Boolean, whether or not the inputs are equal, or matrix array of booleans if + equality_matrix=True is set. Examples -------- - >>> ivy.set_nestable_mode(False) - >>> ivy.nestable_mode - False + With :class:`ivy.Array` inputs: - >>> ivy.unset_nestable_mode() - >>> ivy.nestable_mode + >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) + >>> y = ivy.all_equal(x1, x2) + >>> print(y) True - """ - global nestable_mode_stack - if nestable_mode_stack: - nestable_mode_stack.pop(-1) - mode = nestable_mode_stack[-1] if nestable_mode_stack else True - ivy.__setattr__("nestable_mode", mode, True) - - -ivy.exception_trace_mode = ( - exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" -) + >>> x1 = ivy.array([0, 0]) + >>> x2 = ivy.array([0, 0]) + >>> x3 = ivy.array([1, 0]) + >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) + >>> print(y) + ivy.array([[ True, True, False], + [ True, True, False], + [False, False, True]]) -@handle_exceptions -def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: - """ - Set the mode of whether to show frontend-truncated exception stack traces, ivy- - truncated exception stack traces or full exception stack traces. + With one :class:`ivy.Container` inputs: - Parameter - --------- - mode - str exeption trace mode, one of `ivy`, `full` or `frontend` + >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), + ... b=ivy.array([0, 0, -1, 1, 0])) + >>> x2 = ivy.array([0, 0, -1, 1, 0]) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) + { + a: True, + b: True + } - Examples - -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' + With multiple :class:`ivy.Container` inputs: - >>> ivy.set_exception_trace_mode("full") - >>> ivy.exception_trace_mode - 'full' + >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, 0, 1])) + >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), + ... b=ivy.array([1, 0, -1, -1])) + >>> y = ivy.all_equal(x1, x2, equality_matrix=False) + >>> print(y) + { + a: True, + b: False + } """ - global exception_trace_mode_stack - trace_modes = list(trace_mode_dict.keys()) - ivy.utils.assertions.check_elem_in_list( - mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) - ) - exception_trace_mode_stack.append(mode) - ivy.__setattr__("exception_trace_mode", mode, True) + equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b + if equality_matrix: + num_arrays = len(xs) + mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] + for i, xa in enumerate(xs): + for j_, xb in enumerate(xs[i:]): + j = j_ + i + res = equality_fn(xa, xb) + if ivy.is_native_array(res): + # noinspection PyTypeChecker + res = ivy.to_scalar(res) + # noinspection PyTypeChecker + mat[i][j] = res + # noinspection PyTypeChecker + mat[j][i] = res + return ivy.array(mat) + x0 = xs[0] + for x in xs[1:]: + if not equality_fn(x0, x): + return False + return True @handle_exceptions -def unset_exception_trace_mode() -> None: +def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): """ - Reset the trace mode to the previously set mode. + Return the index and `inspect.Parameter` representation of the specified argument. + In the form of a dict with keys "idx" and "param". - Examples - -------- - >>> ivy.set_exception_trace_mode("ivy") - >>> ivy.exception_trace_mode - 'ivy' + Parameters + ---------- + fn + The function to retrieve the argument information for + name + The name of the argument + idx + the index of the argument in the inputs - >>> ivy.unset_exception_trace_mode() - >>> ivy.exception_trace_mode - 'full' + Returns + ------- + ret + a `dict` containing the idx, and the `inspect.Parameter` for the argument, + which itself contains the parameter name, type, and other helpful information. """ - global exception_trace_mode_stack - if exception_trace_mode_stack: - exception_trace_mode_stack.pop(-1) - mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" - ivy.__setattr__("exception_trace_mode", mode, True) - - -ivy.show_func_wrapper_trace_mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True -) + ivy.utils.assertions.check_all_or_any_fn( + name, + idx, + fn=ivy.exists, + type="any", + limit=[1], + message="exactly one of the keyword arguments name or idx must be provided", + as_array=False, + ) + params = inspect.signature(fn).parameters + if ivy.exists(name): + return {"idx": list(params).index(name), "param": params[name]} + return {"idx": idx, "param": list(params.values())[idx]} @handle_exceptions -def set_show_func_wrapper_trace_mode(mode: bool) -> None: - """ - Set the mode of whether to show the full stack trace with function wrapping traces. - - Parameter - --------- - mode - boolean whether to perform ivy.Array conversions - - Examples - -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False - - >>> ivy.set_show_func_wrapper_trace_mode(True) - >>> ivy.show_func_wrapper_trace_mode - True +def arg_names(receiver): """ - global show_func_wrapper_trace_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - show_func_wrapper_trace_mode_stack.append(mode) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + Get the expected keyword arguments for a function or class constructor. + Parameters + ---------- + receiver + Function or class constructor -@handle_exceptions -def unset_show_func_wrapper_trace_mode() -> None: - """ - Reset the mode of whether to show the full stack trace with function wrapping - traces. + Returns + ------- + ret + List containing the keyword arguments' names for a function or class constructor Examples -------- - >>> ivy.set_show_func_wrapper_trace_mode(False) - >>> ivy.show_func_wrapper_trace_mode - False + >>> x = ivy.arg_names(ivy.tan) + >>> print(x) + ['x', 'out'] - >>> ivy.unset_show_func_wrapper_trace_mode() - >>> ivy.show_func_wrapper_trace_mode - True + >>> x = ivy.arg_names(ivy.optimizers.Adam) + >>> print(x) + ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', + 'stop_gradients', 'compile_on_next_step', 'device'] """ - global show_func_wrapper_trace_mode_stack - if show_func_wrapper_trace_mode_stack: - show_func_wrapper_trace_mode_stack.pop(-1) - mode = ( - show_func_wrapper_trace_mode_stack[-1] - if show_func_wrapper_trace_mode_stack - else True - ) - ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + return list(inspect.signature(receiver).parameters.keys()) @handle_exceptions @@ -653,347 +521,254 @@ def array_equal( @handle_exceptions @handle_nestable -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_array_function -def all_equal( - *xs: Iterable[Any], equality_matrix: bool = False -) -> Union[bool, ivy.Array, ivy.NativeArray]: +def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: """ - Determine whether the inputs are all equal. + Assert that inplace operations are supported for x. Parameters ---------- - xs - inputs to compare. - equality_matrix - Whether to return a matrix of equalities comparing each input with every other. - Default is ``False``. + x + Input variable or array to check for inplace support for. Returns ------- ret - Boolean, whether or not the inputs are equal, or matrix array of booleans if - equality_matrix=True is set. + True if supports, raises IvyBackendException otherwise + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - With :class:`ivy.Array` inputs: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x1 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> x2 = ivy.array([1, 1, 0, 0, 1, -1]) - >>> y = ivy.all_equal(x1, x2) - >>> print(y) + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) True - >>> x1 = ivy.array([0, 0]) - >>> x2 = ivy.array([0, 0]) - >>> x3 = ivy.array([1, 0]) - >>> y = ivy.all_equal(x1, x2, x3, equality_matrix=True) - >>> print(y) - ivy.array([[ True, True, False], - [ True, True, False], - [False, False, True]]) + With :class:`ivy.Array` input and default backend set as `torch`: - With one :class:`ivy.Container` inputs: + >>> ivy.set_backend("torch") + >>> x = ivy.array([1, 2, 3]) + >>> print(x.assert_supports_inplace()) + True - >>> x1 = ivy.Container(a=ivy.array([0, 0, -1, 1, 0]), - ... b=ivy.array([0, 0, -1, 1, 0])) - >>> x2 = ivy.array([0, 0, -1, 1, 0]) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) + With :class:`ivy.Container` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) { a: True, b: True } - With multiple :class:`ivy.Container` inputs: + With :class:`ivy.Container` input and default backend set as `torch`: - >>> x1 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, 0, 1])) - >>> x2 = ivy.Container(a=ivy.array([1, 0, 1, 1]), - ... b=ivy.array([1, 0, -1, -1])) - >>> y = ivy.all_equal(x1, x2, equality_matrix=False) - >>> print(y) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> print(x.assert_supports_inplace()) { a: True, - b: False + b: True } """ - equality_fn = ivy.array_equal if ivy.is_array(xs[0]) else lambda a, b: a == b - if equality_matrix: - num_arrays = len(xs) - mat = [[None for _ in range(num_arrays)] for _ in range(num_arrays)] - for i, xa in enumerate(xs): - for j_, xb in enumerate(xs[i:]): - j = j_ + i - res = equality_fn(xa, xb) - if ivy.is_native_array(res): - # noinspection PyTypeChecker - res = ivy.to_scalar(res) - # noinspection PyTypeChecker - mat[i][j] = res - # noinspection PyTypeChecker - mat[j][i] = res - return ivy.array(mat) - x0 = xs[0] - for x in xs[1:]: - if not equality_fn(x0, x): - return False + ivy.utils.assertions.check_true( + ivy.supports_inplace_updates(x), + "Inplace operations are not supported {} types with {} backend".format( + type(x), ivy.current_backend_str() + ), + ) return True @handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays -@handle_array_function -@handle_device_shifting -def to_numpy( - x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True -) -> np.ndarray: +def cache_fn(func: Callable) -> Callable: """ - Convert an array into a numpy array. + Cache function outputs. + + A decorator to wrap a function, such that computed outputs are cached to avoid + recalculating them later. Parameters ---------- - x - input array - copy - whether to copy the array to a new address or not. - Default is ``True``. + func + The function to wrap, whose output should be cached for later. Returns ------- ret - a numpy array copying all the element of the array ``x``. + The newly cache wrapped function. Examples -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [-1 0 1] - - >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) - >>> y = ivy.to_numpy(x, copy=True) - >>> print(y) - [[-1 0 1] - [-1 0 1] - [ 1 0 -1]] + With positional arguments only: - With :class:`ivy.Container` input: + >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 + >>> cached_sum = ivy.cache_fn(my_sum) + >>> print(cached_sum(3, 5)) + 8 - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([-1, 0, 1], dtype=int32) - } + With keyword arguments: - >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), - ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) - >>> y = ivy.to_numpy(x) - >>> print(y) - { - a: array([[-1., 0., 1.], - [-1., 0., 1.], - [1., 0., -1.]], dtype=float32), - b: array([[-1, 0, 0], - [1, 0, 1], - [1, 1, 1]], dtype=int32) - } + >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc + >>> cached_line_eq = ivy.cache_fn(line_eq) + >>> print(cached_line_eq(3, itc=5, slp=2)) + 11 """ - return current_backend(x).to_numpy(x, copy=copy) + global FN_CACHE + if func not in FN_CACHE: + FN_CACHE[func] = dict() + @wraps(func) + def cached_fn(*args, **kwargs): + key = "".join( + [str(i) + ", " for i in args] + + [" kw, "] + + [str(i) + ", " for i in sorted(kwargs.items())] + ) + cache = FN_CACHE[func] + if key in cache: + return cache[key] + ret = func(*args, **kwargs) + cache[key] = ret + return ret -@handle_exceptions -@handle_nestable -def isscalar(x: Any, /) -> bool: - return np.isscalar(x) + return cached_fn @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: +def clip_matrix_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, + /, + *, + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Convert an array with a single element into a scalar. + Clips (limits) the matrix norm of an array. Parameters ---------- x - Input array with a single element. + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - a scalar copying the element of the array ``x``. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + An array with the matrix norm downscaled to the max norm if needed. Functional Examples ------------------- With :class:`ivy.Array` input: - >>> x = ivy.array([3]) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[0., 1., 2.]]) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) - 3 + ivy.array([[0. , 0.894, 1.79 ]]) - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) + >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) + >>> print(y) + ivy.array([[ 0.0353, -0.424 , 1.31 ], + [ 0. , 2.58 , -0.176 ]]) - >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[[5., 4.], [-2., 6.]], + ... [[3., 7.], [0., -5.]]]) + >>> y = ivy.empty((2, 2, 2)) + >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) >>> print(y) - { - a: -1, - b: 3 - } + ivy.array([[[ 0.339, 0.271], + [-0.135, 0.406]], + [[ 0.168, 0.391], + [ 0. , -0.279]]]) - >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), - ... c=ivy.array([-1])) - >>> y = ivy.to_scalar(x) + >>> x = ivy.array([[0., 1.], + ... [2., 3.]]) + >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) + >>> print(x) + ivy.array([[0., 1.], + [2., 3.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), + ... b=ivy.array([[3., 4., 5.]])) + >>> y = ivy.clip_matrix_norm(x, 2.0) >>> print(y) { - a: 1, - b: 0, - c: -1 + a: ivy.array([[0., 0.894, 1.79]]), + b: ivy.array([[0.849, 1.13, 1.41]]) } """ - return current_backend(x).to_scalar(x) + norms = ivy.matrix_norm(x, ord=p, keepdims=True) + ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) + return ivy.multiply(ratios, x, out=out) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@inputs_to_native_arrays +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: +def clip_vector_norm( + x: Union[ivy.Array, ivy.NativeArray], + max_norm: float, + /, + *, + p: float = 2.0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Create a (possibly nested) list from input array. + Clips (limits) the vector p-norm of an array. Parameters ---------- x - Input array. + Input array containing elements to clip. + max_norm + The maximum value of the array norm. + p + The p-value for computing the p-norm. + Default is 2. + out + optional output array, for writing the result to. + It must have a shape that the inputs broadcast to. Returns ------- ret - A list representation of the input array ``x``. + An array with the vector norm downscaled to the max norm if needed. + + Functional Examples + ------------------ - Examples - -------- With :class:`ivy.Array` input: - >>> x = ivy.array([-1, 0, 1]) - >>> y = ivy.to_list(x) + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.clip_vector_norm(x, 2.0) >>> print(y) - [-1, 0, 1] + ivy.array([0. , 0.894, 1.79 ]) - >>> x = ivy.array([[ 1.1, 2.2, 3.3], - ... [-4.4, -5.5, -6.6]]) - >>> y = ivy.to_list(x) + >>> x = ivy.array([0.5, -0.7, 2.4]) + >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) >>> print(y) - [[1.100000023841858,2.200000047683716,3.299999952316284], - [-4.400000095367432,-5.5,-6.599999904632568]] - - >>> x = ivy.array([[[-1, 0, 1], - ... [ 1, 0, -1]], - ... [[ 1, -1, 0], - ... [ 1, 0, -1]]]) - >>> y = ivy.to_list(x) - >>> print(y) - [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - - With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - - >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [-1, 0, 1] - } - - >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], - ... [-1, 0, 1], - ... [1, 0, -1]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] - } - - >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], - ... [[1, -1, 0],[1, 0, -1]]])) - >>> y = ivy.to_list(x) - >>> print(y) - { - a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - } - """ - return current_backend(x).to_list(x) - - -@handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def clip_vector_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, - /, - *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Clips (limits) the vector p-norm of an array. - - Parameters - ---------- - x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. - It must have a shape that the inputs broadcast to. - - Returns - ------- - ret - An array with the vector norm downscaled to the max norm if needed. - - Functional Examples - ------------------ - - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.clip_vector_norm(x, 2.0) - >>> print(y) - ivy.array([0. , 0.894, 1.79 ]) - - >>> x = ivy.array([0.5, -0.7, 2.4]) - >>> y = ivy.clip_vector_norm(x, 3.0, p=1.0) - >>> print(y) - ivy.array([ 0.417, -0.583, 2. ]) + ivy.array([ 0.417, -0.583, 2. ]) >>> x = ivy.array([[[0., 0.], [1., 3.], [2., 6.]], ... [[3., 9.], [4., 12.], [5., 15.]]]) @@ -1037,85 +812,139 @@ def clip_vector_norm( @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def clip_matrix_norm( - x: Union[ivy.Array, ivy.NativeArray], - max_norm: float, +def container_types(): + """ + Summary. + + Returns + ------- + ret + a key-value structure, and exposes public methods .keys(), .values() and + items(). + """ + # noinspection PyBroadException + try: + return current_backend().container_types() + except ValueError: + return [] + + +@handle_exceptions +def current_backend_str() -> Union[str, None]: + """ + Return framework string. + + Returns + ------- + ret + The framework string. + """ + fw = current_backend() + if not backend_stack: + return "" + return fw.current_backend_str() + + +@handle_exceptions +def default( + x: Any, /, + default_val: Any, *, - p: float = 2.0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + catch_exceptions: bool = False, + rev: bool = False, + with_callable: bool = False, +) -> Any: """ - Clips (limits) the matrix norm of an array. + Return x provided it exists (is not None), else returns default value. Parameters ---------- x - Input array containing elements to clip. - max_norm - The maximum value of the array norm. - p - The p-value for computing the p-norm. - Default is 2. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input which may or may not exist (be None). + default_val + The default value. + catch_exceptions + Whether to catch exceptions from callable x. + Default is ``False``. + rev + Whether to reverse the input x and default_val. + Default is ``False``. + with_callable + Whether either of the arguments might be callable functions. + Default is ``False``. Returns ------- ret - An array with the matrix norm downscaled to the max norm if needed. + x if x exists (is not None), else default. Functional Examples - ------------------- + ------------------ + With :code:`Any` input: - With :class:`ivy.Array` input: + >>> x = None + >>> y = ivy.default(x, "default_string") + >>> print(y) + default_string - >>> x = ivy.array([[0., 1., 2.]]) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = "" + >>> y = ivy.default(x, "default_string") >>> print(y) - ivy.array([[0. , 0.894, 1.79 ]]) - >>> x = ivy.array([[0.1, -1.2, 3.7], [0., 7.3, -0.5]]) - >>> y = ivy.clip_matrix_norm(x, 3.0, p=1.0) + + >>> x = ivy.array([4, 5, 6]) + >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) >>> print(y) - ivy.array([[ 0.0353, -0.424 , 1.31 ], - [ 0. , 2.58 , -0.176 ]]) + ivy.array([1, 2, 3]) - >>> x = ivy.array([[[5., 4.], [-2., 6.]], - ... [[3., 7.], [0., -5.]]]) - >>> y = ivy.empty((2, 2, 2)) - >>> y = ivy.clip_matrix_norm(x, 0.5, p=2.0) + >>> x = lambda: ivy.array([1, 2, 3]) + >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) >>> print(y) - ivy.array([[[ 0.339, 0.271], - [-0.135, 0.406]], - [[ 0.168, 0.391], - [ 0. , -0.279]]]) + ivy.array([1, 2, 3]) - >>> x = ivy.array([[0., 1.], - ... [2., 3.]]) - >>> ivy.clip_matrix_norm(x, 5.0, p=1.0, out=x) - >>> print(x) - ivy.array([[0., 1.], - [2., 3.]]) + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) + >>> print(y) + ivy.array([1, 2, 3]) - With :class:`ivy.Container` input: + >>> x = lambda: None + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) + >>> print(y) + ivy.array([1, 2, 3]) - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.]]), - ... b=ivy.array([[3., 4., 5.]])) - >>> y = ivy.clip_matrix_norm(x, 2.0) + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True) >>> print(y) - { - a: ivy.array([[0., 0.894, 1.79]]), - b: ivy.array([[0.849, 1.13, 1.41]]) - } + ivy.array([1, 2, 3]) + + >>> x = lambda a, b: a + b + >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, + ... catch_exceptions=True, rev=True) + >>> print(y) + ivy.array([1, 2, 3]) """ - norms = ivy.matrix_norm(x, ord=p, keepdims=True) - ratios = ivy.minimum(ivy.stable_divide(max_norm, norms), 1.0) - return ivy.multiply(ratios, x, out=out) + with_callable = catch_exceptions or with_callable + if rev: + tmp = x + x = default_val + default_val = tmp + if with_callable: + x_callable = callable(x) + default_callable = callable(default_val) + else: + x_callable = False + default_callable = False + if catch_exceptions: + # noinspection PyBroadException + try: + x = x() if x_callable else x + except Exception: + return default_val() if default_callable else default_val + else: + x = x() if x_callable else x + return x if exists(x) else default_val() if default_callable else default_val @handle_exceptions @@ -1123,234 +952,267 @@ def clip_matrix_norm( @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def fourier_encode( +def einops_rearrange( x: Union[ivy.Array, ivy.NativeArray], - max_freq: Union[float, ivy.Array, ivy.NativeArray], + pattern: str, /, *, - num_bands: int = 4, - linear: bool = False, - concat: bool = True, - flatten: bool = False, -) -> Union[ivy.Array, ivy.NativeArray, Tuple]: + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Pad an array with fourier encodings. + Perform einops rearrange operation on input array x. Parameters ---------- x - Input array to encode. - max_freq - The maximum frequency of the encoding. - num_bands - The number of frequency bands for the encoding. - Default is 4. - linear - Whether to space the frequency bands linearly as opposed to geometrically. - Default is ``False``. - concat - Whether to concatenate the position, sin and cos values, or return seperately. - Default is ``True``. - flatten - Whether to flatten the position dimension into the batch dimension. - Default is False. + Input array to be re-arranged. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - New array with the final dimension expanded, and the encodings stored in this - channel. + New array with einops.rearrange having been applied. Examples -------- - >>> x = ivy.array([1,2,3]) - >>> y = 1.5 - >>> z = ivy.fourier_encode(x,y) - >>> print(z) - ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00], - [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, - 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, - 1.0000000e+00]]) + With :class:`ivy.Array` instance method: + >>> x = ivy.array([[1, 2, 3], + ... [-4, -5, -6]]) + >>> y = x.einops_rearrange("height width -> width height") + >>> print(y) + ivy.array([[ 1, -4], + [ 2, -5], + [ 3, -6]]) - >>> x = ivy.array([3,10]) - >>> y = 2.5 - >>> z = ivy.fourier_encode(x, y, num_bands=3) - >>> print(z) - ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, - -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], - [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, - 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) - """ - x_in = x - dim = x.shape[-1] - x = ivy.expand_dims(x, axis=-1) - orig_x = x - if linear: - scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) - else: - if ivy.backend == "torch" and isinstance(max_freq, float): - scales = ivy.logspace( - 0.0, - ivy.log(ivy.array(max_freq / 2)) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - else: - scales = ivy.logspace( - 0.0, - ivy.log(max_freq / 2) / math.log(10), - num_bands, - base=10, - device=dev(x), - ) - scales = ivy.astype(scales, ivy.dtype(x)) - scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] - x = x * scales * math.pi - sin_x = ivy.sin(x) - cos_x = ivy.cos(x) - if flatten: - orig_x = x_in - sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) - cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) - if concat: - return ivy.concat([orig_x, sin_x, cos_x], axis=-1) - return sin_x, cos_x + >>> x = ivy.array([[[ 1, 2, 3], + ... [ 4, 5, 6]], + ... [[ 7, 8, 9], + ... [10, 11, 12]]]) + >>> y = x.einops_rearrange("c h w -> c (h w)") + >>> print(y) + ivy.array([[ 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12]]) + >>> x = ivy.array([[1, 2, 3, 4, 5, 6], + ... [7, 8, 9, 10, 11, 12]]) + >>> y = ivy.zeros((4,3)) + >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) + >>> print(y) + ivy.array([[ 1, 2, 3], + [ 4, 5, 6], + [ 7, 8, 9], + [10, 11, 12]]) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def value_is_nan( - x: Union[ivy.Array, ivy.NativeArray, Number], - /, - *, - include_infs: bool = True, -) -> bool: - """ - Determine whether the single valued array or scalar is of nan type. + With :class:`ivy.Container` input: - Parameters - ---------- - x - The input to check Input array. - include_infs - Whether to include infs and -infs in the check. - Default is ``True``. + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> y = ivy.einops_rearrange(x, 'a b -> b a') + >>> print(y) + { + a: ivy.array([[-4.46999979, 3.66000009], + [0.93000001, 24.29000092], + [-3.33999991, 3.6400001]]), + b: ivy.array([[4.96000004, 4.36000013], + [1.51999998, 13.96000004], + [-10.67000008, 0.30000001]]) + } - Returns - ------- - ret - Boolean as to whether the input value is a nan or not. + With varying pattern: - Examples - -------- - >>> x = ivy.array([451]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False + Suppose we have a set of 32 images in "h w c" format (height-width-channel) + and concatenate images along height (vertical axis), 960 = 32 * 30 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') + >>> print(x.shape) + (960, 40, 3) - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - True + # Concatenate images along horizontal axis, 1280 = 32 * 40 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') + >>> print(x.shape) + (30, 1280, 3) - >>> x = ivy.array([float('inf')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - False + # Reorder axes to "b c h w" format for deep learning + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') + >>> print(x.shape) + (32, 3, 30, 40) - >>> x = ivy.array([float('nan')]) - >>> y = ivy.value_is_nan(x, include_infs=False) - >>> print(y) - True + # Flatten each image into a vector, 3600 = 30 * 40 * 3 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') + >>> print(x.shape) + (32, 3600) - >>> x = ivy.array([0]) - >>> y = ivy.value_is_nan(x) - >>> print(y) - False + # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), + # 128 = 32 * 2 * 2 + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', + ... h1=2, w1=2) + >>> print(x.shape) + (128, 15, 20, 3) + + # Space-to-depth operation + >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) + >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, + ... w1=2) + >>> print(x.shape) + (32, 15, 20, 12) """ - x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x - if not x_scalar == x: - return True - if include_infs and (x_scalar == INF or x_scalar == -INF): - return True - return False + ret = einops.rearrange(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret @handle_exceptions @handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@inputs_to_native_arrays @handle_array_function -def has_nans( - x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True -) -> bool: +def einops_reduce( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + reduction: Union[str, Callable], + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: """ - Determine whether the array contains any nans, as well as infs or -infs if - specified. + Perform einops reduce operation on input array x. Parameters ---------- x - Input array. - include_infs - Whether to include ``+infinity`` and ``-infinity`` in the check. - Default is ``True``. + Input array to be reduced. + pattern + Reduction pattern. + reduction + One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Boolean as to whether the array contains nans. + New array with einops.reduce having been applied. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - False + >>> x = ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]) + >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') + >>> print(reduced) + ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) - >>> x = ivy.array([float('nan'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True + With :class:`ivy.Container` input: - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x) - >>> print(y) - True + >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], + ... [3.66, 24.29, 3.64]]), + ... b=ivy.array([[4.96, 1.52, -10.67], + ... [4.36, 13.96, 0.3]])) + >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') + >>> print(reduced) + { + a: ivy.array([-2.29333329, 10.53000069]), + b: ivy.array([-1.39666676, 6.20666695]) + } + """ + ret = einops.reduce(x, pattern, reduction, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret - >>> x = ivy.array([float('inf'), 2, 3]) - >>> y = ivy.has_nans(x, include_infs=False) - >>> print(y) - False + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def einops_repeat( + x: Union[ivy.Array, ivy.NativeArray], + pattern: str, + /, + *, + out: Optional[ivy.Array] = None, + **axes_lengths: Dict[str, int], +) -> ivy.Array: + """ + Perform einops repeat operation on input array x. + + Parameters + ---------- + x + Input array to be repeated. + pattern + Rearrangement pattern. + axes_lengths + Any additional specifications for dimensions. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + New array with einops.repeat having been applied. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4]) + >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) + >>> print(repeated) + ivy.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.has_nans(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[4,5], + ... [1, 3]]), + ... b=ivy.array([[9, 10], + ... [4, 2]])) + >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) + >>> print(repeated) { - a: False, - b: False + a: ivy.array([[4, 5, 4, 5], + [1, 3, 1, 3]]), + b: ivy.array([[9, 10, 9, 10], + [4, 2, 4, 2]]) } """ - return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) + ret = einops.repeat(x._data, pattern, **axes_lengths) + ret = ivy.array(ret, dtype=x.dtype) + if ivy.exists(out): + return ivy.inplace_update(out, ret) + return ret @handle_exceptions @@ -1428,986 +1290,1254 @@ def exists(x: Any) -> bool: @handle_exceptions -def default( - x: Any, +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def fourier_encode( + x: Union[ivy.Array, ivy.NativeArray], + max_freq: Union[float, ivy.Array, ivy.NativeArray], /, - default_val: Any, *, - catch_exceptions: bool = False, - rev: bool = False, - with_callable: bool = False, -) -> Any: + num_bands: int = 4, + linear: bool = False, + concat: bool = True, + flatten: bool = False, +) -> Union[ivy.Array, ivy.NativeArray, Tuple]: """ - Return x provided it exists (is not None), else returns default value. + Pad an array with fourier encodings. Parameters ---------- x - Input which may or may not exist (be None). - default_val - The default value. - catch_exceptions - Whether to catch exceptions from callable x. - Default is ``False``. - rev - Whether to reverse the input x and default_val. - Default is ``False``. - with_callable - Whether either of the arguments might be callable functions. + Input array to encode. + max_freq + The maximum frequency of the encoding. + num_bands + The number of frequency bands for the encoding. + Default is 4. + linear + Whether to space the frequency bands linearly as opposed to geometrically. Default is ``False``. + concat + Whether to concatenate the position, sin and cos values, or return seperately. + Default is ``True``. + flatten + Whether to flatten the position dimension into the batch dimension. + Default is False. Returns ------- ret - x if x exists (is not None), else default. - - Functional Examples - ------------------ - With :code:`Any` input: + New array with the final dimension expanded, and the encodings stored in this + channel. - >>> x = None - >>> y = ivy.default(x, "default_string") - >>> print(y) - default_string + Examples + -------- + >>> x = ivy.array([1,2,3]) + >>> y = 1.5 + >>> z = ivy.fourier_encode(x,y) + >>> print(z) + ivy.array([[ 1.0000000e+00, 1.2246468e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 2.0000000e+00, -2.4492936e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, 1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00], + [ 3.0000000e+00, 3.6739404e-16, 0.0000000e+00, 0.0000000e+00, + 0.0000000e+00, -1.0000000e+00, 1.0000000e+00, 1.0000000e+00, + 1.0000000e+00]]) - >>> x = "" - >>> y = ivy.default(x, "default_string") - >>> print(y) + >>> x = ivy.array([3,10]) + >>> y = 2.5 + >>> z = ivy.fourier_encode(x, y, num_bands=3) + >>> print(z) + ivy.array([[ 3.0000000e+00, 3.6739404e-16, 3.6739404e-16, 3.6739404e-16, + -1.0000000e+00, -1.0000000e+00, -1.0000000e+00], + [ 1.0000000e+01, -1.2246468e-15, -1.2246468e-15, -1.2246468e-15, + 1.0000000e+00, 1.0000000e+00, 1.0000000e+00]]) + """ + x_in = x + dim = x.shape[-1] + x = ivy.expand_dims(x, axis=-1) + orig_x = x + if linear: + scales = ivy.linspace(1.0, max_freq / 2, num_bands, device=dev(x)) + else: + if ivy.backend == "torch" and isinstance(max_freq, float): + scales = ivy.logspace( + 0.0, + ivy.log(ivy.array(max_freq / 2)) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + else: + scales = ivy.logspace( + 0.0, + ivy.log(max_freq / 2) / math.log(10), + num_bands, + base=10, + device=dev(x), + ) + scales = ivy.astype(scales, ivy.dtype(x)) + scales = scales[(*((None,) * (len(x.shape) - len(scales.shape))), Ellipsis)] + x = x * scales * math.pi + sin_x = ivy.sin(x) + cos_x = ivy.cos(x) + if flatten: + orig_x = x_in + sin_x = ivy.reshape(sin_x, [-1, num_bands * dim]) + cos_x = ivy.reshape(cos_x, [-1, num_bands * dim]) + if concat: + return ivy.concat([orig_x, sin_x, cos_x], axis=-1) + return sin_x, cos_x - >>> x = ivy.array([4, 5, 6]) - >>> y = ivy.default(x, ivy.array([1, 2, 3]), rev=True) - >>> print(y) - ivy.array([1, 2, 3]) - >>> x = lambda: ivy.array([1, 2, 3]) - >>> y = ivy.default(x, ivy.array([4, 5, 6]), with_callable=True) - >>> print(y) - ivy.array([1, 2, 3]) +@handle_exceptions +@handle_nestable +def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: + """ + Return the supported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the supported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda: None - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True) - >>> print(y) - ivy.array([1, 2, 3]) - - >>> x = lambda a, b: a + b - >>> y = ivy.default(x, lambda: ivy.array([1, 2, 3]), with_callable=True, - ... catch_exceptions=True, rev=True) - >>> print(y) - ivy.array([1, 2, 3]) - """ - with_callable = catch_exceptions or with_callable - if rev: - tmp = x - x = default_val - default_val = tmp - if with_callable: - x_callable = callable(x) - default_callable = callable(default_val) - else: - x_callable = False - default_callable = False - if catch_exceptions: - # noinspection PyBroadException - try: - x = x() if x_callable else x - except Exception: - return default_val() if default_callable else default_val - else: - x = x() if x_callable else x - return x if exists(x) else default_val() if default_callable else default_val - - -@handle_exceptions -def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: - """ - Return the input shape in ivy.Shape form. - - Parameters - ---------- - shape - The input to be converted + Parameters + ---------- + fn + The function to check for the supported device and dtype attribute + recurse + Whether to recurse into used ivy functions. + Default is ``True``. Returns ------- - ret - the input in ivy.Shape form + ret + Tuple or dict containing the supported devices and dtypes of the function """ - if isinstance(shape, ivy.Shape): - return shape - return ivy.Shape(shape) + ivy.utils.assertions.check_true( + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", + ) + + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_supported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=False), + } + else: + supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) + if recurse: + supported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + supported_devices_dtypes, + _dnd_dict_intersection, + function_supported_devices_and_dtypes, + wrapper=lambda x: x, + ) + + return supported_devices_dtypes @handle_exceptions -def to_native_shape( - shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] -) -> ivy.NativeShape: +@handle_nestable +def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: """ - Return the input shape in its native backend framework form. + Return the unsupported combination of devices and dtypes of the current backend's + function. The function returns a dict containing the unsupported combination of + devices and dtypes of the primary and compositional implementations incase of + partial mixed functions. Parameters ---------- - shape - The input to be converted + fn + The function to check for the unsupported device and dtype attribute + recurse + Whether to recurse into used ivy functions. + Default is ``True``. Returns ------- - ret - the input in its native framework form + ret + Tuple or dict containing the unsupported devices and dtypes of the function """ - native_shape_type = (ivy.NativeShape,) - if ivy.current_backend_str() == "torch": - native_shape_type += (tuple,) - if len(backend_stack) != 0 and isinstance(shape, native_shape_type): - return shape - ivy.utils.assertions.check_isinstance( - shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) - ) - if isinstance(shape, int): - shape = (shape,) - elif isinstance(shape, list): - shape = tuple(shape) - elif is_array(shape): - shape = ivy.to_numpy(shape).tolist() - elif isinstance(shape, ivy.Shape): - shape = shape.shape - ivy.utils.assertions.check_all( - [isinstance(v, int) for v in shape if not is_array(v)], - "shape must take integers only", - as_array=False, - ) ivy.utils.assertions.check_true( - not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" + _is_valid_device_and_dtypes_attributes(fn), + "supported_device_and_dtypes and unsupported_device_and_dtypes " + "attributes cannot both exist in a particular backend", ) - return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) + if hasattr(fn, "partial_mixed_handler"): + return { + "compositional": function_unsupported_devices_and_dtypes( + fn.compos, recurse=recurse + ), + "primary": _get_devices_and_dtypes(fn, complement=True), + } + else: + unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) + if recurse: + unsupported_devices_dtypes = ivy.functional.data_type._nested_get( + fn, + unsupported_devices_dtypes, + _dnd_dict_union, + function_unsupported_devices_and_dtypes, + wrapper=lambda x: x, + ) + return unsupported_devices_dtypes @handle_exceptions +@handle_backend_invalid @handle_nestable -def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def gather( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + batch_dims: int = 0, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Try and return the function, otherwise return None if an exception was raised during - function execution. + Gather slices from params at axis according to indices. Parameters ---------- - fn - Function to try and call and return. - args - list of arguments. - kwargs - dictionay of keyword arguments + params + The array from which to gather values. + indices + The array which indicates the indices that will be gathered along + the specified axis. + axis + optional int, the axis from which to gather from. + Default is ``-1``. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional array, for writing the result to. It must have a shape + that the inputs broadcast to. Returns ------- - Either the function itself or None if an exception was raised - during function execution. + ret + New array with the values gathered at the specified indices along the + specified axis. + + + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - with a function that is executed without any exception: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([4, 5, 6]) - >>> z = ivy.try_else_none(ivy.add, x, y) - >>> print(z.__name__) - add + With :class:`ivy.Array` input: - with a function that is executed with an exception: + >>> x = ivy.array([0., 1., 2.]) + >>> y = ivy.array([1, 2]) + >>> print(ivy.gather(x, y)) + ivy.array([1., 2.]) - >>> x = ivy.array([1, 2, 3]) - >>> y = 'hemant' - >>> z = ivy.try_else_none(ivy.add,x, y) + >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) + >>> y = ivy.array([[0, 1],[1, 2]]) + >>> z = ivy.zeros((2, 2, 2)) + >>> ivy.gather(x, y, out=z) >>> print(z) - None - """ - try: - _ = fn(*args, **kwargs) - return fn - except Exception: - return None - + ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) -@handle_exceptions -def arg_names(receiver): - """ - Get the expected keyword arguments for a function or class constructor. + >>> x = ivy.array([[[0., 1.], [2., 3.]], + ... [[8., 9.], [10., 11.]]]) + >>> y = ivy.array([[0, 1]]) + >>> z = ivy.zeros((1, 2, 2, 2)) + >>> ivy.gather(x, y, axis=0, out=z) + >>> print(z) + ivy.array( + [[[[ 0., 1.], + [ 2., 3.]], + [[ 8., 9.], + [10., 11.]]]]) - Parameters - ---------- - receiver - Function or class constructor + >>> x = ivy.array([[0, 10, 20, 0, 0], + ... [0, 0, 0, 30, 40], + ... [0, 10, 0, 0, 40]]) + >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) + >>> z = ivy.gather(x, y, batch_dims=1) + >>> print(z) + ivy.array([[10, 20], [30, 40],[10, 40]]) - Returns - ------- - ret - List containing the keyword arguments' names for a function or class constructor + With :class:`ivy.Container` input: - Examples - -------- - >>> x = ivy.arg_names(ivy.tan) - >>> print(x) - ['x', 'out'] + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.Container(a = ivy.array([0, 1]), + ... b = ivy.array([1, 2])) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([5., 6.]) + } - >>> x = ivy.arg_names(ivy.optimizers.Adam) - >>> print(x) - ['lr', 'beta1', 'beta2', 'epsilon', 'inplace', - 'stop_gradients', 'compile_on_next_step', 'device'] + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), + ... b = ivy.array([4., 5., 6.])) + >>> y = ivy.array([0, 1]) + >>> print(ivy.gather(x, y)) + { + a: ivy.array([0., 1.]), + b: ivy.array([4., 5.]) + } """ - return list(inspect.signature(receiver).parameters.keys()) + return current_backend(params, indices).gather( + params, indices, axis=axis, batch_dims=batch_dims, out=out + ) @handle_exceptions -def match_kwargs( - kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False -) -> Union[List[Dict], Dict]: +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def gather_nd( + params: Union[ivy.Array, ivy.NativeArray], + indices: Union[ivy.Array, ivy.NativeArray], + /, + *, + batch_dims: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Match keyword arguments to either class or function receivers. + Gather slices from params into a array with shape specified by indices. Parameters ---------- - kwargs - Keyword arguments to match. - receivers - Functions and/or classes to match the keyword arguments to. - allow_duplicates - Whether to allow one keyword argument to be used for multiple receivers. - Default is ``False``. + params + The array from which to gather values. + indices + Index array. + batch_dims + optional int, lets you gather different items from each element of a batch. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - Sequence of keyword arguments split as best as possible. + New array of given shape, with the values gathered at the indices. Examples -------- - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) - >>> print(x) - [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] - - >>> o = ivy.zeros(3) - >>> kwargs = {'out': o, 'bias': ivy.arange(3)} - >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) - >>> print(x) - [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] - """ - split_kwargs = list() - for receiver in receivers: - expected_kwargs = arg_names(receiver) - found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} - if not allow_duplicates: - for k in found_kwargs.keys(): - del kwargs[k] - split_kwargs.append(found_kwargs) - if len(split_kwargs) == 1: - return split_kwargs[0] - return split_kwargs - - -@handle_exceptions -def cache_fn(func: Callable) -> Callable: - """ - Cache function outputs. - - A decorator to wrap a function, such that computed outputs are cached to avoid - recalculating them later. + With :class:`ivy.Array` input: - Parameters - ---------- - func - The function to wrap, whose output should be cached for later. + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) + ivy.array(1.) - Returns - ------- - ret - The newly cache wrapped function. + >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) + >>> y = ivy.array([[0],[1],[1]], dtype='int32') + >>> z = ivy.gather_nd(x,y,batch_dims=1) + ivy.array([0., 3., 5.]) - Examples - -------- - With positional arguments only: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> def my_sum(val1:float, val2:float)->float: return val1 + val2 - >>> cached_sum = ivy.cache_fn(my_sum) - >>> print(cached_sum(3, 5)) - 8 + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) + >>> y = ivy.array([1]) + >>> print(ivy.gather_nd(x, y)) + { + a: ivy.array(1.), + b: ivy.array(5.) + } - With keyword arguments: + With :class:`ivy.Container` input: - >>> def line_eq(x:float, /, *, slp:float=2, itc:float=0)->float: return x*slp+itc - >>> cached_line_eq = ivy.cache_fn(line_eq) - >>> print(cached_line_eq(3, itc=5, slp=2)) - 11 + >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), + ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) + >>> y = ivy.Container(a=ivy.array([1,0]), + ... b=ivy.array([0])) + >>> print(ivy.gather_nd(x, y)) + { + a: ivy.array(30.), + b: ivy.array([0., 100., 200.]) + } """ - global FN_CACHE - if func not in FN_CACHE: - FN_CACHE[func] = dict() - - @wraps(func) - def cached_fn(*args, **kwargs): - key = "".join( - [str(i) + ", " for i in args] - + [" kw, "] - + [str(i) + ", " for i in sorted(kwargs.items())] - ) - cache = FN_CACHE[func] - if key in cache: - return cache[key] - ret = func(*args, **kwargs) - cache[key] = ret - return ret - - return cached_fn + res = current_backend(params, indices).gather_nd( + params, indices, batch_dims=batch_dims + ) + if ivy.exists(out): + return ivy.inplace_update(out, res) + return res @handle_exceptions -def current_backend_str() -> Union[str, None]: +def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: """ - Return framework string. + Get all arrays which are currently alive. Returns ------- ret - The framework string. + All arrays which are alive. + + Examples + -------- + >>> ivy.get_all_arrays_in_memory() + [] + >>> x = ivy.get_all_arrays_in_memory() + >>> x + [] + >>> y = ivy.array([0, 1, 2]) + >>> x + [ivy.array([0, 1, 2])] """ - fw = current_backend() - if not backend_stack: - return "" - return fw.current_backend_str() + all_arrays = list() + for obj in gc.get_objects(): + try: + if ivy.current_backend_str() in ["", "numpy"]: + if ivy.is_ivy_array(obj): + all_arrays.append(obj) + else: + if ivy.is_native_array(obj): + all_arrays.append(obj) + + except Exception: + pass + return all_arrays -@handle_exceptions @handle_nestable -@handle_array_like_without_promotion +@handle_partial_mixed_function +@handle_view_indexing @inputs_to_ivy_arrays @handle_array_function -def einops_rearrange( +@handle_device_shifting +def get_item( x: Union[ivy.Array, ivy.NativeArray], - pattern: str, /, + query: Union[ivy.Array, ivy.NativeArray, Tuple], *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], + copy: Optional[bool] = None, ) -> ivy.Array: """ - Perform einops rearrange operation on input array x. + Gather slices from x according to query array, identical to x[query]. Parameters ---------- x - Input array to be re-arranged. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + array, the array from which to gather values. + query + array, index array, integer indices or boolean mask. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. Returns ------- ret - New array with einops.rearrange having been applied. + New array with the values gathered at the specified indices. - Examples - -------- - With :class:`ivy.Array` instance method: + Functional Examples + ------------------- - >>> x = ivy.array([[1, 2, 3], - ... [-4, -5, -6]]) - >>> y = x.einops_rearrange("height width -> width height") - >>> print(y) - ivy.array([[ 1, -4], - [ 2, -5], - [ 3, -6]]) + >>> x = ivy.array([0, -1, 20]) + >>> query = ivy.array([0, 1]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 0, -1]) - >>> x = ivy.array([[[ 1, 2, 3], - ... [ 4, 5, 6]], - ... [[ 7, 8, 9], - ... [10, 11, 12]]]) - >>> y = x.einops_rearrange("c h w -> c (h w)") - >>> print(y) - ivy.array([[ 1, 2, 3, 4, 5, 6], - [ 7, 8, 9, 10, 11, 12]]) - - >>> x = ivy.array([[1, 2, 3, 4, 5, 6], - ... [7, 8, 9, 10, 11, 12]]) - >>> y = ivy.zeros((4,3)) - >>> x.einops_rearrange("c (h w) -> (c h) w", out=y, h=2, w=3) - >>> print(y) - ivy.array([[ 1, 2, 3], - [ 4, 5, 6], - [ 7, 8, 9], - [10, 11, 12]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> y = ivy.einops_rearrange(x, 'a b -> b a') - >>> print(y) - { - a: ivy.array([[-4.46999979, 3.66000009], - [0.93000001, 24.29000092], - [-3.33999991, 3.6400001]]), - b: ivy.array([[4.96000004, 4.36000013], - [1.51999998, 13.96000004], - [-10.67000008, 0.30000001]]) - } - - With varying pattern: - - Suppose we have a set of 32 images in "h w c" format (height-width-channel) - and concatenate images along height (vertical axis), 960 = 32 * 30 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> (b h) w c') - >>> print(x.shape) - (960, 40, 3) - - # Concatenate images along horizontal axis, 1280 = 32 * 40 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> h (b w) c') - >>> print(x.shape) - (30, 1280, 3) - - # Reorder axes to "b c h w" format for deep learning - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b c h w') - >>> print(x.shape) - (32, 3, 30, 40) - - # Flatten each image into a vector, 3600 = 30 * 40 * 3 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b h w c -> b (c h w)') - >>> print(x.shape) - (32, 3600) - - # Split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), - # 128 = 32 * 2 * 2 - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', - ... h1=2, w1=2) - >>> print(x.shape) - (128, 15, 20, 3) - - # Space-to-depth operation - >>> images = ivy.asarray([ivy.random_normal(shape=(30, 40, 3)) for _ in range(32)]) - >>> x = ivy.einops_rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, - ... w1=2) - >>> print(x.shape) - (32, 15, 20, 12) + >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) + >>> query = ivy.array([[True, False], [False, False], [True, True]]) + >>> print(ivy.get_item(x, query)) + ivy.array([ 4, -2, -10]) """ - ret = einops.rearrange(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) + if ivy.is_array(query) and ivy.is_bool_dtype(query): + if not len(query.shape): + if not query: + return ivy.array([], shape=(0,), dtype=x.dtype) + return ivy.expand_dims(x, axis=0) + query = ivy.nonzero(query, as_tuple=False) + ret = ivy.gather_nd(x, query) + else: + indices, target_shape = _parse_query(query, x.shape) + if indices is None: + return ivy.empty(target_shape, dtype=x.dtype) + ret = ivy.gather_nd(x, indices) + ret = ivy.reshape(ret, target_shape) return ret -@handle_exceptions +@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@inputs_to_native_arrays +@to_native_arrays_and_back @handle_array_function -def einops_reduce( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - reduction: Union[str, Callable], - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: +@handle_device_shifting +def get_num_dims( + x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False +) -> int: """ - Perform einops reduce operation on input array x. + Return the number of dimensions of the array x. Parameters ---------- x - Input array to be reduced. - pattern - Reduction pattern. - reduction - One of available reductions ('min', 'max', 'sum', 'mean', 'prod'), or callable. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array to infer the number of dimensions for. + as_array + Whether to return the shape as a array, default False. Returns ------- ret - New array with einops.reduce having been applied. + Shape of the array - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]) - >>> reduced = ivy.einops_reduce(x, 'a b -> b', 'mean') - >>> print(reduced) - ivy.array([-0.40499985, 12.61000061, 0.1500001 ]) + >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], + ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) + >>> b = ivy.get_num_dims(a, as_array=False) + >>> print(b) + 3 With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[-4.47, 0.93, -3.34], - ... [3.66, 24.29, 3.64]]), - ... b=ivy.array([[4.96, 1.52, -10.67], - ... [4.36, 13.96, 0.3]])) - >>> reduced = ivy.einops_reduce(x, 'a b -> a', 'mean') - >>> print(reduced) + >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) + >>> print(ivy.get_num_dims(a)) { - a: ivy.array([-2.29333329, 10.53000069]), - b: ivy.array([-1.39666676, 6.20666695]) + b: 2 + } + + >>> b = ivy.get_num_dims(a, as_array=True) + >>> print(b) + { + b: ivy.array(2) } """ - ret = einops.reduce(x, pattern, reduction, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + return current_backend(x).get_num_dims(x, as_array=as_array) -# IMPORTANT: assign attribute directly to function instead of wrapper here -einops_reduce.unsupported_dtypes = { - "torch": ("float16",), - "tensorflow": ("complex",), - "paddle": ("complex", "uint8", "int8", "int16", "float16"), -} +def get_referrers_recursive( + item, depth=0, max_depth=None, seen_set=None, local_set=None +): + """ + Summary. + + Parameters + ---------- + item + + depth + (Default value = 0) + max_depth + (Default value = None) + seen_set + (Default value = None) + local_set + (Default value = None`) + """ + seen_set = ivy.default(seen_set, set()) + local_set = ivy.default(local_set, set()) + ret_cont = ivy.Container( + repr=str(item).replace(" ", ""), + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + referrers = [ + ref + for ref in gc.get_referrers(item) + if not ( + isinstance(ref, dict) + and min([k in ref for k in ["depth", "max_depth", "seen_set", "local_set"]]) + ) + ] + local_set.add(str(id(referrers))) + for ref in referrers: + ref_id = str(id(ref)) + if ref_id in local_set or hasattr(ref, "cell_contents"): + continue + seen = ref_id in seen_set + seen_set.add(ref_id) + refs_rec = lambda: get_referrers_recursive( + ref, depth + 1, max_depth, seen_set, local_set + ) + this_repr = "tracked" if seen else str(ref).replace(" ", "") + if not seen and (not max_depth or depth < max_depth): + val = ivy.Container( + repr=this_repr, + alphabetical_keys=False, + keyword_color_dict={"repr": "magenta"}, + ) + refs = refs_rec() + for k, v in refs.items(): + val[k] = v + else: + val = this_repr + ret_cont[str(ref_id)] = val + return ret_cont @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def einops_repeat( - x: Union[ivy.Array, ivy.NativeArray], - pattern: str, - /, - *, - out: Optional[ivy.Array] = None, - **axes_lengths: Dict[str, int], -) -> ivy.Array: +def has_nans( + x: Union[ivy.Array, ivy.NativeArray], /, *, include_infs: bool = True +) -> bool: """ - Perform einops repeat operation on input array x. + Determine whether the array contains any nans, as well as infs or -infs if + specified. Parameters ---------- x - Input array to be repeated. - pattern - Rearrangement pattern. - axes_lengths - Any additional specifications for dimensions. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + Input array. + include_infs + Whether to include ``+infinity`` and ``-infinity`` in the check. + Default is ``True``. Returns ------- ret - New array with einops.repeat having been applied. + Boolean as to whether the array contains nans. - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3, 4]) - >>> repeated = ivy.einops_repeat(x, 'a -> b a', b=2) - >>> print(repeated) - ivy.array([[1, 2, 3, 4], - [1, 2, 3, 4]]) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + False + + >>> x = ivy.array([float('nan'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf'), 2, 3]) + >>> y = ivy.has_nans(x, include_infs=False) + >>> print(y) + False With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[4,5], - ... [1, 3]]), - ... b=ivy.array([[9, 10], - ... [4, 2]])) - >>> repeated = ivy.einops_repeat(x, 'h w -> h (c w)', c=2) - >>> print(repeated) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.has_nans(x) + >>> print(y) { - a: ivy.array([[4, 5, 4, 5], - [1, 3, 1, 3]]), - b: ivy.array([[9, 10, 9, 10], - [4, 2, 4, 2]]) + a: False, + b: False } """ - ret = einops.repeat(x._data, pattern, **axes_lengths) - ret = ivy.array(ret, dtype=x.dtype) - if ivy.exists(out): - return ivy.inplace_update(out, ret) - return ret + return ivy.value_is_nan(ivy.sum(x), include_infs=include_infs) -ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 +@handle_exceptions +def inplace_arrays_supported() -> bool: + """ + Determine whether inplace arrays are supported for the current backend framework. + + Returns + ------- + ret + Boolean, whether or not inplace arrays are supported. + """ + return current_backend().inplace_arrays_supported() @handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays @handle_array_function -def set_min_denominator(val: float) -> None: +@handle_device_shifting +def inplace_decrement( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Set the global minimum denominator used by ivy for numerically stable division. + Perform in-place decrement for the input array. Parameters ---------- + x + The input array to be decremented by the defined value. val - The value to set the global minimum denominator to. - - Examples - -------- - >>> x = ivy.min_denominator - >>> print(x) - 1e-12 + The value of decrement. - >>> ivy.set_min_denominator(1e-13) - >>> y = ivy.min_denominator - >>> print(y) - 1e-13 - """ - global min_denominator_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_denominator_stack.append(val) - ivy.__setattr__("min_denominator", val, True) + Returns + ------- + ret + The array following the in-place decrement. -@handle_exceptions -def unset_min_denominator() -> None: - """ - Reset the global minimum denominator used by ivy for numerically stable division to - the previous value. + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - >>> ivy.set_min_denominator(1e-10) - >>> y = ivy.min_denominator + With :class:`ivy.Array` input: + + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_decrement(x, 1.25) >>> print(y) - 1e-10 + ivy.array([[ 4.05, 5.75, -1.25], + [ 5.55, 6.75, 2.65], + [-1.25, 8.75, 5.05]]) - >>> ivy.unset_min_denominator() - >>> ivy.min_denominator - 1e-12 - """ - global min_denominator_stack - if min_denominator_stack: - min_denominator_stack.pop(-1) - val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 - ivy.__setattr__("min_denominator", val, True) + With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) + >>> y = ivy.inplace_decrement(x, 1.5) + >>> print(y) + { + a: ivy.array([-1., -6.5, 28.5]), + b: ivy.array([-1.5, -26.5, 48.5]) + } -ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([0., 0., 0.]), + b: ivy.array([0., 0., 0.]) + } + + >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) + >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) + >>> z = ivy.inplace_decrement(x, y) + >>> print(z) + { + a: ivy.array([1., 1.5, 3.]), + b: ivy.array([0., 50., 3.5]) + } + """ + return current_backend(x).inplace_decrement(x, val) @handle_exceptions +@handle_backend_invalid +@handle_nestable +@inputs_to_ivy_arrays @handle_array_function -def set_min_base(val: float) -> None: +@handle_device_shifting +def inplace_increment( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], +) -> ivy.Array: """ - Set the global minimum base used by ivy for numerically stable power raising. + Perform in-place increment for the input array. Parameters ---------- + x + The input array to be incremented by the defined value. val - The new value to set the minimum base to. + The value of increment. + + Returns + ------- + ret + The array following the in-place increment. Examples -------- - >>> x = ivy.min_base - >>> print(x) - 1e-05 + With :class:`ivy.Array` input: - >>> ivy.set_min_base(1e-04) - >>> y = ivy.min_base + >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) + >>> y = ivy.inplace_increment(x, 3.) >>> print(y) - 1e-04 - """ - global min_base_stack - ivy.utils.assertions.check_isinstance(val, (int, float)) - min_base_stack.append(val) - ivy.__setattr__("min_base", val, True) - + ivy.array([[ 8.3, 10., 3.], + [ 9.8, 11., 6.9], + [ 3., 13., 9.3]]) -@handle_exceptions -def unset_min_base() -> None: - """ - Reset the global minimum base used by ivy for numerically stable power raising to - the previous value. + With :class:`ivy.Container` input: - Examples - -------- - >>> ivy.set_min_base(1e-07) - >>> y = ivy.min_base + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.inplace_increment(x, 2.5) >>> print(y) - 1e-07 + { + a: ivy.array([2.5, 17.5, 32.5]), + b: ivy.array([2.5, 27.5, 52.5]) + } - >>> ivy.unset_min_base() - >>> ivy.min_base - 1e-05 + + >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) + >>> z = ivy.inplace_increment(x, y) + >>> print(z) + { + a: ivy.array([0., 30., 60.]), + b: ivy.array([0., 50., 100.]) + } """ - global min_base_stack - if min_base_stack: - min_base_stack.pop(-1) - val = min_base_stack[-1] if min_base_stack else 1e-05 - ivy.__setattr__("min_base", val, True) + return current_backend(x).inplace_increment(x, val) @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def stable_divide( - numerator: Union[Number, ivy.Array, ivy.NativeArray], - denominator: Union[Number, ivy.Array, ivy.NativeArray], +# @handle_device_shifting +def inplace_update( + x: Union[ivy.Array, ivy.NativeArray], + val: Union[ivy.Array, ivy.NativeArray], /, *, - min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, -) -> Union[Number, ivy.Array]: + ensure_in_backend: bool = False, + keep_input_dtype: bool = False, +) -> ivy.Array: """ - Divide the numerator by the denominator, with min denominator added to the - denominator for numerical stability. + Perform in-place update for the input array. + + This will always be performed on ivy.Array instances pass in the input, and will + also be performed on the native array classes in the backend when the backend + supports this. If the backend does not natively support inplace updates, and x is an + ivy.NativeArray instance, then an + exception will be thrown. Parameters ---------- - numerator - The numerator of the division. - denominator - The denominator of the division. - min_denominator - The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) - by default. + x + The variable to update. + val + The array to update the variable with. + ensure_in_backend + Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. + In cases where it should be, backends which do not natively support inplace + updates will raise an exception. + keep_input_dtype + Whether or not to preserve `x` data type after the update, otherwise `val` + data type will be applied. Defaults to False. Returns ------- ret - The new item following the numerically stable division. + The array following the in-place update. + + Raises + ------ + IvyException + If backend set doesn't natively support inplace updates and ensure_in_backend is + True, above exception will be raised. + + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the arguments. Examples -------- - With :code:`int` input: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x = ivy.stable_divide(1, 2) + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([0]) + >>> ivy.inplace_update(x, y) >>> print(x) - 0.49999999999975 + ivy.array([0]) - >>> x = ivy.stable_divide(1, 4, min_denominator=1) + With :class:`ivy.Array` input and default backend set as `numpy`: + + >>> ivy.set_backend("numpy") + >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) + >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) + >>> ivy.inplace_update(x, y, keep_input_dtype=True) >>> print(x) - 0.2 + ivy.array([0., 0., 0.]) - With float input: + With :class:`ivy.Container` instances:, and backend set as `torch`: - >>> x = ivy.stable_divide(5.0, 3.33) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) + >>> ivy.inplace_update(x, y) >>> print(x) - 1.5015015015010504 + { + a: ivy.array([1, 1]), + b: ivy.array([2, 2]) + } - With :code:`complex` input: + With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend + set as `torch`: - >>> x = ivy.stable_divide(1+1j, 1-1j) + >>> ivy.set_backend("torch") + >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) + >>> y = ivy.array([1, 2]) + >>> ivy.inplace_update(x, y) >>> print(x) - (5.000444502911705e-13+0.9999999999995j) + { + a: ivy.array([1, 2]), + b: ivy.array([1, 2]) + } + """ + return current_backend(x).inplace_update( + x, + val, + ensure_in_backend=ensure_in_backend, + keep_input_dtype=keep_input_dtype, + ) - With :class:`ivy.Array` input: - >>> x = ivy.asarray([[10., 20., 30.], - ... [40., 50., 60.]]) - >>> y = ivy.stable_divide(x, 10.) - >>> print(y) - ivy.array([[1., 2., 3.], - [4., 5., 6.]]) +@handle_exceptions +def inplace_variables_supported() -> bool: + """ + Determine whether inplace variables are supported for the current backend framework. + Returns + ------- + ret + Boolean, whether or not inplace variables are supported. + """ + return current_backend().inplace_variables_supported() - >>> x = ivy.asarray([1,2,3]) - >>> y = np.array((1., 3., 5.)) - >>> z = ivy.stable_divide(x, y) - >>> print(z) - ivy.array([1. , 0.667, 0.6 ]) - >>> x = ivy.asarray([1., 2., 4.]) - >>> y = ivy.asarray([1., 0.5, 0.25]) - >>> z = ivy.asarray([0.01, 0.02, 0.03]) - >>> w = ivy.stable_divide(x, y, min_denominator=z) - >>> print(w) - ivy.array([ 0.99, 3.85, 14.3 ]) +@handle_exceptions +@handle_backend_invalid +def is_array(x: Any, /, *, exclusive: bool = False) -> bool: + """ + Determine whether the input x is either an Ivy Array or a Native Array. - With :class:`ivy.Container` input: + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. - >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) - >>> y = ivy.stable_divide(x, 0.5) - >>> print(y) - { - a: ivy.array([20., 30.]), - b: ivy.array([40., 50.]) - } + Returns + ------- + ret + Boolean, whether or not x is an array. + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> print(ivy.is_array(x)) + True - >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) - >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) - >>> z = ivy.stable_divide(x, y) - >>> print(z) - { - a: ivy.array([2., 0.8]), - b: ivy.array([0.857, 10.]) - } + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> print(ivy.is_array(x, exclusive=True)) + True + + >>> x = [2, 3] + >>> print(ivy.is_array(x)) + False """ - return numerator / (denominator + default(min_denominator, ivy.min_denominator)) + return ivy.is_ivy_array(x, exclusive=exclusive) or ivy.is_native_array( + x, exclusive=exclusive + ) @handle_exceptions -@handle_nestable -@inputs_to_ivy_arrays -@handle_array_function -def stable_pow( - base: Union[Number, ivy.Array, ivy.NativeArray], - exponent: Union[Number, ivy.Array, ivy.NativeArray], - /, - *, - min_base: float = None, -) -> Any: +@handle_backend_invalid +def is_ivy_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: Optional[bool] = False +) -> bool: """ - Raise the base by the power, with ivy.min_base added to the base when exponent > 1 - for numerical stability. + Determine whether the input x is a valid Ivy Array. Parameters ---------- - base - The base number. - exponent - The exponent number. - min_base - The minimum base to use, use global ivy.min_base by default. + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. Returns ------- ret - The new item following the numerically stable power. + Boolean, whether or not x is a valid Ivy Array. Examples -------- - With :code:`int` input: + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_ivy_array(x) + True - >>> x = ivy.stable_pow(2, 2) - >>> print(x) - ivy.array(4.00004) + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_ivy_array(x, exclusive=True) + False + """ + return isinstance(x, ivy.Array) and ivy.is_native_array(x.data, exclusive=exclusive) - >>> x = ivy.stable_pow(2, 2, min_base=2) - >>> print(x) - ivy.array(16) - With float input: +@handle_exceptions +def is_ivy_container(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Container. - >>> x = ivy.stable_pow(4.0, .5) - >>> print(x) - ivy.array(2.00000262) + Parameters + ---------- + x + The input to check - With :code:`complex` input: + Returns + ------- + ret + Boolean, whether or not x is an ivy container. - >>> x = ivy.stable_pow(3+4j, 2j) - >>> print(x) - ivy.array(-0.15605032-0.01208451j) + Examples + -------- + >>> x = ivy.Container() + >>> print(ivy.is_ivy_container(x)) + True - With :class:`ivy.Array` input: + >>> x = [2, 3] + >>> print(ivy.is_ivy_container(x)) + False + """ + return isinstance(x, ivy.Container) - >>> x = ivy.asarray([[2, 4], - ... [6, 8]]) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) - ivy.array([[ 4.00004, 16.00008], - [36.00012, 64.00016]]) - >>> x = ivy.asarray([2, 4, 6]) - >>> y = ivy.asarray([2, 3, 4]) - >>> z = ivy.stable_pow(x, y) - >>> print(z) - ivy.array([ 4.00004, 64.00048, 1296.00864]) +def is_ivy_nested_array(x: Any, /) -> bool: + """ + Determine whether the input x is an Ivy Nested Array. - With :class:`ivy.Container` input: + Parameters + ---------- + x + The input to check + Returns + ------- + ret + Boolean, whether or not x is an ivy nested array. + """ + return isinstance(x, ivy.NestedArray) - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.stable_pow(x, 2) - >>> print(y) - { - a: ivy.array([4.00004, 16.00008]), - b: ivy.array([36.00012, 64.00016]) - } - >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) - >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) - >>> z = ivy.stable_pow(x, y) - >>> print(z) - { - a: ivy.array([2.00001, 64.00048]), - b: ivy.array([1296.00864, 32768.2048]) - } +@handle_exceptions +@handle_backend_invalid +def is_native_array( + x: Union[ivy.Array, ivy.NativeArray], /, *, exclusive: bool = False +) -> bool: """ - return_dtype = ivy.promote_types( - ivy.default_dtype(item=base), - ivy.default_dtype(item=default(min_base, ivy.min_base)), + Determine whether the input x is an :class:`ivy.NativeArray` instance. + + Parameters + ---------- + x + The input to check + exclusive + Whether to check if the data type is exclusively an array, rather than a + variable or traced array. + + Returns + ------- + ret + Boolean, whether or not x is an :class:`ivy.NativeArray`. + + Examples + -------- + >>> x = ivy.array([0, 1, 2]) + >>> ivy.is_native_array(x) + False + + >>> x = ivy.native_array([9.1, -8.3, 2.8, 3.0]) + >>> ivy.is_native_array(x, exclusive=True) + True + """ + try: + return current_backend(x).is_native_array(x, exclusive=exclusive) + except ValueError: + return False + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@to_native_arrays_and_back +@handle_device_shifting +def isin( + elements: Union[ivy.Array, ivy.NativeArray], + test_elements: Union[ivy.Array, ivy.NativeArray], + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> ivy.Array: + """ + Test if each element of elements is in test_elements. + + Parameters + ---------- + elements + input array + test_elements + values against which to test for each input element + assume_unique + If True, assumes both elements and test_elements contain unique elements, + which can speed up the calculation. Default value is False. + invert + If True, inverts the boolean return array, resulting in True values for + elements not in test_elements. Default value is False. + + Returns + ------- + ret + output a boolean array of the same shape as elements that is True for elements + in test_elements and False otherwise. + + Examples + -------- + >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y) + ivy.array([[False, False, False], [ True, True, True]]) + + >>> x = ivy.array([3, 2, 1, 0]) + >>> y = ivy.array([1, 2, 3]) + >>> ivy.isin(x, y, invert=True) + ivy.array([False, False, False, True]) + """ + return ivy.current_backend(elements, test_elements).isin( + elements, test_elements, assume_unique=assume_unique, invert=invert ) - return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) - ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) - return ret.astype(return_dtype) -stable_pow.unsupported_dtypes = ("bfloat16",) +@handle_exceptions +@handle_nestable +def isscalar(x: Any, /) -> bool: + return np.isscalar(x) @handle_exceptions -def get_all_arrays_in_memory() -> List[Union[ivy.Array, ivy.NativeArray]]: +@handle_backend_invalid +@handle_nestable +@inputs_to_native_arrays +@handle_device_shifting +def itemsize( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> int: """ - Get all arrays which are currently alive. + Return the size of the input array's elements. + + Parameters + ---------- + x + The input array. Returns ------- ret - All arrays which are alive. + An integer specifying the element size in bytes. Examples -------- - >>> ivy.get_all_arrays_in_memory() - [] - >>> x = ivy.get_all_arrays_in_memory() - >>> x - [] - >>> y = ivy.array([0, 1, 2]) - >>> x - [ivy.array([0, 1, 2])] + >>> x = ivy.array([1,2,3], dtype=ivy.float64) + >>> ivy.itemsize(x) + 8 + + >>> x = ivy.array([1,2,3], dtype=ivy.complex128) + >>> ivy.itemsize(x) + 16 """ - all_arrays = list() - for obj in gc.get_objects(): - try: - if ivy.current_backend_str() in ["", "numpy"]: - if ivy.is_ivy_array(obj): - all_arrays.append(obj) - else: - if ivy.is_native_array(obj): - all_arrays.append(obj) + return ivy.current_backend(x).itemsize(x) - except Exception: - pass - return all_arrays + +@handle_exceptions +def match_kwargs( + kwargs: Dict, *receivers: Iterable[Callable], allow_duplicates: bool = False +) -> Union[List[Dict], Dict]: + """ + Match keyword arguments to either class or function receivers. + + Parameters + ---------- + kwargs + Keyword arguments to match. + receivers + Functions and/or classes to match the keyword arguments to. + allow_duplicates + Whether to allow one keyword argument to be used for multiple receivers. + Default is ``False``. + + Returns + ------- + ret + Sequence of keyword arguments split as best as possible. + + Examples + -------- + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.add, ivy.linear) + >>> print(x) + [{'out': ivy.array([0., 0., 0.])}, {'bias': ivy.array([0, 1, 2])}] + + >>> o = ivy.zeros(3) + >>> kwargs = {'out': o, 'bias': ivy.arange(3)} + >>> x = ivy.match_kwargs(kwargs, ivy.linear, ivy.add) + >>> print(x) + [{'out': ivy.array([0., 0., 0.]), 'bias': ivy.array([0, 1, 2])}, {}] + """ + split_kwargs = list() + for receiver in receivers: + expected_kwargs = arg_names(receiver) + found_kwargs = {k: v for k, v in kwargs.items() if k in expected_kwargs} + if not allow_duplicates: + for k in found_kwargs.keys(): + del kwargs[k] + split_kwargs.append(found_kwargs) + if len(split_kwargs) == 1: + return split_kwargs[0] + return split_kwargs + + +@handle_exceptions +@handle_nestable +@handle_array_function +def multiprocessing(context: Optional[str] = None): + """ + Return backend-specific multiprocessing module. + + Parameters + ---------- + context + The context of the multiprocessing, either fork, forkserver or spawn. + Default is ``None``. + + Returns + ------- + ret + Multiprocessing module + """ + return current_backend().multiprocessing(context) @handle_exceptions @@ -2447,365 +2577,286 @@ def print_all_arrays_in_memory(): print(type(arr), arr.shape) -ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 - - @handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back @handle_array_function -def set_queue_timeout(timeout: float): +@handle_device_shifting +def scatter_flat( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], + /, + *, + size: Optional[int] = None, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Set a timeout value (in seconds) for the global queue. - - Set the global queue timeout value (in seconds) Default value without this function - being called is 15 seconds. + Scatter flat updates into a new flat array according to flat indices. Parameters ---------- - timeout - The timeout when waiting for containers to arrive from the queues. - To be set in seconds. - - Examples - -------- - >>> x = ivy.set_queue_timeout(10) - >>> x = ivy.queue_timeout - >>> print(x) - 10.0 - - >>> ivy.set_queue_timeout(30) - >>> y = ivy.queue_timeout - >>> print(y) - 30 - """ - global queue_timeout_stack - ivy.utils.assertions.check_isinstance(timeout, (int, float)) - queue_timeout_stack.append(timeout) - ivy.__setattr__("queue_timeout", timeout, True) + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + size + The size of the result. Default is `None`, in which case tensor + argument out must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + Returns + ------- + ret + New array of given shape, with the values scattered at the indices. -@handle_exceptions -def unset_queue_timeout() -> None: - """ - Reset the global queue timeout value (in seconds) to the previous state. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - >>> ivy.set_queue_timeout(10.0) - >>> y = ivy.queue_timeout - >>> print(y) - 10.0 + With :class:`ivy.Array` input: + >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) + >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) + >>> ivy.scatter_flat(indices, updates, out=out) + >>> print(out) + ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) - >>> ivy.unset_queue_timeout() - >>> ivy.queue_timeout - 15.0 - """ - global queue_timeout_stack - if queue_timeout_stack: - queue_timeout_stack.pop(-1) - timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 - ivy.__setattr__("queue_timeout", timeout, True) + + With :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) -ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" + With :class:`ivy.Container` and :class:`ivy.Array` input: + >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + { + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + } + + + With :class:`ivy.Container` input: + >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), + ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) + >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), + ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) + >>> size = 8 + >>> print(ivy.scatter_flat(indices, updates, size=size)) + { + a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), + b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + } + """ + return current_backend(indices).scatter_flat( + indices, updates, size=size, reduction=reduction, out=out + ) @handle_exceptions -def set_tmp_dir(tmp_dr: str) -> None: +@handle_backend_invalid +@handle_nestable +@inputs_to_native_shapes +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def scatter_nd( + indices: Union[ivy.Array, ivy.NativeArray], + updates: Union[ivy.Array, ivy.NativeArray], + /, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + *, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: """ - Set the directory for saving temporary files. + Scatter updates into a new array according to indices. Parameters ---------- - tmp_dr - The new directory for saving temporary files + indices + Indices for the new values to occupy. + updates + Values for the new array to hold. + shape + The shape of the result. Default is ``None``, in which case tensor + argument must be provided. + reduction + The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + New array of given shape, with the values scattered at the indices. Examples -------- - >>> x = ivy.tmp_dir - >>> print(x) - /tmp + scatter values into an empty array, With :class:`ivy.Array` input: - >>> ivy.set_tmp_dir("/my_tmp") - >>> y = ivy.tmp_dir - >>> print(y) - /my_tmp - """ - global tmp_dir_stack - ivy.utils.assertions.check_isinstance(tmp_dr, str) - tmp_dir_stack.append(tmp_dr) - ivy.__setattr__("tmp_dir", tmp_dr, True) + >>> indices = ivy.array([[4], [3], [1], [7]]) + >>> updates = ivy.array([9, 10, 11, 12]) + >>> shape = ivy.array([8]) + >>> scatter = ivy.scatter_nd(indices, updates, shape) + >>> print(scatter) + ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) + With scatter into an empty array, With :class:`ivy.Container` input: -@handle_exceptions -def unset_tmp_dir() -> None: - """ - Reset the directory for saving temporary files to the previous value. + >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), + ... b=ivy.array([[5],[1],[2]])) + >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), + ... b=ivy.array([20, 30, 40])) + >>> shape = ivy.Container(a=ivy.array([10]), + ... b = ivy.array([10])) + >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') + >>> print(z) + { + a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), + b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) + } - Examples - -------- - >>> ivy.set_tmp_dir("/my_dir") - >>> y = ivy.tmp_dir - >>> print(y) - /my_dir + With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> ivy.unset_tmp_dir() - >>> ivy.tmp_dir - /tmp + >>> indices = ivy.array([[4],[3],[1]]) + >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), + ... b=ivy.array([200, 300, 400])) + >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), + ... b=ivy.array([10, 20, 30, 40, 50])) + >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) + >>> print(z) + { + a: ivy.array([1, 30, 3, 20, 10]), + b: ivy.array([10, 400, 30, 300, 200]) + } """ - global tmp_dir_stack - if tmp_dir_stack: - tmp_dir_stack.pop(-1) - tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" - ivy.__setattr__("tmp_dir", tmp_dr, True) + return current_backend(indices).scatter_nd( + indices, updates, shape=shape, reduction=reduction, out=out + ) @handle_exceptions -def container_types(): +def set_array_mode(mode: bool) -> None: """ - Summary. + Set the mode of whether to convert inputs to ivy.NativeArray, then convert outputs + back to ivy.Array. - Returns - ------- - ret - a key-value structure, and exposes public methods .keys(), .values() and - items(). - """ - # noinspection PyBroadException - try: - return current_backend().container_types() - except ValueError: - return [] + It Stops the conversion of ivy.NativeArray to ivy.Array in the + case when it is set to False. + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions -@handle_exceptions -def inplace_arrays_supported() -> bool: - """ - Determine whether inplace arrays are supported for the current backend framework. + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - Returns - ------- - ret - Boolean, whether or not inplace arrays are supported. + >>> ivy.set_array_mode(True) + >>> ivy.array_mode + True """ - return current_backend().inplace_arrays_supported() + global array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + array_mode_stack.append(mode) + ivy.__setattr__("array_mode", mode, True) @handle_exceptions -def inplace_variables_supported() -> bool: +def set_exception_trace_mode(mode: Literal["ivy", "full", "frontend"]) -> None: """ - Determine whether inplace variables are supported for the current backend framework. + Set the mode of whether to show frontend-truncated exception stack traces, ivy- + truncated exception stack traces or full exception stack traces. - Returns - ------- - ret - Boolean, whether or not inplace variables are supported. - """ - return current_backend().inplace_variables_supported() + Parameter + --------- + mode + str exeption trace mode, one of `ivy`, `full` or `frontend` + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' -@handle_exceptions -@handle_nestable -@inputs_to_native_arrays -@handle_array_function -def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: - """ - Return if in-place operations are supported for x's data type. - - Determine whether in-place operations are supported for x's data type, by the - current backend framework setting. - - Parameters - ---------- - x - Input variable for whose data type we check whether the current backend - framework supports in-place operations. - - Returns - ------- - ret - Value depends on whether in-place operations are supported for - data type of x. - - Raises - ------ - IvyException - If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will - be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. - - Examples - -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) - True - - With :class:`ivy.Container` input and backend set as `torch`: - - >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) - >>> y = ivy.supports_inplace_updates(x) - >>> print(y) - { - a: True, - b: True - } - - With `ivy.Array` input and backend set as "tensorflow": - - >>> x = ivy.array([1., 4.2, 2.2]) - >>> ret = x.supports_inplace_updates() - >>> print(ret) - False + >>> ivy.set_exception_trace_mode("full") + >>> ivy.exception_trace_mode + 'full' """ - if _is_variable(x): - return ivy.inplace_variables_supported() - elif ivy.is_native_array(x): - return ivy.inplace_arrays_supported() - raise ivy.utils.exceptions.IvyException( - "Input x must be either a variable or an array." + global exception_trace_mode_stack + trace_modes = list(trace_mode_dict.keys()) + ivy.utils.assertions.check_elem_in_list( + mode, trace_modes, False, "trace mode must be one of {}".format(trace_modes) ) + exception_trace_mode_stack.append(mode) + ivy.__setattr__("exception_trace_mode", mode, True) @handle_exceptions -@handle_nestable -@inputs_to_native_arrays -@handle_array_function -def assert_supports_inplace(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: +def set_inplace_mode(mode: str = "lenient") -> None: """ - Assert that inplace operations are supported for x. + Set the memory management behavior for in-place updates in Ivy. + + By default, Ivy creates new arrays in the backend for in-place updates. + However, this behavior can be controlled by the user + using the 'inplace_mode' parameter. Parameters ---------- - x - Input variable or array to check for inplace support for. + mode : str + The mode for memory management during in-place updates. + - 'lenient': (Default) In this mode, new arrays will be created during + in-place updates to avoid breaking existing code. + This is the default behavior. + - 'strict': In this mode, an error will be raised if the + 'inplace_update' function is called + in a backend that doesn't support inplace updates natively. Returns ------- - ret - True if supports, raises IvyBackendException otherwise - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + None Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) - True - - With :class:`ivy.Array` input and default backend set as `torch`: - - >>> ivy.set_backend("torch") - >>> x = ivy.array([1, 2, 3]) - >>> print(x.assert_supports_inplace()) - True - - With :class:`ivy.Container` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) - { - a: True, - b: True - } + >>> set_inplace_mode('lenient') + >>> ivy.inplace_mode + 'lenient' - With :class:`ivy.Container` input and default backend set as `torch`: + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> print(x.assert_supports_inplace()) - { - a: True, - b: True - } + Note + ---- + Enabling strict mode can help users have more control over memory management + but may lead to errors if the backend doesn't support inplace updates natively. """ - ivy.utils.assertions.check_true( - ivy.supports_inplace_updates(x), - "Inplace operations are not supported {} types with {} backend".format( - type(x), ivy.current_backend_str() - ), + global inplace_mode_stack + inplace_modes = ["lenient", "strict"] + ivy.utils.assertions.check_elem_in_list( + mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" ) - return True - - -@handle_nestable -@handle_partial_mixed_function -@handle_view_indexing -@inputs_to_ivy_arrays -@handle_array_function -@handle_device_shifting -def get_item( - x: Union[ivy.Array, ivy.NativeArray], - /, - query: Union[ivy.Array, ivy.NativeArray, Tuple], - *, - copy: Optional[bool] = None, -) -> ivy.Array: - """ - Gather slices from x according to query array, identical to x[query]. - - Parameters - ---------- - x - array, the array from which to gather values. - query - array, index array, integer indices or boolean mask. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - - Returns - ------- - ret - New array with the values gathered at the specified indices. - - Functional Examples - ------------------- - - >>> x = ivy.array([0, -1, 20]) - >>> query = ivy.array([0, 1]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 0, -1]) - - >>> x = ivy.array([[4, 5], [20, 128], [-2, -10]]) - >>> query = ivy.array([[True, False], [False, False], [True, True]]) - >>> print(ivy.get_item(x, query)) - ivy.array([ 4, -2, -10]) - """ - if ivy.is_array(query) and ivy.is_bool_dtype(query): - if not len(query.shape): - if not query: - return ivy.array([], shape=(0,), dtype=x.dtype) - return ivy.expand_dims(x, axis=0) - query = ivy.nonzero(query, as_tuple=False) - ret = ivy.gather_nd(x, query) - else: - indices, target_shape = _parse_query(query, x.shape) - if indices is None: - return ivy.empty(target_shape, dtype=x.dtype) - ret = ivy.gather_nd(x, indices) - ret = ivy.reshape(ret, target_shape) - return ret - - -get_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} + inplace_mode_stack.append(mode) + ivy.__setattr__("inplace_mode", mode, True) @handle_nestable @@ -2879,758 +2930,711 @@ def set_item( return ret -set_item.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - ), - "to_skip": ("inputs_to_ivy_arrays",), -} - - -def _parse_query(query, x_shape): - query = (query,) if not isinstance(query, tuple) else query - query_ = tuple([q.to_numpy() if ivy.is_array(q) else q for q in query]) - - # array containing all of x's flat indices - x_ = ivy.arange(0, _numel(x_shape)).reshape(x_shape) +@handle_exceptions +@handle_array_function +def set_min_base(val: float) -> None: + """ + Set the global minimum base used by ivy for numerically stable power raising. - # use numpy's __getitem__ to get the queried indices - x_idxs = ivy.array(x_.to_numpy()[query_]) - target_shape = x_idxs.shape + Parameters + ---------- + val + The new value to set the minimum base to. - if 0 in x_idxs.shape or 0 in x_shape: - return None, target_shape + Examples + -------- + >>> x = ivy.min_base + >>> print(x) + 1e-05 - # convert the flat indices to multi-D indices - x_idxs = ivy.unravel_index(x_idxs, x_shape) - - # stack the multi-D indices to bring them to gather_nd/scatter_nd format - x_idxs = ivy.stack(x_idxs, axis=-1).astype(ivy.int64) - - return x_idxs, target_shape - - -def _numel(shape): - shape = tuple(shape) - return ivy.prod(shape).to_scalar() if shape != () else 1 - - -def _broadcast_to(input, target_shape): - input = ivy.squeeze(input) - if _numel(tuple(input.shape)) == _numel(tuple(target_shape)): - return ivy.reshape(input, target_shape) - else: - input = ivy.expand_dims(input, axis=0) if not len(input.shape) else input - new_dims = () - i_i = len(input.shape) - 1 - for i_t in range(len(target_shape) - 1, -1, -1): - if len(input.shape) + len(new_dims) >= len(target_shape): - break - if i_i < 0 or target_shape[i_t] != input.shape[i_i]: - new_dims += (i_t,) - else: - i_i -= 1 - input = ivy.expand_dims(input, axis=new_dims) - return ivy.broadcast_to(input, target_shape) + >>> ivy.set_min_base(1e-04) + >>> y = ivy.min_base + >>> print(y) + 1e-04 + """ + global min_base_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_base_stack.append(val) + ivy.__setattr__("min_base", val, True) @handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_ivy_arrays @handle_array_function -# @handle_device_shifting -def inplace_update( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], - /, - *, - ensure_in_backend: bool = False, - keep_input_dtype: bool = False, -) -> ivy.Array: +def set_min_denominator(val: float) -> None: """ - Perform in-place update for the input array. - - This will always be performed on ivy.Array instances pass in the input, and will - also be performed on the native array classes in the backend when the backend - supports this. If the backend does not natively support inplace updates, and x is an - ivy.NativeArray instance, then an - exception will be thrown. + Set the global minimum denominator used by ivy for numerically stable division. Parameters ---------- - x - The variable to update. val - The array to update the variable with. - ensure_in_backend - Whether or not to ensure that the `ivy.NativeArray` is also inplace updated. - In cases where it should be, backends which do not natively support inplace - updates will raise an exception. - keep_input_dtype - Whether or not to preserve `x` data type after the update, otherwise `val` - data type will be applied. Defaults to False. - - Returns - ------- - ret - The array following the in-place update. - - Raises - ------ - IvyException - If backend set doesn't natively support inplace updates and ensure_in_backend is - True, above exception will be raised. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the arguments. + The value to set the global minimum denominator to. Examples -------- - With :class:`ivy.Array` input and default backend set as `numpy`: - - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.array([0]) - >>> ivy.inplace_update(x, y) + >>> x = ivy.min_denominator >>> print(x) - ivy.array([0]) + 1e-12 - With :class:`ivy.Array` input and default backend set as `numpy`: + >>> ivy.set_min_denominator(1e-13) + >>> y = ivy.min_denominator + >>> print(y) + 1e-13 + """ + global min_denominator_stack + ivy.utils.assertions.check_isinstance(val, (int, float)) + min_denominator_stack.append(val) + ivy.__setattr__("min_denominator", val, True) - >>> ivy.set_backend("numpy") - >>> x = ivy.array([1, 2, 3], dtype=ivy.float32) - >>> y = ivy.array([0, 0, 0], dtype=ivy.int32) - >>> ivy.inplace_update(x, y, keep_input_dtype=True) - >>> print(x) - ivy.array([0., 0., 0.]) - With :class:`ivy.Container` instances:, and backend set as `torch`: +@handle_exceptions +def set_nestable_mode(mode: bool) -> None: + """ + Set the mode of whether to check if function inputs are ivy.Container. - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.Container(a=ivy.array([1]), b=ivy.array([2])) - >>> ivy.inplace_update(x, y) - >>> print(x) - { - a: ivy.array([1, 1]), - b: ivy.array([2, 2]) - } + Parameter + --------- + mode + boolean whether to check if function inputs are ivy.Container - With mix of :class:`ivy.Array` and :class:`ivy.Container` instances:, and backend - set as `torch`: + Examples + -------- + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False - >>> ivy.set_backend("torch") - >>> x = ivy.Container(a=ivy.array([5, 6]), b=ivy.array([7, 8])) - >>> y = ivy.array([1, 2]) - >>> ivy.inplace_update(x, y) - >>> print(x) - { - a: ivy.array([1, 2]), - b: ivy.array([1, 2]) - } + >>> ivy.set_nestable_mode(True) + >>> ivy.nestable_mode + True """ - return current_backend(x).inplace_update( - x, - val, - ensure_in_backend=ensure_in_backend, - keep_input_dtype=keep_input_dtype, - ) + global nestable_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + nestable_mode_stack.append(mode) + ivy.__setattr__("nestable_mode", mode, True) -inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} +@handle_exceptions +def set_precise_mode(mode: bool) -> None: + """ + Set the mode of whether to use a promotion table that avoids any precision loss or a + compute effecient table that avoids most wider-than-necessary promotions. -ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" + Parameter + --------- + mode + boolean whether to use high precision promtion table + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.set_precise_mode(True) + >>> ivy.precise_mode + True + """ + global precise_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + precise_mode_stack.append(mode) + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) @handle_exceptions -def set_inplace_mode(mode: str = "lenient") -> None: +@handle_array_function +def set_queue_timeout(timeout: float): """ - Set the memory management behavior for in-place updates in Ivy. + Set a timeout value (in seconds) for the global queue. - By default, Ivy creates new arrays in the backend for in-place updates. - However, this behavior can be controlled by the user - using the 'inplace_mode' parameter. + Set the global queue timeout value (in seconds) Default value without this function + being called is 15 seconds. Parameters ---------- - mode : str - The mode for memory management during in-place updates. - - 'lenient': (Default) In this mode, new arrays will be created during - in-place updates to avoid breaking existing code. - This is the default behavior. - - 'strict': In this mode, an error will be raised if the - 'inplace_update' function is called - in a backend that doesn't support inplace updates natively. - - Returns - ------- - None + timeout + The timeout when waiting for containers to arrive from the queues. + To be set in seconds. Examples -------- - >>> set_inplace_mode('lenient') - >>> ivy.inplace_mode - 'lenient' - - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' + >>> x = ivy.set_queue_timeout(10) + >>> x = ivy.queue_timeout + >>> print(x) + 10.0 - Note - ---- - Enabling strict mode can help users have more control over memory management - but may lead to errors if the backend doesn't support inplace updates natively. + >>> ivy.set_queue_timeout(30) + >>> y = ivy.queue_timeout + >>> print(y) + 30 """ - global inplace_mode_stack - inplace_modes = ["lenient", "strict"] - ivy.utils.assertions.check_elem_in_list( - mode, inplace_modes, False, f"inplace mode must be one of {inplace_modes}" - ) - inplace_mode_stack.append(mode) - ivy.__setattr__("inplace_mode", mode, True) + global queue_timeout_stack + ivy.utils.assertions.check_isinstance(timeout, (int, float)) + queue_timeout_stack.append(timeout) + ivy.__setattr__("queue_timeout", timeout, True) @handle_exceptions -def unset_inplace_mode() -> None: +def set_shape_array_mode(mode: bool) -> None: """ - Reset the memory management behavior for in-place updates in Ivy to the previous - state. + Set the mode of returning shape as ivy.Array to the given mode instance. + + Parameter + --------- + mode + boolean whether to return shape as ivy.Array Examples -------- - >>> set_inplace_mode('strict') - >>> ivy.inplace_mode - 'strict' + >>> ivy.set_shape_array_mode(False) + >>> ivy.shape_array_mode + False - >>> unset_inplace_mode() - >>> ivy.inplace_mode - 'lenient' + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True """ - global inplace_mode_stack - if inplace_mode_stack: - inplace_mode_stack.pop(-1) - mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" - ivy.__setattr__("inplace_mode", mode, True) + global shape_array_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + shape_array_mode_stack.append(mode) + ivy.__setattr__("shape_array_mode", mode, True) + + +@handle_exceptions +def set_show_func_wrapper_trace_mode(mode: bool) -> None: + """ + Set the mode of whether to show the full stack trace with function wrapping traces. + + Parameter + --------- + mode + boolean whether to perform ivy.Array conversions + + Examples + -------- + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False + + >>> ivy.set_show_func_wrapper_trace_mode(True) + >>> ivy.show_func_wrapper_trace_mode + True + """ + global show_func_wrapper_trace_mode_stack + ivy.utils.assertions.check_isinstance(mode, bool) + show_func_wrapper_trace_mode_stack.append(mode) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + + +@handle_exceptions +def set_tmp_dir(tmp_dr: str) -> None: + """ + Set the directory for saving temporary files. + + Parameters + ---------- + tmp_dr + The new directory for saving temporary files + + Examples + -------- + >>> x = ivy.tmp_dir + >>> print(x) + /tmp + + >>> ivy.set_tmp_dir("/my_tmp") + >>> y = ivy.tmp_dir + >>> print(y) + /my_tmp + """ + global tmp_dir_stack + ivy.utils.assertions.check_isinstance(tmp_dr, str) + tmp_dir_stack.append(tmp_dr) + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions @handle_backend_invalid @handle_nestable -@inputs_to_ivy_arrays +@handle_array_like_without_promotion +@inputs_to_native_arrays +@outputs_to_ivy_shapes +@outputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def inplace_decrement( +def shape( x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: + /, + *, + as_array: bool = False, +) -> Union[ivy.Shape, ivy.NativeShape]: """ - Perform in-place decrement for the input array. + Return the shape of the array ``x``. Parameters ---------- x - The input array to be decremented by the defined value. - val - The value of decrement. + Input array to infer the shape of. + as_array + Whether to return the shape as an array. + Default is False. Returns ------- ret - The array following the in-place decrement. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Shape of the array ``x``. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_decrement(x, 1.25) - >>> print(y) - ivy.array([[ 4.05, 5.75, -1.25], - [ 5.55, 6.75, 2.65], - [-1.25, 8.75, 5.05]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0.5, -5., 30.]), b=ivy.array([0., -25., 50.])) - >>> y = ivy.inplace_decrement(x, 1.5) + >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) + >>> y = ivy.shape(x) + >>> z = ivy.shape(x, as_array = True) >>> print(y) - { - a: ivy.array([-1., -6.5, 28.5]), - b: ivy.array([-1.5, -26.5, 48.5]) - } - - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_decrement(x, y) - >>> print(z) - { - a: ivy.array([0., 0., 0.]), - b: ivy.array([0., 0., 0.]) - } + (2, 3) - >>> x = ivy.Container(a=ivy.array([3., 7., 10.]), b=ivy.array([0., 75., 5.5])) - >>> y = ivy.Container(a=ivy.array([2., 5.5, 7.]), b=ivy.array([0., 25., 2.])) - >>> z = ivy.inplace_decrement(x, y) >>> print(z) - { - a: ivy.array([1., 1.5, 3.]), - b: ivy.array([0., 50., 3.5]) - } + ivy.array([2, 3]) """ - return current_backend(x).inplace_decrement(x, val) + return current_backend(x).shape(x, as_array=as_array) @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def inplace_increment( - x: Union[ivy.Array, ivy.NativeArray], - val: Union[ivy.Array, ivy.NativeArray], -) -> ivy.Array: +def stable_divide( + numerator: Union[Number, ivy.Array, ivy.NativeArray], + denominator: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + min_denominator: Union[Number, ivy.Array, ivy.NativeArray] = None, +) -> Union[Number, ivy.Array]: """ - Perform in-place increment for the input array. + Divide the numerator by the denominator, with min denominator added to the + denominator for numerical stability. Parameters ---------- - x - The input array to be incremented by the defined value. - val - The value of increment. + numerator + The numerator of the division. + denominator + The denominator of the division. + min_denominator + The minimum denominator to use, use global ivy._MIN_DENOMINATOR (1e-12) + by default. Returns ------- ret - The array following the in-place increment. + The new item following the numerically stable division. Examples -------- + With :code:`int` input: + + >>> x = ivy.stable_divide(1, 2) + >>> print(x) + 0.49999999999975 + + >>> x = ivy.stable_divide(1, 4, min_denominator=1) + >>> print(x) + 0.2 + + With float input: + + >>> x = ivy.stable_divide(5.0, 3.33) + >>> print(x) + 1.5015015015010504 + + With :code:`complex` input: + + >>> x = ivy.stable_divide(1+1j, 1-1j) + >>> print(x) + (5.000444502911705e-13+0.9999999999995j) + With :class:`ivy.Array` input: - >>> x = ivy.array([[5.3, 7., 0.],[6.8, 8, 3.9],[0., 10., 6.3]]) - >>> y = ivy.inplace_increment(x, 3.) + >>> x = ivy.asarray([[10., 20., 30.], + ... [40., 50., 60.]]) + >>> y = ivy.stable_divide(x, 10.) >>> print(y) - ivy.array([[ 8.3, 10., 3.], - [ 9.8, 11., 6.9], - [ 3., 13., 9.3]]) + ivy.array([[1., 2., 3.], + [4., 5., 6.]]) + + + >>> x = ivy.asarray([1,2,3]) + >>> y = np.array((1., 3., 5.)) + >>> z = ivy.stable_divide(x, y) + >>> print(z) + ivy.array([1. , 0.667, 0.6 ]) + + >>> x = ivy.asarray([1., 2., 4.]) + >>> y = ivy.asarray([1., 0.5, 0.25]) + >>> z = ivy.asarray([0.01, 0.02, 0.03]) + >>> w = ivy.stable_divide(x, y, min_denominator=z) + >>> print(w) + ivy.array([ 0.99, 3.85, 14.3 ]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.inplace_increment(x, 2.5) + >>> x = ivy.Container(a=ivy.asarray([10., 15.]), b=ivy.asarray([20., 25.])) + >>> y = ivy.stable_divide(x, 0.5) >>> print(y) { - a: ivy.array([2.5, 17.5, 32.5]), - b: ivy.array([2.5, 27.5, 52.5]) + a: ivy.array([20., 30.]), + b: ivy.array([40., 50.]) } - >>> x = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> y = ivy.Container(a=ivy.array([0., 15., 30.]), b=ivy.array([0., 25., 50.])) - >>> z = ivy.inplace_increment(x, y) + >>> x = ivy.Container(a=ivy.asarray([1., 2.]), b=ivy.asarray([3., 4.])) + >>> y = ivy.Container(a=ivy.asarray([0.5, 2.5]), b=ivy.asarray([3.5, 0.4])) + >>> z = ivy.stable_divide(x, y) >>> print(z) { - a: ivy.array([0., 30., 60.]), - b: ivy.array([0., 50., 100.]) + a: ivy.array([2., 0.8]), + b: ivy.array([0.857, 10.]) } """ - return current_backend(x).inplace_increment(x, val) + return numerator / (denominator + default(min_denominator, ivy.min_denominator)) @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def scatter_flat( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], +def stable_pow( + base: Union[Number, ivy.Array, ivy.NativeArray], + exponent: Union[Number, ivy.Array, ivy.NativeArray], /, *, - size: Optional[int] = None, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: + min_base: float = None, +) -> Any: """ - Scatter flat updates into a new flat array according to flat indices. + Raise the base by the power, with ivy.min_base added to the base when exponent > 1 + for numerical stability. Parameters ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - size - The size of the result. Default is `None`, in which case tensor - argument out must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + base + The base number. + exponent + The exponent number. + min_base + The minimum base to use, use global ivy.min_base by default. Returns ------- ret - New array of given shape, with the values scattered at the indices. - - This function is *nestable*, and therefore also accepts :code:'ivy.Container' - instance in place of the argument. + The new item following the numerically stable power. Examples -------- - With :class:`ivy.Array` input: - >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3]) - >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0]) - >>> ivy.scatter_flat(indices, updates, out=out) - >>> print(out) - ivy.array([8, 7, 5, 4, 0, 0, 0, 0]) + With :code:`int` input: + + >>> x = ivy.stable_pow(2, 2) + >>> print(x) + ivy.array(4.00004) + >>> x = ivy.stable_pow(2, 2, min_base=2) + >>> print(x) + ivy.array(16) + + With float input: + + >>> x = ivy.stable_pow(4.0, .5) + >>> print(x) + ivy.array(2.00000262) + + With :code:`complex` input: + + >>> x = ivy.stable_pow(3+4j, 2j) + >>> print(x) + ivy.array(-0.15605032-0.01208451j) With :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8]) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) - ivy.array([2, 0, 2, 8, 0, 0, 0, 0]) + >>> x = ivy.asarray([[2, 4], + ... [6, 8]]) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) + ivy.array([[ 4.00004, 16.00008], + [36.00012, 64.00016]]) - With :class:`ivy.Container` and :class:`ivy.Array` input: - >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3]) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) + >>> x = ivy.asarray([2, 4, 6]) + >>> y = ivy.asarray([2, 3, 4]) + >>> z = ivy.stable_pow(x, y) + >>> print(z) + ivy.array([ 4.00004, 64.00048, 1296.00864]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.stable_pow(x, 2) + >>> print(y) { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + a: ivy.array([4.00004, 16.00008]), + b: ivy.array([36.00012, 64.00016]) } - - With :class:`ivy.Container` input: - >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), - ... b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3])) - >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), - ... b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3])) - >>> size = 8 - >>> print(ivy.scatter_flat(indices, updates, size=size)) + >>> x = ivy.Container(a=ivy.asarray([2, 4]), b=ivy.asarray([6, 8])) + >>> y = ivy.Container(a=ivy.asarray([1, 3]), b=ivy.asarray([4, 5])) + >>> z = ivy.stable_pow(x, y) + >>> print(z) { - a: ivy.array([2, 0, 2, 8, 0, 0, 0, 0]), - b: ivy.array([2, 7, 2, 3, 0, 0, 0, 0]) + a: ivy.array([2.00001, 64.00048]), + b: ivy.array([1296.00864, 32768.2048]) } """ - return current_backend(indices).scatter_flat( - indices, updates, size=size, reduction=reduction, out=out + return_dtype = ivy.promote_types( + ivy.default_dtype(item=base), + ivy.default_dtype(item=default(min_base, ivy.min_base)), ) + return_dtype = ivy.promote_types(return_dtype, ivy.default_dtype(item=exponent)) + ret = (base + default(min_base, ivy.min_base)) ** ivy.array(exponent) + return ret.astype(return_dtype) @handle_exceptions -@handle_backend_invalid @handle_nestable -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function @handle_device_shifting -def scatter_nd( - indices: Union[ivy.Array, ivy.NativeArray], - updates: Union[ivy.Array, ivy.NativeArray], +def strides( + x: Union[ivy.Array, ivy.NativeArray], /, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - *, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[int]: """ - Scatter updates into a new array according to indices. + Return the input array's strides across each dimension. Parameters ---------- - indices - Indices for the new values to occupy. - updates - Values for the new array to hold. - shape - The shape of the result. Default is ``None``, in which case tensor - argument must be provided. - reduction - The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace' - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + x + The input array. Returns ------- ret - New array of given shape, with the values scattered at the indices. + A tuple containing the strides. Examples -------- - scatter values into an empty array, With :class:`ivy.Array` input: - - >>> indices = ivy.array([[4], [3], [1], [7]]) - >>> updates = ivy.array([9, 10, 11, 12]) - >>> shape = ivy.array([8]) - >>> scatter = ivy.scatter_nd(indices, updates, shape) - >>> print(scatter) - ivy.array([ 0, 11, 0, 10, 9, 0, 0, 12]) - - With scatter into an empty array, With :class:`ivy.Container` input: - - >>> indices = ivy.Container(a=ivy.array([[4],[3],[6]]), - ... b=ivy.array([[5],[1],[2]])) - >>> updates = ivy.Container(a=ivy.array([100, 200, 200]), - ... b=ivy.array([20, 30, 40])) - >>> shape = ivy.Container(a=ivy.array([10]), - ... b = ivy.array([10])) - >>> z = ivy.scatter_nd(indices, updates, shape=shape, reduction='replace') - >>> print(z) - { - a: ivy.array([0, 0, 0, 200, 100, 0, 200, 0, 0, 0]), - b: ivy.array([0, 30, 40, 0, 0, 20, 0, 0, 0, 0]) - } - - With :class:`ivy.Container` and :class:`ivy.Array` input: - - >>> indices = ivy.array([[4],[3],[1]]) - >>> updates = ivy.Container(a=ivy.array([10, 20, 30]), - ... b=ivy.array([200, 300, 400])) - >>> z = ivy.Container(a=ivy.array([1, 2, 3, 4, 5]), - ... b=ivy.array([10, 20, 30, 40, 50])) - >>> ivy.scatter_nd(indices, updates, reduction='replace', out=z) - >>> print(z) - { - a: ivy.array([1, 30, 3, 20, 10]), - b: ivy.array([10, 400, 30, 300, 200]) - } + >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) + >>> ivy.strides(x) + (4, 8) """ - return current_backend(indices).scatter_nd( - indices, updates, shape=shape, reduction=reduction, out=out - ) + if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): + return ivy.to_numpy(x).strides + # if x is an ivy array with a base, + # convert it to a numpy array with the same base: + ret = ivy.to_numpy(x.base) + ivy_numpy = ivy.with_backend("numpy") + for fn, args, kwargs, index in x._manipulation_stack: + ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) + ret = ret[index] if ivy.exists(index) else ret + return ret.to_native().strides @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function -@handle_device_shifting -def gather( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - batch_dims: int = 0, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: +def supports_inplace_updates(x: Union[ivy.Array, ivy.NativeArray], /) -> bool: """ - Gather slices from params at axis according to indices. + Return if in-place operations are supported for x's data type. + + Determine whether in-place operations are supported for x's data type, by the + current backend framework setting. Parameters ---------- - params - The array from which to gather values. - indices - The array which indicates the indices that will be gathered along - the specified axis. - axis - optional int, the axis from which to gather from. - Default is ``-1``. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional array, for writing the result to. It must have a shape - that the inputs broadcast to. + x + Input variable for whose data type we check whether the current backend + framework supports in-place operations. Returns ------- ret - New array with the values gathered at the specified indices along the - specified axis. + Value depends on whether in-place operations are supported for + data type of x. + Raises + ------ + IvyException + If x isn't a class instance of ivy.Array or ivy.NativeArray, an exception will + be raised. - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + This function is *nestable*, and therefore also accepts :code:'ivy.Container' + instance in place of the argument. Examples -------- - With :class:`ivy.Array` input: + With :class:`ivy.Array` input and default backend set as `numpy`: - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.array([1, 2]) - >>> print(ivy.gather(x, y)) - ivy.array([1., 2.]) + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) + True - >>> x = ivy.array([[0., 1., 2.],[3., 4., 5.]]) - >>> y = ivy.array([[0, 1],[1, 2]]) - >>> z = ivy.zeros((2, 2, 2)) - >>> ivy.gather(x, y, out=z) - >>> print(z) - ivy.array([[[0., 1.],[1., 2.]],[[3., 4.],[4., 5.]]]) - - >>> x = ivy.array([[[0., 1.], [2., 3.]], - ... [[8., 9.], [10., 11.]]]) - >>> y = ivy.array([[0, 1]]) - >>> z = ivy.zeros((1, 2, 2, 2)) - >>> ivy.gather(x, y, axis=0, out=z) - >>> print(z) - ivy.array( - [[[[ 0., 1.], - [ 2., 3.]], - [[ 8., 9.], - [10., 11.]]]]) - - >>> x = ivy.array([[0, 10, 20, 0, 0], - ... [0, 0, 0, 30, 40], - ... [0, 10, 0, 0, 40]]) - >>> y = ivy.array([[1, 2],[3, 4],[1, 4]]) - >>> z = ivy.gather(x, y, batch_dims=1) - >>> print(z) - ivy.array([[10, 20], [30, 40],[10, 40]]) - - With :class:`ivy.Container` input: + With :class:`ivy.Container` input and backend set as `torch`: - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.Container(a = ivy.array([0, 1]), - ... b = ivy.array([1, 2])) - >>> print(ivy.gather(x, y)) + >>> x = ivy.Container(a=ivy.array([5., 6.]), b=ivy.array([7., 8.])) + >>> y = ivy.supports_inplace_updates(x) + >>> print(y) { - a: ivy.array([0., 1.]), - b: ivy.array([5., 6.]) + a: True, + b: True } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With `ivy.Array` input and backend set as "tensorflow": - >>> x = ivy.Container(a = ivy.array([0., 1., 2.]), - ... b = ivy.array([4., 5., 6.])) - >>> y = ivy.array([0, 1]) - >>> print(ivy.gather(x, y)) - { - a: ivy.array([0., 1.]), - b: ivy.array([4., 5.]) - } + >>> x = ivy.array([1., 4.2, 2.2]) + >>> ret = x.supports_inplace_updates() + >>> print(ret) + False """ - return current_backend(params, indices).gather( - params, indices, axis=axis, batch_dims=batch_dims, out=out + if _is_variable(x): + return ivy.inplace_variables_supported() + elif ivy.is_native_array(x): + return ivy.inplace_arrays_supported() + raise ivy.utils.exceptions.IvyException( + "Input x must be either a variable or an array." ) +@handle_exceptions +def to_ivy_shape(shape: Union[ivy.Shape, ivy.NativeShape]) -> ivy.Shape: + """ + Return the input shape in ivy.Shape form. + + Parameters + ---------- + shape + The input to be converted + + Returns + ------- + ret + the input in ivy.Shape form + """ + if isinstance(shape, ivy.Shape): + return shape + return ivy.Shape(shape) + + @handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function @handle_device_shifting -def gather_nd( - params: Union[ivy.Array, ivy.NativeArray], - indices: Union[ivy.Array, ivy.NativeArray], - /, - *, - batch_dims: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: +def to_list(x: Union[ivy.Array, ivy.NativeArray], /) -> List: """ - Gather slices from params into a array with shape specified by indices. + Create a (possibly nested) list from input array. Parameters ---------- - params - The array from which to gather values. - indices - Index array. - batch_dims - optional int, lets you gather different items from each element of a batch. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + x + Input array. Returns ------- ret - New array of given shape, with the values gathered at the indices. + A list representation of the input array ``x``. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6.]) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) - ivy.array(1.) + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_list(x) + >>> print(y) + [-1, 0, 1] - >>> x = ivy.array([[0., 1.], [2., 3.], [4., 5.]]) - >>> y = ivy.array([[0],[1],[1]], dtype='int32') - >>> z = ivy.gather_nd(x,y,batch_dims=1) - ivy.array([0., 3., 5.]) + >>> x = ivy.array([[ 1.1, 2.2, 3.3], + ... [-4.4, -5.5, -6.6]]) + >>> y = ivy.to_list(x) + >>> print(y) + [[1.100000023841858,2.200000047683716,3.299999952316284], + [-4.400000095367432,-5.5,-6.599999904632568]] - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([[[-1, 0, 1], + ... [ 1, 0, -1]], + ... [[ 1, -1, 0], + ... [ 1, 0, -1]]]) + >>> y = ivy.to_list(x) + >>> print(y) + [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]),b=ivy.array([4., 5., 6.])) - >>> y = ivy.array([1]) - >>> print(ivy.gather_nd(x, y)) + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: + + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_list(x) + >>> print(y) { - a: ivy.array(1.), - b: ivy.array(5.) + a: [-1, 0, 1] } - With :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([[-1, 0, 1], + ... [-1, 0, 1], + ... [1, 0, -1]])) + >>> y = ivy.to_list(x) + >>> print(y) + { + a: [[-1, 0, 1], [-1, 0, 1], [1,0,-1]] + } - >>> x = ivy.Container(a=ivy.array([[0., 10., 20.],[30.,40.,50.]]), - ... b=ivy.array([[0., 100., 200.],[300.,400.,500.]])) - >>> y = ivy.Container(a=ivy.array([1,0]), - ... b=ivy.array([0])) - >>> print(ivy.gather_nd(x, y)) + >>> x = ivy.Container(a=ivy.array([[[-1, 0, 1],[1, 0, -1]], + ... [[1, -1, 0],[1, 0, -1]]])) + >>> y = ivy.to_list(x) + >>> print(y) { - a: ivy.array(30.), - b: ivy.array([0., 100., 200.]) + a: [[[-1, 0, 1], [1, 0, -1]], [[1, -1, 0], [1, 0, -1]]] } """ - res = current_backend(params, indices).gather_nd( - params, indices, batch_dims=batch_dims - ) - if ivy.exists(out): - return ivy.inplace_update(out, res) - return res + return current_backend(x).to_list(x) @handle_exceptions -@handle_nestable -@handle_array_function -def multiprocessing(context: Optional[str] = None): +def to_native_shape( + shape: Union[ivy.Array, ivy.Shape, ivy.NativeShape, tuple, int, list] +) -> ivy.NativeShape: """ - Return backend-specific multiprocessing module. + Return the input shape in its native backend framework form. Parameters ---------- - context - The context of the multiprocessing, either fork, forkserver or spawn. - Default is ``None``. + shape + The input to be converted Returns ------- - ret - Multiprocessing module + ret + the input in its native framework form """ - return current_backend().multiprocessing(context) + native_shape_type = (ivy.NativeShape,) + if ivy.current_backend_str() == "torch": + native_shape_type += (tuple,) + if len(backend_stack) != 0 and isinstance(shape, native_shape_type): + return shape + ivy.utils.assertions.check_isinstance( + shape, (int, list, tuple, ivy.Array, ivy.NativeArray, ivy.Shape) + ) + if isinstance(shape, int): + shape = (shape,) + elif isinstance(shape, list): + shape = tuple(shape) + elif is_array(shape): + shape = ivy.to_numpy(shape).tolist() + elif isinstance(shape, ivy.Shape): + shape = shape.shape + ivy.utils.assertions.check_all( + [isinstance(v, int) for v in shape if not is_array(v)], + "shape must take integers only", + as_array=False, + ) + ivy.utils.assertions.check_true( + not is_array(shape) or ivy.is_int_dtype(shape), "shape must take integers only" + ) + return ivy.NativeShape(shape) if len(backend_stack) != 0 else ivy.Shape(shape) @handle_exceptions @@ -3638,422 +3642,490 @@ def multiprocessing(context: Optional[str] = None): @handle_nestable @handle_array_like_without_promotion @inputs_to_native_arrays -@outputs_to_ivy_shapes -@outputs_to_ivy_arrays @handle_array_function @handle_device_shifting -def shape( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - as_array: bool = False, -) -> Union[ivy.Shape, ivy.NativeShape]: +def to_numpy( + x: Union[ivy.Array, ivy.NativeArray], /, *, copy: bool = True +) -> np.ndarray: """ - Return the shape of the array ``x``. + Convert an array into a numpy array. Parameters ---------- x - Input array to infer the shape of. - as_array - Whether to return the shape as an array. - Default is False. + input array + copy + whether to copy the array to a new address or not. + Default is ``True``. Returns ------- ret - Shape of the array ``x``. + a numpy array copying all the element of the array ``x``. Examples -------- - >>> x = ivy.array([[-1, 0, 1], [1, 0, -1]]) - >>> y = ivy.shape(x) - >>> z = ivy.shape(x, as_array = True) + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([-1, 0, 1]) + >>> y = ivy.to_numpy(x, copy=True) >>> print(y) - (2, 3) + [-1 0 1] - >>> print(z) - ivy.array([2, 3]) - """ - return current_backend(x).shape(x, as_array=as_array) + >>> x = ivy.array([[-1, 0, 1],[-1, 0, 1], [1,0,-1]]) + >>> y = ivy.to_numpy(x, copy=True) + >>> print(y) + [[-1 0 1] + [-1 0 1] + [ 1 0 -1]] + With :class:`ivy.Container` input: -ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False - - -@handle_exceptions -def set_shape_array_mode(mode: bool) -> None: - """ - Set the mode of returning shape as ivy.Array to the given mode instance. - - Parameter - --------- - mode - boolean whether to return shape as ivy.Array - - Examples - -------- - >>> ivy.set_shape_array_mode(False) - >>> ivy.shape_array_mode - False + >>> x = ivy.Container(a=ivy.array([-1, 0, 1])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([-1, 0, 1], dtype=int32) + } - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True + >>> x = ivy.Container(a=ivy.array([[-1.0, 0., 1.], [-1, 0, 1], [1, 0, -1]]), + ... b=ivy.array([[-1, 0, 0], [1, 0, 1], [1, 1, 1]])) + >>> y = ivy.to_numpy(x) + >>> print(y) + { + a: array([[-1., 0., 1.], + [-1., 0., 1.], + [1., 0., -1.]], dtype=float32), + b: array([[-1, 0, 0], + [1, 0, 1], + [1, 1, 1]], dtype=int32) + } """ - global shape_array_mode_stack - ivy.utils.assertions.check_isinstance(mode, bool) - shape_array_mode_stack.append(mode) - ivy.__setattr__("shape_array_mode", mode, True) + return current_backend(x).to_numpy(x, copy=copy) @handle_exceptions -def unset_shape_array_mode() -> None: - """ - Reset the mode of returning shape as ivy.Array to the previous state. - - Examples - -------- - >>> ivy.set_shape_array_mode(True) - >>> ivy.shape_array_mode - True - - >>> ivy.unset_shape_array_mode() - >>> ivy.shape_array_mode - False - """ - global shape_array_mode_stack - if shape_array_mode_stack: - shape_array_mode_stack.pop(-1) - mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False - ivy.__setattr__("shape_array_mode", mode, True) - - @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@to_native_arrays_and_back +@inputs_to_native_arrays @handle_array_function @handle_device_shifting -def get_num_dims( - x: Union[ivy.Array, ivy.NativeArray], /, *, as_array: bool = False -) -> int: +def to_scalar(x: Union[ivy.Array, ivy.NativeArray], /) -> Number: """ - Return the number of dimensions of the array x. + Convert an array with a single element into a scalar. Parameters ---------- x - Input array to infer the number of dimensions for. - as_array - Whether to return the shape as a array, default False. + Input array with a single element. Returns ------- ret - Shape of the array + a scalar copying the element of the array ``x``. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` instances in place of any of the arguments. - Examples - -------- + Functional Examples + ------------------- + With :class:`ivy.Array` input: - >>> a = ivy.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]], - ... [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) - >>> b = ivy.get_num_dims(a, as_array=False) - >>> print(b) + >>> x = ivy.array([3]) + >>> y = ivy.to_scalar(x) + >>> print(y) 3 - With :class:`ivy.Container` input: + With a mix of :class:`ivy.Container` and :class:`ivy.Array` input: - >>> a = ivy.Container(b = ivy.asarray([[0.,1.,1.],[1.,0.,0.],[8.,2.,3.]])) - >>> print(ivy.get_num_dims(a)) + >>> x = ivy.Container(a=ivy.array([-1]), b=ivy.array([3])) + >>> y = ivy.to_scalar(x) + >>> print(y) { - b: 2 + a: -1, + b: 3 } - >>> b = ivy.get_num_dims(a, as_array=True) - >>> print(b) + >>> x = ivy.Container(a=ivy.array([1]), b=ivy.array([0]), + ... c=ivy.array([-1])) + >>> y = ivy.to_scalar(x) + >>> print(y) { - b: ivy.array(2) + a: 1, + b: 0, + c: -1 } """ - return current_backend(x).get_num_dims(x, as_array=as_array) + return current_backend(x).to_scalar(x) @handle_exceptions -def arg_info(fn: Callable, *, name: Optional[str] = None, idx: Optional[int] = None): +@handle_nestable +def try_else_none(fn: Callable, *args: Any, **kwargs: Any) -> Union[Callable, None]: """ - Return the index and `inspect.Parameter` representation of the specified argument. - In the form of a dict with keys "idx" and "param". + Try and return the function, otherwise return None if an exception was raised during + function execution. Parameters ---------- fn - The function to retrieve the argument information for - name - The name of the argument - idx - the index of the argument in the inputs + Function to try and call and return. + args + list of arguments. + kwargs + dictionay of keyword arguments Returns ------- - ret - a `dict` containing the idx, and the `inspect.Parameter` for the argument, - which itself contains the parameter name, type, and other helpful information. - """ - ivy.utils.assertions.check_all_or_any_fn( - name, - idx, - fn=ivy.exists, - type="any", - limit=[1], - message="exactly one of the keyword arguments name or idx must be provided", - as_array=False, - ) - params = inspect.signature(fn).parameters - if ivy.exists(name): - return {"idx": list(params).index(name), "param": params[name]} - return {"idx": idx, "param": list(params.values())[idx]} + Either the function itself or None if an exception was raised + during function execution. + Examples + -------- + with a function that is executed without any exception: -def _valid_attrib_combinations(fn, backend, dnd_dict, first_attr_name, other_attr_name): - attr_list = () - if hasattr(fn, other_attr_name): - attr_list = getattr(fn, other_attr_name) - if isinstance(attr_list, dict): - attr_list = attr_list.get(backend, ()) - ivy.utils.assertions.check_false( - dnd_dict and attr_list, - f"Cannot specify both {first_attr_name} and {other_attr_name} " - "cannot both be defined for the same function", - ) + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.array([4, 5, 6]) + >>> z = ivy.try_else_none(ivy.add, x, y) + >>> print(z.__name__) + add + with a function that is executed with an exception: -def _is_valid_device_and_dtypes_attributes(fn: Callable) -> bool: - fn_unsupported_dnd = {} - fn_supported_dnd = {} - backend = ivy.current_backend_str() - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_unsupported_dnd.__get__().values())[0], dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, {}) - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype - # if it's a nested dict, unwrap for the current backend - if isinstance(list(fn_supported_dnd.__get__().values())[0], dict): - fn_supported_dnd = fn_supported_dnd.get(backend, {}) + >>> x = ivy.array([1, 2, 3]) + >>> y = 'hemant' + >>> z = ivy.try_else_none(ivy.add,x, y) + >>> print(z) + None + """ + try: + _ = fn(*args, **kwargs) + return fn + except Exception: + return None - ivy.utils.assertions.check_false( - fn_unsupported_dnd and fn_supported_dnd, - "unsupported_device_and_dtype and supported_device_and_dtype cannot" - " both be defined for the same function", - ) - us = "unsupported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_devices") - _valid_attrib_combinations(fn, backend, fn_unsupported_dnd, us, "supported_dtypes") +@handle_exceptions +def unset_array_mode() -> None: + """ + Reset the mode of converting inputs to ivy.NativeArray, then converting outputs back + to ivy.Array to the previous state. - ss = "supported_device_and_dtype" - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_device") - _valid_attrib_combinations(fn, backend, fn_supported_dnd, ss, "unsupported_dtypes") + Examples + -------- + >>> ivy.set_array_mode(False) + >>> ivy.array_mode + False - return True + >>> ivy.unset_shape_array_mode() + >>> ivy.array_mode + True + """ + global array_mode_stack + if array_mode_stack: + array_mode_stack.pop(-1) + mode = array_mode_stack[-1] if array_mode_stack else True + ivy.__setattr__("array_mode", mode, True) -def _all_dnd_combinations(): - all_comb = {} - for device in ivy.all_devices: - all_comb[device] = ivy.all_dtypes - return all_comb +@handle_exceptions +def unset_exception_trace_mode() -> None: + """ + Reset the trace mode to the previously set mode. + Examples + -------- + >>> ivy.set_exception_trace_mode("ivy") + >>> ivy.exception_trace_mode + 'ivy' -def _dnd_dict_intersection(a, b): - res = {} - for device in a: - if device in b: - intersection = set.intersection(set(a[device]), set(b[device])) - if intersection: - res[device] = tuple(intersection) - return res + >>> ivy.unset_exception_trace_mode() + >>> ivy.exception_trace_mode + 'full' + """ + global exception_trace_mode_stack + if exception_trace_mode_stack: + exception_trace_mode_stack.pop(-1) + mode = exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" + ivy.__setattr__("exception_trace_mode", mode, True) -def _dnd_dict_difference(a, b): - res = a - for device in list(a): - if device in b: - difference = set.difference(set(a[device]), set(b[device])) - if difference: - res[device] = tuple(difference) - else: - del res[device] - return res +@handle_exceptions +def unset_inplace_mode() -> None: + """ + Reset the memory management behavior for in-place updates in Ivy to the previous + state. + Examples + -------- + >>> set_inplace_mode('strict') + >>> ivy.inplace_mode + 'strict' -def _dnd_dict_union(a, b): - res = {} - for device in set(list(a) + list(b)): - u1 = set(a.get(device, ())) - u2 = set(b.get(device, ())) - res[device] = tuple(set.union(u1, u2)) + >>> unset_inplace_mode() + >>> ivy.inplace_mode + 'lenient' + """ + global inplace_mode_stack + if inplace_mode_stack: + inplace_mode_stack.pop(-1) + mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" + ivy.__setattr__("inplace_mode", mode, True) - return res +@handle_exceptions +def unset_min_base() -> None: + """ + Reset the global minimum base used by ivy for numerically stable power raising to + the previous value. -def _get_devices_and_dtypes(fn, recurse=False, complement=True): - supported_devices = ivy.function_supported_devices(fn, recurse=recurse) - supported_dtypes = ivy.function_supported_dtypes(fn, recurse=recurse) + Examples + -------- + >>> ivy.set_min_base(1e-07) + >>> y = ivy.min_base + >>> print(y) + 1e-07 - if hasattr(fn, "partial_mixed_handler"): - supported_devices = supported_devices["primary"] - supported_dtypes = supported_dtypes["primary"] + >>> ivy.unset_min_base() + >>> ivy.min_base + 1e-05 + """ + global min_base_stack + if min_base_stack: + min_base_stack.pop(-1) + val = min_base_stack[-1] if min_base_stack else 1e-05 + ivy.__setattr__("min_base", val, True) - supported = {} - # Generate a base supported set from other attributes - for device in supported_devices: - supported[device] = supported_dtypes - is_backend_fn = "backend" in fn.__module__ - is_frontend_fn = "frontend" in fn.__module__ - is_einops_fn = "einops" in fn.__name__ - if not is_backend_fn and not is_frontend_fn and not is_einops_fn: - if complement: - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported +@handle_exceptions +def unset_min_denominator() -> None: + """ + Reset the global minimum denominator used by ivy for numerically stable division to + the previous value. - backend = ivy.current_backend_str() + Examples + -------- + >>> ivy.set_min_denominator(1e-10) + >>> y = ivy.min_denominator + >>> print(y) + 1e-10 - # Their values are formatted like either - # 1. fn.supported_device_and_dtype = {"cpu":("float16",)} - if hasattr(fn, "supported_device_and_dtype"): - fn_supported_dnd = fn.supported_device_and_dtype.__get__() + >>> ivy.unset_min_denominator() + >>> ivy.min_denominator + 1e-12 + """ + global min_denominator_stack + if min_denominator_stack: + min_denominator_stack.pop(-1) + val = min_denominator_stack[-1] if min_denominator_stack else 1e-12 + ivy.__setattr__("min_denominator", val, True) - if "einops" in fn.__name__ and isinstance(fn_supported_dnd, dict): - fn_supported_dnd = fn_supported_dnd.get(backend, supported) - ivy.utils.assertions.check_isinstance(list(fn_supported_dnd.values())[0], tuple) - # dict intersection - supported = _dnd_dict_intersection(supported, fn_supported_dnd) +@handle_exceptions +def unset_nestable_mode() -> None: + """ + Reset the mode of whether to check if function inputs are ivy.Container to the + previous state. - if hasattr(fn, "unsupported_device_and_dtype"): - fn_unsupported_dnd = fn.unsupported_device_and_dtype.__get__() + Examples + -------- + >>> ivy.set_nestable_mode(False) + >>> ivy.nestable_mode + False - if "einops" in fn.__name__ and isinstance(fn_unsupported_dnd, dict): - fn_unsupported_dnd = fn_unsupported_dnd.get(backend, supported) + >>> ivy.unset_nestable_mode() + >>> ivy.nestable_mode + True + """ + global nestable_mode_stack + if nestable_mode_stack: + nestable_mode_stack.pop(-1) + mode = nestable_mode_stack[-1] if nestable_mode_stack else True + ivy.__setattr__("nestable_mode", mode, True) - ivy.utils.assertions.check_isinstance( - list(fn_unsupported_dnd.values())[0], tuple - ) - # dict difference - supported = _dnd_dict_difference(supported, fn_unsupported_dnd) - if complement: - # dict difference - all_comb = _all_dnd_combinations() - supported = _dnd_dict_difference(all_comb, supported) - return supported +@handle_exceptions +def unset_precise_mode() -> None: + """ + Reset the mode of whether to use a promotion table that avoids any precision loss or + a compute effecient table that avoids most wider-than-necessary promotions. + + Examples + -------- + >>> ivy.set_precise_mode(False) + >>> ivy.precise_mode + False + + >>> ivy.unset_precise_mode() + >>> ivy.precise_mode + True + """ + global precise_mode_stack + if precise_mode_stack: + precise_mode_stack.pop(-1) + mode = precise_mode_stack[-1] if precise_mode_stack else True + ivy.__setattr__("precise_mode", mode, True) + _update_promotion_table(precise=mode) @handle_exceptions -@handle_nestable -def function_supported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: +def unset_queue_timeout() -> None: """ - Return the supported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the supported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + Reset the global queue timeout value (in seconds) to the previous state. - Parameters - ---------- - fn - The function to check for the supported device and dtype attribute - recurse - Whether to recurse into used ivy functions. - Default is ``True``. + Examples + -------- + >>> ivy.set_queue_timeout(10.0) + >>> y = ivy.queue_timeout + >>> print(y) + 10.0 - Returns - ------- - ret - Tuple or dict containing the supported devices and dtypes of the function + >>> ivy.unset_queue_timeout() + >>> ivy.queue_timeout + 15.0 """ - ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", - ) + global queue_timeout_stack + if queue_timeout_stack: + queue_timeout_stack.pop(-1) + timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 + ivy.__setattr__("queue_timeout", timeout, True) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_supported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=False), - } - else: - supported_devices_dtypes = _get_devices_and_dtypes(fn, complement=False) - if recurse: - supported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - supported_devices_dtypes, - _dnd_dict_intersection, - function_supported_devices_and_dtypes, - wrapper=lambda x: x, - ) - return supported_devices_dtypes +@handle_exceptions +def unset_shape_array_mode() -> None: + """ + Reset the mode of returning shape as ivy.Array to the previous state. + + Examples + -------- + >>> ivy.set_shape_array_mode(True) + >>> ivy.shape_array_mode + True + + >>> ivy.unset_shape_array_mode() + >>> ivy.shape_array_mode + False + """ + global shape_array_mode_stack + if shape_array_mode_stack: + shape_array_mode_stack.pop(-1) + mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False + ivy.__setattr__("shape_array_mode", mode, True) + + +@handle_exceptions +def unset_show_func_wrapper_trace_mode() -> None: + """ + Reset the mode of whether to show the full stack trace with function wrapping + traces. + + Examples + -------- + >>> ivy.set_show_func_wrapper_trace_mode(False) + >>> ivy.show_func_wrapper_trace_mode + False + + >>> ivy.unset_show_func_wrapper_trace_mode() + >>> ivy.show_func_wrapper_trace_mode + True + """ + global show_func_wrapper_trace_mode_stack + if show_func_wrapper_trace_mode_stack: + show_func_wrapper_trace_mode_stack.pop(-1) + mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True + ) + ivy.__setattr__("show_func_wrapper_trace_mode", mode, True) + + +@handle_exceptions +def unset_tmp_dir() -> None: + """ + Reset the directory for saving temporary files to the previous value. + + Examples + -------- + >>> ivy.set_tmp_dir("/my_dir") + >>> y = ivy.tmp_dir + >>> print(y) + /my_dir + + >>> ivy.unset_tmp_dir() + >>> ivy.tmp_dir + /tmp + """ + global tmp_dir_stack + if tmp_dir_stack: + tmp_dir_stack.pop(-1) + tmp_dr = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" + ivy.__setattr__("tmp_dir", tmp_dr, True) @handle_exceptions @handle_nestable -def function_unsupported_devices_and_dtypes(fn: Callable, recurse: bool = True) -> Dict: +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def value_is_nan( + x: Union[ivy.Array, ivy.NativeArray, Number], + /, + *, + include_infs: bool = True, +) -> bool: """ - Return the unsupported combination of devices and dtypes of the current backend's - function. The function returns a dict containing the unsupported combination of - devices and dtypes of the primary and compositional implementations incase of - partial mixed functions. + Determine whether the single valued array or scalar is of nan type. Parameters ---------- - fn - The function to check for the unsupported device and dtype attribute - recurse - Whether to recurse into used ivy functions. + x + The input to check Input array. + include_infs + Whether to include infs and -infs in the check. Default is ``True``. Returns ------- ret - Tuple or dict containing the unsupported devices and dtypes of the function + Boolean as to whether the input value is a nan or not. + + Examples + -------- + >>> x = ivy.array([451]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False + + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + True + + >>> x = ivy.array([float('inf')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + False + + >>> x = ivy.array([float('nan')]) + >>> y = ivy.value_is_nan(x, include_infs=False) + >>> print(y) + True + + >>> x = ivy.array([0]) + >>> y = ivy.value_is_nan(x) + >>> print(y) + False """ - ivy.utils.assertions.check_true( - _is_valid_device_and_dtypes_attributes(fn), - "supported_device_and_dtypes and unsupported_device_and_dtypes " - "attributes cannot both exist in a particular backend", - ) - if hasattr(fn, "partial_mixed_handler"): - return { - "compositional": function_unsupported_devices_and_dtypes( - fn.compos, recurse=recurse - ), - "primary": _get_devices_and_dtypes(fn, complement=True), - } - else: - unsupported_devices_dtypes = _get_devices_and_dtypes(fn, complement=True) - if recurse: - unsupported_devices_dtypes = ivy.functional.data_type._nested_get( - fn, - unsupported_devices_dtypes, - _dnd_dict_union, - function_unsupported_devices_and_dtypes, - wrapper=lambda x: x, - ) - return unsupported_devices_dtypes + x_scalar = ivy.to_scalar(x) if ivy.is_array(x) else x + if not x_scalar == x: + return True + if include_infs and (x_scalar == INF or x_scalar == -INF): + return True + return False @handle_exceptions @@ -4115,142 +4187,48 @@ def vmap( return current_backend().vmap(func, in_axes, out_axes) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@to_native_arrays_and_back -@handle_device_shifting -def isin( - elements: Union[ivy.Array, ivy.NativeArray], - test_elements: Union[ivy.Array, ivy.NativeArray], - /, - *, - assume_unique: bool = False, - invert: bool = False, -) -> ivy.Array: - """ - Test if each element of elements is in test_elements. - - Parameters - ---------- - elements - input array - test_elements - values against which to test for each input element - assume_unique - If True, assumes both elements and test_elements contain unique elements, - which can speed up the calculation. Default value is False. - invert - If True, inverts the boolean return array, resulting in True values for - elements not in test_elements. Default value is False. - - Returns - ------- - ret - output a boolean array of the same shape as elements that is True for elements - in test_elements and False otherwise. - - Examples - -------- - >>> x = ivy.array([[10, 7, 4], [3, 2, 1]]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y) - ivy.array([[False, False, False], [ True, True, True]]) - - >>> x = ivy.array([3, 2, 1, 0]) - >>> y = ivy.array([1, 2, 3]) - >>> ivy.isin(x, y, invert=True) - ivy.array([False, False, False, True]) - """ - return ivy.current_backend(elements, test_elements).isin( - elements, test_elements, assume_unique=assume_unique, invert=invert - ) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@inputs_to_native_arrays -@handle_device_shifting -def itemsize( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> int: - """ - Return the size of the input array's elements. - - Parameters - ---------- - x - The input array. - - Returns - ------- - ret - An integer specifying the element size in bytes. - - Examples - -------- - >>> x = ivy.array([1,2,3], dtype=ivy.float64) - >>> ivy.itemsize(x) - 8 - - >>> x = ivy.array([1,2,3], dtype=ivy.complex128) - >>> ivy.itemsize(x) - 16 - """ - return ivy.current_backend(x).itemsize(x) - - -@handle_exceptions -@handle_nestable -@handle_device_shifting -def strides( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> Tuple[int]: - """ - Return the input array's strides across each dimension. - - Parameters - ---------- - x - The input array. - - Returns - ------- - ret - A tuple containing the strides. - - Examples - -------- - >>> x = ivy.array([[1, 5, 9], [2, 6, 10]]) - >>> ivy.strides(x) - (4, 8) - """ - if ivy.is_native_array(x) or (ivy.is_ivy_array(x) and x.base is None): - return ivy.to_numpy(x).strides - # if x is an ivy array with a base, - # convert it to a numpy array with the same base: - ret = ivy.to_numpy(x.base) - ivy_numpy = ivy.with_backend("numpy") - for fn, args, kwargs, index in x._manipulation_stack: - ret = ivy_numpy.__dict__[fn](ret, *args, **kwargs) - ret = ret[index] if ivy.exists(index) else ret - return ret.to_native().strides - - -def is_ivy_nested_array(x: Any, /) -> bool: - """ - Determine whether the input x is an Ivy Nested Array. - - Parameters - ---------- - x - The input to check - Returns - ------- - ret - Boolean, whether or not x is an ivy nested array. - """ - return isinstance(x, ivy.NestedArray) +trace_mode_dict["frontend"] = "ivy/functional/frontends" +trace_mode_dict["ivy"] = "ivy/" +trace_mode_dict["full"] = "" +trace_mode_dict["none"] = "" +ivy.precise_mode = precise_mode_stack[-1] if precise_mode_stack else True +ivy.array_mode = array_mode_stack[-1] if array_mode_stack else True +ivy.nestable_mode = nestable_mode_stack[-1] if nestable_mode_stack else True +ivy.exception_trace_mode = ( + exception_trace_mode_stack[-1] if exception_trace_mode_stack else "full" +) +ivy.show_func_wrapper_trace_mode = ( + show_func_wrapper_trace_mode_stack[-1] + if show_func_wrapper_trace_mode_stack + else True +) +# IMPORTANT: assign attribute directly to function instead of wrapper here +einops_reduce.unsupported_dtypes = { + "torch": ("float16",), + "tensorflow": ("complex",), + "paddle": ("complex", "uint8", "int8", "int16", "float16"), +} +ivy.min_denominator = min_denominator_stack[-1] if min_denominator_stack else 1e-12 +ivy.min_base = min_base_stack[-1] if min_base_stack else 1e-05 +stable_pow.unsupported_dtypes = ("bfloat16",) +ivy.queue_timeout = queue_timeout_stack[-1] if queue_timeout_stack else 15.0 +ivy.tmp_dir = tmp_dir_stack[-1] if tmp_dir_stack else "/tmp" +get_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +set_item.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} +inplace_update.unsupported_dtypes = {"torch": ("bfloat16",)} +ivy.inplace_mode = inplace_mode_stack[-1] if inplace_mode_stack else "lenient" +ivy.shape_array_mode = shape_array_mode_stack[-1] if shape_array_mode_stack else False diff --git a/ivy/functional/ivy/gradients.py b/ivy/functional/ivy/gradients.py index fc0c9b2a1f411..72e3961ffc7e9 100644 --- a/ivy/functional/ivy/gradients.py +++ b/ivy/functional/ivy/gradients.py @@ -22,18 +22,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # - - -def _get_duplicate_index_chains(xs): - """Generate a list of duplicate index chains for a given nested structure.""" - duplicate_index_chains = () - if isinstance(xs, ivy.Container): - duplicate_index_chains = xs.cont_duplicate_array_keychains() - elif isinstance(xs, (list, tuple, dict)): - duplicate_index_chains = ivy.duplicate_array_index_chains(xs) - return duplicate_index_chains +# --- Helpers --- # +# --------------- # def _arrays_to_float_variables(xs, xs_grad_idxs=None): @@ -60,6 +50,89 @@ def inner_fn(x): return ivy.nested_map(xs, map_fn, include_derived=True, shallow=False) +def _get_duplicate_index_chains(xs): + """Generate a list of duplicate index chains for a given nested structure.""" + duplicate_index_chains = () + if isinstance(xs, ivy.Container): + duplicate_index_chains = xs.cont_duplicate_array_keychains() + elif isinstance(xs, (list, tuple, dict)): + duplicate_index_chains = ivy.duplicate_array_index_chains(xs) + return duplicate_index_chains + + +def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): + """Extract all relevant results from the output nested structure of a function.""" + + def map_fn(x_): + if ivy.is_array(x_): + x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ + if create_var: + x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ + if len(x_.shape) == 0: + return ivy.to_native(x_) + if reshape: + if x_.size == 1: + if reshape: + return ivy.to_native(ivy.reshape(x_, [])) + return ivy.to_native(x_) + else: + return ivy.to_ivy(x_) + else: + return ivy.to_native(x_) + return x_ + + if ivy.is_array(x): + return [], map_fn(x) + + x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) + arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) + if _check_if_empty(arr_idxs): + return arr_idxs, [] + else: + if idxs is not None: + arr_idxs = [ + arr_idx + for arr_idx in arr_idxs + if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) + ] + arr_values = ivy.multi_index_nest(x, arr_idxs) + arr_idxs = _idxs_to_str(arr_idxs) + return arr_idxs, arr_values + + +def _get_native_y(y): + """Convert all outputs to native arrays.""" + array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) + y_final = [] + if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: + y_final = ivy.multi_index_nest(y, array_idxs) + return y_final + + +def _get_required_float_variables(xs, xs_grad_idxs): + """ + Convert all required arrays to float variables for gradient calculation. + + Also, returns a list of duplicate index chains for the nested + structure. + """ + if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: + xs_grad_idxs = None + duplicate_index_chains = _get_duplicate_index_chains(xs) + xs = _to_ivy(xs) + xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) + xs = _set_duplicates(xs, duplicate_index_chains) + xs_required = _get_required_native_variables(xs, xs_grad_idxs) + required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) + return ( + xs, + xs_grad_idxs, + xs_required, + required_duplicate_index_chains, + duplicate_index_chains, + ) + + def _get_required_native_variables(xs, xs_grad_idxs): """Extract all required native variables from a nested structure.""" # To make sure that only the required arrays are converted to native arrays @@ -102,68 +175,46 @@ def map_fn(x): return xs -def _get_required_float_variables(xs, xs_grad_idxs): - """ - Convert all required arrays to float variables for gradient calculation. - - Also, returns a list of duplicate index chains for the nested - structure. - """ - if (ivy.is_ivy_container(xs) or ivy.is_array(xs)) and xs_grad_idxs == [[0]]: - xs_grad_idxs = None - duplicate_index_chains = _get_duplicate_index_chains(xs) - xs = _to_ivy(xs) - xs = _arrays_to_float_variables(xs, xs_grad_idxs=xs_grad_idxs) - xs = _set_duplicates(xs, duplicate_index_chains) - xs_required = _get_required_native_variables(xs, xs_grad_idxs) - required_duplicate_index_chains = _get_duplicate_index_chains(xs_required) - return ( - xs, - xs_grad_idxs, - xs_required, - required_duplicate_index_chains, - duplicate_index_chains, +def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): + """Get the relevant outputs from the function return value.""" + if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ + [0] + ]: + ret_grad_idxs = None + ret_idxs, ret_values = _get_native_variables_and_indices( + func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape ) + if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): + return func_ret, {} + if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: + y = ret_values[0] + else: + y = ret_values + return ret_grad_idxs, y, ret_idxs -def _get_native_variables_and_indices(x, reshape=True, idxs=None, create_var=False): - """Extract all relevant results from the output nested structure of a function.""" +def _is_variable(x, exclusive=False, to_ignore=None) -> bool: + x = ivy.to_native(x, nested=True, to_ignore=to_ignore) + return ivy.nested_map( + x, + lambda x: current_backend(x).is_variable(x, exclusive=exclusive), + include_derived=True, + shallow=False, + to_ignore=to_ignore, + ) - def map_fn(x_): - if ivy.is_array(x_): - x_ = ivy.to_ivy(x_) if ivy.is_native_array(x_) else x_ - if create_var: - x_ = _variable(x_) if not _is_variable(x_, exclusive=True) else x_ - if len(x_.shape) == 0: - return ivy.to_native(x_) - if reshape: - if x_.size == 1: - if reshape: - return ivy.to_native(ivy.reshape(x_, [])) - return ivy.to_native(x_) - else: - return ivy.to_ivy(x_) - else: - return ivy.to_native(x_) - return x_ - if ivy.is_array(x): - return [], map_fn(x) +def _process_func_ret_and_grads(func_ret, grads, retain_grads): + """ + Stop gradients propagation. - x = ivy.nested_map(x, map_fn, include_derived=True, shallow=False) - arr_idxs = ivy.nested_argwhere(x, lambda x: ivy.is_native_array(x)) - if _check_if_empty(arr_idxs): - return arr_idxs, [] - else: - if idxs is not None: - arr_idxs = [ - arr_idx - for arr_idx in arr_idxs - if "_".join(str(x) for x in arr_idx) in _idxs_to_str(idxs) - ] - arr_values = ivy.multi_index_nest(x, arr_idxs) - arr_idxs = _idxs_to_str(arr_idxs) - return arr_idxs, arr_values + Set the gradients of non-finite values to zero, and stopping + gradient propagation of the function results. + """ + grads = _non_finite_to_zero(grads) + func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) + grads = _to_ivy(grads) + return func_ret, grads def _set_duplicates(xs, duplicate_index_chains): @@ -192,33 +243,6 @@ def _set_duplicates(xs, duplicate_index_chains): return xs -def _get_y_and_ret_idxs(func_ret, ret_grad_idxs, create_var=False, reshape=True): - """Get the relevant outputs from the function return value.""" - if (ivy.is_ivy_container(func_ret) or ivy.is_array(func_ret)) and ret_grad_idxs == [ - [0] - ]: - ret_grad_idxs = None - ret_idxs, ret_values = _get_native_variables_and_indices( - func_ret, idxs=ret_grad_idxs, create_var=create_var, reshape=reshape - ) - if ret_values is None or (isinstance(ret_values, list) and len(ret_values) == 0): - return func_ret, {} - if isinstance(ret_values, list) and len(ret_values) == 1 and ret_grad_idxs is None: - y = ret_values[0] - else: - y = ret_values - return ret_grad_idxs, y, ret_idxs - - -def _get_native_y(y): - """Convert all outputs to native arrays.""" - array_idxs = ivy.nested_argwhere(y, lambda x: ivy.is_native_array(x)) - y_final = [] - if isinstance(array_idxs, list) and np.asarray(array_idxs, "object").size > 0: - y_final = ivy.multi_index_nest(y, array_idxs) - return y_final - - def _stop_grad_and_index(func_ret, retain_grads, grads): """Stop gradient propagation of the function results.""" if not retain_grads: @@ -232,46 +256,6 @@ def _stop_grad_and_index(func_ret, retain_grads, grads): return func_ret, grads -def _process_func_ret_and_grads(func_ret, grads, retain_grads): - """ - Stop gradients propagation. - - Set the gradients of non-finite values to zero, and stopping - gradient propagation of the function results. - """ - grads = _non_finite_to_zero(grads) - func_ret, grads = _stop_grad_and_index(func_ret, retain_grads, grads) - grads = _to_ivy(grads) - return func_ret, grads - - -_check_if_empty = ( - lambda idxs: not isinstance(idxs, list) - or np.asarray(idxs, dtype="object").size == 0 -) - - -_idxs_to_str = lambda idxs: [ - "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) -] - - -_to_ivy = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) - - -_non_finite_to_zero = lambda xs: ivy.nested_map( - xs, - lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, - include_derived=True, - shallow=False, -) - - # Private Variable Helpers # # -------------------------# @@ -284,17 +268,6 @@ def _variable(x): return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) -def _is_variable(x, exclusive=False, to_ignore=None) -> bool: - x = ivy.to_native(x, nested=True, to_ignore=to_ignore) - return ivy.nested_map( - x, - lambda x: current_backend(x).is_variable(x, exclusive=exclusive), - include_derived=True, - shallow=False, - to_ignore=to_ignore, - ) - - def _variable_data( x: Union[ivy.Array, ivy.NativeArray] ) -> Union[ivy.Array, ivy.NativeArray]: @@ -318,128 +291,367 @@ def _variable_data( return ivy.nested_map(ret, ivy.to_ivy, include_derived=True) +# --- Main --- # +# ------------ # + + +# Optimizer Steps # + + @handle_exceptions -@handle_backend_invalid -@handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def stop_gradient( - x: Union[ivy.Array, ivy.NativeArray], +def adam_step( + dcdw: Union[ivy.Array, ivy.NativeArray], + mw: Union[ivy.Array, ivy.NativeArray], + vw: Union[ivy.Array, ivy.NativeArray], + step: Union[int, float], /, *, - preserve_type: bool = True, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Stop gradient computation. + Compute adam step delta, given the derivatives of some cost c with respect to + weights ws, using ADAM update. `[reference] + + `_ Parameters ---------- - x - Array for which to stop the gradient. - preserve_type - Whether to preserve gradient computation on ivy.Array instances. Default is - True. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + mw + running average of the gradients + vw + running average of second moments of the gradients + step + training step + beta1 + gradient forgetting factor (Default value = 0.9) + beta2 + second moment of gradient forgetting factor (Default value = 0.999) + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7) out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the effective grad of adam_step to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The same array x, but with no gradient information. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The adam step delta. Examples -------- With :class:`ivy.Array` inputs: - >>> x = ivy.array([1., 2., 3.]) - >>> y = ivy.stop_gradient(x, preserve_type=True) - >>> print(y) - ivy.array([1., 2., 3.]) + >>> dcdw = ivy.array([1, 2, 3]) + >>> mw = ivy.ones(3) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) + >>> print(adam_step_delta) + (ivy.array([0.2020105 , 0.22187898, 0.24144873]), + ivy.array([0.99999998, 1.09999998, 1.19999998]), + ivy.array([1.00000001, 1.00300001, 1.00800001])) - >>> x = ivy.zeros((2, 3)) - >>> ivy.stop_gradient(x, preserve_type=False, out=x) - >>> print(x) - ivy.array([[0., 0., 0.], - [0., 0., 0.]]) + >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) + >>> mw = ivy.zeros((2,3)) + >>> vw = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.86 + >>> beta2 = 0.95 + >>> epsilon = 1e-6 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + (ivy.array([[ 1., 1., -1.], + [ 1., 1., 1.]]), + ivy.array([[ 0.14, 0.56, -0.42], + [ 0.28, 0.42, 0.07]]), + ivy.array([[0.05 , 0.8 , 0.45 ], + [0.2 , 0.45 , 0.0125]])) - With one :class:`ivy.Container` inputs: + >>> dcdw = ivy.array([0.1, -0.7, 2]) + >>> mw = ivy.ones(1) + >>> vw = ivy.ones(1) + >>> step = ivy.array(3.6) + >>> out = ivy.zeros_like(dcdw) + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) + >>> print(out) + ivy.array([0.17294501, 0.15770318, 0.20863818]) - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.stop_gradient(x, preserve_type=False) - >>> print(y) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + With one :class:`ivy.Container` input: + + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.array([1., 4., 9.]) + >>> vw = ivy.array([0.,]) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), + b: ivy.array([2.02, 4.82, 8.17]) + }, { + a: ivy.array([0.87, 3.61, 8.09]), + b: ivy.array([1.26, 4., 8.48]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> ivy.stop_gradient(x, preserve_type=True, out=x) - >>> print(x) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> vw = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.87 + >>> beta2 = 0.976 + >>> epsilon = 1e-5 + >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, + ... epsilon=epsilon) + >>> print(adam_step_delta) + ({ + a: ivy.array([0., 0.626, 0.626]), + b: ivy.array([0.626, 0.626, 0.626]) + }, { + a: ivy.array([0., 0.13, 0.26]), + b: ivy.array([0.39, 0.52, 0.65]) + }, { + a: ivy.array([0., 0.024, 0.096]), + b: ivy.array([0.216, 0.384, 0.6]) + }) """ - return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) - - -# AutoGrad # + step = float(step) + mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) + dcdw_sqrd = dcdw**2 + vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) + vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 + beta1_pow = beta1**step + beta2_pow = beta2**step + alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) + return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw @handle_exceptions -@handle_device_shifting -def execute_with_gradients( - func, - xs: Union[ivy.Array, ivy.NativeArray], +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def adam_update( + w: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], + lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, - retain_grads: bool = False, - xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], - ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], -) -> Tuple[ivy.Array, ivy.Array]: + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, + stop_gradients: bool = True, + out: Optional[ivy.Array] = None, +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Call function func with input of xs variables, and return the function result - func_ret and the gradients of each output variable w.r.t each input variable, + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, using ADAM update. `[reference] + + `_ Parameters ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - xs - Variables for which to compute the function gradients with respective to. This - can be a single array or an arbitrary nest of arrays. - retain_grads - Whether to retain the gradients of the returned values. (Default value = False) - xs_grad_idxs - Indices of the input arrays to compute gradients with respect to. If None, - gradients are returned with respect to all input arrays. If ``xs`` is an - ``ivy.Array`` or ``ivy.Container``, the default value is ``None``, otherwise the - default value is ``[[0]]``. - ret_grad_idxs - Indices of the returned arrays for which to return computed gradients. If None, - gradients are returned for all returned arrays. If the returned object from the - ``func`` is an ``ivy.Array`` or ``ivy.Container``, the default value is ``None`` - otherwise the default value is ``[[0]]``. + w + Weights of the function to be updated. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. + out + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - the function result func_ret and a dictionary of gradients of each output - variable w.r.t each input variable. + The new function weights ws_new, and also new mw and vw, following the adam + updates. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = 1 + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(updated_weights) + (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), + ivy.array([0.05, 0.02, 0.01]), + ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) + + >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) + >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = 2 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> out = ivy.zeros_like(w) + >>> stop_gradients = True + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, beta2=beta2, + ... epsilon=epsilon, out=out, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ( + ivy.array([[0.92558873, 1.92558754, 2.92558718], + [3.92558694, 1.92558682, 3.92558861], + [5.92558861, 3.92558694, 1.92558718]]), + ivy.array([[0.01, 0.02, 0.03], + [0.04, 0.05, 0.01], + [0.01, 0.05, 0.03]]), + ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], + [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], + [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) + ) + + With one :class:`ivy.Container` input: + + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([0.5, 0.2, 0.4]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(0.01) + >>> step = 2 + >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(updated_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), + ... b=ivy.array([0.3,0.2,0.2])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = 3 + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> stop_gradients = False + >>> lr = ivy.array(0.001) + >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, + ... beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... stop_gradients=stop_gradients) + >>> print(updated_weights) + ({ + a: ivy.array([0.99936122, 1.99936116, 2.99936128]), + b: ivy.array([3.99936128, 4.99936104, 5.99936104]) + }, { + a: ivy.array([0.01, 0.03, 0.03]), + b: ivy.array([0.03, 0.02, 0.02]) + }, { + a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), + b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) + }) + """ + effective_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + return ( + ivy.optimizer_update( + w, effective_grads, lr, stop_gradients=stop_gradients, out=out + ), + mw, + vw, + ) + + +# AutoGrad # + + +@handle_exceptions +@handle_device_shifting +def execute_with_gradients( + func, + xs: Union[ivy.Array, ivy.NativeArray], + /, + *, + retain_grads: bool = False, + xs_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], + ret_grad_idxs: Optional[Sequence[Sequence[Union[str, int]]]] = [[0]], +) -> Tuple[ivy.Array, ivy.Array]: + """ + Call function func with input of xs variables, and return the function result + func_ret and the gradients of each output variable w.r.t each input variable, + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + xs + Variables for which to compute the function gradients with respective to. This + can be a single array or an arbitrary nest of arrays. + retain_grads + Whether to retain the gradients of the returned values. (Default value = False) + xs_grad_idxs + Indices of the input arrays to compute gradients with respect to. If None, + gradients are returned with respect to all input arrays. If ``xs`` is an + ``ivy.Array`` or ``ivy.Container``, the default value is ``None``, otherwise the + default value is ``[[0]]``. + ret_grad_idxs + Indices of the returned arrays for which to return computed gradients. If None, + gradients are returned for all returned arrays. If the returned object from the + ``func`` is an ``ivy.Array`` or ``ivy.Container``, the default value is ``None`` + otherwise the default value is ``[[0]]``. + + Returns + ------- + ret + the function result func_ret and a dictionary of gradients of each output + variable w.r.t each input variable. Examples -------- @@ -483,77 +695,6 @@ def execute_with_gradients( ) -execute_with_gradients.computes_gradients = True - - -@handle_exceptions -def value_and_grad(func: Callable) -> Callable: - """ - Create a function that evaluates both func and the gradient of func. - - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - - Returns - ------- - ret - A function that returns both func and the gradient of func. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> grad_fn = ivy.value_and_grad(func) - >>> value_grad = grad_fn(x) - >>> print(value_grad) - (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], - [0.93333334, 0.43333334, 2.0666666 ]])) - """ - return current_backend(None).value_and_grad(func) - - -value_and_grad.computes_gradients = True - - -@handle_exceptions -def jac(func: Callable) -> Callable: - """ - Call function func, and return func's Jacobian partial derivatives. - - Parameters - ---------- - func - Function for which we compute the gradients of the output with respect to xs - input. - - Returns - ------- - ret - the Jacobian function - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) - >>> func = lambda x: ivy.mean(ivy.square(x)) - >>> jac_fn = ivy.jac(func) - >>> jacobian = jac_fn(x) - >>> print(jacobian) - ivy.array([[1.53 , 0.7 , 1.67 ], - ... [0.933, 0.433, 2.07 ]]) - """ - return current_backend(None).jac(func) - - -jac.computes_gradients = True - - @handle_exceptions def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: """ @@ -585,380 +726,297 @@ def grad(func: Callable, argnums: Union[int, Sequence[int]] = 0) -> Callable: return current_backend(None).grad(func, argnums=argnums) -grad.computes_gradients = True - - -# Optimizer Steps # - - @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def adam_step( +def gradient_descent_update( + w: Union[ivy.Array, ivy.NativeArray], dcdw: Union[ivy.Array, ivy.NativeArray], - mw: Union[ivy.Array, ivy.NativeArray], - vw: Union[ivy.Array, ivy.NativeArray], - step: Union[int, float], + lr: Union[float, ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, + stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Compute adam step delta, given the derivatives of some cost c with respect to - weights ws, using ADAM update. `[reference] - - `_ + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws]. Parameters ---------- + w + Weights of the function to be updated. dcdw Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - mw - running average of the gradients - vw - running average of second moments of the gradients - step - training step - beta1 - gradient forgetting factor (Default value = 0.9) - beta2 - second moment of gradient forgetting factor (Default value = 0.999) - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7) + lr + Learning rate(s), the rate(s) at which the weights should be updated relative to + the gradient. + stop_gradients + Whether to stop the gradients of the variables after each gradient step. + Default is ``True``. out - optional output array, for writing the effective grad of adam_step to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The adam step delta. + The new weights, following the gradient descent updates. Examples -------- With :class:`ivy.Array` inputs: - >>> dcdw = ivy.array([1, 2, 3]) - >>> mw = ivy.ones(3) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step) - >>> print(adam_step_delta) - (ivy.array([0.2020105 , 0.22187898, 0.24144873]), - ivy.array([0.99999998, 1.09999998, 1.19999998]), - ivy.array([1.00000001, 1.00300001, 1.00800001])) - - >>> dcdw = ivy.array([[1., 4., -3.], [2., 3., 0.5]]) - >>> mw = ivy.zeros((2,3)) - >>> vw = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.86 - >>> beta2 = 0.95 - >>> epsilon = 1e-6 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - (ivy.array([[ 1., 1., -1.], - [ 1., 1., 1.]]), - ivy.array([[ 0.14, 0.56, -0.42], - [ 0.28, 0.42, 0.07]]), - ivy.array([[0.05 , 0.8 , 0.45 ], - [0.2 , 0.45 , 0.0125]])) + >>> w = ivy.array([[1., 2, 3], + ... [4, 6, 1], + ... [1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1], + ... [0.3, 0.6, 0.4], + ... [0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) + >>> print(new_weights) + ivy.array([[ 0.95, 1.98, 2.99], + ... [ 3.97, 5.94, 0.96], + ... [ 0.96, -0.07, 6.98]]) - >>> dcdw = ivy.array([0.1, -0.7, 2]) - >>> mw = ivy.ones(1) - >>> vw = ivy.ones(1) - >>> step = ivy.array(3.6) - >>> out = ivy.zeros_like(dcdw) - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, out=out) + >>> w = ivy.array([1., 2., 3.]) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> out = ivy.zeros_like(w) + >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) >>> print(out) - ivy.array([0.17294501, 0.15770318, 0.20863818]) + ivy.array([0.85, 1.94, 2.97]) - With one :class:`ivy.Container` input: + With one :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.array([1., 4., 9.]) - >>> vw = ivy.array([0.,]) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([6.49e+04, 1.74e+01, 1.95e+01]), - b: ivy.array([2.02, 4.82, 8.17]) - }, { - a: ivy.array([0.87, 3.61, 8.09]), - b: ivy.array([1.26, 4., 8.48]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.array([0.5, 0.2, 0.1]) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([3.33, 5.66, 1.95]) + } With multiple :class:`ivy.Container` inputs: - >>> dcdw = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> mw = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> vw = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.87 - >>> beta2 = 0.976 - >>> epsilon = 1e-5 - >>> adam_step_delta = ivy.adam_step(dcdw, mw, vw, step, beta1=beta1, beta2=beta2, - ... epsilon=epsilon) - >>> print(adam_step_delta) - ({ - a: ivy.array([0., 0.626, 0.626]), - b: ivy.array([0.626, 0.626, 0.626]) - }, { - a: ivy.array([0., 0.13, 0.26]), - b: ivy.array([0.39, 0.52, 0.65]) - }, { - a: ivy.array([0., 0.024, 0.096]), - b: ivy.array([0.216, 0.384, 0.6]) - }) + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), + ... b=ivy.array([3.48, 5.72, 1.98])) + >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), + ... b=ivy.array([2., 3.42, 1.69])) + >>> lr = ivy.array(0.3) + >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) + >>> print(w_new) + { + a: ivy.array([0.85, 1.94, 2.97]), + b: ivy.array([2.88, 4.69, 1.47]) + } """ - step = float(step) - mw = ivy.add(beta1 * mw, (1 - beta1) * dcdw) - dcdw_sqrd = dcdw**2 - vw = ivy.add(ivy.multiply(beta2, vw), (1 - beta2) * dcdw_sqrd) - vw_sqrt = ivy.maximum(vw, 0.0) ** 0.5 - beta1_pow = beta1**step - beta2_pow = beta2**step - alpha = (1 - beta2_pow) ** 0.5 / (1 - beta1_pow + epsilon) - return ivy.divide(alpha * mw, vw_sqrt + epsilon, out=out), mw, vw + return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) -adam_step.out_index = 0 +@handle_exceptions +def jac(func: Callable) -> Callable: + """ + Call function func, and return func's Jacobian partial derivatives. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + Returns + ------- + ret + the Jacobian function -# Optimizer Updates # + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> jac_fn = ivy.jac(func) + >>> jacobian = jac_fn(x) + >>> print(jacobian) + ivy.array([[1.53 , 0.7 , 1.67 ], + ... [0.933, 0.433, 2.07 ]]) + """ + return current_backend(None).jac(func) @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def optimizer_update( +def lamb_update( w: Union[ivy.Array, ivy.NativeArray], - effective_grad: Union[ivy.Array, ivy.NativeArray], + dcdw: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], + mw_tm1: Union[ivy.Array, ivy.NativeArray], + vw_tm1: Union[ivy.Array, ivy.NativeArray], + step: int, /, *, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-7, + max_trust_ratio: Union[int, float] = 10, + decay_lambda: float = 0, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> ivy.Array: +) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: """ - Update weights ws of some function, given the true or effective derivatives of some - cost c with respect to ws, [dc/dw for w in ws]. + Update weights ws of some function, given the derivatives of some cost c with + respect to ws, [dc/dw for w in ws], by applying LAMB method. Parameters ---------- w Weights of the function to be updated. - effective_grad - Effective gradients of the cost c with respect to the weights ws, - [dc/dw for w in ws]. + dcdw + Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. + mw_tm1 + running average of the gradients, from the previous time-step. + vw_tm1 + running average of second moments of the gradients, from the previous time-step. + step + training step. + beta1 + gradient forgetting factor (Default value = 0.9). + beta2 + second moment of gradient forgetting factor (Default value = 0.999). + epsilon + divisor during adam update, preventing division by zero (Default value = 1e-7). + max_trust_ratio + The maximum value for the trust ratio. (Default value = 10) + decay_lambda + The factor used for weight decay. (Default value = 0). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the new function weights ws_new to. It must + have a shape that the inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the optimizer updates. + The new function weights ws_new, following the LAMB updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - ivy.array([1., 2., 3.]) - - >>> w = ivy.array([1., 2., 3.]) - >>> effective_grad = ivy.zeros(3) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... out=None, stop_gradients=True) - >>> print(ws_new) - ivy.array([1., 2., 3.]) + >>> w = ivy.array([1., 2, 3]) + >>> dcdw = ivy.array([0.5,0.2,0.1]) + >>> lr = ivy.array(0.1) + >>> vw_tm1 = ivy.zeros(1) + >>> mw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) + >>> print(new_weights) + (ivy.array([0.784, 1.78 , 2.78 ]), + ... ivy.array([0.05, 0.02, 0.01]), + ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) - >>> w = ivy.array([[1., 2.], [4., 5.]]) + >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) + >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) + >>> lr = ivy.array(0.1) + >>> mw_tm1 = ivy.zeros((3,3)) + >>> vw_tm1 = ivy.zeros(3) + >>> step = ivy.array(1) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) - >>> lr = ivy.array([3e-4, 1e-2]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) - >>> print(out) - ivy.array([[0.999, 1.95], - [4., 4.92]]) - - >>> w = ivy.array([1., 2., 3.]) - >>> out = ivy.zeros_like(w) - >>> effective_grad = ivy.array([4., 5., 6.]) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False, out=out) + >>> stop_gradients = True + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, out=out, + ... stop_gradients=stop_gradients) >>> print(out) - ivy.array([0.999, 2. , 3. ]) + ivy.array([[ 0.639, 1.64 , 2.64 ], + ... [ 3.64 , 5.64 , 0.639], + ... [ 0.639, -0.361, 6.64 ]]) - With one :class:`ivy.Container` input: + With one :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.array([0., 0., 0.]) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } + >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) + >>> dcdw = ivy.array([3., 4., 5.]) + >>> mw_tm1 = ivy.array([0., 0., 0.]) + >>> vw_tm1 = ivy.array([0.]) + >>> lr = ivy.array(1.) + >>> step = ivy.array([2]) + >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) + >>> print(new_weights) + ({ + a: ivy.array([1., 2., 3.]), + b: ivy.array([4., 5., 6.]) + }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = 3e-4 - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) - >>> print(w) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), - ... b=ivy.array([0., 0., 0.])) - >>> lr = ivy.array([3e-4]) - >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, - ... stop_gradients=False) - >>> print(ws_new) - { - a: ivy.array([0., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - """ - deltas = effective_grad * lr - w = ivy.subtract(w, deltas, out=out) - if stop_gradients: - return ivy.stop_gradient(w, preserve_type=True, out=out) - return w - - -@handle_exceptions -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def gradient_descent_update( - w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], - /, - *, - stop_gradients: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws]. - - Parameters - ---------- - w - Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The new weights, following the gradient descent updates. - - Examples - -------- - With :class:`ivy.Array` inputs: + >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), + ... b=ivy.array([3.,4.,2.])) + >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), + ... b=ivy.array([0.6,0.4,0.7])) + >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), + ... b=ivy.array([0.,0.,0.])) - >>> w = ivy.array([[1., 2, 3], - ... [4, 6, 1], - ... [1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1], - ... [0.3, 0.6, 0.4], - ... [0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> new_weights = ivy.gradient_descent_update(w, dcdw, lr, stop_gradients=True) + >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), + ... b=ivy.array([0.,])) + >>> step = ivy.array([3.4]) + >>> beta1 = 0.9 + >>> beta2 = 0.999 + >>> epsilon = 1e-7 + >>> max_trust_ratio = 10 + >>> decay_lambda = 0 + >>> stop_gradients = True + >>> lr = ivy.array(0.5) + >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, + ... beta2=beta2, epsilon=epsilon, + ... max_trust_ratio=max_trust_ratio, + ... decay_lambda=decay_lambda, + ... stop_gradients=stop_gradients) >>> print(new_weights) - ivy.array([[ 0.95, 1.98, 2.99], - ... [ 3.97, 5.94, 0.96], - ... [ 0.96, -0.07, 6.98]]) - - >>> w = ivy.array([1., 2., 3.]) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> out = ivy.zeros_like(w) - >>> ivy.gradient_descent_update(w, dcdw, lr, out=out) - >>> print(out) - ivy.array([0.85, 1.94, 2.97]) - - With one :class:`ivy.Container` inputs: - - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.array([0.5, 0.2, 0.1]) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([3.33, 5.66, 1.95]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), - ... b=ivy.array([3.48, 5.72, 1.98])) - >>> dcdw = ivy.Container(a=ivy.array([0.5, 0.2, 0.1]), - ... b=ivy.array([2., 3.42, 1.69])) - >>> lr = ivy.array(0.3) - >>> w_new = ivy.gradient_descent_update(w, dcdw, lr) - >>> print(w_new) - { - a: ivy.array([0.85, 1.94, 2.97]), - b: ivy.array([2.88, 4.69, 1.47]) - } + ({ + a: ivy.array([-0.708, 1.29, 3.29]), + b: ivy.array([1.45, 2.45, 0.445]) + }, { + a: ivy.array([0.02, 0.03, 0.06]), + b: ivy.array([0.06, 0.04, 0.07]) + }, { + a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), + b: ivy.array([0.00036, 0.00016, 0.00049]) + }) """ - return ivy.optimizer_update(w, dcdw, lr, stop_gradients=stop_gradients, out=out) + r1 = ivy.vector_norm(w) + eff_grads, mw, vw = ivy.adam_step( + dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon + ) + if decay_lambda > 0: + r2 = ivy.vector_norm(eff_grads + decay_lambda * w) + else: + r2 = ivy.vector_norm(eff_grads) + r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) + lr = r * lr + return ( + ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), + mw, + vw, + ) @handle_exceptions @@ -1012,338 +1070,264 @@ def lars_update( ) +# Optimizer Updates # + + @handle_exceptions @handle_array_like_without_promotion @inputs_to_ivy_arrays @handle_array_function -def adam_update( +def optimizer_update( w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], + effective_grad: Union[ivy.Array, ivy.NativeArray], lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, stop_gradients: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, using ADAM update. `[reference] - - `_ + Update weights ws of some function, given the true or effective derivatives of some + cost c with respect to ws, [dc/dw for w in ws]. Parameters ---------- w Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. + effective_grad + Effective gradients of the cost c with respect to the weights ws, + [dc/dw for w in ws]. lr Learning rate(s), the rate(s) at which the weights should be updated relative to the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). stop_gradients Whether to stop the gradients of the variables after each gradient step. Default is ``True``. out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new function weights ws_new, and also new mw and vw, following the adam - updates. + The new function weights ws_new, following the optimizer updates. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = 1 - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(updated_weights) - (ivy.array([0.90000075, 1.90000164, 2.9000032 ]), - ivy.array([0.05, 0.02, 0.01]), - ivy.array([2.50000012e-04, 4.00000063e-05, 1.00000016e-05])) + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 2, 3],[4, 2, 4],[6, 4, 2]]) - >>> dcdw = ivy.array([[0.1, 0.2, 0.3],[0.4, 0.5, 0.1],[0.1, 0.5, 0.3]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = 2 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 + >>> w = ivy.array([1., 2., 3.]) + >>> effective_grad = ivy.zeros(3) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... out=None, stop_gradients=True) + >>> print(ws_new) + ivy.array([1., 2., 3.]) + + >>> w = ivy.array([[1., 2.], [4., 5.]]) >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, beta2=beta2, - ... epsilon=epsilon, out=out, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ( - ivy.array([[0.92558873, 1.92558754, 2.92558718], - [3.92558694, 1.92558682, 3.92558861], - [5.92558861, 3.92558694, 1.92558718]]), - ivy.array([[0.01, 0.02, 0.03], - [0.04, 0.05, 0.01], - [0.01, 0.05, 0.03]]), - ivy.array([[1.00000016e-05, 4.00000063e-05, 9.00000086e-05], - [1.60000025e-04, 2.50000012e-04, 1.00000016e-05], - [1.00000016e-05, 2.50000012e-04, 9.00000086e-05]]) - ) + >>> effective_grad = ivy.array([[4., 5.], [7., 8.]]) + >>> lr = ivy.array([3e-4, 1e-2]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=out) + >>> print(out) + ivy.array([[0.999, 1.95], + [4., 4.92]]) + + >>> w = ivy.array([1., 2., 3.]) + >>> out = ivy.zeros_like(w) + >>> effective_grad = ivy.array([4., 5., 6.]) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False, out=out) + >>> print(out) + ivy.array([0.999, 2. , 3. ]) With one :class:`ivy.Container` input: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([0.5, 0.2, 0.4]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(0.01) - >>> step = 2 - >>> updated_weights = ivy.adam_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(updated_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.05, 0.02, 0.04]), ivy.array([0.01024, 0.01003, 0.01015])) + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.array([0., 0., 0.]) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), ... b=ivy.array([3., 4., 5.])) - >>> dcdw = ivy.Container(a=ivy.array([0.1,0.3,0.3]), - ... b=ivy.array([0.3,0.2,0.2])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = 3 - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> stop_gradients = False - >>> lr = ivy.array(0.001) - >>> updated_weights = ivy.adam_update(w, dcdw, lr, mw_tm1, vw_tm1, step, - ... beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... stop_gradients=stop_gradients) - >>> print(updated_weights) - ({ - a: ivy.array([0.99936122, 1.99936116, 2.99936128]), - b: ivy.array([3.99936128, 4.99936104, 5.99936104]) - }, { - a: ivy.array([0.01, 0.03, 0.03]), - b: ivy.array([0.03, 0.02, 0.02]) - }, { - a: ivy.array([1.00000016e-05, 9.00000086e-05, 9.00000086e-05]), - b: ivy.array([9.00000086e-05, 4.00000063e-05, 4.00000063e-05]) - }) - """ - effective_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - return ( - ivy.optimizer_update( - w, effective_grads, lr, stop_gradients=stop_gradients, out=out - ), - mw, - vw, - ) - + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = 3e-4 + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, out=w) + >>> print(w) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } -adam_update.out_index = 0 + >>> w = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> effective_grad = ivy.Container(a=ivy.array([0., 0., 0.]), + ... b=ivy.array([0., 0., 0.])) + >>> lr = ivy.array([3e-4]) + >>> ws_new = ivy.optimizer_update(w, effective_grad, lr, + ... stop_gradients=False) + >>> print(ws_new) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + """ + deltas = effective_grad * lr + w = ivy.subtract(w, deltas, out=out) + if stop_gradients: + return ivy.stop_gradient(w, preserve_type=True, out=out) + return w @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def lamb_update( - w: Union[ivy.Array, ivy.NativeArray], - dcdw: Union[ivy.Array, ivy.NativeArray], - lr: Union[float, ivy.Array, ivy.NativeArray], - mw_tm1: Union[ivy.Array, ivy.NativeArray], - vw_tm1: Union[ivy.Array, ivy.NativeArray], - step: int, +@handle_device_shifting +def stop_gradient( + x: Union[ivy.Array, ivy.NativeArray], /, *, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-7, - max_trust_ratio: Union[int, float] = 10, - decay_lambda: float = 0, - stop_gradients: bool = True, + preserve_type: bool = True, out: Optional[ivy.Array] = None, -) -> Tuple[ivy.Array, ivy.Array, ivy.Array]: +) -> ivy.Array: """ - Update weights ws of some function, given the derivatives of some cost c with - respect to ws, [dc/dw for w in ws], by applying LAMB method. + Stop gradient computation. Parameters ---------- - w - Weights of the function to be updated. - dcdw - Derivates of the cost c with respect to the weights ws, [dc/dw for w in ws]. - lr - Learning rate(s), the rate(s) at which the weights should be updated relative to - the gradient. - mw_tm1 - running average of the gradients, from the previous time-step. - vw_tm1 - running average of second moments of the gradients, from the previous time-step. - step - training step. - beta1 - gradient forgetting factor (Default value = 0.9). - beta2 - second moment of gradient forgetting factor (Default value = 0.999). - epsilon - divisor during adam update, preventing division by zero (Default value = 1e-7). - max_trust_ratio - The maximum value for the trust ratio. (Default value = 10) - decay_lambda - The factor used for weight decay. (Default value = 0). - stop_gradients - Whether to stop the gradients of the variables after each gradient step. - Default is ``True``. + x + Array for which to stop the gradient. + preserve_type + Whether to preserve gradient computation on ivy.Array instances. Default is + True. out - optional output array, for writing the new function weights ws_new to. It must - have a shape that the inputs broadcast to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - The new function weights ws_new, following the LAMB updates. + The same array x, but with no gradient information. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` inputs: - >>> w = ivy.array([1., 2, 3]) - >>> dcdw = ivy.array([0.5,0.2,0.1]) - >>> lr = ivy.array(0.1) - >>> vw_tm1 = ivy.zeros(1) - >>> mw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step) - >>> print(new_weights) - (ivy.array([0.784, 1.78 , 2.78 ]), - ... ivy.array([0.05, 0.02, 0.01]), - ... ivy.array([2.5e-04, 4.0e-05, 1.0e-05])) + >>> x = ivy.array([1., 2., 3.]) + >>> y = ivy.stop_gradient(x, preserve_type=True) + >>> print(y) + ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 2, 3],[4, 6, 1],[1, 0, 7]]) - >>> dcdw = ivy.array([[0.5, 0.2, 0.1],[0.3, 0.6, 0.4],[0.4, 0.7, 0.2]]) - >>> lr = ivy.array(0.1) - >>> mw_tm1 = ivy.zeros((3,3)) - >>> vw_tm1 = ivy.zeros(3) - >>> step = ivy.array(1) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 - >>> out = ivy.zeros_like(w) - >>> stop_gradients = True - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, out=out, - ... stop_gradients=stop_gradients) - >>> print(out) - ivy.array([[ 0.639, 1.64 , 2.64 ], - ... [ 3.64 , 5.64 , 0.639], - ... [ 0.639, -0.361, 6.64 ]]) + >>> x = ivy.zeros((2, 3)) + >>> ivy.stop_gradient(x, preserve_type=False, out=x) + >>> print(x) + ivy.array([[0., 0., 0.], + [0., 0., 0.]]) With one :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1., 2., 3.]), b=ivy.array([4., 5., 6.])) - >>> dcdw = ivy.array([3., 4., 5.]) - >>> mw_tm1 = ivy.array([0., 0., 0.]) - >>> vw_tm1 = ivy.array([0.]) - >>> lr = ivy.array(1.) - >>> step = ivy.array([2]) - >>> new_weights = ivy.lamb_update(w, dcdw, mw_tm1, vw_tm1, lr, step) - >>> print(new_weights) - ({ - a: ivy.array([1., 2., 3.]), - b: ivy.array([4., 5., 6.]) - }, ivy.array([0.3, 0.4, 0.5]), ivy.array([1.01, 1.01, 1.02])) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.stop_gradient(x, preserve_type=False) + >>> print(y) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } With multiple :class:`ivy.Container` inputs: - >>> w = ivy.Container(a=ivy.array([1.,3.,5.]), - ... b=ivy.array([3.,4.,2.])) - >>> dcdw = ivy.Container(a=ivy.array([0.2,0.3,0.6]), - ... b=ivy.array([0.6,0.4,0.7])) - >>> mw_tm1 = ivy.Container(a=ivy.array([0.,0.,0.]), - ... b=ivy.array([0.,0.,0.])) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> ivy.stop_gradient(x, preserve_type=True, out=x) + >>> print(x) + { + a: ivy.array([0., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + """ + return current_backend(x).stop_gradient(x, preserve_type=preserve_type, out=out) - >>> vw_tm1 = ivy.Container(a=ivy.array([0.,]), - ... b=ivy.array([0.,])) - >>> step = ivy.array([3.4]) - >>> beta1 = 0.9 - >>> beta2 = 0.999 - >>> epsilon = 1e-7 - >>> max_trust_ratio = 10 - >>> decay_lambda = 0 - >>> stop_gradients = True - >>> lr = ivy.array(0.5) - >>> new_weights = ivy.lamb_update(w, dcdw, lr, mw_tm1, vw_tm1, step, beta1=beta1, - ... beta2=beta2, epsilon=epsilon, - ... max_trust_ratio=max_trust_ratio, - ... decay_lambda=decay_lambda, - ... stop_gradients=stop_gradients) - >>> print(new_weights) - ({ - a: ivy.array([-0.708, 1.29, 3.29]), - b: ivy.array([1.45, 2.45, 0.445]) - }, { - a: ivy.array([0.02, 0.03, 0.06]), - b: ivy.array([0.06, 0.04, 0.07]) - }, { - a: ivy.array([4.0e-05, 9.0e-05, 3.6e-04]), - b: ivy.array([0.00036, 0.00016, 0.00049]) - }) + +@handle_exceptions +def value_and_grad(func: Callable) -> Callable: """ - r1 = ivy.vector_norm(w) - eff_grads, mw, vw = ivy.adam_step( - dcdw, mw_tm1, vw_tm1, step, beta1=beta1, beta2=beta2, epsilon=epsilon - ) - if decay_lambda > 0: - r2 = ivy.vector_norm(eff_grads + decay_lambda * w) - else: - r2 = ivy.vector_norm(eff_grads) - r = ivy.minimum(ivy.stable_divide(r1, r2), ivy.array(max_trust_ratio)) - lr = r * lr - return ( - ivy.optimizer_update(w, eff_grads, lr, stop_gradients=stop_gradients, out=out), - mw, - vw, - ) + Create a function that evaluates both func and the gradient of func. + + Parameters + ---------- + func + Function for which we compute the gradients of the output with respect to xs + input. + + Returns + ------- + ret + A function that returns both func and the gradient of func. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[4.6, 2.1, 5], [2.8, 1.3, 6.2]]) + >>> func = lambda x: ivy.mean(ivy.square(x)) + >>> grad_fn = ivy.value_and_grad(func) + >>> value_grad = grad_fn(x) + >>> print(value_grad) + (ivy.array(16.42333412), ivy.array([[1.5333333 , 0.69999999, 1.66666675], + [0.93333334, 0.43333334, 2.0666666 ]])) + """ + return current_backend(None).value_and_grad(func) +execute_with_gradients.computes_gradients = True +value_and_grad.computes_gradients = True +jac.computes_gradients = True +grad.computes_gradients = True +adam_step.out_index = 0 +adam_update.out_index = 0 lamb_update.out_index = 0 +_check_if_empty = ( + lambda idxs: not isinstance(idxs, list) + or np.asarray(idxs, dtype="object").size == 0 +) +_idxs_to_str = lambda idxs: [ + "_".join(list(map(lambda x: str(x), idxs[i]))) for i in range(len(idxs)) +] +_to_ivy = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.to_ivy(x) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) +_non_finite_to_zero = lambda xs: ivy.nested_map( + xs, + lambda x: ivy.where(ivy.isfinite(x), x, 0.0) if ivy.is_array(x) else x, + include_derived=True, + shallow=False, +) diff --git a/ivy/functional/ivy/layers.py b/ivy/functional/ivy/layers.py index c8096d4d1fe26..1a239557e8b53 100644 --- a/ivy/functional/ivy/layers.py +++ b/ivy/functional/ivy/layers.py @@ -20,6 +20,111 @@ ) from ivy.utils.exceptions import handle_exceptions + +# --- Helpers --- # +# --------------- # + + +def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): + kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) + if padding == "SAME": + dim_size = dim_size * stride_size + else: + dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) + return dim_size + + +def _depth_max_pooling_helper( + x_shape, kernel, strides, dims, data_format="channel_last" +): + # Determine depth pooling. + # We assume that the kernel and the data have the same data_format. + depth_pooling = False + CHANNEL_LAST = "channel_last" + channel_idx = -1 if data_format == CHANNEL_LAST else 1 + if len(kernel) == dims + 2: + spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] + if kernel[channel_idx] != 1: + depth_pooling = True + if any(i != 1 for i in spatial_kernel): + raise NotImplementedError( + "MaxPooling supports exactly one of pooling across" + " depth or pooling across width/height." + ) + if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to equal the depth" + " stride" + ) + if x_shape[channel_idx] % kernel[channel_idx] != 0: + raise NotImplementedError( + "Depthwise max pooling requires the depth window to evenly divide" + " the input depth" + ) + kernel = [kernel[channel_idx], *[1] * (dims - 1)] + strides = [strides[channel_idx], *[1] * (dims - 1)] + else: + kernel = spatial_kernel + if len(strides) == dims + 2: + strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] + return kernel, strides, depth_pooling + + +def _get_num_padded_values(i, p, n, k, s): + """ + Get number of padded values in a specific window. + + Parameters + ---------- + i window index + p total amount of padding + n input size + k kernel size + s stride + + Returns + ------- + number of padded values in a particular window represented by i + """ + current_index = s * i + left_padding = p // 2 + return max(0, left_padding - current_index) + max( + 0, current_index + k - n - left_padding + ) + + +def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): + if dims == 1: + if data_format == "channel_first": + return "NCW" + else: + return "NWC" + if dims == 2: + if data_format == "channel_first": + return "NCHW" + else: + return "NHWC" + elif dims == 3: + if data_format == "channel_first": + return "NCDHW" + else: + return "NDHWC" + + +# Helpers # + + +def _handle_padding(x, strides, filters, padding): + if isinstance(padding, str) and padding.upper() == "SAME": + if x % strides == 0: + pad = max(filters - strides, 0) + else: + pad = max(filters - (x % strides), 0) + else: + pad = 0 + return pad + + # Extra # # ------# @@ -72,34 +177,134 @@ def _in_projection( ) -# Linear # +def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): + if isinstance(kernel, int): + kernel = (kernel,) * dims + elif len(kernel) == 1: + kernel = (kernel[0],) * dims + elif (len(kernel) != dims) and (len(kernel) != dims + 2): + raise ValueError( + "The kernel should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" + ) + + if isinstance(strides, int): + strides = (strides,) * dims + elif len(strides) == 1: + strides = (strides[0],) * dims + elif (len(strides) != dims) and (len(strides) != dims + 2): + raise ValueError( + "The stride should be an integer, or a tuple of length" + f" {list(set((1, dims, dims+2)))}" + ) + + if isinstance(padding, int): + padding = [(padding,) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == 1: + padding = [(padding[0],) * 2] * dims + elif isinstance(padding, tuple) and len(padding) == dims: + padding = [(padding[i],) * 2 for i in range(dims)] + elif isinstance(padding, list) and len(padding) == dims: + if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): + raise ValueError("Explicit padding must be a list of tuple of two integers") + if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: + raise ValueError( + f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" + ) + + if isinstance(dilation, int): + dilation = (dilation,) * dims + elif len(dilation) == 1: + dilation = (dilation[0],) * dims + elif len(dilation) != dims: + raise ValueError( + f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" + ) + if min(dilation) < 1: + raise ValueError("All values of `dilation` must be positive") + + # Other errors + if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: + raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") + assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" + + # Account for dilation when padding > kernel/2. Not the case in torch by default. + new_kernel = tuple( + [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] + ) + if isinstance(padding, list) and len(padding) == len(new_kernel): + ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) + + return kernel, strides, padding, dilation + + +# --- Main --- # +# ------------ # + + @handle_exceptions -@handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@inputs_to_native_shapes @handle_array_function -def linear( +def conv( x: Union[ivy.Array, ivy.NativeArray], - weight: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], /, *, + transpose: bool = False, + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. - The operation also supports batching of the weight matrices. This is useful if a - batch of different network parameters are to be represented. + """ + Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D + input x respectively and filters arrays. Parameters ---------- x - The input x to compute linear transformation on. - *[outer_batch_shape,inner_batch_shape,in_features]* - weight - The weight matrix. *[outer_batch_shape,out_features,in_features]* + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. + filters + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + transpose + True for computing transpose convolution, and False for dilated convolution. + When True, `x_dilations` must be 1 (the default). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output (Default value = None) + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) bias - The bias vector, default is ``None``. *[outer_batch_shape,out_features]* + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -107,179 +312,92 @@ def linear( Returns ------- ret - Result array of the linear transformation. - *[outer_batch_shape,inner_batch_shape,out_features]* - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1., 2., 3.]) - >>> w = ivy.array([[1., 0., 0.]]) - >>> y = ivy.linear(x, w) - >>> print(y) - ivy.array([1.]) - - >>> x = ivy.array([[0.666, -0.4269, 1.911]]) - >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) - >>> y = ivy.zeros((1, 2)) - >>> ivy.linear(x, w, out=y) - >>> print(y) - ivy.array([[0.666, 1.91 ]]) - - >>> x = ivy.array([[1.546, 5.234, 6.487], - ... [0.157, 5.753, 4.52], - ... [5.165, 3.159, 7.101]]) - >>> w = ivy.array([[1.545, 2.547, 3.124], - ... [5.852, 8.753, 6.963]]) - >>> b = ivy.array([-1., 1.]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.linear(x, w, bias=b, out=y) - >>> print(y) - ivy.array([[ 34.98495483, 101.0293808 ], - [ 28.0159359 , 83.74752808], - [ 37.20942307, 108.3205719 ]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [4., 5., 6.]]), - ... b=ivy.array([1.1, 2.2, 3.3])) - >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], - ... [-1., 1., 2.]]), - ... b=ivy.array([[0., -1., 1.], - ... [0., 1., 1.]])) - >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) - >>> y = ivy.linear(x, w, bias=b) - >>> print(y) - { - a: ivy.array([[15., 6.], - [33., 12.]]), - b: ivy.array([2.1, 6.5]) - } - - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], - ... [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], - ... [7., 13., 17.]])) - >>> w = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.]]) - >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), - ... b=ivy.array([1., 1., 0.])) - >>> ivy.linear(x, w, bias=b, out=x) - >>> print(x) - { - a: ivy.array([[16.4, 35.2, 54.], - [155., 352., 549.]]), - b: ivy.array([[15.1, 32., 47.9], - [85., 196., 306.]]) - } - + The result of the transpose or dilated convolution operation. """ - outer_batch_shape = list(weight.shape[:-2]) - num_outer_batch_dims = len(outer_batch_shape) - inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) - num_inner_batch_dims = len(inner_batch_shape) - num_out_feats, num_in_feats = list(weight.shape[-2:]) - - # OBS x IBS x OF - y = ivy.matmul( - x, - ivy.swapaxes( - ivy.reshape( - weight, - outer_batch_shape - + [1] * max(num_inner_batch_dims - 1, 0) - + [num_out_feats, num_in_feats], - ), - -1, - -2, - ), - ) - - if ivy.exists(bias): - # OBS x [1]*len(IBS) x OF - bias_broadcast = ivy.reshape( - bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] + if transpose: + return conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, + ) + else: + return conv_general_dilated( + x, + filters, + strides, + padding, + dims=dims, + data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - - # OBS x IBS x OF - y = y + bias_broadcast - - if ivy.exists(out): - return ivy.inplace_update(out, y) - return y - - -linear.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} -# Dropout # +# Convolutions # @handle_exceptions +@handle_backend_invalid @handle_nestable -@handle_partial_mixed_function @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@handle_out_argument +@to_native_arrays_and_back @handle_array_function -def dropout( +@handle_device_shifting +def conv1d( x: Union[ivy.Array, ivy.NativeArray], - prob: float, + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - scale: bool = True, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - training: bool = True, - seed: Optional[int] = None, - noise_shape: Optional[Sequence[int]] = None, + data_format: str = "NWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int]] = 1, + dilations: Union[int, Tuple[int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Randomly setting a fraction of input tensor to zeroes with probability. - - `prob` at each update during training time to prevent possible overfitting. - The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that - overall sum is unchanged at training time and inference time. + Compute a 1-D convolution given 3-D input x and filters arrays. Parameters ---------- x - The input array x to perform dropout on. - prob - The probability of zeroing out each array element, float between 0 and 1. - scale - Whether to scale the output by `1/(1-prob)`. Default is ``True``. - dtype - output array data type. If dtype is None, the output array data type - must be inferred from x. Default is ``None``. - training - Turn on dropout if training, turn off otherwise. Default is ``True``. - seed - Set a default seed for random number generating (for reproducibility). Default - is ``None``. - noise_shape - a sequence representing the shape of the binary dropout mask that will be - multiplied with the input. A shape dimension set to None means that a different - mask value will be applied to each element of the input across that dimension. A - dimension set to 1 means the same mask value will be applied to all elements of - the input across that dimension. + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + filters + Convolution filters *[fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + data_format + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -287,7 +405,7 @@ def dropout( Returns ------- ret - Result array after dropout is performed. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -297,175 +415,95 @@ def dropout( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - ivy.array([[ 1.42857146, 2.85714293, 4.28571415], - [ 0. , 7.14285755, 8.5714283 ], - [10. , 11.4285717 , 0. ], - [14.2857151 , 0. , 17.1428566 ]]) - - - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - ivy.array([[ 0. , 5.19999981], - [ 0. , 0. ], - [ 0. , 17.39999962]]) + >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC + >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO + >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) + >>> print(result) + ivy.array([[[0.], [3.], [0.]]]) - >>> x = ivy.array([[1., 2., 3.], - ... [4., 5., 6.], - ... [7., 8., 9.], - ... [10., 11., 12.]]) - >>> y = ivy.dropout(x,0.3,scale=False) - >>> print(y) - ivy.array([[ 1., 2., 3.], - [ 4., 5., 0.], - [ 7., 0., 9.], - [10., 11., 0.]]) + With :class:`ivy.NativeArray` input: - >>> x = ivy.array([[1.5, 2.6], - ... [4.9, 6.6], - ... [7.2, 8.7]]) - >>> y = ivy.dropout(x,0.5,scale=False) - >>> print(y) - ivy.array([[0., 2.6], - [0., 0. ], - [0., 8.7]]) + >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) + >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) + >>> result = ivy.conv1d(x, filters, (2,),'VALID') + >>> print(result) + ivy.array([[[3., 1.], + ... [7., 5.]]]) - With :class:`ivy.Container` input: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], + ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), + ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) + >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) + >>> result = ivy.conv1d(x, filters, 3, 'VALID') + >>> print(result) { - a: ivy.array([[0., 0., 4.28571415], - [5.71428585, 7.14285755, 0.]]), - b: ivy.array([0., 11.4285717, 12.8571434]) - } - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - { - a: ivy.array([[0., 4.4000001, 6.5999999], - [22., 44., 0.]]), - b: ivy.array([[2.49000001, 0.55599999, 8.21000004], - [14., 0., 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), - ... b=ivy.array([7., 8., 9.])) - >>> y = ivy.dropout(x,0.3) - >>> print(y) - { - a: ivy.array([[0., 0., 3.], - [4., 5., 0.]]), - b: ivy.array([0., 8., 9.]) - } - - >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), - ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) - >>> y = ivy.dropout(x,0.5) - >>> print(y) - { - a: ivy.array([[0., 2.2, 3.3], - [11., 22., 0.]]), - b: ivy.array([[1.245, 0.278, 4.105], - [7., 0., 0.]]) + a: ivy.array([[[6., 7.9, 1.2], + ... [15.6, 11.7, 6.1]]]), + ... b: ivy.array([[[15.4, 14.3, 8.8]]]) } """ - if prob == 0 or not training: - if dtype is not None: - x = ivy.astype(x, dtype) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) - if noise_shape is None: - noise_shape = x.shape - else: - noise_shape = list(noise_shape) - for i, v in enumerate(noise_shape): - if v is None: - noise_shape[i] = x.shape[i] - mask = ivy.where( - ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) - < prob, - 0.0, - 1.0, + return current_backend(x).conv1d( + x, + filters, + strides, + padding, + data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - x = x * mask - if scale: - x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) - return x if not ivy.exists(out) else ivy.inplace_update(out, x) - - -dropout.mixed_backend_wrappers = { - "to_add": ( - "handle_backend_invalid", - "handle_out_argument", - "inputs_to_native_arrays", - "outputs_to_ivy_arrays", - "handle_device_shifting", - ), - "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), -} - - -# Attention # @handle_exceptions +@handle_backend_invalid +@handle_nestable @handle_array_like_without_promotion +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back @handle_array_function -def scaled_dot_product_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Union[ivy.Array, ivy.NativeArray], - value: Union[ivy.Array, ivy.NativeArray], +@handle_device_shifting +def conv1d_transpose( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int]], + padding: str, /, *, - scale: Optional[float] = None, - mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - dropout_p: Optional[float] = 0.0, - is_causal: Optional[bool] = False, - training: Optional[bool] = False, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NWC", + dilations: Union[int, Tuple[int]] = 1, + bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Apply scaled dot product attention to inputs x using optional mask. + Compute a 1-D transpose convolution given 3-D input x and filters arrays. Parameters ---------- - query - The queries input array. The shape of queries input array should be in - *[batch_shape,num_queries,feat_dim]*. The queries input array should have the - same size as keys and values. - key - The keys input array. The shape of keys input array should be in - *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same - size as queries and values. - value - The values input array. The shape of values input should be in - *[batch_shape,num_keys,feat_dim]*. The values input array should have the same - size as queries and keys. - scale - The scale float value. - The scale float value is used to scale the query-key pairs before softmax. - mask - The mask input array. The mask to apply to the query-key values. Default is - None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. - dropout_p - Specifies the dropout probablity, if greater than 0.0, dropout is applied - is_causal - If true, assumes causal attention masking - and errors if both `mask` and `is_causal` are set. - training - If True, dropout is used, otherwise dropout is not activated. + x + Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + filters + Convolution filters *[fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) + data_format + The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" + corresponds to input with shape (batch_size, width, channels), while "NCW" + corresponds to input with shape (batch_size, channels, width). + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -473,9 +511,7 @@ def scaled_dot_product_attention( Returns ------- ret - The output following application of scaled dot-product attention. - The output array is the weighted sum produced by the attention score and value. - The shape of output array is *[batch_shape,num_queries,feat_dim]* . + The result of the transpose convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -485,403 +521,227 @@ def scaled_dot_product_attention( -------- With :class:`ivy.Array` input: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) - - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 6) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 64) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) + >>> y = ivy.zeros((1, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) + >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 32) - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + With :class:`ivy.NativeArray` input: - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) + >>> filters = ivy.native_array( + ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 512, 32) - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) + With one :class:`ivy.Container` input: - ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) + >>> x = ivy.full((1, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) + >>> print(y.shape) + { + a: ivy.Shape(1, 10, 1), + b: ivy.Shape(1, 10, 1) + } - >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True, - ... out=out) - >>> print(out) + With multiple :class:`ivy.Container` inputs: - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) - - With :class:`ivy.Container` input: - - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) - >>> print(result) + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a: ivy.array([[[5.19999981, 1.], - ... [2.59249449, 2.68226194], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[0.2, 1.], - ... [2.19603825, 2.9960382], - ... [4.4000001, 5.5999999]]]) + a: { + c: ivy.Shape(1, 28, 3), + d: ivy.Shape(1, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 3), + d: ivy.Shape(1, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 3), + d: ivy.Shape(6, 6, 3) + } } + """ + return current_backend(x).conv1d_transpose( + x, + filters, + strides, + padding, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + bias=bias, + out=out, + ) - >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), - ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) - >>> mask = ivy.Container( - ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), - ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) - ... ) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) - >>> print(result) - { - a: ivy.array([[[4.26894283, 5.40236187], - ... [4.39999437, 5.59999037], - ... [4.4000001, 5.5999999]]]), - b: ivy.array([[[4.35046196, 5.54282808], - ... [4.39989519, 5.5998764], - ... [4.4000001, 5.5999999]]]) - } - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def conv2d( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int]] = 1, + dilations: Union[int, Tuple[int, int]] = 1, + bias: Optional[ivy.Array] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D convolution given 4-D input x and filters arrays. - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... dropout_p=0.1, - ... is_causal=True, - ... training=True) + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + filters + Convolution filters *[fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + data_format + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", + input data formats, while "channel_last" corresponds to "HWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) + dilations + The dilation factor for each dimension of input. (Default value = 1) + bias + Bias array of shape *[d_out]*. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the convolution operation. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]]) + >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], + ... [[[0.]],[[1.]], [[0.]]], + ... [[[0.]],[[1.]], [[0.]]]]) + >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) >>> print(result) + ivy.array([[ + [[2.],[4.],[6.]], + [[3.],[6.],[9.]], + [[2.],[4.],[6.]] + ]]) - ivy.array([[[0.40000001, 1.29999995], - ... [2.19994521, 3.09994531], - ... [4.30000019, 5.30000019]]]) + With one :class:`ivy.Container` input: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) - >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) - >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> out = ivy.zeros(shape=(1, 3, 2)) - >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) - >>> print(out) - ivy.array([[[4.03946018, 5.0280633 ], - ... [4.29981947, 5.29981089], - ... [4.30000019, 5.30000019]]]) + >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]], + ... [[1.], [2.0],[3.]]]])) + >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) + >>> print(result) + { + a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) + } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With multiple :class:`ivy.Container` inputs: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) - >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), + ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[1, 1, 1], + ... [0, 1, 1], + ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') >>> print(result) { - a: ivy.array([[[0.40000001, 1.29999995], - ... [2.06345534, 2.9634552], - ... [4.30000019, 5.30000019]]]), - b: ivy.array([[[0.40000001, 1.29999995], - ... [2.19336844, 3.09336829], - ... [4.30000019, 5.30000019]]]) + a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), + b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), + c:ivy.array([[[[2.], [0.], [0.]], + [[1.], [3.], [0.]], + [[0.], [1.], [2.]] + ]]) } + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) - >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), - ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) - >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) - >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) - >>> result = ivy.scaled_dot_product_attention(q, - ... k, - ... v, - ... scale=1, - ... mask=mask, - ... dropout_p=0.1, - ... training=True) + >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), + ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) + >>> filters = ivy.array([[2, 0, 1], + ... [1, 3, 1], + ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) + >>> result = ivy.conv2d(x, filters, 2, 'SAME') >>> print(result) { - a: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]), - b: ivy.array([[[2.30000019, 3.23333359], - ... [2.30000019, 3.23333359], - ... [2.30000019, 3.23333359]]]) + a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), + b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) } """ - ivy.assertions.check_all( - (not is_causal) or (is_causal and mask is None), - "is_causal and attn_mask cannot be set at the same time", - ) - embed_dim = query.shape[-1] - scale = 1 / (embed_dim**0.5) if not scale else scale - sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale - sim = ivy.dropout(sim, dropout_p, training=training) - if ivy.exists(mask): - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, - ) - elif is_causal: - L = query.shape[-2] # Source sequence length - S = key.shape[-2] # Target sequence length - mask = ivy.tril(ivy.ones((L, S)), k=0) - mask = ivy.astype(mask, ivy.bool) - sim = ivy.where( - ivy.logical_not(mask), - -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, - sim, - ) - attn = ivy.softmax(sim, axis=-1) - result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) - return result if not ivy.exists(out) else ivy.inplace_update(out, result) - - -@handle_exceptions -@handle_nestable -@handle_out_argument -# @handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def multi_head_attention( - query: Union[ivy.Array, ivy.NativeArray], - key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - /, - *, - num_heads: Optional[int] = 8, - scale: Optional[float] = None, - attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - is_causal: Optional[bool] = False, - return_attention_weights: Optional[bool] = False, - average_attention_weights: Optional[bool] = True, - dropout: Optional[float] = 0.0, - training: Optional[bool] = False, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Union[ivy.Array, ivy.NativeArray]: - """ - Apply multi-head attention to inputs x. This is an implementation of multi-headed - attention as described in the paper "Attention is all you Need" (Vaswani et al., - 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each - timestep in `query` attends to the corresponding sequence in `key`, and returns a - fixed-width vector. This layer first projects `query`, `key` and `value`. These are - (effectively) a list of tensors of length `num_attention_heads`, where the - corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, - , key_dim)`, `(batch_size, , - value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are - softmaxed to obtain attention probabilities. The value tensors are then interpolated - by these probabilities, then concatenated back to a single tensor. Finally, the - result tensor with the last dimension as value_dim can take an linear projection and - return. - - Parameters - ---------- - query - query embeddings *[batch_shape,num_queries,query_dim]*. - key - key embeddings *[batch_shape,num_queries,key_dim]*. - value - value embeddings *[batch_shape,num_queries,value_dim]*. - num_heads - The number of attention heads to use. - scale - The value by which to scale the query-key similarity measure before softmax. - attention_mask - The mask to apply to the query-key values. Default is ``None``. - *[batch_shape,num_queries,num_keys]*. - in_proj_weights - The weights used to project query, key and value *[3*E, E]. - q_proj_weights - The weights used to project query if in_proj_weights is None *[new_E, E]. - k_proj_weights - The weights used to project key if in_proj_weights is None *[new_E, E]. - v_proj_weights - The weights used to project value if in_proj_weights is None *[new_E, E]. - out_proj_weights - The weights used to project the output. - in_proj_bias - The bias used when projecting with query, key and value. - out_proj_bias - The bias used when projecting the output. - is_causal - If True, Uses a causal attention mask and ignores provided attention_mask. - return_attention_weights - If True, returns attention_weights alongside the output - as a tuple (output, attenion_weights). Defaults to `False`. - average_attention_weights - If true, indicates that the returned ``attention_weights`` should be averaged - across heads. Otherwise, ``attention_weights`` are provided separately per head. - Note that this flag only has an effect when ``return_attention_weights=True``. - Default: ``True`` (i.e. average weights across heads) - dropout - Specifies the dropout probablity, dropout is applied to attention_weights. - training - If True, dropout is used, otherwise dropout is not activated. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The output following application of multi-head attention. - *[batch_shape,num_queries,out_feat_dim]* if input is batched - otherwise *[num_queries, out_feat_dim] - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - """ - num_dims = query.ndim - ivy.assertions.check_all( - num_dims > 1 and num_dims < 4, - "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" - f" input), got {num_dims}", - ) - if key is None and value is None: - key = value = query - if num_dims == 2: - query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] - if ivy.exists(in_proj_weights): - q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) - elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): - if ivy.exists(in_proj_bias): - b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) - else: - b_q = b_k = b_v = None - q, k, v = ( - ivy.linear(query, q_proj_weights, bias=b_q), - ivy.linear(key, k_proj_weights, bias=b_k), - ivy.linear(value, v_proj_weights, bias=b_v), - ) - else: - q, k, v = query, key, value - batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] - k_seq_length = k.shape[1] - ivy.assertions.check_true( - emb_dim % num_heads == 0, "features must be divisible by number of heads" - ) - dims_per_head = emb_dim // num_heads - # isolate heads - q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 3, 1) - ) - v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( - (0, 2, 1, 3) - ) - # perform bmm - attn_scores = ivy.matmul(q, k) - # scale - scale = 1 / (dims_per_head**0.5) if not scale else scale - attn_scores *= scale - # apply attention mask - if ivy.exists(attention_mask) or is_causal: - if is_causal: - # create causal mask - attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) - attention_mask = attention_mask.astype("bool") - attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) - # perform softmax - attn_weights = ivy.softmax(attn_scores, axis=-1) - # perform dropout - attn_weights = ivy.dropout(attn_weights, dropout, training=training) - # bmm with values - attention_out = ivy.matmul(attn_weights, v) - attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( - (batch_size, q_seq_length, -1) + return current_backend(x).conv2d( + x, + filters, + strides, + padding, + data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, + dilations=dilations, + bias=bias, + out=out, ) - # proj out if out_proj_weight exists - if ivy.exists(out_proj_weights): - attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) - # if input was unbatched, unbatchify the output - if num_dims == 2: - attention_out = attention_out.squeeze(axis=0) - if return_attention_weights: - if average_attention_weights: - attn_weights = attn_weights.mean(axis=1) - if num_dims == 2: - attn_weights = attn_weights.squeeze(axis=0) - return attention_out, attn_weights - else: - return attention_out - - -# Convolutions # @handle_exceptions @@ -889,46 +749,47 @@ def multi_head_attention( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv1d( +def conv2d_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int]], + padding: str, /, *, - data_format: str = "NWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int]] = 1, - dilations: Union[int, Tuple[int]] = 1, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D convolution given 3-D input x and filters arrays. + Compute a 2-D transpose convolution given 4-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. filters - Convolution filters *[fw,d_in,d_out]*. + Convolution filters *[fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). + The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" + corresponds to inputs with shape (batch_size, height, width, channels), while + "NCHW" corresponds to input with shape (batch_size, channels, height, width). filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - input data formats, while "channel_last" corresponds to "WIO", "HWIO", "DHWIO". - x_dilations + Either "channel_first" or "channel_last". "channel_first" corresponds to + "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . + x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations The dilation factor for each dimension of input. (Default value = 1) @@ -941,7 +802,7 @@ def conv1d( Returns ------- ret - The result of the convolution operation. + The result of the transpose convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -950,44 +811,72 @@ def conv1d( Examples -------- With :class:`ivy.Array` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) + ivy.Shape(1, 56, 56, 6) - >>> x = ivy.asarray([[[0.], [3.], [0.]]]) #NWC - >>> filters = ivy.array([[[0.]], [[1.]], [[0.]]]) #WIO - >>> result = ivy.conv1d(x, filters, (1,), 'SAME', data_format='NWC',dilations= (1,)) - >>> print(result) - ivy.array([[[0.], [3.], [0.]]]) - - With :class:`ivy.NativeArray` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) + >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) + >>> print(x.shape) + ivy.Shape(1, 128, 128, 64) - >>> x = ivy.native_array([[[1., 3.], [2., 4.], [5., 7]]]) - >>> filters = ivy.native_array([[[0., 1.], [1., 0.]]]) - >>> result = ivy.conv1d(x, filters, (2,),'VALID') - >>> print(result) - ivy.array([[[3., 1.], - ... [7., 5.]]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) + >>> y = ivy.zeros((1, 258, 258, 32)) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) + >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) + >>> print(y.shape) + ivy.Shape(1, 258, 258, 32) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + With one :class:`ivy.Container` inputs: + >>> x = ivy.full((1, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) + >>> filters = ivy.Container(a=a, b=b) + >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) + >>> print(y.shape) + { + a: ivy.Shape(1, 10, 10, 1), + b: ivy.Shape(1, 10, 10, 1) + } - >>> x = ivy.Container(a=ivy.array([[[1.2, 3.1, 4.8], [5.9, 2.2, 3.3], - ... [10.8, 7.6, 4.9], [6.1, 2.2, 9.5]]]), - ... b=ivy.array([[[8.8, 7.7, 6.6], [1.1, 2.2, 3.5]]])) - >>> filters = ivy.array([[[1., 0., 1.], [0., 1., 0.], [1., 1., 0.]]]) - >>> result = ivy.conv1d(x, filters, 3, 'VALID') - >>> print(result) + With multiple :class:`ivy.Container` inputs: + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') + >>> print(y.shape) { - a: ivy.array([[[6., 7.9, 1.2], - ... [15.6, 11.7, 6.1]]]), - ... b: ivy.array([[[15.4, 14.3, 8.8]]]) + a: { + c: ivy.Shape(1, 28, 28, 3), + d: ivy.Shape(1, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 56, 56, 3), + d: ivy.Shape(1, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 3) + } } """ - return current_backend(x).conv1d( + return current_backend(x).conv2d_transpose( x, filters, strides, padding, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -999,47 +888,52 @@ def conv1d( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv1d_transpose( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int]], - padding: str, +def conv3d( + x: Union[ivy.Array, ivy.NativeArray, ivy.Container], + filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], + strides: Union[int, Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NWC", - dilations: Union[int, Tuple[int]] = 1, + data_format: str = "NDHWC", + filter_format: str = "channel_last", + x_dilations: Union[int, Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int, int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D transpose convolution given 3-D input x and filters arrays. + Compute a 3-D convolution given 5-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*. + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. data_format - The ordering of the dimensions in the input, one of "NWC" or "NCW". "NWC" - corresponds to input with shape (batch_size, width, channels), while "NCW" - corresponds to input with shape (batch_size, channels, width). + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). + filter_format + Either "channel_first" or "channel_last". "channel_first" corresponds + to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". + x_dilations + The dilation factor for each dimension of input. (Default value = 1) dilations The dilation factor for each dimension of input. (Default value = 1) bias - Bias array of shape *[d_out]*. + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1047,7 +941,7 @@ def conv1d_transpose( Returns ------- ret - The result of the transpose convolution operation. + The result of the convolution operation. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1057,139 +951,106 @@ def conv1d_transpose( -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 6]) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 6) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 64, 64]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 64) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 64]) - >>> y = ivy.zeros((1, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 64, 32]) - >>> ivy.conv1d_transpose(x, filters, 1, 'VALID', out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 32) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[1,256,128])) - >>> filters = ivy.native_array( - ... ivy.random_normal(mean=0, std=1, shape=[3, 128, 32])) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 512, 32) + >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], + ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) + >>> filters = ivy.array([[[0.,1.,0.], + ... [0.,1.,0.], + ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) + >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) + >>> print(result) + ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], + [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) With one :class:`ivy.Container` input: - >>> x = ivy.full((1, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv1d_transpose(x, filters, 1, 'VALID', dilations=2) - >>> print(y.shape) + >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) + >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 2, 'SAME') + >>> print(result) { - a: ivy.Shape(1, 10, 1), - b: ivy.Shape(1, 10, 1) + a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) } - With multiple :class:`ivy.Container` inputs: + With multiple :class:`ivy.Container` input: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv1d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) + >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 3, 5, 5, 1]), + ... b = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 5, 32 ,32, 1]), + ... c = ivy.random_normal(mean = 0, std = 1, + ... shape = [1, 32, 32, 32, 1])) + >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) + >>> result = ivy.conv3d(x, filters, 1, 'SAME') + >>> print(result.cont_shapes) { - a: { - c: ivy.Shape(1, 28, 3), - d: ivy.Shape(1, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 3), - d: ivy.Shape(1, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 3), - d: ivy.Shape(6, 6, 3) - } + a: ivy.Shape(1, 3, 5, 5, 3), + b: ivy.Shape(1, 5, 32, 32, 3), + c: ivy.Shape(1, 32, 32, 32, 3) } """ - return current_backend(x).conv1d_transpose( + return current_backend(x).conv3d( x, filters, strides, padding, - output_shape=output_shape, data_format=data_format, + filter_format=filter_format, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, ) +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv2d( +def conv3d_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int, int, int]], + padding: str, /, *, - data_format: str = "NHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int]] = 1, - dilations: Union[int, Tuple[int, int]] = 1, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "NDHWC", + dilations: Union[int, Tuple[int, int, int]] = 1, bias: Optional[ivy.Array] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D convolution given 4-D input x and filters arrays. + Compute a 3-D transpose convolution given 5-D input x and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in,d_out]*. strides The stride of the sliding window for each dimension of input. padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + output_shape + Shape of the output (Default value = None) data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIHW", - input data formats, while "channel_last" corresponds to "HWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) + The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" + corresponds to inputs with shape (batch_size, depth, height, width, channels), + while "NCDHW" corresponds to input with shape (batch_size, channels, depth, + height, width). dilations The dilation factor for each dimension of input. (Default value = 1) bias - Bias array of shape *[d_out]*. + Bias array of shape *[d_out]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1197,83 +1058,85 @@ def conv2d( Returns ------- ret - The result of the convolution operation. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + The result of the transpose convolution operation. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]]) - >>> filters = ivy.array([[[[0.]],[[1.]],[[0.]]], - ... [[[0.]],[[1.]], [[0.]]], - ... [[[0.]],[[1.]], [[0.]]]]) - >>> result = ivy.conv2d(x, filters, 1, 'SAME', data_format='NHWC', dilations=1) - >>> print(result) - ivy.array([[ - [[2.],[4.],[6.]], - [[3.],[6.],[9.]], - [[2.],[4.],[6.]] - ]]) + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) + ivy.Shape(1, 6, 56, 56, 6) - With one :class:`ivy.Container` input: + >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) + >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) + >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') + >>> print(y.shape) + ivy.Shape(1, 9, 258, 258, 32) - >>> x = ivy.Container(a=ivy.array([[[[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]], - ... [[1.], [2.0],[3.]]]])) - >>> filters = ivy.eye(3, 3).reshape((3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv2d(x, filters, 2, 'SAME', data_format='NHWC', dilations= 1) - >>> print(result) + With :class:`ivy.Container` inputs: + + >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) + >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) + >>> x = ivy.Container(a=a, b=b) + >>> filters = ivy.Container(c=c, d=d) + >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') + >>> print(y.shape) { - a:ivy.array([[[[3.], [3.]], [[1.], [5.]]]]) + a: { + c: ivy.Shape(1, 6, 28, 28, 3), + d: ivy.Shape(1, 6, 28, 28, 3) + }, + b: { + c: ivy.Shape(1, 6, 56, 56, 3), + d: ivy.Shape(1, 6, 56, 56, 3) + }, + c: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + }, + d: { + c: ivy.Shape(6, 6, 6, 6, 3), + d: ivy.Shape(6, 6, 6, 6, 3) + } } - With multiple :class:`ivy.Container` inputs: + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(4, 4).reshape((1, 4, 4, 1)), - ... c = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[1, 1, 1], - ... [0, 1, 1], - ... [0, 0, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) + >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) { - a:ivy.array([[[[2.], [0.]], [[1.], [2.]]]]), - b:ivy.array([[[[3.], [0.]], [[1.], [2.]]]]), - c:ivy.array([[[[2.], [0.], [0.]], - [[1.], [3.], [0.]], - [[0.], [1.], [2.]] - ]]) + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) } - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.Container(a = ivy.eye(3, 3).reshape((1, 3, 3, 1)), - ... b = ivy.eye(5, 5).reshape((1, 5, 5, 1))) - >>> filters = ivy.array([[2, 0, 1], - ... [1, 3, 1], - ... [0, 1, 1]], dtype = ivy.float32).reshape((3, 3, 1, 1)) - >>> result = ivy.conv2d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) + >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) + >>> filters = ivy.Container(a = a, b = b) + >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) + >>> print(y.shape) { - a:ivy.array([[[[4.],[0.]],[[1.],[5.]]]]), - b:ivy.array([[[[4.],[0.],[0.]],[[1.],[6.],[0.]],[[0.],[1.],[5.]]]]) + a: ivy.Shape(1, 8, 8, 8, 1), + b: ivy.Shape(1, 8, 8, 8, 1) } """ - return current_backend(x).conv2d( + return current_backend(x).conv3d_transpose( x, filters, strides, padding, + output_shape=output_shape, data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1285,50 +1148,58 @@ def conv2d( @handle_nestable @handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def conv2d_transpose( +def conv_general_dilated( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: str, + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: Union[str, int, Sequence[Tuple[int, int]]], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, - bias: Optional[ivy.Array] = None, + dims: int = 2, + data_format: str = "channel_last", + filter_format: str = "channel_last", + feature_group_count: int = 1, + x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D transpose convolution given 4-D input x and filters arrays. + Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively + and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in,d_out]*. + Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. strides The stride of the sliding window for each dimension of input. padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) + either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no + padding), or a sequence of n (low, high) integer pairs that give the padding to + apply before and after each spatial dimension. + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. data_format - The ordering of the dimensions in the input, one of "NHWC" or "NCHW". "NHWC" - corresponds to inputs with shape (batch_size, height, width, channels), while - "NCHW" corresponds to input with shape (batch_size, channels, height, width). + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to - "OIDHW" input data formats, while "channel_last" corresponds to "DHWIO" . + Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", + "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + (Default value = 1) x_dilations The dilation factor for each dimension of input. (Default value = 1) dilations - The dilation factor for each dimension of input. (Default value = 1) + The dilation factor for each dimension of filter. (Default value = 1) bias Bias array of shape *[d_out]*. out @@ -1339,80 +1210,17 @@ def conv2d_transpose( ------- ret The result of the transpose convolution operation. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 6]) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) - ivy.Shape(1, 56, 56, 6) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 128, 128, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[1, 1, 64, 64]) - >>> ivy.conv2d_transpose(x,filters,1,'VALID',out=x) - >>> print(x.shape) - ivy.Shape(1, 128, 128, 64) - - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 256, 256, 64]) - >>> y = ivy.zeros((1, 258, 258, 32)) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 64, 32]) - >>> ivy.conv2d_transpose(x,filters,[1, 1, 1],'VALID',out=y) - >>> print(y.shape) - ivy.Shape(1, 258, 258, 32) - - With one :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 1, 1]) - >>> filters = ivy.Container(a=a, b=b) - >>> y = ivy.conv2d_transpose(x,filters,1,'VALID',dilations=2) - >>> print(y.shape) - { - a: ivy.Shape(1, 10, 10, 1), - b: ivy.Shape(1, 10, 10, 1) - } - - With multiple :class:`ivy.Container` inputs: - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv2d_transpose(x,filters,2,'SAME') - >>> print(y.shape) - { - a: { - c: ivy.Shape(1, 28, 28, 3), - d: ivy.Shape(1, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 56, 56, 3), - d: ivy.Shape(1, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 3) - } - } """ - return current_backend(x).conv2d_transpose( + return current_backend(x).conv_general_dilated( x, filters, strides, padding, - output_shape=output_shape, + dims=dims, data_format=data_format, + filter_format=filter_format, + feature_group_count=feature_group_count, + x_dilations=x_dilations, dilations=dilations, bias=bias, out=out, @@ -1424,29 +1232,106 @@ def conv2d_transpose( @handle_nestable @handle_array_like_without_promotion @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def depthwise_conv2d( +def conv_general_transpose( x: Union[ivy.Array, ivy.NativeArray], filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], + strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], + padding: str, /, *, - data_format: str = "NHWC", - dilations: Union[int, Tuple[int, int]] = 1, + dims: int = 2, + output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + data_format: str = "channel_last", + dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, + feature_group_count: int = 1, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 2-D depthwise convolution given 4-D input ``x`` and filters arrays. + Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x + respectively and filters arrays. Parameters ---------- x - Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. filters - Convolution filters *[fh,fw,d_in]*. (d_in must be the same as d from x) + Convolution filters *[fd,fh,fw,d_in,d_out]*. + strides + The stride of the sliding window for each dimension of input. + padding + Either ‘SAME’ (padding so that the output's shape is the same as the + input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). + dims + Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. + output_shape + Shape of the output. + data_format + Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", + "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, + while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. + dilations + The dilation factor for each dimension of input. (Default value = 1) + feature_group_count + split input into groups, d_in should be divisible by the number of groups. + bias + Bias array of shape *[d_out]*. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The result of the transpose convolution operation. + """ + return current_backend(x).conv_general_transpose( + x, + filters, + strides, + padding, + dims=dims, + output_shape=output_shape, + data_format=data_format, + dilations=dilations, + feature_group_count=feature_group_count, + bias=bias, + out=out, + ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def depthwise_conv2d( + x: Union[ivy.Array, ivy.NativeArray], + filters: Union[ivy.Array, ivy.NativeArray], + strides: Union[int, Tuple[int, int]], + padding: Union[str, Sequence[Tuple[int, int]]], + /, + *, + data_format: str = "NHWC", + dilations: Union[int, Tuple[int, int]] = 1, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute a 2-D depthwise convolution given 4-D input ``x`` and filters arrays. + + Parameters + ---------- + x + Input image *[batch_size,h,w,d_in]* or *[batch_size,d_in,h,w]*. + filters + Convolution filters *[fh,fw,d_in]*. (d_in must be the same as d from x) strides The stride of the sliding window for each dimension of input. padding @@ -1558,57 +1443,56 @@ def depthwise_conv2d( ) +# Dropout # + + @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv3d( - x: Union[ivy.Array, ivy.NativeArray, ivy.Container], - filters: Union[ivy.Array, ivy.NativeArray, ivy.Container], - strides: Union[int, Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], +def dropout( + x: Union[ivy.Array, ivy.NativeArray], + prob: float, /, *, - data_format: str = "NDHWC", - filter_format: str = "channel_last", - x_dilations: Union[int, Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int, int, int]] = 1, - bias: Optional[ivy.Array] = None, + scale: bool = True, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + training: bool = True, + seed: Optional[int] = None, + noise_shape: Optional[Sequence[int]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 3-D convolution given 5-D input x and filters arrays. + Randomly setting a fraction of input tensor to zeroes with probability. + + `prob` at each update during training time to prevent possible overfitting. + The inputs not set to 0 are scaled up `1 / (1 - prob)` by default, so that + overall sum is unchanged at training time and inference time. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds - to "OIDHW",input data formats, while "channel_last" corresponds to "DHWIO". - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]* + The input array x to perform dropout on. + prob + The probability of zeroing out each array element, float between 0 and 1. + scale + Whether to scale the output by `1/(1-prob)`. Default is ``True``. + dtype + output array data type. If dtype is None, the output array data type + must be inferred from x. Default is ``None``. + training + Turn on dropout if training, turn off otherwise. Default is ``True``. + seed + Set a default seed for random number generating (for reproducibility). Default + is ``None``. + noise_shape + a sequence representing the shape of the binary dropout mask that will be + multiplied with the input. A shape dimension set to None means that a different + mask value will be applied to each element of the input across that dimension. A + dimension set to 1 means the same mask value will be applied to all elements of + the input across that dimension. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1616,7 +1500,7 @@ def conv3d( Returns ------- ret - The result of the convolution operation. + Result array after dropout is performed. Both the description and the type hints above assumes an array input for simplicity, but this function is *nestable*, and therefore also accepts :class:`ivy.Container` @@ -1626,106 +1510,142 @@ def conv3d( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([[[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]], - ... [[1., 2. ,1.], [1., 2. ,1.], [1., 2. ,1.]]]).reshape((1, 3, 3, 3, 1)) - >>> filters = ivy.array([[[0.,1.,0.], - ... [0.,1.,0.], - ... [0.,1.,0.]]]).reshape((1,3,3,1,1)) - >>> result = ivy.conv3d(x, filters, 1, 'SAME', data_format='NDHWC', dilations=1) - >>> print(result) - ivy.array([[[[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]], - [[[2.],[4.],[2.]],[[3.],[6.],[3.]],[[2.],[4.],[2.]]]]]) + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3) + >>> print(y) + ivy.array([[ 1.42857146, 2.85714293, 4.28571415], + [ 0. , 7.14285755, 8.5714283 ], + [10. , 11.4285717 , 0. ], + [14.2857151 , 0. , 17.1428566 ]]) - With one :class:`ivy.Container` input: - >>> x = ivy.Container(a = ivy.ones((1, 3, 3, 3, 1)).astype(ivy.float32)) - >>> filters = ivy.ones((3, 3, 3, 1, 1)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 2, 'SAME') - >>> print(result) + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + ivy.array([[ 0. , 5.19999981], + [ 0. , 0. ], + [ 0. , 17.39999962]]) + + >>> x = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.], + ... [10., 11., 12.]]) + >>> y = ivy.dropout(x,0.3,scale=False) + >>> print(y) + ivy.array([[ 1., 2., 3.], + [ 4., 5., 0.], + [ 7., 0., 9.], + [10., 11., 0.]]) + + >>> x = ivy.array([[1.5, 2.6], + ... [4.9, 6.6], + ... [7.2, 8.7]]) + >>> y = ivy.dropout(x,0.5,scale=False) + >>> print(y) + ivy.array([[0., 2.6], + [0., 0. ], + [0., 8.7]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) { - a: ivy.array([[[[[8.],[8.]],[[8.],[8.]]],[[[8.],[8.]],[[8.],[8.]]]]]) + a: ivy.array([[0., 0., 4.28571415], + [5.71428585, 7.14285755, 0.]]), + b: ivy.array([0., 11.4285717, 12.8571434]) } - With multiple :class:`ivy.Container` input: + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 4.4000001, 6.5999999], + [22., 44., 0.]]), + b: ivy.array([[2.49000001, 0.55599999, 8.21000004], + [14., 0., 0.]]) + } - >>> x = ivy.Container( a = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 3, 5, 5, 1]), - ... b = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 5, 32 ,32, 1]), - ... c = ivy.random_normal(mean = 0, std = 1, - ... shape = [1, 32, 32, 32, 1])) - >>> filters = ivy.ones((3, 5, 5, 1, 3)).astype(ivy.float32) - >>> result = ivy.conv3d(x, filters, 1, 'SAME') - >>> print(result.cont_shapes) + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], [4., 5., 6.]]), + ... b=ivy.array([7., 8., 9.])) + >>> y = ivy.dropout(x,0.3) + >>> print(y) { - a: ivy.Shape(1, 3, 5, 5, 3), - b: ivy.Shape(1, 5, 32, 32, 3), - c: ivy.Shape(1, 32, 32, 32, 3) + a: ivy.array([[0., 0., 3.], + [4., 5., 0.]]), + b: ivy.array([0., 8., 9.]) } - """ - return current_backend(x).conv3d( - x, - filters, - strides, - padding, - data_format=data_format, - filter_format=filter_format, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, - ) + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], [7., 13., 17.]])) + >>> y = ivy.dropout(x,0.5) + >>> print(y) + { + a: ivy.array([[0., 2.2, 3.3], + [11., 22., 0.]]), + b: ivy.array([[1.245, 0.278, 4.105], + [7., 0., 0.]]) + } + """ + if prob == 0 or not training: + if dtype is not None: + x = ivy.astype(x, dtype) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + if noise_shape is None: + noise_shape = x.shape + else: + noise_shape = list(noise_shape) + for i, v in enumerate(noise_shape): + if v is None: + noise_shape[i] = x.shape[i] + mask = ivy.where( + ivy.random_uniform(shape=noise_shape, device=ivy.dev(x), dtype=dtype, seed=seed) + < prob, + 0.0, + 1.0, + ) + x = x * mask + if scale: + x = ivy.multiply(x, 1.0 / (1.0 - prob), out=out) + return x if not ivy.exists(out) else ivy.inplace_update(out, x) + +# Linear # @handle_exceptions -@handle_backend_invalid @handle_nestable +@handle_partial_mixed_function @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv3d_transpose( +def linear( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int, int, int]], - padding: str, + weight: Union[ivy.Array, ivy.NativeArray], /, *, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "NDHWC", - dilations: Union[int, Tuple[int, int, int]] = 1, - bias: Optional[ivy.Array] = None, + bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Compute a 3-D transpose convolution given 5-D input x and filters arrays. + """Apply a linear transformation to the incoming data: y = x * t(weight) + bias. + The operation also supports batching of the weight matrices. This is useful if a + batch of different network parameters are to be represented. Parameters ---------- x - Input volume *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - output_shape - Shape of the output (Default value = None) - data_format - The ordering of the dimensions in the input, one of "NDHWC" or "NCDHW". "NDHWC" - corresponds to inputs with shape (batch_size, depth, height, width, channels), - while "NCDHW" corresponds to input with shape (batch_size, channels, depth, - height, width). - dilations - The dilation factor for each dimension of input. (Default value = 1) + The input x to compute linear transformation on. + *[outer_batch_shape,inner_batch_shape,in_features]* + weight + The weight matrix. *[outer_batch_shape,out_features,in_features]* bias - Bias array of shape *[d_out]* + The bias vector, default is ``None``. *[outer_batch_shape,out_features]* out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1733,228 +1653,303 @@ def conv3d_transpose( Returns ------- ret - The result of the transpose convolution operation. + Result array of the linear transformation. + *[outer_batch_shape,inner_batch_shape,out_features]* + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 3, 6]) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) - ivy.Shape(1, 6, 56, 56, 6) + >>> x = ivy.array([1., 2., 3.]) + >>> w = ivy.array([[1., 0., 0.]]) + >>> y = ivy.linear(x, w) + >>> print(y) + ivy.array([1.]) - >>> x = ivy.random_normal(mean=0, std=1, shape=[1, 7, 256, 256, 64]) - >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 64, 32]) - >>> y = ivy.conv3d_transpose(x, filters, [1, 1, 1], 'VALID') - >>> print(y.shape) - ivy.Shape(1, 9, 258, 258, 32) + >>> x = ivy.array([[0.666, -0.4269, 1.911]]) + >>> w = ivy.array([[1., 0., 0.], [0., 0., 1.]]) + >>> y = ivy.zeros((1, 2)) + >>> ivy.linear(x, w, out=y) + >>> print(y) + ivy.array([[0.666, 1.91 ]]) - With :class:`ivy.Container` inputs: + >>> x = ivy.array([[1.546, 5.234, 6.487], + ... [0.157, 5.753, 4.52], + ... [5.165, 3.159, 7.101]]) + >>> w = ivy.array([[1.545, 2.547, 3.124], + ... [5.852, 8.753, 6.963]]) + >>> b = ivy.array([-1., 1.]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.linear(x, w, bias=b, out=y) + >>> print(y) + ivy.array([[ 34.98495483, 101.0293808 ], + [ 28.0159359 , 83.74752808], + [ 37.20942307, 108.3205719 ]]) - >>> a = ivy.random_normal(mean=0, std=1, shape=[1, 3, 14, 14, 3]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[1, 3, 28, 28, 3]) - >>> c = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> d = ivy.random_normal(mean=0, std=1, shape=[6, 3, 3, 3, 3]) - >>> x = ivy.Container(a=a, b=b) - >>> filters = ivy.Container(c=c, d=d) - >>> y = ivy.conv3d_transpose(x, filters, 2, 'SAME') - >>> print(y.shape) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [4., 5., 6.]]), + ... b=ivy.array([1.1, 2.2, 3.3])) + >>> w = ivy.Container(a=ivy.array([[1., 2., 3.], + ... [-1., 1., 2.]]), + ... b=ivy.array([[0., -1., 1.], + ... [0., 1., 1.]])) + >>> b = ivy.Container(a=ivy.array([1., -1.]), b=ivy.array([1., 1.])) + >>> y = ivy.linear(x, w, bias=b) + >>> print(y) { - a: { - c: ivy.Shape(1, 6, 28, 28, 3), - d: ivy.Shape(1, 6, 28, 28, 3) - }, - b: { - c: ivy.Shape(1, 6, 56, 56, 3), - d: ivy.Shape(1, 6, 56, 56, 3) - }, - c: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - }, - d: { - c: ivy.Shape(6, 6, 6, 6, 3), - d: ivy.Shape(6, 6, 6, 6, 3) - } + a: ivy.array([[15., 6.], + [33., 12.]]), + b: ivy.array([2.1, 6.5]) } With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - >>> x = ivy.full((1, 6, 6, 6, 1), 2.7) - >>> a = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> b = ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1]) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) + >>> x = ivy.Container(a=ivy.array([[1.1, 2.2, 3.3], + ... [11., 22., 33.]]), + ... b=ivy.array([[1.245, 0.278, 4.105], + ... [7., 13., 17.]])) + >>> w = ivy.array([[1., 2., 3.], + ... [4., 5., 6.], + ... [7., 8., 9.]]) + >>> b = ivy.Container(a=ivy.array([1., 0., -1.]), + ... b=ivy.array([1., 1., 0.])) + >>> ivy.linear(x, w, bias=b, out=x) + >>> print(x) { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) + a: ivy.array([[16.4, 35.2, 54.], + [155., 352., 549.]]), + b: ivy.array([[15.1, 32., 47.9], + [85., 196., 306.]]) } - - >>> x = ivy.full((1, 6, 6, 6, 1), 1.23) - >>> a = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> b = ivy.array(ivy.random_normal(mean=0, std=1, shape=[3, 3, 3, 1, 1])) - >>> filters = ivy.Container(a = a, b = b) - >>> y = ivy.conv3d_transpose(x, filters, 1, 'VALID', dilations=1) - >>> print(y.shape) - { - a: ivy.Shape(1, 8, 8, 8, 1), - b: ivy.Shape(1, 8, 8, 8, 1) - } """ - return current_backend(x).conv3d_transpose( + outer_batch_shape = list(weight.shape[:-2]) + num_outer_batch_dims = len(outer_batch_shape) + inner_batch_shape = list(x.shape[num_outer_batch_dims:-1]) + num_inner_batch_dims = len(inner_batch_shape) + num_out_feats, num_in_feats = list(weight.shape[-2:]) + + # OBS x IBS x OF + y = ivy.matmul( x, - filters, - strides, - padding, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - bias=bias, - out=out, + ivy.swapaxes( + ivy.reshape( + weight, + outer_batch_shape + + [1] * max(num_inner_batch_dims - 1, 0) + + [num_out_feats, num_in_feats], + ), + -1, + -2, + ), ) + if ivy.exists(bias): + # OBS x [1]*len(IBS) x OF + bias_broadcast = ivy.reshape( + bias, outer_batch_shape + [1] * num_inner_batch_dims + [num_out_feats] + ) + + # OBS x IBS x OF + y = y + bias_broadcast + + if ivy.exists(out): + return ivy.inplace_update(out, y) + return y + + +# LSTM # + @handle_exceptions -@handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv_general_dilated( +def lstm_update( x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, int, Sequence[Tuple[int, int]]], + init_h: Union[ivy.Array, ivy.NativeArray], + init_c: Union[ivy.Array, ivy.NativeArray], + kernel: Union[ivy.Array, ivy.NativeArray], + recurrent_kernel: Union[ivy.Array, ivy.NativeArray], /, *, - dims: int = 2, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[ivy.Array, ivy.Array]: """ - Compute a 1-D, 2-D, and 3-D convolution given 3-D, 4-D and 5-D input x respectively - and filters arrays. + Perform long-short term memory update by unrolling time dimension of input array. Parameters ---------- x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of filter. (Default value = 1) + input tensor of LSTM layer *[batch_shape, t, in]*. + init_h + initial state tensor for the cell output *[batch_shape, out]*. + init_c + initial state tensor for the cell hidden state *[batch_shape, out]*. + kernel + weights for cell kernel *[in, 4 x out]*. + recurrent_kernel + weights for cell recurrent kernel *[out, 4 x out]*. bias - Bias array of shape *[d_out]*. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + bias for cell kernel *[4 x out]*. (Default value = None) + recurrent_bias + bias for cell recurrent kernel *[4 x out]*. (Default value = None) Returns ------- ret - The result of the transpose convolution operation. + hidden state for all timesteps *[batch_shape,t,out]* and cell state for last + timestep *[batch_shape,out]* """ - return current_backend(x).conv_general_dilated( - x, - filters, - strides, - padding, - dims=dims, - data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, + # get shapes + x_shape = list(x.shape) + batch_shape = x_shape[:-2] + timesteps = x_shape[-2] + input_channels = x_shape[-1] + x_flat = ivy.reshape(x, (-1, input_channels)) + + # input kernel + Wi = kernel + Wi_x = ivy.reshape( + ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), + batch_shape + [timesteps, -1], ) + Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) + + # recurrent kernel + Wh = recurrent_kernel + + # lstm states + ht = init_h + ct = init_c + + # lstm outputs + hts_list = list() + + # unrolled time dimension with lstm steps + for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( + ivy.unstack(Wii_x, axis=-2), + ivy.unstack(Wif_x, axis=-2), + ivy.unstack(Wig_x, axis=-2), + ivy.unstack(Wio_x, axis=-2), + ): + htm1 = ht + ctm1 = ct + + Wh_htm1 = ivy.matmul(htm1, Wh) + ( + recurrent_bias if recurrent_bias is not None else 0 + ) + Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( + Wh_htm1, num_or_size_splits=4, axis=-1 + ) + + it = ivy.sigmoid(Wii_xt + Whi_htm1) + ft = ivy.sigmoid(Wif_xt + Whf_htm1) + gt = ivy.tanh(Wig_xt + Whg_htm1) + ot = ivy.sigmoid(Wio_xt + Who_htm1) + ct = ft * ctm1 + it * gt + ht = ot * ivy.tanh(ct) + + hts_list.append(ivy.expand_dims(ht, axis=-2)) + + return ivy.concat(hts_list, axis=-2), ct @handle_exceptions -@handle_backend_invalid @handle_nestable -@handle_array_like_without_promotion @handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back +# @handle_array_like_without_promotion +@inputs_to_ivy_arrays @handle_array_function -@handle_device_shifting -def conv_general_transpose( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: str, +def multi_head_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + value: Optional[Union[ivy.Array, ivy.NativeArray]] = None, /, *, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - feature_group_count: int = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + num_heads: Optional[int] = 8, + scale: Optional[float] = None, + attention_mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + q_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + k_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + v_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_weights: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + in_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + out_proj_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + is_causal: Optional[bool] = False, + return_attention_weights: Optional[bool] = False, + average_attention_weights: Optional[bool] = True, + dropout: Optional[float] = 0.0, + training: Optional[bool] = False, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Union[ivy.Array, ivy.NativeArray]: """ - Compute a 1-D, 2-D, and 3-D transpose convolution given 3-D, 4-D and 5-D input x - respectively and filters arrays. + Apply multi-head attention to inputs x. This is an implementation of multi-headed + attention as described in the paper "Attention is all you Need" (Vaswani et al., + 2017). If `query`, `key`, `value` are the same, then this is self-attention. Each + timestep in `query` attends to the corresponding sequence in `key`, and returns a + fixed-width vector. This layer first projects `query`, `key` and `value`. These are + (effectively) a list of tensors of length `num_attention_heads`, where the + corresponding shapes are `(batch_size, , key_dim)`, `(batch_size, + , key_dim)`, `(batch_size, , + value_dim)`. Then, the query and key tensors are dot-producted and scaled. These are + softmaxed to obtain attention probabilities. The value tensors are then interpolated + by these probabilities, then concatenated back to a single tensor. Finally, the + result tensor with the last dimension as value_dim can take an linear projection and + return. Parameters ---------- - x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - Either ‘SAME’ (padding so that the output's shape is the same as the - input's), or ‘VALID’ (padding so that the output's shape is `output_shape`). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output. - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - dilations - The dilation factor for each dimension of input. (Default value = 1) - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - bias - Bias array of shape *[d_out]*. + query + query embeddings *[batch_shape,num_queries,query_dim]*. + key + key embeddings *[batch_shape,num_queries,key_dim]*. + value + value embeddings *[batch_shape,num_queries,value_dim]*. + num_heads + The number of attention heads to use. + scale + The value by which to scale the query-key similarity measure before softmax. + attention_mask + The mask to apply to the query-key values. Default is ``None``. + *[batch_shape,num_queries,num_keys]*. + in_proj_weights + The weights used to project query, key and value *[3*E, E]. + q_proj_weights + The weights used to project query if in_proj_weights is None *[new_E, E]. + k_proj_weights + The weights used to project key if in_proj_weights is None *[new_E, E]. + v_proj_weights + The weights used to project value if in_proj_weights is None *[new_E, E]. + out_proj_weights + The weights used to project the output. + in_proj_bias + The bias used when projecting with query, key and value. + out_proj_bias + The bias used when projecting the output. + is_causal + If True, Uses a causal attention mask and ignores provided attention_mask. + return_attention_weights + If True, returns attention_weights alongside the output + as a tuple (output, attenion_weights). Defaults to `False`. + average_attention_weights + If true, indicates that the returned ``attention_weights`` should be averaged + across heads. Otherwise, ``attention_weights`` are provided separately per head. + Note that this flag only has an effect when ``return_attention_weights=True``. + Default: ``True`` (i.e. average weights across heads) + dropout + Specifies the dropout probablity, dropout is applied to attention_weights. + training + If True, dropout is used, otherwise dropout is not activated. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -1962,86 +1957,140 @@ def conv_general_transpose( Returns ------- ret - The result of the transpose convolution operation. - """ - return current_backend(x).conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - + The output following application of multi-head attention. + *[batch_shape,num_queries,out_feat_dim]* if input is batched + otherwise *[num_queries, out_feat_dim] + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + """ + num_dims = query.ndim + ivy.assertions.check_all( + num_dims > 1 and num_dims < 4, + "Number of dimensions should be 2 (for unbatched input) or 3 (for batched" + f" input), got {num_dims}", + ) + if key is None and value is None: + key = value = query + if num_dims == 2: + query, key, value = [ivy.expand_dims(x, axis=0) for x in [query, key, value]] + if ivy.exists(in_proj_weights): + q, k, v = _in_projection(query, key, value, w=in_proj_weights, b=in_proj_bias) + elif all([ivy.exists(x) for x in [q_proj_weights, k_proj_weights, v_proj_weights]]): + if ivy.exists(in_proj_bias): + b_q, b_k, b_v = ivy.split(in_proj_bias, num_or_size_splits=3) + else: + b_q = b_k = b_v = None + q, k, v = ( + ivy.linear(query, q_proj_weights, bias=b_q), + ivy.linear(key, k_proj_weights, bias=b_k), + ivy.linear(value, v_proj_weights, bias=b_v), + ) + else: + q, k, v = query, key, value + batch_size, q_seq_length, emb_dim = q.shape[0], q.shape[1], q.shape[-1] + k_seq_length = k.shape[1] + ivy.assertions.check_true( + emb_dim % num_heads == 0, "features must be divisible by number of heads" + ) + dims_per_head = emb_dim // num_heads + # isolate heads + q = q.reshape((batch_size, q_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + k = k.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 3, 1) + ) + v = v.reshape((batch_size, k_seq_length, num_heads, dims_per_head)).permute_dims( + (0, 2, 1, 3) + ) + # perform bmm + attn_scores = ivy.matmul(q, k) + # scale + scale = 1 / (dims_per_head**0.5) if not scale else scale + attn_scores *= scale + # apply attention mask + if ivy.exists(attention_mask) or is_causal: + if is_causal: + # create causal mask + attention_mask = ivy.tril(ivy.ones((q_seq_length, k_seq_length))) + attention_mask = attention_mask.astype("bool") + attn_scores = ivy.where(attention_mask, attn_scores, -ivy.inf) + # perform softmax + attn_weights = ivy.softmax(attn_scores, axis=-1) + # perform dropout + attn_weights = ivy.dropout(attn_weights, dropout, training=training) + # bmm with values + attention_out = ivy.matmul(attn_weights, v) + attention_out = attention_out.permute_dims((0, 2, 1, 3)).reshape( + (batch_size, q_seq_length, -1) + ) + # proj out if out_proj_weight exists + if ivy.exists(out_proj_weights): + attention_out = ivy.linear(attention_out, out_proj_weights, bias=out_proj_bias) + # if input was unbatched, unbatchify the output + if num_dims == 2: + attention_out = attention_out.squeeze(axis=0) + if return_attention_weights: + if average_attention_weights: + attn_weights = attn_weights.mean(axis=1) + if num_dims == 2: + attn_weights = attn_weights.squeeze(axis=0) + return attention_out, attn_weights + else: + return attention_out + + +# Attention # + @handle_exceptions @handle_array_like_without_promotion -@handle_out_argument -@inputs_to_native_shapes @handle_array_function -def conv( - x: Union[ivy.Array, ivy.NativeArray], - filters: Union[ivy.Array, ivy.NativeArray], - strides: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]], - padding: Union[str, Sequence[Tuple[int, int]]], +def scaled_dot_product_attention( + query: Union[ivy.Array, ivy.NativeArray], + key: Union[ivy.Array, ivy.NativeArray], + value: Union[ivy.Array, ivy.NativeArray], /, *, - transpose: bool = False, - dims: int = 2, - output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - data_format: str = "channel_last", - filter_format: str = "channel_last", - feature_group_count: int = 1, - x_dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - dilations: Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]] = 1, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + scale: Optional[float] = None, + mask: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + dropout_p: Optional[float] = 0.0, + is_causal: Optional[bool] = False, + training: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Compute a 1-D, 2-D, and 3-D transpose or dilated convolution given 3-D, 4-D and 5-D - input x respectively and filters arrays. + Apply scaled dot product attention to inputs x using optional mask. Parameters ---------- - x - Input image *[batch_size,d,h,w,d_in]* or *[batch_size,d_in,d,h,w]*. - filters - Convolution filters *[fd,fh,fw,d_in/feature_group_count,d_out]*. - strides - The stride of the sliding window for each dimension of input. - padding - either the string ‘SAME’ (padding with zeros evenly), the string ‘VALID’ (no - padding), or a sequence of n (low, high) integer pairs that give the padding to - apply before and after each spatial dimension. - transpose - True for computing transpose convolution, and False for dilated convolution. - When True, `x_dilations` must be 1 (the default). - dims - Either 1, 2, or 3 corresponding to 1-D, 2-D, and 3-D convolution. - output_shape - Shape of the output (Default value = None) - data_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "NCW", - "NCHW", "NCDHW" input data formatS for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "NWC", "NHWC", "NDHWC" respectively. - filter_format - Either "channel_first" or "channel_last". "channel_first" corresponds to "OIW", - "OIHW", "OIDHW" input data formats for 1-D, 2-D, 3-D convolution respectively, - while "channel_last" corresponds to "WIO", "HWIO", "DHWIO" respectively. - feature_group_count - split input into groups, d_in should be divisible by the number of groups. - (Default value = 1) - x_dilations - The dilation factor for each dimension of input. (Default value = 1) - dilations - The dilation factor for each dimension of input. (Default value = 1) - bias - Bias array of shape *[d_out]*. + query + The queries input array. The shape of queries input array should be in + *[batch_shape,num_queries,feat_dim]*. The queries input array should have the + same size as keys and values. + key + The keys input array. The shape of keys input array should be in + *[batch_shape,num_keys,feat_dim]*. The keys input array should have the same + size as queries and values. + value + The values input array. The shape of values input should be in + *[batch_shape,num_keys,feat_dim]*. The values input array should have the same + size as queries and keys. + scale + The scale float value. + The scale float value is used to scale the query-key pairs before softmax. + mask + The mask input array. The mask to apply to the query-key values. Default is + None. The shape of mask input should be in *[batch_shape,num_queries,num_keys]*. + dropout_p + Specifies the dropout probablity, if greater than 0.0, dropout is applied + is_causal + If true, assumes causal attention masking + and errors if both `mask` and `is_causal` are set. + training + If True, dropout is used, otherwise dropout is not activated. out optional output array, for writing the result to. It must have a shape that the inputs broadcast to. @@ -2049,294 +2098,252 @@ def conv( Returns ------- ret - The result of the transpose or dilated convolution operation. - """ - if transpose: - return conv_general_transpose( - x, - filters, - strides, - padding, - dims=dims, - output_shape=output_shape, - data_format=data_format, - dilations=dilations, - feature_group_count=feature_group_count, - bias=bias, - out=out, - ) - else: - return conv_general_dilated( - x, - filters, - strides, - padding, - dims=dims, - data_format=data_format, - filter_format=filter_format, - feature_group_count=feature_group_count, - x_dilations=x_dilations, - dilations=dilations, - bias=bias, - out=out, - ) + The output following application of scaled dot-product attention. + The output array is the weighted sum produced by the attention score and value. + The shape of output array is *[batch_shape,num_queries,feat_dim]* . + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. -# LSTM # + Examples + -------- + With :class:`ivy.Array` input: + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def lstm_update( - x: Union[ivy.Array, ivy.NativeArray], - init_h: Union[ivy.Array, ivy.NativeArray], - init_c: Union[ivy.Array, ivy.NativeArray], - kernel: Union[ivy.Array, ivy.NativeArray], - recurrent_kernel: Union[ivy.Array, ivy.NativeArray], - /, - *, - bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - recurrent_bias: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[ivy.Array, ivy.Array]: - """ - Perform long-short term memory update by unrolling time dimension of input array. + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - Parameters - ---------- - x - input tensor of LSTM layer *[batch_shape, t, in]*. - init_h - initial state tensor for the cell output *[batch_shape, out]*. - init_c - initial state tensor for the cell hidden state *[batch_shape, out]*. - kernel - weights for cell kernel *[in, 4 x out]*. - recurrent_kernel - weights for cell recurrent kernel *[out, 4 x out]*. - bias - bias for cell kernel *[4 x out]*. (Default value = None) - recurrent_bias - bias for cell recurrent kernel *[4 x out]*. (Default value = None) - - Returns - ------- - ret - hidden state for all timesteps *[batch_shape,t,out]* and cell state for last - timestep *[batch_shape,out]* - """ - # get shapes - x_shape = list(x.shape) - batch_shape = x_shape[:-2] - timesteps = x_shape[-2] - input_channels = x_shape[-1] - x_flat = ivy.reshape(x, (-1, input_channels)) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - # input kernel - Wi = kernel - Wi_x = ivy.reshape( - ivy.matmul(x_flat, Wi) + (bias if bias is not None else 0), - batch_shape + [timesteps, -1], - ) - Wii_x, Wif_x, Wig_x, Wio_x = ivy.split(Wi_x, num_or_size_splits=4, axis=-1) + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - # recurrent kernel - Wh = recurrent_kernel + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) - # lstm states - ht = init_h - ct = init_c + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - # lstm outputs - hts_list = list() + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0],[0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) - # unrolled time dimension with lstm steps - for Wii_xt, Wif_xt, Wig_xt, Wio_xt in zip( - ivy.unstack(Wii_x, axis=-2), - ivy.unstack(Wif_x, axis=-2), - ivy.unstack(Wig_x, axis=-2), - ivy.unstack(Wio_x, axis=-2), - ): - htm1 = ht - ctm1 = ct + ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) - Wh_htm1 = ivy.matmul(htm1, Wh) + ( - recurrent_bias if recurrent_bias is not None else 0 - ) - Whi_htm1, Whf_htm1, Whg_htm1, Who_htm1 = ivy.split( - Wh_htm1, num_or_size_splits=4, axis=-1 - ) + >>> q = ivy.native_array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True, + ... out=out) + >>> print(out) - it = ivy.sigmoid(Wii_xt + Whi_htm1) - ft = ivy.sigmoid(Wif_xt + Whf_htm1) - gt = ivy.tanh(Wig_xt + Whg_htm1) - ot = ivy.sigmoid(Wio_xt + Who_htm1) - ct = ft * ctm1 + it * gt - ht = ot * ivy.tanh(ct) + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) - hts_list.append(ivy.expand_dims(ht, axis=-2)) + With :class:`ivy.Container` input: - return ivy.concat(hts_list, axis=-2), ct + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) + { + a: ivy.array([[[5.19999981, 1.], + ... [2.59249449, 2.68226194], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[0.2, 1.], + ... [2.19603825, 2.9960382], + ... [4.4000001, 5.5999999]]]) + } + >>> q = ivy.Container(a=ivy.array([[[0.2, 1.], [2.7, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[1.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.Container(a=ivy.array([[[5.2, 1.], [2.1, 3.], [4.4, 5.6]]]), + ... b=ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]])) + >>> mask = ivy.Container( + ... a=ivy.array([[[1.0, 1.0, 1.0],[1.0, 1.0, 1.0],[1.0, 1.0, 1.0]]]), + ... b=ivy.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,1.0]]]) + ... ) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,mask=mask) + >>> print(result) + { + a: ivy.array([[[4.26894283, 5.40236187], + ... [4.39999437, 5.59999037], + ... [4.4000001, 5.5999999]]]), + b: ivy.array([[[4.35046196, 5.54282808], + ... [4.39989519, 5.5998764], + ... [4.4000001, 5.5999999]]]) + } -# Helpers # + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3],[4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... dropout_p=0.1, + ... is_causal=True, + ... training=True) + >>> print(result) -def _handle_padding(x, strides, filters, padding): - if isinstance(padding, str) and padding.upper() == "SAME": - if x % strides == 0: - pad = max(filters - strides, 0) - else: - pad = max(filters - (x % strides), 0) - else: - pad = 0 - return pad + ivy.array([[[0.40000001, 1.29999995], + ... [2.19994521, 3.09994531], + ... [4.30000019, 5.30000019]]]) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.], [4.4, 5.6]]]) + >>> k = ivy.native_array([[[0.6, 1.5], [2.4, 3.3], [4.2, 5.1]]]) + >>> v = ivy.native_array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> out = ivy.zeros(shape=(1, 3, 2)) + >>> ivy.scaled_dot_product_attention(q,k,v,scale=1,out=out) + >>> print(out) + ivy.array([[[4.03946018, 5.0280633 ], + ... [4.29981947, 5.29981089], + ... [4.30000019, 5.30000019]]]) -def _validate_max_pool_params(kernel, strides, padding, dilation, ceil_mode, dims): - if isinstance(kernel, int): - kernel = (kernel,) * dims - elif len(kernel) == 1: - kernel = (kernel[0],) * dims - elif (len(kernel) != dims) and (len(kernel) != dims + 2): - raise ValueError( - "The kernel should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - if isinstance(strides, int): - strides = (strides,) * dims - elif len(strides) == 1: - strides = (strides[0],) * dims - elif (len(strides) != dims) and (len(strides) != dims + 2): - raise ValueError( - "The stride should be an integer, or a tuple of length" - f" {list(set((1, dims, dims+2)))}" - ) + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3], [4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6], [4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1], [4.3, 5.3]]]) + >>> result = ivy.scaled_dot_product_attention(q,k,v,scale=1,is_causal=True) + >>> print(result) + { + a: ivy.array([[[0.40000001, 1.29999995], + ... [2.06345534, 2.9634552], + ... [4.30000019, 5.30000019]]]), + b: ivy.array([[[0.40000001, 1.29999995], + ... [2.19336844, 3.09336829], + ... [4.30000019, 5.30000019]]]) + } + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - if isinstance(padding, int): - padding = [(padding,) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == 1: - padding = [(padding[0],) * 2] * dims - elif isinstance(padding, tuple) and len(padding) == dims: - padding = [(padding[i],) * 2 for i in range(dims)] - elif isinstance(padding, list) and len(padding) == dims: - if not all([isinstance(p, tuple) and len(p) == 2 for p in padding]): - raise ValueError("Explicit padding must be a list of tuple of two integers") - if isinstance(padding, str) and padding.upper() not in ["VALID", "SAME"]: - raise ValueError( - f"Invalid padding arg {padding}Must be one of: 'VALID' or 'SAME'" + >>> q = ivy.array([[[0.2, 1.], [2.2, 3.],[4.4, 5.6]]]) + >>> k = ivy.Container(a=ivy.array([[[4.2, 1.], [2.2, 3.3],[4.4, 5.6]]]), + ... b=ivy.array([[[3.2, 1.], [2.2, 3.6],[4.0, 5.6]]])) + >>> v = ivy.array([[[0.4, 1.3], [2.2, 3.1],[4.3, 5.3]]]) + >>> mask = ivy.native_array([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]) + >>> result = ivy.scaled_dot_product_attention(q, + ... k, + ... v, + ... scale=1, + ... mask=mask, + ... dropout_p=0.1, + ... training=True) + >>> print(result) + { + a: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]), + b: ivy.array([[[2.30000019, 3.23333359], + ... [2.30000019, 3.23333359], + ... [2.30000019, 3.23333359]]]) + } + """ + ivy.assertions.check_all( + (not is_causal) or (is_causal and mask is None), + "is_causal and attn_mask cannot be set at the same time", + ) + embed_dim = query.shape[-1] + scale = 1 / (embed_dim**0.5) if not scale else scale + sim = ivy.einsum("... q f, ... k f -> ... q k", query, key) * scale + sim = ivy.dropout(sim, dropout_p, training=training) + if ivy.exists(mask): + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, ) - - if isinstance(dilation, int): - dilation = (dilation,) * dims - elif len(dilation) == 1: - dilation = (dilation[0],) * dims - elif len(dilation) != dims: - raise ValueError( - f"Dilation must be an integer or a tuple of length {list(set((1, dims)))}" + elif is_causal: + L = query.shape[-2] # Source sequence length + S = key.shape[-2] # Target sequence length + mask = ivy.tril(ivy.ones((L, S)), k=0) + mask = ivy.astype(mask, ivy.bool) + sim = ivy.where( + ivy.logical_not(mask), + -ivy.ones_like(sim) * ivy.finfo(ivy.dtype(sim)).max, + sim, ) - if min(dilation) < 1: - raise ValueError("All values of `dilation` must be positive") - - # Other errors - if isinstance(padding, str) and (padding.upper() == "VALID") and ceil_mode: - raise ValueError("When 'padding' is 'VALID', 'ceil_mode' must be False") - assert len(kernel) == len(strides), f"len({kernel}) must equal len({strides})" - - # Account for dilation when padding > kernel/2. Not the case in torch by default. - new_kernel = tuple( - [dilation[i] * (kernel[i] - 1) + 1 for i in range(1, len(kernel))] - ) - if isinstance(padding, list) and len(padding) == len(new_kernel): - ivy.utils.assertions.check_kernel_padding_size(new_kernel, padding) - - return kernel, strides, padding, dilation - - -def _depth_max_pooling_helper( - x_shape, kernel, strides, dims, data_format="channel_last" -): - # Determine depth pooling. - # We assume that the kernel and the data have the same data_format. - depth_pooling = False - CHANNEL_LAST = "channel_last" - channel_idx = -1 if data_format == CHANNEL_LAST else 1 - if len(kernel) == dims + 2: - spatial_kernel = kernel[1:-1] if data_format == CHANNEL_LAST else kernel[2:] - if kernel[channel_idx] != 1: - depth_pooling = True - if any(i != 1 for i in spatial_kernel): - raise NotImplementedError( - "MaxPooling supports exactly one of pooling across" - " depth or pooling across width/height." - ) - if len(strides) != dims + 2 or strides[channel_idx] != kernel[channel_idx]: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to equal the depth" - " stride" - ) - if x_shape[channel_idx] % kernel[channel_idx] != 0: - raise NotImplementedError( - "Depthwise max pooling requires the depth window to evenly divide" - " the input depth" - ) - kernel = [kernel[channel_idx], *[1] * (dims - 1)] - strides = [strides[channel_idx], *[1] * (dims - 1)] - else: - kernel = spatial_kernel - if len(strides) == dims + 2: - strides = strides[1:-1] if data_format == CHANNEL_LAST else strides[2:] - return kernel, strides, depth_pooling - - -def _deconv_length(dim_size, stride_size, kernel_size, padding, dilation=1): - kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) - if padding == "SAME": - dim_size = dim_size * stride_size - else: - dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0) - return dim_size - - -def _get_x_data_format(dims: int = 2, data_format: str = "channel_first"): - if dims == 1: - if data_format == "channel_first": - return "NCW" - else: - return "NWC" - if dims == 2: - if data_format == "channel_first": - return "NCHW" - else: - return "NHWC" - elif dims == 3: - if data_format == "channel_first": - return "NCDHW" - else: - return "NDHWC" - - -def _get_num_padded_values(i, p, n, k, s): - """ - Get number of padded values in a specific window. + attn = ivy.softmax(sim, axis=-1) + result = ivy.einsum("... qk, ...kf -> ...qf", attn, value) + return result if not ivy.exists(out) else ivy.inplace_update(out, result) - Parameters - ---------- - i window index - p total amount of padding - n input size - k kernel size - s stride - Returns - ------- - number of padded values in a particular window represented by i - """ - current_index = s * i - left_padding = p // 2 - return max(0, left_padding - current_index) + max( - 0, current_index + k - n - left_padding - ) +linear.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} +dropout.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "handle_out_argument", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), +} diff --git a/ivy/functional/ivy/linear_algebra.py b/ivy/functional/ivy/linear_algebra.py index 612315f0b987a..8190c0d4b57d5 100644 --- a/ivy/functional/ivy/linear_algebra.py +++ b/ivy/functional/ivy/linear_algebra.py @@ -345,6 +345,95 @@ def det( return current_backend(x).det(x, out=out) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def diag( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + k: int = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the specified diagonals of the input array, or an array with the input + array's elements as diagonals. + + Parameters + ---------- + x + An array with rank >= 1. + k + An integer that controls which diagonal to consider. + Positive value means superdiagonal, + 0 refers to the main diagonal, + and negative value means subdiagonal. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + If x is a 1-D array, the function returns a 2-D square array with the elements + of input as diagonals. + If x is a 2-D array, the function returns a 1-D array with the diagonal elements + of x. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------ + + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x) + ivy.array([0, 4, 8]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=1) + ivy.array([1, 5]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(x, k=-1) + ivy.array([3, 7]) + + >>> x = ivy.array([[0, 1, 2], + >>> [3, 4, 5], + >>> [6, 7, 8]]) + >>> ivy.diag(ivy.diag(x)) + ivy.array([[0, 0, 0], + [0, 4, 0], + [0, 0, 8]]) + """ + return current_backend(x).diag(x, k=k, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -929,6 +1018,41 @@ def inv( return current_backend(x).inv(x, adjoint=adjoint, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_device_shifting +def lu_factor( + A: Union[ivy.Array, ivy.NativeArray], + /, + *, + pivot: bool = True, + out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Parameters + ---------- + A + tensor of shape (*, m, n) where * is zero or more batch dimensions. + + pivot + Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. pivot = False not supported on CPU. Default: True. + + out + tuple of two tensors to write the output to. Ignored if None. Default: None. + + Returns + ------- + ret + A named tuple (LU, pivots). + """ + return current_backend(A).lu_factor(A, pivot=pivot, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2315,6 +2439,51 @@ def tensordot( return current_backend(x1, x2).tensordot(x1, x2, axes=axes, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def tensorsolve( + x1: Union[ivy.Array, ivy.NativeArray], + x2: Union[ivy.Array, ivy.NativeArray], + /, + *, + axes: Union[int, Tuple[List[int], List[int]]] = 2, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + ndim1 = ivy.get_num_dims(x1) + ndim2 = ivy.get_num_dims(x2) + + if axes is not None: + allaxes = list(range(0, ndim1)) + for k in axes: + allaxes.remove(k) + allaxes.insert(ndim1, k) + + x1 = ivy.matrix_transpose(x1, allaxes) + + old_shape = x1.shape[-(ndim1 - ndim2) :] + + prod = 1 + for k in old_shape: + prod *= k + + if ivy.shape(ivy.flatten(x1))[0] != prod**2: + raise ivy.utils.exceptions.IvyException( + "Input arrays must satisfy the requirement " + "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" + ) + + x1 = ivy.reshape(x1, (prod, prod)) + x2 = ivy.flatten(x2) + res = ivy.solve(x1, x2) + res = ivy.reshape(res, old_shape) + return res + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2443,6 +2612,80 @@ def trace( return current_backend(x).trace(x, offset=offset, axis1=axis1, axis2=axis2, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def vander( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + N: Optional[int] = None, + increasing: bool = False, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Generate a Vandermonde matrix. The columns of the output matrix are elementwise + powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, + the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a + geometric progression in each row is named for Alexandre-Theophile Vandermonde. + + Parameters + ---------- + x + 1-D input array. + N + Number of columns in the output. If N is not specified, + a square array is returned (N = len(x)) + increasing + Order of the powers of the columns. If True, the powers increase + from left to right, if False (the default) they are reversed. + out + optional output array, for writing the result to. + + Returns + ------- + ret + Vandermonde matrix. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x) + ivy.array( + [[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3) + ivy.array( + [[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]] + ) + + >>> x = ivy.array([1, 2, 3, 5]) + >>> ivy.vander(x, N=3, increasing=True) + ivy.array( + [[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]] + ) + """ + return current_backend(x).vander(x, N=N, increasing=increasing, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2657,169 +2900,6 @@ def vector_norm( ) -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def diag( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - k: int = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the specified diagonals of the input array, or an array with the input - array's elements as diagonals. - - Parameters - ---------- - x - An array with rank >= 1. - k - An integer that controls which diagonal to consider. - Positive value means superdiagonal, - 0 refers to the main diagonal, - and negative value means subdiagonal. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - If x is a 1-D array, the function returns a 2-D square array with the elements - of input as diagonals. - If x is a 2-D array, the function returns a 1-D array with the diagonal elements - of x. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------ - - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x) - ivy.array([0, 4, 8]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=1) - ivy.array([1, 5]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(x, k=-1) - ivy.array([3, 7]) - - >>> x = ivy.array([[0, 1, 2], - >>> [3, 4, 5], - >>> [6, 7, 8]]) - >>> ivy.diag(ivy.diag(x)) - ivy.array([[0, 0, 0], - [0, 4, 0], - [0, 0, 8]]) - """ - return current_backend(x).diag(x, k=k, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def vander( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - N: Optional[int] = None, - increasing: bool = False, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Generate a Vandermonde matrix. The columns of the output matrix are elementwise - powers of the input vector x^{(N-1)}, x^{(N-2)}, ..., x^0x. If increasing is True, - the order of the columns is reversed x^0, x^1, ..., x^{(N-1)}. Such a matrix with a - geometric progression in each row is named for Alexandre-Theophile Vandermonde. - - Parameters - ---------- - x - 1-D input array. - N - Number of columns in the output. If N is not specified, - a square array is returned (N = len(x)) - increasing - Order of the powers of the columns. If True, the powers increase - from left to right, if False (the default) they are reversed. - out - optional output array, for writing the result to. - - Returns - ------- - ret - Vandermonde matrix. - - Examples - -------- - With :class:`ivy.Array` inputs: - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x) - ivy.array( - [[ 1, 1, 1, 1], - [ 8, 4, 2, 1], - [ 27, 9, 3, 1], - [125, 25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3) - ivy.array( - [[ 1, 1, 1], - [ 4, 2, 1], - [ 9, 3, 1], - [25, 5, 1]] - ) - - >>> x = ivy.array([1, 2, 3, 5]) - >>> ivy.vander(x, N=3, increasing=True) - ivy.array( - [[ 1, 1, 1], - [ 1, 2, 4], - [ 1, 3, 9], - [ 1, 5, 25]] - ) - """ - return current_backend(x).vander(x, N=N, increasing=increasing, out=out) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -2854,84 +2934,3 @@ def vector_to_skew_symmetric_matrix( instances in place of any of the arguments. """ return current_backend(vector).vector_to_skew_symmetric_matrix(vector, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_device_shifting -def lu_factor( - A: Union[ivy.Array, ivy.NativeArray], - /, - *, - pivot: bool = True, - out: Optional[Union[ivy.Array, ivy.NativeArray]] = None, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Parameters - ---------- - A - tensor of shape (*, m, n) where * is zero or more batch dimensions. - - pivot - Whether to compute the LU decomposition with partial pivoting, or the regular LU - decomposition. pivot = False not supported on CPU. Default: True. - - out - tuple of two tensors to write the output to. Ignored if None. Default: None. - - Returns - ------- - ret - A named tuple (LU, pivots). - """ - return current_backend(A).lu_factor(A, pivot=pivot, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def tensorsolve( - x1: Union[ivy.Array, ivy.NativeArray], - x2: Union[ivy.Array, ivy.NativeArray], - /, - *, - axes: Union[int, Tuple[List[int], List[int]]] = 2, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - ndim1 = ivy.get_num_dims(x1) - ndim2 = ivy.get_num_dims(x2) - - if axes is not None: - allaxes = list(range(0, ndim1)) - for k in axes: - allaxes.remove(k) - allaxes.insert(ndim1, k) - - x1 = ivy.matrix_transpose(x1, allaxes) - - old_shape = x1.shape[-(ndim1 - ndim2) :] - - prod = 1 - for k in old_shape: - prod *= k - - if ivy.shape(ivy.flatten(x1))[0] != prod**2: - raise ivy.utils.exceptions.IvyException( - "Input arrays must satisfy the requirement " - "prod(x1.shape[x2.ndim:]) == prod(x1.shape[:x2.ndim])" - ) - - x1 = ivy.reshape(x1, (prod, prod)) - x2 = ivy.flatten(x2) - res = ivy.solve(x1, x2) - res = ivy.reshape(res, old_shape) - return res - # return current_backend(x1, x2).tensorsolve(x1, x2, axes=axes, out=out) diff --git a/ivy/functional/ivy/losses.py b/ivy/functional/ivy/losses.py index d866cd8bc2487..d803eb6fca839 100644 --- a/ivy/functional/ivy/losses.py +++ b/ivy/functional/ivy/losses.py @@ -12,8 +12,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # +# --- Helpers --- # +# --------------- # def _reduce_loss(red, loss, axis, out): @@ -25,64 +25,8 @@ def _reduce_loss(red, loss, axis, out): return ivy.negative(loss, out=out) -# Extra # -# ------# - - -@handle_exceptions -@handle_nestable -@handle_array_like_without_promotion -@inputs_to_ivy_arrays -@handle_array_function -def cross_entropy( - true: Union[ivy.Array, ivy.NativeArray], - pred: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - epsilon: float = 1e-7, - reduction: str = "sum", - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Compute cross-entropy between predicted and true discrete distributions. - - Parameters - ---------- - true - input array containing true labels. - pred - input array containing the predicted labels. - axis - the axis along which to compute the cross-entropy. If axis is ``-1``, - the cross-entropy will be computed along the last dimension. Default: ``-1``. - epsilon - a float in [0.0, 1.0] specifying the amount of smoothing when calculating - the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. - - Returns - ------- - ret - The cross-entropy loss between the given distributions - - Examples - -------- - >>> x = ivy.array([0, 0, 1, 0]) - >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) - >>> print(ivy.cross_entropy(x, y)) - ivy.array(1.3862944) - - >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) - >>> print(ivy.cross_entropy(x, z)) - ivy.array(0.35667497) - """ - ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) - pred = ivy.clip(pred, epsilon, 1 - epsilon) - log_pred = ivy.log(pred) - return _reduce_loss(reduction, log_pred * true, axis, out) +# --- Main --- # +# ------------ # @handle_exceptions @@ -268,6 +212,66 @@ def binary_cross_entropy( return _reduce_loss(reduction, loss, axis, out) +# Extra # +# ------# + + +@handle_exceptions +@handle_nestable +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +def cross_entropy( + true: Union[ivy.Array, ivy.NativeArray], + pred: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + epsilon: float = 1e-7, + reduction: str = "sum", + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Compute cross-entropy between predicted and true discrete distributions. + + Parameters + ---------- + true + input array containing true labels. + pred + input array containing the predicted labels. + axis + the axis along which to compute the cross-entropy. If axis is ``-1``, + the cross-entropy will be computed along the last dimension. Default: ``-1``. + epsilon + a float in [0.0, 1.0] specifying the amount of smoothing when calculating + the loss. If epsilon is ``0``, no smoothing will be applied. Default: ``1e-7``. + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + The cross-entropy loss between the given distributions + + Examples + -------- + >>> x = ivy.array([0, 0, 1, 0]) + >>> y = ivy.array([0.25, 0.25, 0.25, 0.25]) + >>> print(ivy.cross_entropy(x, y)) + ivy.array(1.3862944) + + >>> z = ivy.array([0.1, 0.1, 0.7, 0.1]) + >>> print(ivy.cross_entropy(x, z)) + ivy.array(0.35667497) + """ + ivy.utils.assertions.check_elem_in_list(reduction, ["none", "sum", "mean"]) + pred = ivy.clip(pred, epsilon, 1 - epsilon) + log_pred = ivy.log(pred) + return _reduce_loss(reduction, log_pred * true, axis, out) + + @handle_exceptions @handle_nestable @handle_array_like_without_promotion diff --git a/ivy/functional/ivy/manipulation.py b/ivy/functional/ivy/manipulation.py index 71d5e65b61c89..af74e3523eb5f 100644 --- a/ivy/functional/ivy/manipulation.py +++ b/ivy/functional/ivy/manipulation.py @@ -20,6 +20,10 @@ from ivy.utils.exceptions import handle_exceptions +# --- Helpers --- # +# --------------- # + + def _calculate_out_shape(axis, array_shape): if type(axis) not in (tuple, list): axis = (axis,) @@ -33,6 +37,141 @@ def _calculate_out_shape(axis, array_shape): return out_shape +# --- Main --- # +# ------------ # + + +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def clip( + x: Union[ivy.Array, ivy.NativeArray], + x_min: Union[Number, ivy.Array, ivy.NativeArray], + x_max: Union[Number, ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Clips (limits) the values in an array. + + Given an interval, values outside the interval are clipped to the interval edges + (element-wise). For example, if an interval of [0, 1] is specified, values smaller + than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller + or equal to maximum value to return correct results. + + Parameters + ---------- + x + Input array containing elements to clip. + x_min + Minimum value. + x_max + Maximum value. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + An array with the elements of x, but where values < x_min are replaced with + x_min, and those > x_max with x_max. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> y = ivy.zeros_like(x) + >>> ivy.clip(x, 2., 7., out=y) + >>> print(y) + ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) + >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.NativeArray` input: + + >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) + >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) + >>> y = ivy.clip(x, x_min, x_max) + >>> print(y) + ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> y = ivy.clip(x, 1., 5.) + >>> print(y) + { + a: ivy.array([1., 1., 2.]), + b: ivy.array([3., 4., 5.]) + } + + With multiple :class:`ivy.Container` inputs: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> x_min = ivy.Container(a=0, b=-3) + >>> x_max = ivy.Container(a=1, b=-1) + >>> y = ivy.clip(x, x_min,x_max) + >>> print(y) + { + a: ivy.array([0., 1., 1.]), + b: ivy.array([-1., -1., -1.]) + } + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + >>> x_min = ivy.array([3., 0., 1]) + >>> x_max = ivy.array([5., 4., 3.]) + >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([3., 4., 5.])) + >>> z = ivy.clip(y, x_min, x_max) + >>> print(z) + { + a: ivy.array([3., 1., 2.]), + b: ivy.array([3., 4., 3.]) + } + """ + return current_backend(x).clip(x, x_min, x_max, out=out) + + # Array API Standard # # -------------------# @@ -97,6 +236,103 @@ def concat( return current_backend(xs[0]).concat(xs, axis=axis, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def constant_pad( + x: Union[ivy.Array, ivy.NativeArray], + /, + pad_width: Iterable[Tuple[int]], + *, + value: Number = 0, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Pad an array with a constant value. + + Parameters + ---------- + x + Input array to pad. + pad_width + Number of values padded to the edges of each axis. + Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of + axes of x. + value + The constant value to pad the array with. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Padded array of rank equal to x with shape increased according to pad_width. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Functional Examples + ------------------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3, 4, 5]) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) + >>> print(y) + ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) + >>> print(y) + ivy.array([[0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 1, 2, 0, 0, 0], + [0, 0, 3, 4, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0]]) + + >>> x = ivy.array([[2.], [3.]]) + >>> y = ivy.zeros((4, 3)) + >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) + >>> print(y) + ivy.array([[5., 5., 5.], + [5., 2., 5.], + [5., 3., 5.], + [5., 5., 5.]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), + ... b = ivy.array([3., 4., 5.])) + >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) + >>> print(y) + { + a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), + b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) + } + """ + return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -435,40 +671,113 @@ def permute_dims( @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion -@handle_view @handle_out_argument @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def reshape( +def repeat( x: Union[ivy.Array, ivy.NativeArray], /, - shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + repeats: Union[int, Iterable[int]], *, - copy: Optional[bool] = None, - order: str = "C", - allowzero: bool = True, + axis: int = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Give a new shape to an array without changing its data. + Repeat values along a given dimension. Parameters ---------- x - Input array to be reshaped. - shape - a new shape compatible with the original shape. One shape dimension - can be -1. In this case, the value is inferred from the length of the array and - remaining dimensions. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - order - Read the elements of x using this index order, and place the elements into - the reshaped array using this index order. + Input array. + repeats + The number of repetitions for each element. repeats is broadcast to fit the + shape of the given axis. + axis + The axis along which to repeat values. By default, use the flattened input + array, and return a flat output array. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + The repeated output array. + + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([3, 4, 5]) + >>> y = ivy.repeat(x, 2) + >>> print(y) + ivy.array([3, 3, 4, 4, 5, 5]) + + >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) + >>> y = ivy.repeat(x, [1, 2], axis=0) + >>> print(y) + ivy.array([[1, 2, 3], + [4, 5, 6], + [4, 5, 6]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), + ... b=ivy.array([0., 1., 2.])) + >>> y = ivy.repeat(x, 2, axis=0) + >>> print(y) + { + a: ivy.array([0., 0., 1., 1., 2., 2.]), + b: ivy.array([0., 0., 1., 1., 2., 2.]) + } + """ + return current_backend(x).repeat(x, repeats, axis=axis, out=out) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def reshape( + x: Union[ivy.Array, ivy.NativeArray], + /, + shape: Union[ivy.Shape, ivy.NativeShape, Sequence[int]], + *, + copy: Optional[bool] = None, + order: str = "C", + allowzero: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Give a new shape to an array without changing its data. + + Parameters + ---------- + x + Input array to be reshaped. + shape + a new shape compatible with the original shape. One shape dimension + can be -1. In this case, the value is inferred from the length of the array and + remaining dimensions. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + order + Read the elements of x using this index order, and place the elements into + the reshaped array using this index order. ``C`` means to read / write the elements using C-like index order, with the last axis index changing fastest, back to the first axis index changing slowest. @@ -671,6 +980,93 @@ def roll( return current_backend(x).roll(x, shift, axis=axis, out=out) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_view +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def split( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + copy: Optional[bool] = None, + num_or_size_splits: Optional[ + Union[int, Sequence[int], ivy.Array, ivy.NativeArray] + ] = None, + axis: int = 0, + with_remainder: bool = False, +) -> List[ivy.Array]: + """ + Split an array into multiple sub-arrays. + + Parameters + ---------- + x + array to be divided into sub-arrays. + copy + boolean indicating whether or not to copy the input array. + If True, the function must always copy. + If False, the function must never copy. + In case copy is False we avoid copying by returning a view of the input array. + num_or_size_splits + Number of equal arrays to divide the array into along the given axis if an + integer. The size of each split element if a sequence of integers or 1-D array. + Default is to divide into as many 1-dimensional arrays as the axis dimension. + axis + The axis along which to split, default is ``0``. + with_remainder + If the tensor does not split evenly, then store the last remainder entry. + Default is ``False``. + + Returns + ------- + ret + A list of sub-arrays. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1, 2, 3]) + >>> y = ivy.split(x) + >>> print(y) + [ivy.array([1]),ivy.array([2]),ivy.array([3])] + + >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) + >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) + >>> print(y) + [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] + + >>> x = ivy.array([4, 6, 5, 3]) + >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) + >>> print(y) + ivy.array([[4], [6, 5, 3]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([10, 45, 2])) + >>> y = ivy.split(x) + >>> print(y) + { + a:(list[3],shape=[1]) + } + """ + return current_backend(x).split( + x, + copy=copy, + num_or_size_splits=num_or_size_splits, + axis=axis, + with_remainder=with_remainder, + ) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -851,394 +1247,6 @@ def stack( return res -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def clip( - x: Union[ivy.Array, ivy.NativeArray], - x_min: Union[Number, ivy.Array, ivy.NativeArray], - x_max: Union[Number, ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Clips (limits) the values in an array. - - Given an interval, values outside the interval are clipped to the interval edges - (element-wise). For example, if an interval of [0, 1] is specified, values smaller - than 0 become 0, and values larger than 1 become 1. Minimum value needs to smaller - or equal to maximum value to return correct results. - - Parameters - ---------- - x - Input array containing elements to clip. - x_min - Minimum value. - x_max - Maximum value. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - An array with the elements of x, but where values < x_min are replaced with - x_min, and those > x_max with x_max. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - ivy.array([1., 1., 2., 3., 4., 5., 5., 5., 5., 5.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> y = ivy.zeros_like(x) - >>> ivy.clip(x, 2., 7., out=y) - >>> print(y) - ivy.array([2., 2., 2., 3., 4., 5., 6., 7., 7., 7.]) - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 3., 1., 0., 2., 3., 4., 0., 4., 4.]) - >>> x_max = ivy.array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With a mix of :class:`ivy.Array` and :class:`ivy.NativeArray` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.native_array([3., 3., 1., 0., 2., 3., 4., 2., 4., 4.]) - >>> x_max = ivy.native_array([5., 4., 3., 3., 5., 7., 8., 3., 8., 8.]) - >>> y = ivy.clip(x, x_min, x_max) - >>> print(y) - ivy.array([3., 3., 2., 3., 4., 5., 6., 3., 8., 8.]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> y = ivy.clip(x, 1., 5.) - >>> print(y) - { - a: ivy.array([1., 1., 2.]), - b: ivy.array([3., 4., 5.]) - } - - With multiple :class:`ivy.Container` inputs: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> x_min = ivy.Container(a=0, b=-3) - >>> x_max = ivy.Container(a=1, b=-1) - >>> y = ivy.clip(x, x_min,x_max) - >>> print(y) - { - a: ivy.array([0., 1., 1.]), - b: ivy.array([-1., -1., -1.]) - } - - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: - - >>> x = ivy.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - >>> x_min = ivy.array([3., 0., 1]) - >>> x_max = ivy.array([5., 4., 3.]) - >>> y = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([3., 4., 5.])) - >>> z = ivy.clip(y, x_min, x_max) - >>> print(z) - { - a: ivy.array([3., 1., 2.]), - b: ivy.array([3., 4., 3.]) - } - """ - return current_backend(x).clip(x, x_min, x_max, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def constant_pad( - x: Union[ivy.Array, ivy.NativeArray], - /, - pad_width: Iterable[Tuple[int]], - *, - value: Number = 0, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Pad an array with a constant value. - - Parameters - ---------- - x - Input array to pad. - pad_width - Number of values padded to the edges of each axis. - Specified as ((before_1, after_1), … (before_N, after_N)), where N is number of - axes of x. - value - The constant value to pad the array with. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Padded array of rank equal to x with shape increased according to pad_width. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Functional Examples - ------------------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3, 4, 5]) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]]) - >>> print(y) - ivy.array([0, 0, 1, 2, 3, 4, 5, 0, 0, 0]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width=[(2, 3), (2, 3)]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> y = ivy.constant_pad(x, pad_width = [[3, 2], [2, 3]]) - >>> print(y) - ivy.array([[0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 1, 2, 0, 0, 0], - [0, 0, 3, 4, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0]]) - - >>> x = ivy.array([[2.], [3.]]) - >>> y = ivy.zeros((4, 3)) - >>> ivy.constant_pad(x, pad_width = [(1, 1), (1, 1)], value = 5.0, out= y) - >>> print(y) - ivy.array([[5., 5., 5.], - [5., 2., 5.], - [5., 3., 5.], - [5., 5., 5.]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a = ivy.array([1., 2., 3.]), - ... b = ivy.array([3., 4., 5.])) - >>> y = ivy.constant_pad(x, pad_width = [[2, 3]], value = 5.0) - >>> print(y) - { - a: ivy.array([5., 5., 1., 2., 3., 5., 5., 5.]), - b: ivy.array([5., 5., 3., 4., 5., 5., 5., 5.]) - } - """ - return current_backend(x).constant_pad(x, pad_width=pad_width, value=value, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def repeat( - x: Union[ivy.Array, ivy.NativeArray], - /, - repeats: Union[int, Iterable[int]], - *, - axis: int = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Repeat values along a given dimension. - - Parameters - ---------- - x - Input array. - repeats - The number of repetitions for each element. repeats is broadcast to fit the - shape of the given axis. - axis - The axis along which to repeat values. By default, use the flattened input - array, and return a flat output array. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - The repeated output array. - - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([3, 4, 5]) - >>> y = ivy.repeat(x, 2) - >>> print(y) - ivy.array([3, 3, 4, 4, 5, 5]) - - >>> x = ivy.array([[1, 2, 3], [4, 5, 6]]) - >>> y = ivy.repeat(x, [1, 2], axis=0) - >>> print(y) - ivy.array([[1, 2, 3], - [4, 5, 6], - [4, 5, 6]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), - ... b=ivy.array([0., 1., 2.])) - >>> y = ivy.repeat(x, 2, axis=0) - >>> print(y) - { - a: ivy.array([0., 0., 1., 1., 2., 2.]), - b: ivy.array([0., 0., 1., 1., 2., 2.]) - } - """ - return current_backend(x).repeat(x, repeats, axis=axis, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_view -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def split( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - copy: Optional[bool] = None, - num_or_size_splits: Optional[ - Union[int, Sequence[int], ivy.Array, ivy.NativeArray] - ] = None, - axis: int = 0, - with_remainder: bool = False, -) -> List[ivy.Array]: - """ - Split an array into multiple sub-arrays. - - Parameters - ---------- - x - array to be divided into sub-arrays. - copy - boolean indicating whether or not to copy the input array. - If True, the function must always copy. - If False, the function must never copy. - In case copy is False we avoid copying by returning a view of the input array. - num_or_size_splits - Number of equal arrays to divide the array into along the given axis if an - integer. The size of each split element if a sequence of integers or 1-D array. - Default is to divide into as many 1-dimensional arrays as the axis dimension. - axis - The axis along which to split, default is ``0``. - with_remainder - If the tensor does not split evenly, then store the last remainder entry. - Default is ``False``. - - Returns - ------- - ret - A list of sub-arrays. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 2, 3]) - >>> y = ivy.split(x) - >>> print(y) - [ivy.array([1]),ivy.array([2]),ivy.array([3])] - - >>> x = ivy.array([[3, 2, 1], [4, 5, 6]]) - >>> y = ivy.split(x, num_or_size_splits=2, axis=1, with_remainder=True) - >>> print(y) - [ivy.array([[3,2],[4,5]]),ivy.array([[1],[6]])] - - >>> x = ivy.array([4, 6, 5, 3]) - >>> y = x.split(num_or_size_splits=[1, 3], axis=0, with_remainder=False) - >>> print(y) - ivy.array([[4], [6, 5, 3]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([10, 45, 2])) - >>> y = ivy.split(x) - >>> print(y) - { - a:(list[3],shape=[1]) - } - """ - return current_backend(x).split( - x, - copy=copy, - num_or_size_splits=num_or_size_splits, - axis=axis, - with_remainder=with_remainder, - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable diff --git a/ivy/functional/ivy/meta.py b/ivy/functional/ivy/meta.py index ab885bb3b207c..66d5452e48832 100644 --- a/ivy/functional/ivy/meta.py +++ b/ivy/functional/ivy/meta.py @@ -7,6 +7,11 @@ # local from typing import Optional, Union, Callable, Tuple, Any + +# --- Helpers --- # +# --------------- # + + # Extra # # ------# @@ -194,6 +199,70 @@ def _train_task( return final_cost, variables, all_grads +def _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + batched, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, +): + if batched: + return _train_tasks_batched( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) + return _train_tasks_with_for_loop( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + order, + average_across_steps, + inner_v, + keep_innver_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + stop_gradients, + ) + + def _train_tasks_batched( batch, inner_batch_fn, @@ -336,68 +405,8 @@ def _train_tasks_with_for_loop( return total_cost / num_tasks -def _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - batched, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, -): - if batched: - return _train_tasks_batched( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) - return _train_tasks_with_for_loop( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - order, - average_across_steps, - inner_v, - keep_innver_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - stop_gradients, - ) +# --- Main --- # +# ------------ # # Public # @@ -526,7 +535,136 @@ def fomaml_step( return cost, grads -fomaml_step.computes_gradients = True +# Second Order + + +@handle_exceptions +@handle_array_function +def maml_step( + batch: ivy.Container, + inner_cost_fn: Callable, + outer_cost_fn: Callable, + variables: ivy.Container, + inner_grad_steps: int, + inner_learning_rate: float, + /, + *, + inner_optimization_step: Callable = gradient_descent_update, + inner_batch_fn: Optional[Callable] = None, + outer_batch_fn: Optional[Callable] = None, + average_across_steps: bool = False, + batched: bool = True, + inner_v: Optional[ivy.Container] = None, + keep_inner_v: bool = True, + outer_v: Optional[ivy.Container] = None, + keep_outer_v: bool = True, + return_inner_v: Union[str, bool] = False, + num_tasks: Optional[int] = None, + stop_gradients: bool = True, +) -> Tuple[ivy.Array, ivy.Container, Any]: + """ + Perform step of vanilla second order MAML. + + Parameters + ---------- + batch + The input batch + inner_cost_fn + callable for the inner loop cost function, receiving sub-batch, inner vars and + outer vars + outer_cost_fn + callable for the outer loop cost function, receving task-specific sub-batch, + inner vars and outer vars. If None, the cost from the inner loop will also be + optimized in the outer loop. + variables + Variables to be optimized during the meta step + inner_grad_steps + Number of gradient steps to perform during the inner loop. + inner_learning_rate + The learning rate of the inner loop. + inner_optimization_step + The function used for the inner loop optimization. + Default is ivy.gradient_descent_update. + inner_batch_fn + Function to apply to the task sub-batch, before passing to the inner_cost_fn. + Default is ``None``. + outer_batch_fn + Function to apply to the task sub-batch, before passing to the outer_cost_fn. + Default is ``None``. + average_across_steps + Whether to average the inner loop steps for the outer loop update. + Default is ``False``. + batched + Whether to batch along the time dimension, and run the meta steps in batch. + Default is ``True``. + inner_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_inner_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + outer_v + Nested variable keys to be optimized during the inner loop, with same keys and + boolean values. (Default value = None) + keep_outer_v + If True, the key chains in inner_v will be kept, otherwise they will be removed. + Default is ``True``. + return_inner_v + Either 'first', 'all', or False. 'first' means the variables for the first task + inner loop will also be returned. variables for all tasks will be returned with + 'all'. Default is ``False``. + num_tasks + Number of unique tasks to inner-loop optimize for the meta step. Determined from + batch by default. + stop_gradients + Whether to stop the gradients of the cost. Default is ``True``. + + Returns + ------- + ret + The cost and the gradients with respect to the outer loop variables. + """ + if num_tasks is None: + num_tasks = batch.cont_shape[0] + unique_outer = outer_v is not None + func_ret, grads = ivy.execute_with_gradients( + lambda v: _train_tasks( + batch, + inner_batch_fn, + outer_batch_fn, + inner_cost_fn, + outer_cost_fn, + variables.cont_set_at_key_chains(v) if unique_outer else v, + inner_grad_steps, + inner_learning_rate, + inner_optimization_step, + 2, + average_across_steps, + batched, + inner_v, + keep_inner_v, + outer_v, + keep_outer_v, + return_inner_v, + num_tasks, + False, + ), + ( + variables.cont_at_key_chains(outer_v, ignore_none=True) + if keep_outer_v + else variables.cont_prune_key_chains(outer_v, ignore_none=True) + ), + ) + if isinstance(func_ret, tuple): + grads = grads["0"] if "0" in grads else grads + cost = func_ret[0] + rest = func_ret[1] + else: + cost = func_ret + rest = () + if stop_gradients: + cost = ivy.stop_gradient(cost, preserve_type=False) + return cost, grads.sum(axis=0), rest @handle_exceptions @@ -665,139 +803,6 @@ def reptile_step( return cost, grads +fomaml_step.computes_gradients = True reptile_step.computes_gradients = True - - -# Second Order - - -@handle_exceptions -@handle_array_function -def maml_step( - batch: ivy.Container, - inner_cost_fn: Callable, - outer_cost_fn: Callable, - variables: ivy.Container, - inner_grad_steps: int, - inner_learning_rate: float, - /, - *, - inner_optimization_step: Callable = gradient_descent_update, - inner_batch_fn: Optional[Callable] = None, - outer_batch_fn: Optional[Callable] = None, - average_across_steps: bool = False, - batched: bool = True, - inner_v: Optional[ivy.Container] = None, - keep_inner_v: bool = True, - outer_v: Optional[ivy.Container] = None, - keep_outer_v: bool = True, - return_inner_v: Union[str, bool] = False, - num_tasks: Optional[int] = None, - stop_gradients: bool = True, -) -> Tuple[ivy.Array, ivy.Container, Any]: - """ - Perform step of vanilla second order MAML. - - Parameters - ---------- - batch - The input batch - inner_cost_fn - callable for the inner loop cost function, receiving sub-batch, inner vars and - outer vars - outer_cost_fn - callable for the outer loop cost function, receving task-specific sub-batch, - inner vars and outer vars. If None, the cost from the inner loop will also be - optimized in the outer loop. - variables - Variables to be optimized during the meta step - inner_grad_steps - Number of gradient steps to perform during the inner loop. - inner_learning_rate - The learning rate of the inner loop. - inner_optimization_step - The function used for the inner loop optimization. - Default is ivy.gradient_descent_update. - inner_batch_fn - Function to apply to the task sub-batch, before passing to the inner_cost_fn. - Default is ``None``. - outer_batch_fn - Function to apply to the task sub-batch, before passing to the outer_cost_fn. - Default is ``None``. - average_across_steps - Whether to average the inner loop steps for the outer loop update. - Default is ``False``. - batched - Whether to batch along the time dimension, and run the meta steps in batch. - Default is ``True``. - inner_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_inner_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - outer_v - Nested variable keys to be optimized during the inner loop, with same keys and - boolean values. (Default value = None) - keep_outer_v - If True, the key chains in inner_v will be kept, otherwise they will be removed. - Default is ``True``. - return_inner_v - Either 'first', 'all', or False. 'first' means the variables for the first task - inner loop will also be returned. variables for all tasks will be returned with - 'all'. Default is ``False``. - num_tasks - Number of unique tasks to inner-loop optimize for the meta step. Determined from - batch by default. - stop_gradients - Whether to stop the gradients of the cost. Default is ``True``. - - Returns - ------- - ret - The cost and the gradients with respect to the outer loop variables. - """ - if num_tasks is None: - num_tasks = batch.cont_shape[0] - unique_outer = outer_v is not None - func_ret, grads = ivy.execute_with_gradients( - lambda v: _train_tasks( - batch, - inner_batch_fn, - outer_batch_fn, - inner_cost_fn, - outer_cost_fn, - variables.cont_set_at_key_chains(v) if unique_outer else v, - inner_grad_steps, - inner_learning_rate, - inner_optimization_step, - 2, - average_across_steps, - batched, - inner_v, - keep_inner_v, - outer_v, - keep_outer_v, - return_inner_v, - num_tasks, - False, - ), - ( - variables.cont_at_key_chains(outer_v, ignore_none=True) - if keep_outer_v - else variables.cont_prune_key_chains(outer_v, ignore_none=True) - ), - ) - if isinstance(func_ret, tuple): - grads = grads["0"] if "0" in grads else grads - cost = func_ret[0] - rest = func_ret[1] - else: - cost = func_ret - rest = () - if stop_gradients: - cost = ivy.stop_gradient(cost, preserve_type=False) - return cost, grads.sum(axis=0), rest - - maml_step.computes_gradients = True diff --git a/ivy/functional/ivy/nest.py b/ivy/functional/ivy/nest.py index 05e47e2358582..c7a052fd7cdb6 100644 --- a/ivy/functional/ivy/nest.py +++ b/ivy/functional/ivy/nest.py @@ -11,6 +11,246 @@ from ivy.utils.exceptions import handle_exceptions +@handle_exceptions +def all_nested_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, + /, + include_nests: bool = False, + _index: Optional[Union[int, Sequence[int]]] = None, + _base: bool = True, + extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return indices of all the elements in nest. + + Parameters + ---------- + nest + The nest to check the leaves of. + include_nests + Whether to also include indices of the nests themselves, not only + leaves. Default is ``False``. + _index + The indices detected so far. None at the beginning. Used internally, + do not set manually. + _base + Whether the current function call is the first function call in the + recursive stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + out + Optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. + + Returns + ------- + ret + A set of indices of all elements in nest + + Both the description and the type hints above assumes an array input + for simplicity, but this function is nestable, and therefore also + accepts :class:ivy.Container instances in place of the arguments. + + Examples + -------- + With :class:`Dict` input: + + >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} + >>> y = ivy.all_nested_indices(x) + >>> print(y) + [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] + + With :class:`ivy.Array` input: + + >>> x = ivy.array([0., 1., 2., 3., 4.]) + >>> y = ivy.all_nested_indices(x, False, out=x) + >>> print(y) + [[]] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.all_nested_indices(x, True) + >>> print(y) + [['a'], ['b']] + """ + _index = list() if _index is None else _index + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + ind = ivy.argwhere(ivy.ones_like(nest)) + indices = list() + for i in range(len(ind)): + indices.append(_index + ind.to_list()[i]) + return indices + _indices = [ + all_nested_indices( + item, include_nests, _index + [i], False, extra_nest_types + ) + for i, item in enumerate(nest) + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + elif isinstance(nest, dict): + _indices = [ + all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) + for k, v in nest.items() + ] + _indices = [idx for idxs in _indices if idxs for idx in idxs] + if include_nests: + _indices.append(_index) + else: + return [_index] + return [index for index in _indices if index] + + +@handle_exceptions +def copy_nest( + nest: Union[ivy.Array, ivy.NativeArray, Iterable], + /, + include_derived: bool = False, + to_mutable: bool = False, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[ivy.Array, ivy.NativeArray, Iterable]: + """ + Copy a nest deeply, but without copying leaves of the nest, only the nest lists, + tuples and dicts are copied. + + Parameters + ---------- + nest + The nest to copy. + include_derived + Whether to also recursive for classes derived from tuple, list and dict. + Default is ``False``. + to_mutable + Whether to convert the nest to a mutable form, changing all tuples to lists. + Default is ``False``. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + The copied nest. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + ivy.array([[1., 2., 3.], + [7., 8., 9.]]) + + With :code:`Iterable` input: + + >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + >>> copied_nest = ivy.copy_nest(nest, include_derived = True) + >>> print(copied_nest) + [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] + + >>> nest = ([23, 25, 1337], [63, 98, 6]) + >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) + >>> print(copied_nest) + [[23, 25, 1337], [63, 98, 6]] + + >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} + >>> copied_nest = ivy.copy_nest(nest) + >>> print(copied_nest) + {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + class_instance = type(nest) + check_fn = ( + (lambda x_, t: isinstance(nest, t)) + if include_derived + else (lambda x_, t: type(nest) is t) + ) + if check_fn(nest, tuple): + ret_list = [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + if to_mutable: + return ret_list + if hasattr(nest, "_fields"): + return class_instance(**dict(zip(nest._fields, ret_list))) + return class_instance(tuple(ret_list)) + elif check_fn(nest, list) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + return copy.deepcopy(nest) + return class_instance( + [ + copy_nest( + i, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for i in nest + ] + ) + elif check_fn(nest, dict): + class_instance = type(nest) + dict_ = { + k: copy_nest( + v, + include_derived=include_derived, + to_mutable=to_mutable, + extra_nest_types=extra_nest_types, + ) + for k, v in nest.items() + } + if isinstance(nest, OrderedDict): + return class_instance(**dict_) + return class_instance(dict_) + return nest + + +@handle_exceptions +def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): + """ + Group all unique index chains in a nest. This function is useful for finding all + unique index chains in a nest, and then duplicating the values at those index chains + for functional frameworks. + + Parameters + ---------- + nest + nest to get duplicate index chains for. + + Returns + ------- + list of index chains to duplicate. + """ + all_index_chains = ivy.nested_argwhere(nest, lambda _: True) + duplicates = [] + duplicate_index_chains = {} + for index_chain in all_index_chains: + val = ivy.index_nest(nest, index_chain) + if ivy.is_array(val): + for i in range(len(duplicates)): + if val is duplicates[i]: + duplicate_index_chains[i].append(index_chain) + break + else: + duplicates.append(val) + duplicate_index_chains[len(duplicates) - 1] = [index_chain] + return list(duplicate_index_chains.values()) + + # Extra # # ------# @@ -92,140 +332,146 @@ def index_nest( @handle_exceptions -def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: +def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: + if len(index) == 1: + idx = index[0] + if isinstance(nest, list): + nest.insert(idx, value) + else: + nest[index[0]] = value + else: + insert_into_nest_at_index(nest[index[0]], index[1:], value) + + +@handle_exceptions +def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: """ - Prune a nested object at a specified index. + Insert a value into the nested item at specified indices with specified values. Parameters ---------- nest - The nested object to prune. - index - A tuple of indices for the index at which to prune. + The nested object to insert into. + indices + A tuple of tuples of indices for the indices at which to insert + values. + values + The new values for inserting. """ - if len(index) == 1: - del nest[index[0]] - else: - prune_nest_at_index(nest[index[0]], index[1:]) + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + [ + insert_into_nest_at_index(nest, index, value) + for index, value in zip(indices, values) + ] + + +# noinspection PyShadowingBuiltins @handle_exceptions -def set_nest_at_index( - nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], - index: Sequence[Union[str, int]], - value: Any, - /, - shallow: bool = True, - _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: +def map( + fn: Callable, + constant: Optional[Dict[str, Any]] = None, + unique: Optional[Dict[str, Iterable[Any]]] = None, + mean: bool = False, +) -> List: """ - Set the value of a nested item at a specified index. + Apply a function on each item of an iterable x. Parameters ---------- - nest - The nested object to update. - index - A tuple of indices for the index at which to update. - value - The new value for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - _result - Placeholder for the result of the update. do not set this paramter. + fn + The function to map onto x. + constant + keyword arguments which remain constant between each function call. + Default is ``None``. + unique + keyword arguments which are unique for each function call. Default is ``None``. + mean + Whether to compute the mean across the return values, and return this mean. + Default is ``False``. Returns ------- ret - nest with changed value at the given index. + x following the application of fn to each of its iterated items. Examples -------- - With :class:`ivy.Array` inputs: + With :code:`int` inputs: - >>> x = ivy.array([[1., 2.], [3., 4.]]) - >>> y = (1, 1) - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([[1., 2.], [3., 5.]]) + >>> def special_square(x : float) -> float : return np.square(x) + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x' : [1,2,3]}, + ... mean = False) + >>> print(results) + [1, 4, 9] - >>> x = ivy.array([1., 2., 3., 4.]) - >>> y = [1] - >>> z = 5. - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - ivy.array([1., 5., 3., 4.]) + >>> results = ivy.map(fn = special_square, + ... constant = None, + ... unique = {'x':[0,1,2]}, + ... mean = True) + >>> print(results) + 1.6666666666666667 - With :code:`Dict` input: + >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = False) + >>> print(results) + [array([1,1]), + array([1,2]), + array([1,3])] - >>> x = {1 : [1, [2, 3]], 2: (4, 5)} - >>> y = (1, 1) - >>> z = 2 - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - {1: [1, 2], 2: (4, 5)} + >>> results = ivy.map(fn = special_pow, + ... constant = {'y':[0,1]}, + ... unique = {'x':[1,2,3]}, + ... mean = True) + >>> print(results) + [1. 2.] - With :code:`List` inputs: + With float inputs: - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> y = (2, 1, 0) - >>> z = 'H' - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - [['a','b','c'],['d','e','f'],['g',['H','i']]] + >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':10., 'b':1.}, + ... unique = {'x':[0.,1.,2.]}, + ... mean = False) + >>> print(results) + [1.0, 11.0, 21.0] - With :class:`ivy.Container` input: + With :class:`ivy.Array` inputs: - >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) - >>> y = ('b',) - >>> z = ivy.array([3., 4.]) - >>> ivy.set_nest_at_index(x, y, z) - >>> print(x) - { - a: ivy.array([1., 2.]), - b: ivy.array([3., 4.]) - } + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = False) + >>> print(results) + [ivy.array([0., 10., 100.]), + ivy.array([1., 10., 101.])] + + >>> results = ivy.map(fn = linear_model, + ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, + ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, + ... mean = True) + >>> print(results) + ivy.array([ 0.5, 10. , 100. ]) """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if _result is None: - if shallow: - _result = nest_type(nest) - else: - _result = copy_nest(nest, include_derived=True) - _result = list(_result) if is_tuple else _result - if len(index) == 1: - if shallow: - try: - nest[index[0]] = value - except TypeError: - pass - _result[index[0]] = value - else: - _result[index[0]] = set_nest_at_index( - nest[index[0]], index[1:], value, shallow, _result[index[0]] + c = ivy.default(constant, {}) + u = ivy.default(unique, {}) + rets = [ + r + for r in _map( + lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() ) - try: - _result = nest_type(_result) - except TypeError: - _result = nest_type(*_result) - return _result - + ] + if mean: + rets = sum(rets) / len(rets) -@handle_exceptions -def insert_into_nest_at_index(nest: Iterable, index: Tuple, value, /) -> None: - if len(index) == 1: - idx = index[0] - if isinstance(nest, list): - nest.insert(idx, value) - else: - nest[index[0]] = value - else: - insert_into_nest_at_index(nest[index[0]], index[1:], value) + return rets @handle_exceptions @@ -337,162 +583,71 @@ def map_nest_at_index( @handle_exceptions -def multi_index_nest( - nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], - indices: Iterable[Iterable[int]], +def map_nest_at_indices( + nest: Iterable, + indices: Tuple, + fn: Callable, /, -) -> Iterable[Any]: + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: """ - Repeatedly index a nested object, using a tuple of tuples of indices or keys in the - case of dicts. + Map a function to the values of a nested item at the specified indices. Parameters ---------- nest - The nested object to slice. + The nested object to update. indices - A tuple of tuples of indices to apply. + A tuple of tuples of indices for the indices at which to update. + fn + The function to perform on the nest at the given index. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. Returns ------- ret - The result elements through indexing the nested object. + nest with applicable of fn on given indices. Examples -------- - With :code:`Tuple` inputs: - - >>> x = (1, 2) - >>> y = [[0]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [1] - - With :class:`ivy.Array` inputs: + With :code:`List` inputs: - >>> x = ivy.array([[1., 2.], - ... [3., 4.]]) - >>> y = [[0],[1]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [ivy.array([1., 2.], ivy.array([3., 4.])] + >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] + >>> indices = [[0, 4], [1, 5]] + >>> function = lambda x : x + 'b' + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] - With :class:`ivy.Container` input: + With :code:`Tuple` inputs: - >>> x = ivy.Container(a=ivy.array([1,2]), - ... b=[30,40]) - >>> y = ('a', ('b', 0)) - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [ivy.array([1, 2]), 30] + >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) + >>> indices = ((0, 2),(1, 0),(1, 2)) + >>> function = abs + >>> ivy.map_nest_at_indices(nest, indices, function) + >>> print(nest) + ([-9, 8, 27], [9, -4, 5, 7]) With :code:`Dict` input: - >>> x = {'a': 0, 'b': [1, [2, 3]], 'c': (4, 5)} - >>> y = (('b', 1), 'a') - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - [[2, 3], 0] - - With :code:`List` inputs: - - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> y = [[2, 1, 0], [0, 1]] - >>> z = ivy.multi_index_nest(x, y) - >>> print(z) - ['h', 'b'] - """ - return [index_nest(nest, index) for index in indices] - - -@handle_exceptions -def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: - """ - Prune a nested object at specified indices. - - Parameters - ---------- - nest - The nested object to prune. - indices - A tuple of tuples of indices for the indices at which to prune. - """ - # Delete first deeper elements and elements with larger index - indices_sorted = sorted( - indices, - key=str, - reverse=True, - ) - [prune_nest_at_index(nest, index) for index in indices_sorted] - - -@handle_exceptions -def set_nest_at_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], - indices: Union[List[int], Tuple[int], Iterable[int]], - values: Union[List[int], Tuple[int], Iterable[int]], - /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: - """ - Set the value of a nested item at specified indices with specified values. - - Parameters - ---------- - nest - The nested object to update. - indices - A tuple of tuples of indices for the indices at which to update. - values - The new values for updating. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - - Returns - ------- - ret - nest with updated values at the given indices. - - Examples - -------- - With :code:`List` inputs: - - >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] - >>> indices = [[0, 4], [1, 3]] - >>> values = [111, 'x'] - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] - - With :code:`Tuple` inputs: - - >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] - >>> indices = ((0, 1),(1, 2)) - >>> values = ('ivy', 'x') - >>> ivy.set_nest_at_indices(nest, indices, values) - >>> print(nest) - (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) - - With :code:`Dict` input: - - >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} - >>> indices = (('a', 1), ('b', 2), ('c', 0)) - >>> values = (11., 22., 33.) - >>> ivy.set_nest_at_indices(nest, indices, values) + >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} + >>> indices = (('a', 2), ('b', 0), ('c', 1)) + >>> function = lambda x : x + 1 + >>> ivy.map_nest_at_indices(nest, indices, function) >>> print(nest) - {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} + {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} With :class:`ivy.Array` inputs: - >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) - >>> indices = ((0, 1),(1, 2)) - >>> values = (11., 22.) - >>> ivy.set_nest_at_indices(nest, indices, values) + >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) + >>> indices = ((0, 1),(1, 1),(1, 2)) + >>> function = lambda x : x ** 2 + >>> ivy.map_nest_at_indices(nest, indices, function) >>> print(nest) - ivy.array([[1., 11., 3.], [4., 5., 22.]]) + ivy.array([[ -9., 64., -17.], + [ 11., 9., 25.]]) """ is_tuple = isinstance(nest, tuple) nest_type = type(nest) if is_tuple else lambda x: x @@ -501,10 +656,8 @@ def set_nest_at_indices( else: result = copy_nest(nest, include_derived=True) result = list(result) if is_tuple else result - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - for index, value in zip(indices, values): - result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) + for i, index in enumerate(indices): + result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) try: result = nest_type(result) except TypeError: @@ -513,125 +666,87 @@ def set_nest_at_indices( @handle_exceptions -def insert_into_nest_at_indices(nest: Iterable, indices: Tuple, values, /) -> None: - """ - Insert a value into the nested item at specified indices with specified values. - - Parameters - ---------- - nest - The nested object to insert into. - indices - A tuple of tuples of indices for the indices at which to insert - values. - values - The new values for inserting. - """ - if not isinstance(values, (list, tuple)): - values = [values] * len(indices) - [ - insert_into_nest_at_index(nest, index, value) - for index, value in zip(indices, values) - ] - - -@handle_exceptions -def map_nest_at_indices( - nest: Iterable, - indices: Tuple, - fn: Callable, +def multi_index_nest( + nest: Union[List, Dict, Tuple, ivy.Array, ivy.NativeArray, ivy.Container], + indices: Iterable[Iterable[int]], /, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: +) -> Iterable[Any]: """ - Map a function to the values of a nested item at the specified indices. + Repeatedly index a nested object, using a tuple of tuples of indices or keys in the + case of dicts. Parameters ---------- nest - The nested object to update. + The nested object to slice. indices - A tuple of tuples of indices for the indices at which to update. - fn - The function to perform on the nest at the given index. - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. + A tuple of tuples of indices to apply. Returns ------- ret - nest with applicable of fn on given indices. + The result elements through indexing the nested object. Examples -------- - With :code:`List` inputs: + With :code:`Tuple` inputs: - >>> nest = [['a', 'c', 'e', 'd', 'u', 'k'], ['m', 'n', 'f', 'p', 'q', 't']] - >>> indices = [[0, 4], [1, 5]] - >>> function = lambda x : x + 'b' - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - [['a', 'c', 'e', 'd', 'ub', 'k'], ['m', 'n', 'f', 'p', 'q', 'tb']] + >>> x = (1, 2) + >>> y = [[0]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [1] - With :code:`Tuple` inputs: + With :class:`ivy.Array` inputs: - >>> nest = ([-9, 8, -27],[9, -4, -5, 7]) - >>> indices = ((0, 2),(1, 0),(1, 2)) - >>> function = abs - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ([-9, 8, 27], [9, -4, 5, 7]) + >>> x = ivy.array([[1., 2.], + ... [3., 4.]]) + >>> y = [[0],[1]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [ivy.array([1., 2.], ivy.array([3., 4.])] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1,2]), + ... b=[30,40]) + >>> y = ('a', ('b', 0)) + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [ivy.array([1, 2]), 30] With :code:`Dict` input: - >>> nest = {'a': [8., 16., 22.], 'b': [10., 44., 81.], 'c': [9., 75., 37.]} - >>> indices = (('a', 2), ('b', 0), ('c', 1)) - >>> function = lambda x : x + 1 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - {'a': [8.0, 16.0, 23.0], 'b': [11.0, 44.0, 81.0], 'c': [9.0, 76.0, 37.0]} + >>> x = {'a': 0, 'b': [1, [2, 3]], 'c': (4, 5)} + >>> y = (('b', 1), 'a') + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + [[2, 3], 0] - With :class:`ivy.Array` inputs: + With :code:`List` inputs: - >>> nest = ivy.array([[-9., 8., -17.],[11., -3., 5.]]) - >>> indices = ((0, 1),(1, 1),(1, 2)) - >>> function = lambda x : x ** 2 - >>> ivy.map_nest_at_indices(nest, indices, function) - >>> print(nest) - ivy.array([[ -9., 64., -17.], - [ 11., 9., 25.]]) + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> y = [[2, 1, 0], [0, 1]] + >>> z = ivy.multi_index_nest(x, y) + >>> print(z) + ['h', 'b'] """ - is_tuple = isinstance(nest, tuple) - nest_type = type(nest) if is_tuple else lambda x: x - if shallow: - result = nest_type(nest) - else: - result = copy_nest(nest, include_derived=True) - result = list(result) if is_tuple else result - for i, index in enumerate(indices): - result = map_nest_at_index(nest, index, fn, _result=result, shallow=shallow) - try: - result = nest_type(result) - except TypeError: - result = nest_type(*result) - return result + return [index_nest(nest, index) for index in indices] @handle_exceptions -def nested_argwhere( +def nested_any( nest: Iterable, fn: Callable, check_nests: bool = False, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - _index: Optional[List] = None, _base: bool = True, - stop_after_n_found: Optional[int] = None, extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[Iterable, bool]: +) -> bool: """ - Check the leaf nodes of nested x via function fn, and returns all nest indices where - the method evaluates as True. + Check the leaf nodes of nest x via function fn, and returns True if any evaluate to + True, else False. Parameters ---------- @@ -642,9 +757,66 @@ def nested_argwhere( check_nests Whether to also check the nests for the condition, not only nest leaves. Default is ``False``. - to_ignore - Types to ignore when deciding whether to go deeper into the nest or not - _index + _base + Whether the current function call is the first function call in the recursive + stack. Used internally, do not set manually. + extra_nest_types + Types to recursively check when deciding whether to go deeper into the + nest or not + + Returns + ------- + ret + A boolean, whether the function evaluates to true for any leaf node. + """ + extra_nest_types = ivy.default(extra_nest_types, ()) + if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): + if isinstance(nest, (ivy.Array, ivy.NativeArray)): + if ivy.any(fn(nest)): + return True + for i, item in enumerate(nest): + if nested_any(item, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif isinstance(nest, dict): + for k, v in nest.items(): + if nested_any(v, fn, check_nests, False, extra_nest_types): + return True + if check_nests and fn(nest): + return True + elif fn(nest): + return True + return False + + +@handle_exceptions +def nested_argwhere( + nest: Iterable, + fn: Callable, + check_nests: bool = False, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + _index: Optional[List] = None, + _base: bool = True, + stop_after_n_found: Optional[int] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, +) -> Union[Iterable, bool]: + """ + Check the leaf nodes of nested x via function fn, and returns all nest indices where + the method evaluates as True. + + Parameters + ---------- + nest + The nest to check the leaves of. + fn + The conditon function, returning True or False. + check_nests + Whether to also check the nests for the condition, not only nest leaves. + Default is ``False``. + to_ignore + Types to ignore when deciding whether to go deeper into the nest or not + _index The indices detected so far. None at the beginning. Used internally, do not set manually. _base @@ -799,330 +971,125 @@ def nested_argwhere( @handle_exceptions -def all_nested_indices( - nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray, ivy.Container] = None, +def nested_map( + x: Union[ivy.Array, ivy.NativeArray, Iterable], /, - include_nests: bool = False, - _index: Optional[Union[int, Sequence[int]]] = None, - _base: bool = True, - extra_nest_types: Optional[Union[ivy.Dtype, Sequence[ivy.Dtype]]] = None, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: + fn: Callable, + include_derived: Optional[Union[Dict[type, bool], bool]] = None, + to_ignore: Optional[Union[type, Tuple[type]]] = None, + to_mutable: bool = False, + max_depth: Optional[int] = None, + _depth: int = 0, + _tuple_check_fn: Optional[Callable] = None, + _list_check_fn: Optional[Callable] = None, + _dict_check_fn: Optional[Callable] = None, + extra_nest_types: Optional[Union[type, Tuple[type]]] = None, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: """ - Return indices of all the elements in nest. + Apply a function on x in a nested manner, whereby all dicts, lists and tuples are + traversed to their lowest leaves before applying the method and returning x. If x is + not nested, the method is applied to x directly. Parameters ---------- - nest - The nest to check the leaves of. - include_nests - Whether to also include indices of the nests themselves, not only - leaves. Default is ``False``. - _index - The indices detected so far. None at the beginning. Used internally, - do not set manually. - _base - Whether the current function call is the first function call in the - recursive stack. Used internally, do not set manually. + x + The item to apply the mapped function to. + fn + The function to map onto x. + include_derived + Whether to also recursive for classes derived from tuple, list and dict. + Default is ``False``. + to_ignore + Types to ignore when deciding whether to go deeper into the nest or not + to_mutable + Whether to convert the nest to a mutable form, changing all tuples to lists. + Default is ``False``. + max_depth + The maximum nested depth to reach. Default is 1. Increase this if the nest is + deeper. + _depth + Placeholder for tracking the recursive depth, do not set this parameter. + _tuple_check_fn + Placeholder for the tuple check function, do not set this parameter. + _list_check_fn + Placeholder for the list check function, do not set this parameter. + _dict_check_fn + Placeholder for the dict check function, do not set this parameter. extra_nest_types Types to recursively check when deciding whether to go deeper into the nest or not - out - Optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. Returns ------- ret - A set of indices of all elements in nest - - Both the description and the type hints above assumes an array input - for simplicity, but this function is nestable, and therefore also - accepts :class:ivy.Container instances in place of the arguments. + x following the applicable of fn to it's nested leaves, or x itself if x is not + nested. Examples -------- - With :class:`Dict` input: - - >>> x = {'a': 2., 'b': [6., [15., 9.]], 'c': (7., 56.)} - >>> y = ivy.all_nested_indices(x) - >>> print(y) - [['a'], ['b', 0], ['b', 1, 0], ['b', 1, 1], ['c', 0], ['c', 1]] + With :class:`Tuple` inputs: - With :class:`ivy.Array` input: + >>> x = ([[1., 2.], [3., 4.]]) + >>> function = lambda a : a * 2 + >>> ivy.nested_map(x, function) + [[2.0, 4.0], [6.0, 8.0]] + >>> print(x) + [[2.0, 4.0], [6.0, 8.0]] - >>> x = ivy.array([0., 1., 2., 3., 4.]) - >>> y = ivy.all_nested_indices(x, False, out=x) - >>> print(y) - [[]] + With :code:`Dict` input: - With :class:`ivy.Container` input: + >>> x = {1 : [1, [2, 3]], 2: (4, 5)} + >>> function = lambda a : a + 1 + >>> ivy.nested_map(x, function) + {1 : [2, [3, 4]], 2: (5, 6)} + >>> print(x) + {1 : [2, [3, 4]], 2: (5, 6)} - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.all_nested_indices(x, True) - >>> print(y) - [['a'], ['b']] - """ - _index = list() if _index is None else _index - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - ind = ivy.argwhere(ivy.ones_like(nest)) - indices = list() - for i in range(len(ind)): - indices.append(_index + ind.to_list()[i]) - return indices - _indices = [ - all_nested_indices( - item, include_nests, _index + [i], False, extra_nest_types - ) - for i, item in enumerate(nest) - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - elif isinstance(nest, dict): - _indices = [ - all_nested_indices(v, include_nests, _index + [k], False, extra_nest_types) - for k, v in nest.items() - ] - _indices = [idx for idxs in _indices if idxs for idx in idxs] - if include_nests: - _indices.append(_index) - else: - return [_index] - return [index for index in _indices if index] + With :code:`List` inputs: + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> function = lambda a: a + 'H' + >>> ivy.nested_map(x, function) + [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] + >>> print(x) + [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] -# noinspection PyShadowingBuiltins + With :class:`ivy.Container` input: + >>> x = ivy.Container( + ... a=ivy.array([[1, 2, 3], [9, 8, 7]]) , b=ivy.array([[4, 5, 6], [12, 13, 14]]) + ... ) + >>> function = lambda a : a + 1 + >>> ivy.nested_map(x, function) + { + a: ivy.array([[2, 3, 4], + [10, 9, 8]]), + b: ivy.array([[5, 6, 7], + [13, 14, 15]]) + } + >>> print(x) + { + a: ivy.array([[2, 3, 4], + [10, 9, 8]]), + b: ivy.array([[5, 6, 7], + [13, 14, 15]]) + } -@handle_exceptions -def map( - fn: Callable, - constant: Optional[Dict[str, Any]] = None, - unique: Optional[Dict[str, Iterable[Any]]] = None, - mean: bool = False, -) -> List: - """ - Apply a function on each item of an iterable x. + >>> nest = ([1, 2], [3, 4], [5, 6], {"a": 1, "b": 2, "c": 3}) + >>> function = lambda a : a * 2 + >>> ivy.nested_map(nest, function, to_ignore=list) + ([1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6], {'a': 2, 'b': 4, 'c': 6}) - Parameters - ---------- - fn - The function to map onto x. - constant - keyword arguments which remain constant between each function call. - Default is ``None``. - unique - keyword arguments which are unique for each function call. Default is ``None``. - mean - Whether to compute the mean across the return values, and return this mean. - Default is ``False``. - - Returns - ------- - ret - x following the application of fn to each of its iterated items. - - Examples - -------- - With :code:`int` inputs: - - >>> def special_square(x : float) -> float : return np.square(x) - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x' : [1,2,3]}, - ... mean = False) - >>> print(results) - [1, 4, 9] - - >>> results = ivy.map(fn = special_square, - ... constant = None, - ... unique = {'x':[0,1,2]}, - ... mean = True) - >>> print(results) - 1.6666666666666667 - - >>> def special_pow(x:float,y:float) ->float : return np.power(x,y) - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = False) - >>> print(results) - [array([1,1]), - array([1,2]), - array([1,3])] - - >>> results = ivy.map(fn = special_pow, - ... constant = {'y':[0,1]}, - ... unique = {'x':[1,2,3]}, - ... mean = True) - >>> print(results) - [1. 2.] - - With float inputs: - - >>> def linear_model(w:float, x:float, b:float) -> float: return w*x + b - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':10., 'b':1.}, - ... unique = {'x':[0.,1.,2.]}, - ... mean = False) - >>> print(results) - [1.0, 11.0, 21.0] - - With :class:`ivy.Array` inputs: - - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = False) - >>> print(results) - [ivy.array([0., 10., 100.]), - ivy.array([1., 10., 101.])] - - >>> results = ivy.map(fn = linear_model, - ... constant = {'w':ivy.array([1.,0.,1.]), 'b':ivy.array([0.,10.,100.])}, - ... unique = {'x':[ivy.array([0.,1.,0.]), ivy.array([1.,1.,1.])]}, - ... mean = True) - >>> print(results) - ivy.array([ 0.5, 10. , 100. ]) - """ - c = ivy.default(constant, {}) - u = ivy.default(unique, {}) - rets = [ - r - for r in _map( - lambda *uv: fn(**dict(**c, **dict(zip(u.keys(), uv)))), *u.values() - ) - ] - if mean: - rets = sum(rets) / len(rets) - - return rets - - -@handle_exceptions -def nested_map( - x: Union[ivy.Array, ivy.NativeArray, Iterable], - /, - fn: Callable, - include_derived: Optional[Union[Dict[type, bool], bool]] = None, - to_ignore: Optional[Union[type, Tuple[type]]] = None, - to_mutable: bool = False, - max_depth: Optional[int] = None, - _depth: int = 0, - _tuple_check_fn: Optional[Callable] = None, - _list_check_fn: Optional[Callable] = None, - _dict_check_fn: Optional[Callable] = None, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, - shallow: bool = True, -) -> Union[ivy.Array, ivy.NativeArray, Iterable, Dict]: - """ - Apply a function on x in a nested manner, whereby all dicts, lists and tuples are - traversed to their lowest leaves before applying the method and returning x. If x is - not nested, the method is applied to x directly. - - Parameters - ---------- - x - The item to apply the mapped function to. - fn - The function to map onto x. - include_derived - Whether to also recursive for classes derived from tuple, list and dict. - Default is ``False``. - to_ignore - Types to ignore when deciding whether to go deeper into the nest or not - to_mutable - Whether to convert the nest to a mutable form, changing all tuples to lists. - Default is ``False``. - max_depth - The maximum nested depth to reach. Default is 1. Increase this if the nest is - deeper. - _depth - Placeholder for tracking the recursive depth, do not set this parameter. - _tuple_check_fn - Placeholder for the tuple check function, do not set this parameter. - _list_check_fn - Placeholder for the list check function, do not set this parameter. - _dict_check_fn - Placeholder for the dict check function, do not set this parameter. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - shallow - Whether to inplace update the input nest or not - Only works if nest is a mutable type. Default is ``True``. - - Returns - ------- - ret - x following the applicable of fn to it's nested leaves, or x itself if x is not - nested. - - Examples - -------- - With :class:`Tuple` inputs: - - >>> x = ([[1., 2.], [3., 4.]]) - >>> function = lambda a : a * 2 - >>> ivy.nested_map(x, function) - [[2.0, 4.0], [6.0, 8.0]] - >>> print(x) - [[2.0, 4.0], [6.0, 8.0]] - - With :code:`Dict` input: - - >>> x = {1 : [1, [2, 3]], 2: (4, 5)} - >>> function = lambda a : a + 1 - >>> ivy.nested_map(x, function) - {1 : [2, [3, 4]], 2: (5, 6)} - >>> print(x) - {1 : [2, [3, 4]], 2: (5, 6)} - - With :code:`List` inputs: - - >>> x = [['a', 'b', 'c'], - ... ['d', 'e', 'f'], - ... ['g', ['h', 'i']]] - >>> function = lambda a: a + 'H' - >>> ivy.nested_map(x, function) - [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] - >>> print(x) - [['aH','bH','cH'],['dH','eH','fH'],['gH',['hH','iH']]] - - With :class:`ivy.Container` input: - - >>> x = ivy.Container( - ... a=ivy.array([[1, 2, 3], [9, 8, 7]]) , b=ivy.array([[4, 5, 6], [12, 13, 14]]) - ... ) - >>> function = lambda a : a + 1 - >>> ivy.nested_map(x, function) - { - a: ivy.array([[2, 3, 4], - [10, 9, 8]]), - b: ivy.array([[5, 6, 7], - [13, 14, 15]]) - } - >>> print(x) - { - a: ivy.array([[2, 3, 4], - [10, 9, 8]]), - b: ivy.array([[5, 6, 7], - [13, 14, 15]]) - } - - >>> nest = ([1, 2], [3, 4], [5, 6], {"a": 1, "b": 2, "c": 3}) - >>> function = lambda a : a * 2 - >>> ivy.nested_map(nest, function, to_ignore=list) - ([1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6], {'a': 2, 'b': 4, 'c': 6}) - - >>> nest = [[1, 2], [3, [4, 5]], [[6], [7, 8, [9, 10]]]] - >>> function = lambda a : a * 2 - >>> ivy.nested_map(nest, function, max_depth = 3) - [[2, 4], [6, [8, 10]], [[12], [14, 16, [9, 10]]]] + >>> nest = [[1, 2], [3, [4, 5]], [[6], [7, 8, [9, 10]]]] + >>> function = lambda a : a * 2 + >>> ivy.nested_map(nest, function, max_depth = 3) + [[2, 4], [6, [8, 10]], [[12], [14, 16, [9, 10]]]] >>> nest = ([23, 25, 1337], [63, 98, 6]) >>> function = lambda a : a + 1 @@ -1261,211 +1228,46 @@ def nested_map( @handle_exceptions -def nested_any( - nest: Iterable, - fn: Callable, - check_nests: bool = False, - _base: bool = True, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> bool: +def nested_multi_map( + func: Callable, + nests: List[Iterable], + index_chains=None, + to_apply=True, + prune_unapplied=False, + index_chain="", + config=None, + to_ivy=True, +): """ - Check the leaf nodes of nest x via function fn, and returns True if any evaluate to - True, else False. + Apply function to all array values from a collection of identically structured ivy + arrays. Parameters ---------- + func + Function to apply to each nest entry. nest - The nest to check the leaves of. - fn - The conditon function, returning True or False. - check_nests - Whether to also check the nests for the condition, not only nest leaves. - Default is ``False``. - _base - Whether the current function call is the first function call in the recursive - stack. Used internally, do not set manually. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - + nests to map. + index_chains + The key-chains to apply or not apply the method to. Default is ``None``. + to_apply + If True, the method will be applied to index_chains, otherwise index_chains will + be skipped. Default is ``True``. + prune_unapplied + Whether to prune index_chains for which the function was not applied, + otherwise the leftmost nest value is used. Default is ``False``. + index_chain + Chain of keys for this dict entry (Default value = '') + config + The configuration for the nests. Default is the same as nest0. + to_ivy + convert the output to ivy_arrays. Default is ``True`` Returns ------- - ret - A boolean, whether the function evaluates to true for any leaf node. - """ - extra_nest_types = ivy.default(extra_nest_types, ()) - if isinstance(nest, (tuple, list)) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - if ivy.any(fn(nest)): - return True - for i, item in enumerate(nest): - if nested_any(item, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif isinstance(nest, dict): - for k, v in nest.items(): - if nested_any(v, fn, check_nests, False, extra_nest_types): - return True - if check_nests and fn(nest): - return True - elif fn(nest): - return True - return False - - -@handle_exceptions -def copy_nest( - nest: Union[ivy.Array, ivy.NativeArray, Iterable], - /, - include_derived: bool = False, - to_mutable: bool = False, - extra_nest_types: Optional[Union[type, Tuple[type]]] = None, -) -> Union[ivy.Array, ivy.NativeArray, Iterable]: - """ - Copy a nest deeply, but without copying leaves of the nest, only the nest lists, - tuples and dicts are copied. - - Parameters - ---------- - nest - The nest to copy. - include_derived - Whether to also recursive for classes derived from tuple, list and dict. - Default is ``False``. - to_mutable - Whether to convert the nest to a mutable form, changing all tuples to lists. - Default is ``False``. - extra_nest_types - Types to recursively check when deciding whether to go deeper into the - nest or not - - Returns - ------- - ret - The copied nest. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> nest = ivy.array([[1.,2.,3.],[7.,8.,9.]]) - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - ivy.array([[1., 2., 3.], - [7., 8., 9.]]) - - With :code:`Iterable` input: - - >>> nest = [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - >>> copied_nest = ivy.copy_nest(nest, include_derived = True) - >>> print(copied_nest) - [[1, 2, 3, 4, 5], [23, 24, 25, 26, 27]] - - >>> nest = ([23, 25, 1337], [63, 98, 6]) - >>> copied_nest = ivy.copy_nest(nest, to_mutable = True) - >>> print(copied_nest) - [[23, 25, 1337], [63, 98, 6]] - - >>> nest = {'first': [23., 24., 25], 'second': [46., 48., 50]} - >>> copied_nest = ivy.copy_nest(nest) - >>> print(copied_nest) - {'first': [23.0, 24.0, 25], 'second': [46.0, 48.0, 50]} - """ - extra_nest_types = ivy.default(extra_nest_types, ()) - class_instance = type(nest) - check_fn = ( - (lambda x_, t: isinstance(nest, t)) - if include_derived - else (lambda x_, t: type(nest) is t) - ) - if check_fn(nest, tuple): - ret_list = [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - if to_mutable: - return ret_list - if hasattr(nest, "_fields"): - return class_instance(**dict(zip(nest._fields, ret_list))) - return class_instance(tuple(ret_list)) - elif check_fn(nest, list) or isinstance(nest, extra_nest_types): - if isinstance(nest, (ivy.Array, ivy.NativeArray)): - return copy.deepcopy(nest) - return class_instance( - [ - copy_nest( - i, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for i in nest - ] - ) - elif check_fn(nest, dict): - class_instance = type(nest) - dict_ = { - k: copy_nest( - v, - include_derived=include_derived, - to_mutable=to_mutable, - extra_nest_types=extra_nest_types, - ) - for k, v in nest.items() - } - if isinstance(nest, OrderedDict): - return class_instance(**dict_) - return class_instance(dict_) - return nest - - -@handle_exceptions -def nested_multi_map( - func: Callable, - nests: List[Iterable], - index_chains=None, - to_apply=True, - prune_unapplied=False, - index_chain="", - config=None, - to_ivy=True, -): - """ - Apply function to all array values from a collection of identically structured ivy - arrays. - - Parameters - ---------- - func - Function to apply to each nest entry. - nest - nests to map. - index_chains - The key-chains to apply or not apply the method to. Default is ``None``. - to_apply - If True, the method will be applied to index_chains, otherwise index_chains will - be skipped. Default is ``True``. - prune_unapplied - Whether to prune index_chains for which the function was not applied, - otherwise the leftmost nest value is used. Default is ``False``. - index_chain - Chain of keys for this dict entry (Default value = '') - config - The configuration for the nests. Default is the same as nest0. - to_ivy - convert the output to ivy_arrays. Default is ``True`` - Returns - ------- - nest containing the result of the function. The structure of the output is the - same as the input with the result of the function applied to each applicable - leaf and the value at that leaf in the first nest for a non-applicable leaf if - prune_unapplied is False else unapplied leaves are pruned. + nest containing the result of the function. The structure of the output is the + same as the input with the result of the function applied to each applicable + leaf and the value at that leaf in the first nest for a non-applicable leaf if + prune_unapplied is False else unapplied leaves are pruned. """ nest0 = None for nest in nests: @@ -1583,38 +1385,6 @@ def _found_in_index_chains(this_index_chain, index_chains): ) -@handle_exceptions -def duplicate_array_index_chains(nest: Union[ivy.Array, ivy.NativeArray, Iterable]): - """ - Group all unique index chains in a nest. This function is useful for finding all - unique index chains in a nest, and then duplicating the values at those index chains - for functional frameworks. - - Parameters - ---------- - nest - nest to get duplicate index chains for. - - Returns - ------- - list of index chains to duplicate. - """ - all_index_chains = ivy.nested_argwhere(nest, lambda _: True) - duplicates = [] - duplicate_index_chains = {} - for index_chain in all_index_chains: - val = ivy.index_nest(nest, index_chain) - if ivy.is_array(val): - for i in range(len(duplicates)): - if val is duplicates[i]: - duplicate_index_chains[i].append(index_chain) - break - else: - duplicates.append(val) - duplicate_index_chains[len(duplicates) - 1] = [index_chain] - return list(duplicate_index_chains.values()) - - def prune_empty(nest): """ Prune empty nests from a nest. @@ -1650,3 +1420,233 @@ def prune_empty(nest): if not valid and not (ivy.is_array(nest) or isinstance(nest, (int, float, str))): return None return nest + + +@handle_exceptions +def prune_nest_at_index(nest: Iterable, index: Tuple, /) -> None: + """ + Prune a nested object at a specified index. + + Parameters + ---------- + nest + The nested object to prune. + index + A tuple of indices for the index at which to prune. + """ + if len(index) == 1: + del nest[index[0]] + else: + prune_nest_at_index(nest[index[0]], index[1:]) + + +@handle_exceptions +def prune_nest_at_indices(nest: Iterable, indices: Tuple, /) -> None: + """ + Prune a nested object at specified indices. + + Parameters + ---------- + nest + The nested object to prune. + indices + A tuple of tuples of indices for the indices at which to prune. + """ + # Delete first deeper elements and elements with larger index + indices_sorted = sorted( + indices, + key=str, + reverse=True, + ) + [prune_nest_at_index(nest, index) for index in indices_sorted] + + +@handle_exceptions +def set_nest_at_index( + nest: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple], + index: Sequence[Union[str, int]], + value: Any, + /, + shallow: bool = True, + _result: Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple] = None, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: + """ + Set the value of a nested item at a specified index. + + Parameters + ---------- + nest + The nested object to update. + index + A tuple of indices for the index at which to update. + value + The new value for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + _result + Placeholder for the result of the update. do not set this paramter. + + Returns + ------- + ret + nest with changed value at the given index. + + Examples + -------- + With :class:`ivy.Array` inputs: + + >>> x = ivy.array([[1., 2.], [3., 4.]]) + >>> y = (1, 1) + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([[1., 2.], [3., 5.]]) + + >>> x = ivy.array([1., 2., 3., 4.]) + >>> y = [1] + >>> z = 5. + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + ivy.array([1., 5., 3., 4.]) + + With :code:`Dict` input: + + >>> x = {1 : [1, [2, 3]], 2: (4, 5)} + >>> y = (1, 1) + >>> z = 2 + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + {1: [1, 2], 2: (4, 5)} + + With :code:`List` inputs: + + >>> x = [['a', 'b', 'c'], + ... ['d', 'e', 'f'], + ... ['g', ['h', 'i']]] + >>> y = (2, 1, 0) + >>> z = 'H' + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + [['a','b','c'],['d','e','f'],['g',['H','i']]] + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1., 2.]) , b=ivy.array([4., 5.])) + >>> y = ('b',) + >>> z = ivy.array([3., 4.]) + >>> ivy.set_nest_at_index(x, y, z) + >>> print(x) + { + a: ivy.array([1., 2.]), + b: ivy.array([3., 4.]) + } + """ + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if _result is None: + if shallow: + _result = nest_type(nest) + else: + _result = copy_nest(nest, include_derived=True) + _result = list(_result) if is_tuple else _result + if len(index) == 1: + if shallow: + try: + nest[index[0]] = value + except TypeError: + pass + _result[index[0]] = value + else: + _result[index[0]] = set_nest_at_index( + nest[index[0]], index[1:], value, shallow, _result[index[0]] + ) + try: + _result = nest_type(_result) + except TypeError: + _result = nest_type(*_result) + return _result + + +@handle_exceptions +def set_nest_at_indices( + nest: Union[List, Tuple, Dict, ivy.Array, ivy.NativeArray], + indices: Union[List[int], Tuple[int], Iterable[int]], + values: Union[List[int], Tuple[int], Iterable[int]], + /, + shallow: bool = True, +) -> Union[ivy.Array, ivy.NativeArray, ivy.Container, Dict, List, Tuple]: + """ + Set the value of a nested item at specified indices with specified values. + + Parameters + ---------- + nest + The nested object to update. + indices + A tuple of tuples of indices for the indices at which to update. + values + The new values for updating. + shallow + Whether to inplace update the input nest or not + Only works if nest is a mutable type. Default is ``True``. + + Returns + ------- + ret + nest with updated values at the given indices. + + Examples + -------- + With :code:`List` inputs: + + >>> nest = [[1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']] + >>> indices = [[0, 4], [1, 3]] + >>> values = [111, 'x'] + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + [[1, 2, 3, 4, 111, 6], ['a', 'b', 'c', 'x', 'e', 'f']] + + With :code:`Tuple` inputs: + + >>> nest = [['abc', 'xyz', 'pqr'],[1, 4, 'a', 'b']] + >>> indices = ((0, 1),(1, 2)) + >>> values = ('ivy', 'x') + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + (['abc', 'ivy', 'pqr'], [1, 4, 'x', 'b']) + + With :code:`Dict` input: + + >>> nest = {'a': [1., 2., 3.], 'b': [4., 5., 6.], 'c': [0.]} + >>> indices = (('a', 1), ('b', 2), ('c', 0)) + >>> values = (11., 22., 33.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + {'a': [1.0, 11.0, 3.0], 'b': [4.0, 5.0, 22.0], 'c': [33.0]} + + With :class:`ivy.Array` inputs: + + >>> nest = ivy.array([[1., 2., 3.],[4., 5., 6.]]) + >>> indices = ((0, 1),(1, 2)) + >>> values = (11., 22.) + >>> ivy.set_nest_at_indices(nest, indices, values) + >>> print(nest) + ivy.array([[1., 11., 3.], [4., 5., 22.]]) + """ + is_tuple = isinstance(nest, tuple) + nest_type = type(nest) if is_tuple else lambda x: x + if shallow: + result = nest_type(nest) + else: + result = copy_nest(nest, include_derived=True) + result = list(result) if is_tuple else result + if not isinstance(values, (list, tuple)): + values = [values] * len(indices) + for index, value in zip(indices, values): + result = set_nest_at_index(nest, index, value, _result=result, shallow=shallow) + try: + result = nest_type(result) + except TypeError: + result = nest_type(*result) + return result diff --git a/ivy/functional/ivy/random.py b/ivy/functional/ivy/random.py index 7ed96ae869d1c..70e8c21a5db45 100644 --- a/ivy/functional/ivy/random.py +++ b/ivy/functional/ivy/random.py @@ -20,8 +20,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# ------- # +# --- Helpers --- # +# --------------- # def _check_bounds_and_get_shape(low, high, shape): @@ -51,6 +51,17 @@ def _check_bounds_and_get_shape(low, high, shape): return ivy.Shape(()) +def _check_shapes_broadcastable(out, inp): + if out is not None: + ivy.utils.assertions.check_shapes_broadcastable(out, inp) + + +def _check_valid_scale(std): + ivy.utils.assertions.check_greater( + std, 0, allow_equal=True, message="std must be non-negative" + ) + + def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_all_or_any_fn( low, @@ -73,66 +84,51 @@ def _randint_check_dtype_and_bound(low, high, dtype): ivy.utils.assertions.check_less(low, high) -def _check_valid_scale(std): - ivy.utils.assertions.check_greater( - std, 0, allow_equal=True, message="std must be non-negative" - ) - - -def _check_shapes_broadcastable(out, inp): - if out is not None: - ivy.utils.assertions.check_shapes_broadcastable(out, inp) - - -# Extra # -# ------# +# --- Main --- # +# ------------ # @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument -@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function -@infer_dtype @handle_device_shifting @infer_device -def random_uniform( +def multinomial( + population_size: int, + num_samples: int, + /, *, - low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, - high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, - shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, + batch_size: int = 1, + probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, + replace: bool = True, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a uniform distribution. Samples are uniformly distributed over - the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In - other words, any value within the given interval is equally likely to be drawn by - uniform. + Draws samples from a multinomial distribution. Specifically, returns a tensor where + each row contains num_samples indices sampled from the multinomial probability + distribution located in the corresponding row of tensor input. Parameters ---------- - low - Lower boundary of the output interval. All values generated will be greater than - or equal to ``low``. If array, must have same shape as ``high``. - high - Upper boundary of the output interval. All the values generated will be less - than ``high``. If array, must have same shape as ``low``. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. - Can only be specified when ``low`` and ``high`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. + population_size + The size of the population from which to draw samples. + num_samples + Number of independent samples to draw from the population. + batch_size + Number of tensors to generate. Default is 1. + probs + The unnormalized probabilities for all elements in population, + default is uniform *[batch_shape, population_size]* + replace + Whether to replace samples once they've been drawn. Default is ``True``. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default floating-point data type. Default ``None`` + (Default value = None) seed A python integer. Used to create a random seed distribution out @@ -142,67 +138,145 @@ def random_uniform( Returns ------- ret - Drawn samples from the parameterized uniform distribution. + Drawn samples indices from the multinomial distribution. - Functional Examples - ------------------- + Examples + -------- + >>> y = ivy.multinomial(10, 5) + >>> print(y) + ivy.array([[1, 8, 7, 8, 3]]) - >>> ivy.random_uniform() - ivy.array(0.26431865) + >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) + >>> print(y) + ivy.array([[3, 9, 7, 5, 1], + [1, 0, 8, 6, 7]]) - >>> ivy.random_uniform(shape=3) - ivy.array([0.475, 0.878, 0.861]) + >>> y = ivy.multinomial(10, 5, replace=False) + >>> print(y) + ivy.array([[2, 6, 4, 7, 0]]) - >>> ivy.random_uniform(shape=(2,3)) - ivy.array([[0.929 , 0.545 , 0.789 ], - [0.519 , 0.0435, 0.381 ]]) + With :class:`ivy.Array` input: - >>> ivy.random_uniform(low=3.0, high=6.0) - ivy.array(3.4608004) + >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) + >>> print(y) + ivy.array([5, 2, 7, 6, 9]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) - ivy.array([[1.85], - [1.81]]) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) + >>> print(y) + ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) - >>> z = ivy.zeros(()) - >>> ivy.random_uniform(low=1.0, high=2.0, out=z) - ivy.array(1.8458502) + >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), + ... replace=False) + >>> print(y) + ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') - ivy.array([[1.81, 1.8 ], - [1.32, 1.43]]) + With :class:`ivy.NativeArray` input: - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', - ... dtype='int32') - ivy.array([[1, 1], - [1, 1]]) + >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) + >>> print(y) + ivy.array([5, 7, 4, 2, 1]) - >>> z = ivy.zeros((1,2)) - >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', - ... dtype='float64', out=z) - ivy.array([[1.34, 1.02]]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) + >>> print(y) + ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) - >>> x = ivy.array([4.8, 5.6]) - >>> y = ivy.array([9.8, 7.4]) - >>> ivy.random_uniform(low=x, high=y) - ivy.array([0.475, 0.878]) + >>> y = ivy.multinomial(10, 5, batch_size=2, + ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), + ... replace=False) + >>> print(y) + ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) + """ + return ivy.current_backend().multinomial( + population_size, + num_samples, + batch_size=batch_size, + probs=probs, + replace=replace, + device=device, + seed=seed, + out=out, + ) - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) - ivy.array([6.67270088, 7.31128597]) - >>> ivy.random_uniform(low=x, high=y, device='cpu') - ivy.array([6.88, 6.75]) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_out_argument +@inputs_to_native_shapes +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +@infer_device +def randint( + low: Union[int, ivy.NativeArray, ivy.Array], + high: Union[int, ivy.NativeArray, ivy.Array], + /, + *, + shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, + device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + seed: Optional[int] = None, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return an array filled with random integers generated uniformly between low + (inclusive) and high (exclusive). - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') - ivy.array([8.62, 6.47]) + Parameters + ---------- + low + Lowest integer that can be drawn from the distribution. + high + One above the highest integer that can be drawn from the distribution. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn + Can only be specified when ``mean`` and ``std`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. + device + device on which to create the array. 'cuda:0', + 'cuda:1', 'cpu' etc. (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default integer data type. Default ``None`` + seed + A python integer. Used to create a random seed distribution + out + optional output array, for writing the result to. It must have a shape + that the inputs broadcast to. - >>> z = ivy.zeros((2,)) - >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) - ivy.array([5. , 7.3]) + Returns + ------- + ret + Returns an array with the given shape filled with integers from + the uniform distribution in the “half-open” interval [low, high) + + Examples + -------- + >>> y = ivy.randint(0, 9, shape=(1,1)) + >>> print(y) + ivy.array([[5]]) + + >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) + >>> print(y) + ivy.array([[ 8, 16], + [12, 9]]) + + >>> x = ivy.array([1, 2, 3]) + >>> ivy.randint(0, 10, shape=(3,), out=x) + >>> print(x) + ivy.array([2, 6, 7]) + + >>> y = ivy.zeros((3, 3)) + >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) + >>> print(y) + ivy.array([[ 7, 7, 5], + [12, 8, 8], + [ 8, 11, 3]]) """ - return ivy.current_backend().random_uniform( - low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed + return ivy.current_backend().randint( + low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out ) @@ -320,47 +394,55 @@ def random_normal( ) +# Extra # +# ------# + + @handle_exceptions @handle_backend_invalid @handle_nestable @handle_out_argument +@inputs_to_native_shapes @to_native_arrays_and_back @handle_array_function +@infer_dtype @handle_device_shifting @infer_device -def multinomial( - population_size: int, - num_samples: int, - /, +def random_uniform( *, - batch_size: int = 1, - probs: Optional[Union[ivy.Array, ivy.NativeArray]] = None, - replace: bool = True, + low: Union[float, ivy.NativeArray, ivy.Array] = 0.0, + high: Union[float, ivy.NativeArray, ivy.Array] = 1.0, + shape: Optional[Union[ivy.Array, ivy.Shape, ivy.NativeShape]] = None, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, seed: Optional[int] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Draws samples from a multinomial distribution. Specifically, returns a tensor where - each row contains num_samples indices sampled from the multinomial probability - distribution located in the corresponding row of tensor input. + Draws samples from a uniform distribution. Samples are uniformly distributed over + the half-open interval ``[low, high)`` (includes ``low``, but excludes ``high``). In + other words, any value within the given interval is equally likely to be drawn by + uniform. Parameters ---------- - population_size - The size of the population from which to draw samples. - num_samples - Number of independent samples to draw from the population. - batch_size - Number of tensors to generate. Default is 1. - probs - The unnormalized probabilities for all elements in population, - default is uniform *[batch_shape, population_size]* - replace - Whether to replace samples once they've been drawn. Default is ``True``. + low + Lower boundary of the output interval. All values generated will be greater than + or equal to ``low``. If array, must have same shape as ``high``. + high + Upper boundary of the output interval. All the values generated will be less + than ``high``. If array, must have same shape as ``low``. + shape + If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn. + Can only be specified when ``low`` and ``high`` are numeric values, else + exception will be raised. + Default is ``None``, where a single value is returned. device device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. - (Default value = None) + (Default value = None). + dtype + output array data type. If ``dtype`` is ``None``, the output array data + type will be the default floating-point data type. Default ``None`` seed A python integer. Used to create a random seed distribution out @@ -370,145 +452,67 @@ def multinomial( Returns ------- ret - Drawn samples indices from the multinomial distribution. - - Examples - -------- - >>> y = ivy.multinomial(10, 5) - >>> print(y) - ivy.array([[1, 8, 7, 8, 3]]) - - >>> y = ivy.multinomial(10, 5, batch_size=2, seed=42) - >>> print(y) - ivy.array([[3, 9, 7, 5, 1], - [1, 0, 8, 6, 7]]) - - >>> y = ivy.multinomial(10, 5, replace=False) - >>> print(y) - ivy.array([[2, 6, 4, 7, 0]]) - - With :class:`ivy.Array` input: + Drawn samples from the parameterized uniform distribution. - >>> y = ivy.multinomial(10, 5, probs=ivy.array([1/10]*10)) - >>> print(y) - ivy.array([5, 2, 7, 6, 9]) + Functional Examples + ------------------- - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7])) - >>> print(y) - ivy.array([[0, 4, 3, 4, 5], [1, 1, 0, 3, 2]]) + >>> ivy.random_uniform() + ivy.array(0.26431865) - >>> y = ivy.multinomial(7, 5, batch_size=2, probs=ivy.array([[1/7]*7, [1/7]*7]), - ... replace=False) - >>> print(y) - ivy.array([[2, 6, 1, 0, 3], [1, 0, 2, 5, 6]]) + >>> ivy.random_uniform(shape=3) + ivy.array([0.475, 0.878, 0.861]) - With :class:`ivy.NativeArray` input: + >>> ivy.random_uniform(shape=(2,3)) + ivy.array([[0.929 , 0.545 , 0.789 ], + [0.519 , 0.0435, 0.381 ]]) - >>> y = ivy.multinomial(10, 5, probs=ivy.native_array([1/10]*10)) - >>> print(y) - ivy.array([5, 7, 4, 2, 1]) + >>> ivy.random_uniform(low=3.0, high=6.0) + ivy.array(3.4608004) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10])) - >>> print(y) - ivy.array([[8, 0, 4, 1, 7], [2, 3, 4, 9, 3]]) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,1)) + ivy.array([[1.85], + [1.81]]) - >>> y = ivy.multinomial(10, 5, batch_size=2, - ... probs=ivy.native_array([[1/10]*10, [1/10]*10]), - ... replace=False) - >>> print(y) - ivy.array([[0, 2, 6, 9, 1], [6, 7, 2, 4, 3]]) - """ - return ivy.current_backend().multinomial( - population_size, - num_samples, - batch_size=batch_size, - probs=probs, - replace=replace, - device=device, - seed=seed, - out=out, - ) + >>> z = ivy.zeros(()) + >>> ivy.random_uniform(low=1.0, high=2.0, out=z) + ivy.array(1.8458502) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu') + ivy.array([[1.81, 1.8 ], + [1.32, 1.43]]) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_out_argument -@inputs_to_native_shapes -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -@infer_device -def randint( - low: Union[int, ivy.NativeArray, ivy.Array], - high: Union[int, ivy.NativeArray, ivy.Array], - /, - *, - shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None, - device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - seed: Optional[int] = None, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return an array filled with random integers generated uniformly between low - (inclusive) and high (exclusive). + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(2,2), device='cpu', + ... dtype='int32') + ivy.array([[1, 1], + [1, 1]]) - Parameters - ---------- - low - Lowest integer that can be drawn from the distribution. - high - One above the highest integer that can be drawn from the distribution. - shape - If the given shape is, e.g ``(m, n, k)``, then ``m * n * k`` samples are drawn - Can only be specified when ``mean`` and ``std`` are numeric values, else - exception will be raised. - Default is ``None``, where a single value is returned. - device - device on which to create the array. 'cuda:0', - 'cuda:1', 'cpu' etc. (Default value = None). - dtype - output array data type. If ``dtype`` is ``None``, the output array data - type will be the default integer data type. Default ``None`` - seed - A python integer. Used to create a random seed distribution - out - optional output array, for writing the result to. It must have a shape - that the inputs broadcast to. + >>> z = ivy.zeros((1,2)) + >>> ivy.random_uniform(low=1.0, high=2.0, shape=(1,2), device='cpu', + ... dtype='float64', out=z) + ivy.array([[1.34, 1.02]]) - Returns - ------- - ret - Returns an array with the given shape filled with integers from - the uniform distribution in the “half-open” interval [low, high) + >>> x = ivy.array([4.8, 5.6]) + >>> y = ivy.array([9.8, 7.4]) + >>> ivy.random_uniform(low=x, high=y) + ivy.array([0.475, 0.878]) - Examples - -------- - >>> y = ivy.randint(0, 9, shape=(1,1)) - >>> print(y) - ivy.array([[5]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, out=z, seed=42) + ivy.array([6.67270088, 7.31128597]) - >>> y = ivy.randint(2, 20, shape=(2, 2), device='cpu', seed=42) - >>> print(y) - ivy.array([[ 8, 16], - [12, 9]]) + >>> ivy.random_uniform(low=x, high=y, device='cpu') + ivy.array([6.88, 6.75]) - >>> x = ivy.array([1, 2, 3]) - >>> ivy.randint(0, 10, shape=(3,), out=x) - >>> print(x) - ivy.array([2, 6, 7]) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64') + ivy.array([8.62, 6.47]) - >>> y = ivy.zeros((3, 3)) - >>> ivy.randint(3, 15, shape=(3, 3), device='cpu', out=y) - >>> print(y) - ivy.array([[ 7, 7, 5], - [12, 8, 8], - [ 8, 11, 3]]) + >>> z = ivy.zeros((2,)) + >>> ivy.random_uniform(low=x, high=y, device='cpu', dtype='float64', out=z) + ivy.array([5. , 7.3]) """ - return ivy.current_backend().randint( - low, high, shape=shape, device=device, dtype=dtype, seed=seed, out=out + return ivy.current_backend().random_uniform( + low=low, high=high, shape=shape, device=device, dtype=dtype, out=out, seed=seed ) diff --git a/ivy/functional/ivy/searching.py b/ivy/functional/ivy/searching.py index 573c7fc1068ed..c5391a424a431 100644 --- a/ivy/functional/ivy/searching.py +++ b/ivy/functional/ivy/searching.py @@ -236,6 +236,81 @@ def argmin( ) +# Extra # +# ------# + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def argwhere( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return the indices of all non-zero elements of the input array. + + Parameters + ---------- + x + input array, for which indices are desired. + out + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. + + Returns + ------- + ret + Indices of non-zero elements. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([[1, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> res = ivy.argwhere(x) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + >>> x = ivy.array([[0, 2], [3, 4]]) + >>> y = ivy.zeros((3, 2), dtype=ivy.int64) + >>> res = ivy.argwhere(x, out=y) + >>> print(res) + ivy.array([[0, 1], [1, 0], [1, 1]]) + + With a :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0], [1]]), + b: ivy.array([[0], [1]]) + } + + >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) + >>> res = ivy.argwhere(x) + >>> print(res) + { + a: ivy.array([[0]]), + b: ivy.array([[0], [1]]) + } + """ + return current_backend(x).argwhere(x, out=out) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -454,78 +529,3 @@ def where( } """ return current_backend(x1).where(condition, x1, x2, out=out) - - -# Extra # -# ------# - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def argwhere( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return the indices of all non-zero elements of the input array. - - Parameters - ---------- - x - input array, for which indices are desired. - out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. - - Returns - ------- - ret - Indices of non-zero elements. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([[1, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 0], [0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> res = ivy.argwhere(x) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - >>> x = ivy.array([[0, 2], [3, 4]]) - >>> y = ivy.zeros((3, 2), dtype=ivy.int64) - >>> res = ivy.argwhere(x, out=y) - >>> print(res) - ivy.array([[0, 1], [1, 0], [1, 1]]) - - With a :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([1, 2]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0], [1]]), - b: ivy.array([[0], [1]]) - } - - >>> x = ivy.Container(a=ivy.array([1, 0]), b=ivy.array([3, 4])) - >>> res = ivy.argwhere(x) - >>> print(res) - { - a: ivy.array([[0]]), - b: ivy.array([[0], [1]]) - } - """ - return current_backend(x).argwhere(x, out=out) diff --git a/ivy/functional/ivy/set.py b/ivy/functional/ivy/set.py index 1a3e69831ede7..2c27798531508 100644 --- a/ivy/functional/ivy/set.py +++ b/ivy/functional/ivy/set.py @@ -148,6 +148,111 @@ def unique_all( return ivy.current_backend(x).unique_all(x, axis=axis, by_value=by_value) +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def unique_counts( + x: Union[ivy.Array, ivy.NativeArray], + /, +) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: + """ + Return the unique elements of an input array ``x`` and the corresponding counts for + each unique element in ``x``. + + .. admonition:: Data-dependent output shape + :class: important + + The shapes of two of the output arrays for this function depend on the data + values in the input array; hence, array libraries which build computation graphs + (e.g., JAX, Dask, etc.) may find this function difficult to implement without + knowing array values. Accordingly, such libraries may choose to omit this + function. See :ref:`data-dependent-output-shapes` section for more details. + + .. note:: + Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). + For input arrays having floating-point data types, value-based equality implies + the following behavior. + + - As ``nan`` values compare as ``False``, ``nan`` values should be considered + distinct. + - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be + considered distinct, and the corresponding unique element will be + implementation-dependent (e.g., an implementation could choose to return + ``-0`` if ``-0`` occurs before ``+0``). + + Parameters + ---------- + x + input array. If ``x`` has more than one dimension, the function must flatten + ``x`` and return the unique elements of the flattened array. + + Returns + ------- + ret + a namedtuple ``(values, counts)`` whose + + - first element must have the field name ``values`` and must be an + array containing the unique elements of ``x``. + The array must have the same data type as ``x``. + - second element must have the field name ``counts`` and must be an array + containing the number of times each unique element occurs in ``x``. + The returned array must have same shape as ``values`` and must + have the default array index data type. + + .. note:: + The order of unique elements is not specified and may vary between + implementations. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([1,2,1,3,4,1,3]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) + + >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) + + >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) + >>> y = ivy.unique_counts(x) + >>> print(y) + Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, + 2.29999995]), + counts=ivy.array([3, 1, 1, 1, 1])) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), + ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) + >>> y = ivy.unique_counts(x) + >>> print(y) + { + a: (list[2],shape=[4]), + b: (list[2],shape=[4]) + } + """ + return ivy.current_backend(x).unique_counts(x) + + @handle_exceptions @handle_backend_invalid @handle_nestable @@ -348,108 +453,3 @@ def unique_values( array([0., 1., 2., 3., 4., 5., nan, -0.]) """ return ivy.current_backend(x).unique_values(x, out=out) - - -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def unique_counts( - x: Union[ivy.Array, ivy.NativeArray], - /, -) -> Tuple[Union[ivy.Array, ivy.NativeArray], Union[ivy.Array, ivy.NativeArray]]: - """ - Return the unique elements of an input array ``x`` and the corresponding counts for - each unique element in ``x``. - - .. admonition:: Data-dependent output shape - :class: important - - The shapes of two of the output arrays for this function depend on the data - values in the input array; hence, array libraries which build computation graphs - (e.g., JAX, Dask, etc.) may find this function difficult to implement without - knowing array values. Accordingly, such libraries may choose to omit this - function. See :ref:`data-dependent-output-shapes` section for more details. - - .. note:: - Uniqueness should be determined based on value equality (i.e., ``x_i == x_j``). - For input arrays having floating-point data types, value-based equality implies - the following behavior. - - - As ``nan`` values compare as ``False``, ``nan`` values should be considered - distinct. - - As ``-0`` and ``+0`` compare as ``True``, signed zeros should not be - considered distinct, and the corresponding unique element will be - implementation-dependent (e.g., an implementation could choose to return - ``-0`` if ``-0`` occurs before ``+0``). - - Parameters - ---------- - x - input array. If ``x`` has more than one dimension, the function must flatten - ``x`` and return the unique elements of the flattened array. - - Returns - ------- - ret - a namedtuple ``(values, counts)`` whose - - - first element must have the field name ``values`` and must be an - array containing the unique elements of ``x``. - The array must have the same data type as ``x``. - - second element must have the field name ``counts`` and must be an array - containing the number of times each unique element occurs in ``x``. - The returned array must have same shape as ``values`` and must - have the default array index data type. - - .. note:: - The order of unique elements is not specified and may vary between - implementations. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1,2,1,3,4,1,3]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4]), counts=ivy.array([3, 1, 2, 1])) - - >>> x = ivy.asarray([[1,2,3,4],[2,3,4,5],[3,4,5,6]]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([1, 2, 3, 4, 5, 6]), counts=ivy.array([1, 2, 3, 3, 2, 1])) - - >>> x = ivy.array([0.2,0.3,0.4,0.2,1.4,2.3,0.2]) - >>> y = ivy.unique_counts(x) - >>> print(y) - Results(values=ivy.array([0.2 , 0.30000001, 0.40000001, 1.39999998, - 2.29999995]), - counts=ivy.array([3, 1, 1, 1, 1])) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([0., 1., 3. , 2. , 1. , 0.]), - ... b=ivy.array([1, 2, 1, 3, 4, 1, 3])) - >>> y = ivy.unique_counts(x) - >>> print(y) - { - a: (list[2],shape=[4]), - b: (list[2],shape=[4]) - } - """ - return ivy.current_backend(x).unique_counts(x) diff --git a/ivy/functional/ivy/sorting.py b/ivy/functional/ivy/sorting.py index a3423a308e41a..085db2f176276 100644 --- a/ivy/functional/ivy/sorting.py +++ b/ivy/functional/ivy/sorting.py @@ -140,113 +140,6 @@ def argsort( ) -@handle_exceptions -@handle_backend_invalid -@handle_nestable -@handle_array_like_without_promotion -@handle_out_argument -@to_native_arrays_and_back -@handle_array_function -@handle_device_shifting -def sort( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, - out: Optional[ivy.Array] = None, -) -> ivy.Array: - """ - Return a sorted copy of an array. - - Parameters - ---------- - x - input array - axis - axis along which to sort. If set to ``-1``, the function must sort along the - last axis. Default: ``-1``. - descending - direction The direction in which to sort the values - stable - sort stability. If ``True``, - the returned indices must maintain the relative order of ``x`` values which - compare as equal. If ``False``, the returned indices may or may not maintain the - relative order of ``x`` values which compare as equal (i.e., the relative order - of ``x`` values which compare as equal is implementation-dependent). - Default: ``True``. - out - optional output array, for writing the result to. It must have the same shape - as ``x``. - - Returns - ------- - ret - An array with the same dtype and shape as ``x``, with the elements sorted - along the given `axis`. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments - - - Examples - -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([7, 8, 6]) - >>> y = ivy.sort(x) - >>> print(y) - ivy.array([6, 7, 8]) - - >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) - >>> y = ivy.sort(x, axis=1, descending=True, stable=False) - >>> print(y) - ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) - - >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) - >>> y = ivy.zeros(5) - >>> ivy.sort(x, descending=True, stable=False, out=y) - >>> print(y) - ivy.array([3.2, 2.5, 1.5, 0.7]) - - >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) - >>> ivy.sort(x, out=x) - >>> print(x) - ivy.array([[ 1.1, 2.2, 3.3], - [-6.6, -5.5, -4.4]]) - - With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) - >>> y = ivy.sort(x, descending=True) - >>> print(y) - { - a: ivy.array([8, 6, 6]), - b: ivy.array([[9., 0.7], [0.4, 0.]]) - } - - >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) - >>> y = ivy.sort(x, descending=False, stable=False) - >>> print(y) - { - a: ivy.array([0.7, 1., 3.]), - b: ivy.array([[0.9, 4.], [0.2, 0.6]]) - } - """ - return ivy.current_backend(x).sort( - x, axis=axis, descending=descending, stable=stable, out=out - ) - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -369,3 +262,110 @@ def searchsorted( out=out, ret_dtype=ret_dtype, ) + + +@handle_exceptions +@handle_backend_invalid +@handle_nestable +@handle_array_like_without_promotion +@handle_out_argument +@to_native_arrays_and_back +@handle_array_function +@handle_device_shifting +def sort( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, + out: Optional[ivy.Array] = None, +) -> ivy.Array: + """ + Return a sorted copy of an array. + + Parameters + ---------- + x + input array + axis + axis along which to sort. If set to ``-1``, the function must sort along the + last axis. Default: ``-1``. + descending + direction The direction in which to sort the values + stable + sort stability. If ``True``, + the returned indices must maintain the relative order of ``x`` values which + compare as equal. If ``False``, the returned indices may or may not maintain the + relative order of ``x`` values which compare as equal (i.e., the relative order + of ``x`` values which compare as equal is implementation-dependent). + Default: ``True``. + out + optional output array, for writing the result to. It must have the same shape + as ``x``. + + Returns + ------- + ret + An array with the same dtype and shape as ``x``, with the elements sorted + along the given `axis`. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments + + + Examples + -------- + With :class:`ivy.Array` input: + + >>> x = ivy.array([7, 8, 6]) + >>> y = ivy.sort(x) + >>> print(y) + ivy.array([6, 7, 8]) + + >>> x = ivy.array([[[8.9,0], [19,5]],[[6,0.3], [19,0.5]]]) + >>> y = ivy.sort(x, axis=1, descending=True, stable=False) + >>> print(y) + ivy.array([[[19. , 5. ],[ 8.9, 0. ]],[[19. , 0.5],[ 6. , 0.3]]]) + + >>> x = ivy.array([1.5, 3.2, 0.7, 2.5]) + >>> y = ivy.zeros(5) + >>> ivy.sort(x, descending=True, stable=False, out=y) + >>> print(y) + ivy.array([3.2, 2.5, 1.5, 0.7]) + + >>> x = ivy.array([[1.1, 2.2, 3.3],[-4.4, -5.5, -6.6]]) + >>> ivy.sort(x, out=x) + >>> print(x) + ivy.array([[ 1.1, 2.2, 3.3], + [-6.6, -5.5, -4.4]]) + + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([8, 6, 6]),b=ivy.array([[9, 0.7], [0.4, 0]])) + >>> y = ivy.sort(x, descending=True) + >>> print(y) + { + a: ivy.array([8, 6, 6]), + b: ivy.array([[9., 0.7], [0.4, 0.]]) + } + + >>> x = ivy.Container(a=ivy.array([3, 0.7, 1]),b=ivy.array([[4, 0.9], [0.6, 0.2]])) + >>> y = ivy.sort(x, descending=False, stable=False) + >>> print(y) + { + a: ivy.array([0.7, 1., 3.]), + b: ivy.array([[0.9, 4.], [0.2, 0.6]]) + } + """ + return ivy.current_backend(x).sort( + x, axis=axis, descending=descending, stable=stable, out=out + ) diff --git a/ivy/functional/ivy/statistical.py b/ivy/functional/ivy/statistical.py index aaea746d500d0..81d7c3b8f17d9 100644 --- a/ivy/functional/ivy/statistical.py +++ b/ivy/functional/ivy/statistical.py @@ -16,8 +16,8 @@ from ivy.utils.exceptions import handle_exceptions -# Helpers # -# --------# +# --- Helpers --- # +# --------------- # def _get_promoted_type_of_operands(operands): @@ -31,10 +31,11 @@ def _get_promoted_type_of_operands(operands): return ivy.as_native_dtype(dtype) -# Array API Standard # -# -------------------# +# --- Main --- # +# ------------ # +@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -42,105 +43,138 @@ def _get_promoted_type_of_operands(operands): @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def min( +def cumprod( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the minimum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the minimum value is zero, the - minimum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the maximum possible value - for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, - return ``+infinity``). - - **Special Cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` - (i.e., ``NaN`` values propagate). + Return the cumulative product of the elements along a given axis. Parameters ---------- x - Input array. Should have a real-valued data type. + Input array. axis - axis or axes along which minimum values must be computed. By default, the - minimum value must be computed over the entire array. If a tuple of integers, - minimum values must be computed over multiple axes. Default: ``None``. - - keepdims - optional boolean, if ``True``, the reduced axes (dimensions) must be included - in the result as singleton dimensions, and, accordingly, the result must be - compatible with the input array (see :ref:`broadcasting`). Otherwise, - if ``False``, the reduced axes (dimensions) must not be included in the result. - Default: ``False``. + int , axis along which the cumulative product is computed. By default 0. + exclusive + optional bool, Whether to perform the cumprod exclusively. Defaults is False. + reverse + Whether to perform the cumprod from last to first element in the selected + axis. Default is ``False`` (from first to last element) out - optional output array, for writing the result to. + optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - if the minimum value was computed over the entire array, a zero-dimensional - array containing the minimum value; otherwise, a non-zero-dimensional array - containing the minimum values. The returned array must have the same data type - as ``x``. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Input array with cumulatively multiplied elements along axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.min(x) - >>> print(z) - ivy.array(1) - - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array([0, 0, 0]) - >>> y = ivy.min(x, out=z) - >>> print(z) - ivy.array(0) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x) + >>> print(y) + ivy.array([2, 6, 24]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x, axis=0, keepdims=True) + >>> x = ivy.array([2, 3, 4]) + >>> y = ivy.cumprod(x, exclusive=True) >>> print(y) - ivy.array([[0, 1, 2]]) + ivy.array([1, 2, 6]) - >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.min(x) + >>> x = ivy.array([[2, 3], + [5, 7], + [11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) >>> print(y) - ivy.array(0) + ivy.array([[ 1., 2.], + [ 1., 5.], + [ 1., 11.]]) + + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) + >>> print(x) + ivy.array([[1, 1], + [2, 3], + [10, 21]]) + + >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) + >>> y = ivy.zeros((3, 2)) + >>> x.cumprod(axis=0, exclusive=True, out=y) + >>> print(x) + ivy.array([[1., 1.], + [2., 3.], + [10., 21.]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) - >>> z = ivy.min(x) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x) + >>> print(y) { - a: ivy.array(1), - b: ivy.array(2) + a: ivy.array([2, 6, 24]), + b: ivy.array([3, 12, 60]) + } + + >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) + >>> y = ivy.cumprod(x, exclusive=True) + >>> print(y) + { + a: ivy.array([1, 2, 6]), + b: ivy.array([1, 3, 12]) + } + + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) + >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) + >>> print(y) + { + a: ivy.array([[1, 2], + [1, 5], + [1, 11]]), + b: ivy.array([[1, 3], + [1, 4], + [1, 5]]) + } + + >>> x = ivy.Container(a=ivy.array([[2, 3], + [5, 7], + [11, 13]]), + b=ivy.array([[3, 4], + [4, 5], + [5, 6]])) + >>> x.cumprod(axis=0, exclusive=True, out=x) + >>> print(x) + { + a: ivy.array([[1, 1], + [2, 3], + [10, 21]]), + b: ivy.array([[1, 1], + [3, 4], + [15, 42]]) } """ - return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).cumprod( + x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out + ) + + +# Extra # +# ------# @handle_exceptions @@ -151,107 +185,141 @@ def min( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def max( +def cumsum( x: Union[ivy.Array, ivy.NativeArray], - /, + axis: int = 0, + exclusive: bool = False, + reverse: bool = False, *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the maximum value of the input array ``x``. - - .. note:: - When the number of elements over which to compute the maximum value is zero, the - maximum value is implementation-defined. Specification-compliant libraries may - choose to raise an error, return a sentinel value (e.g., if ``x`` is a - floating-point input array, return ``NaN``), or return the minimum possible - value for the input array ``x`` data type (e.g., if ``x`` is a floating-point - array, return ``-infinity``). - - **Special Cases** - - For floating-point operands, - - - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values - propagate). + Return the cumulative sum of the elements along a given axis. Parameters ---------- x - input array. Should have a numeric data type. + Input array. axis - axis or axes along which maximum values must be computed. By default, the - maximum value must be computed over the entire array. If a tuple of integers, - maximum values must be computed over multiple axes. Default: ``None``. - keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + Axis along which the cumulative sum is computed. Default is ``0``. + exclusive + Whether to perform cumsum exclusively. Default is ``False``. + reverse + Whether to perform the cumsum from last to first element in the selected + axis. Default is ``False`` (from first to last element) + dtype + Data type of the returned array. Default is ``None``. + If None, if the default data type corresponding to the data type “kind” + (integer or floating-point) of x has a smaller range of values than the + data type of x (e.g., x has data type int64 and the default data type + is int32, or x has data type uint64 and the default data type is int64), + the returned array must have the same data type as x. + If x has a floating-point data type, the returned array must have the + default floating-point data type. + If x has a signed integer data type (e.g., int16), the returned array + must have the default integer data type. + If x has an unsigned integer data type (e.g., uint16), the returned + array must have an unsigned integer data type having the same number of + bits as the default integer data type (e.g., if the default integer data + type is int32, the returned array must have a uint32 data type). + If the data type (either specified or resolved) differs from the data type + of x, the input array should be cast to the specified data type before + computing the product. out - optional output array, for writing the result to. + Optional output array, for writing the result to. It must have a shape that the + inputs broadcast to. Returns ------- ret - if the maximum value was computed over the entire array, a zero-dimensional - array containing the maximum value; otherwise, a non-zero-dimensional array - containing the maximum values. The returned array must have the same data type - as ``x``. - - - This method conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ - in the standard. - - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Array which holds the result of applying cumsum at each + original array elements along the specified axis. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.max(x) - >>> print(z) - ivy.array(3) + >>> x = ivy.array([1, 5, 2, 0]) + >>> y = ivy.cumsum(x, exclusive= True, reverse=False) + >>> print(y) + ivy.array([0, 1, 6, 8]) - >>> x = ivy.array([0, 1, 2]) - >>> z = ivy.array(0) - >>> y = ivy.max(x, out=z) - >>> print(z) - ivy.array(2) + >>> x = ivy.array([[6, 4, 2], + ... [1, 3, 0]]) + >>> y = ivy.zeros((2,3)) + >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) + >>> print(y) + ivy.array([[7, 7, 2], + [1, 3, 0]]) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.max(x, axis=0, keepdims=True) + >>> x = ivy.array([[1, 5, 2], + ... [4, 3, 0]]) + >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) >>> print(y) - ivy.array([[4, 6, 10]]) + ivy.array([[4, 3, 0], + [0, 0, 0]]) + + >>> x = ivy.array([[2, 4, 5], + ... [3, 6, 5], + ... [1, 3, 10]]) + >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) + >>> print(x) + ivy.array([[11, 9, 5], + [14, 11, 5], + [14, 13, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.max(x) + >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), + ... b=ivy.array([[3, 5, 7]])) + >>> y = ivy.cumsum(x, axis= 0) >>> print(y) { - a: ivy.array(2.), - b: ivy.array(5.) + a: ivy.array([[1, 3, 5]]), + b: ivy.array([[3, 5, 7]]) } - >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), - ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) - >>> z = ivy.max(x, axis=1) - >>> print(z) + >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), + ... b=ivy.array([[3, 5, 8], + ... [5, 6, 5]]), + ... c=ivy.array([[2, 4, 1], + ... [3, 6, 9], + ... [0, 2, 3]])) + >>> y = ivy.Container(a = ivy.zeros((1, 3)), + ... b = ivy.zeros((2, 3)), + ... c = ivy.zeros((3,3))) + >>> ivy.cumsum(x,axis=1,reverse=True, out=y) + >>> print(y) { - a: ivy.array([3, 2]), - b: ivy.array([4, 2]) + a: ivy.array([[8, 7, 4]]), + b: ivy.array([[16, 13, 8], + [16, 11, 5]]), + c: ivy.array([[7, 5, 1], + [18, 15, 9], + [5, 5, 3]]) + } + + >>> x = ivy.Container(a=ivy.array([[0], + ... [5]]), + ... b=ivy.array([[6, 8, 7], + ... [4, 2, 3]]), + ... c=ivy.array([[1, 2], + ... [3, 4], + ... [6, 4]])) + >>> ivy.cumsum(x,axis=0,out=x) + >>> print(x) + { + a: ivy.array([[0], + [5]]), + b: ivy.array([[6, 8, 7], + [10, 10, 10]]), + c: ivy.array([[1, 2], + [4, 6], + [10, 10]]) } """ - return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) @handle_exceptions @@ -262,108 +330,126 @@ def max( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def mean( - x: Union[ivy.Array, ivy.NativeArray], - /, - *, - axis: Optional[Union[int, Sequence[int]]] = None, - keepdims: bool = False, +def einsum( + equation: str, + *operands: Union[ivy.Array, ivy.NativeArray], out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the arithmetic mean of the input array ``x``. - - **Special Cases** - - Let ``N`` equal the number of elements over which to compute the arithmetic mean. - - If ``N`` is ``0``, the arithmetic mean is ``NaN``. - - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values - propagate). + Sum the product of the elements of the input operands along dimensions specified + using a notation based on the Einstein summation convention. Parameters ---------- - x - input array. Should have a floating-point data type. - axis - axis or axes along which arithmetic means must be computed. By default, the mean - must be computed over the entire array. If a Sequence of integers, arithmetic - means must be computed over multiple axes. Default: ``None``. - keepdims - bool, if ``True``, the reduced axes (dimensions) must be included in the result - as singleton dimensions, and, accordingly, the result must be compatible with - the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced - axes (dimensions) must not be included in the result. Default: ``False``. + equation + A str describing the contraction, in the same format as numpy.einsum. + operands + seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes + should be consistent with equation. out optional output array, for writing the result to. Returns ------- ret - array, if the arithmetic mean was computed over the entire array, a - zero-dimensional array containing the arithmetic mean; otherwise, a - non-zero-dimensional array containing the arithmetic means. The returned - array must have the same data type as ``x``. - .. note:: - While this specification recommends that this function only accept input - arrays having a floating-point data type, specification-compliant array - libraries may choose to accept input arrays having an integer data type. - While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have the - default floating-point data type. - - - This function conforms to the `Array API Standard - `_. This docstring is an extension of the - `docstring `_ in the standard. + The array with sums computed. - Both the description and the type hints above assumes an array input for - simplicity, but this function is *nestable*, and therefore also accepts - :class:`ivy.Container` instances in place of any of the arguments. + Functional Examples + ------------------- - Examples - -------- With :class:`ivy.Array` input: - >>> x = ivy.array([3., 4., 5.]) - >>> y = ivy.mean(x) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> y = ivy.einsum('ii', x) >>> print(y) - ivy.array(4.) + ivy.array(12) - >>> x = ivy.array([0., 1., 2.]) - >>> y = ivy.array(0.) - >>> ivy.mean(x, out=y) - >>> print(y) - ivy.array(1.) + >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + >>> z = ivy.einsum('ij -> j', x) + >>> print(z) + ivy.array([ 9, 12, 15]) - >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) - >>> y = ivy.array([0., 0.]) - >>> ivy.mean(x, axis=1, out=y) - >>> print(y) - ivy.array([-1.4, 1.4]) + >>> A = ivy.array([0, 1, 2]) + >>> B = ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]) + >>> C = ivy.einsum('i,ij->i', A, B) + >>> print(C) + ivy.array([ 0, 22, 76]) + >>> A = ivy.array([[1, 1, 1], + ... [2, 2, 2], + ... [5, 5, 5]]) + >>> B = ivy.array([[0, 1, 0], + ... [1, 1, 0], + ... [1, 1, 1]]) + >>> C = ivy.einsum('ij,jk->ik', A, B) + >>> print(C) + ivy.array([[ 2, 3, 1], + [ 4, 6, 2], + [10, 15, 5]]) - With :class:`ivy.Container` input: + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i->', A) + >>> print(C) + ivy.array(45) - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.mean(x) - >>> print(y) + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->i', A, B) + >>> print(C) + ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' + >>> print(C) + ivy.array(510) + + >>> A = ivy.arange(10) + >>> B = ivy.arange(5, 15) + >>> C = ivy.einsum('i,j->ij', A, B) + >>> print(C) + ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], + [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], + [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], + [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], + [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], + [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], + [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], + [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], + [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) + + With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + + >>> x = ivy.array([0, 1, 2]) + >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], + ... [ 4, 5, 6, 7], + ... [ 8, 9, 10, 11]]), + ... b=ivy.array([[ 0, 1, 2], + ... [ 4, 5, 6], + ... [ 8, 9, 10]])) + >>> z = ivy.einsum('i,ij->i', x, y) + >>> print(z) { - a: ivy.array(0.), - b: ivy.array(0.90000004) + a: ivy.array([0, 22, 76]), + b: ivy.array([0, 15, 54]) } - >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), - ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) - >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) - >>> ivy.mean(x, axis=0, out=y) + With :class:`ivy.Container` input: + + >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), + ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) + >>> y = ivy.einsum('ii', x) >>> print(y) { - a: ivy.array([1.5, 2.5, 3.5]), - b: ivy.array([4.5, 5.5, 6.5]) + a: ivy.array(2), + b: ivy.array(15) } """ - return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) + return current_backend(operands[0]).einsum(equation, *operands, out=out) @handle_exceptions @@ -374,72 +460,61 @@ def mean( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def prod( +def max( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the product of input array x elements. + Calculate the maximum value of the input array ``x``. - **Special Cases** + .. note:: + When the number of elements over which to compute the maximum value is zero, the + maximum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the minimum possible + value for the input array ``x`` data type (e.g., if ``x`` is a floating-point + array, return ``-infinity``). - Let ``N`` equal the number of elements over which to compute the product. + **Special Cases** - - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). + For floating-point operands, - For both both real-valued and complex floating-point operands, special - cases must be handled as the operation is implemented by successive application - of :func:`ivy.multiply`: + - If ``x_i`` is ``NaN``, the maximum value is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- x input array. Should have a numeric data type. axis - axis or axes along which products must be computed. By default, the product must - be computed over the entire array. If a tuple of integers, products must be - computed over multiple axes. Default: ``None``. + axis or axes along which maximum values must be computed. By default, the + maximum value must be computed over the entire array. If a tuple of integers, + maximum values must be computed over multiple axes. Default: ``None``. keepdims - bool, if True, the reduced axes (dimensions) must be included in the result as + if ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the - input array (see Broadcasting). Otherwise, if False, the reduced axes + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. - dtype - data type of the returned array. If None, - if the default data type corresponding to the data type “kind” (integer or - floating-point) of x has a smaller range of values than the data type of x - (e.g., x has data type int64 and the default data type is int32, or x has data - type uint64 and the default data type is int64), the returned array must have - the same data type as x. if x has a floating-point data type, the returned array - must have the default floating-point data type. if x has a signed integer data - type (e.g., int16), the returned array must have the default integer data type. - if x has an unsigned integer data type (e.g., uint16), the returned array must - have an unsigned integer data type having the same number of bits as the default - integer data type (e.g., if the default integer data type is int32, the returned - array must have a uint32 data type). If the data type (either specified or - resolved) differs from the data type of x, the input array should be cast to the - specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - array, if the product was computed over the entire array, a zero-dimensional - array containing the product; otherwise, a non-zero-dimensional array containing - the products. The returned array must have a data type as described by the dtype - parameter above. + if the maximum value was computed over the entire array, a zero-dimensional + array containing the maximum value; otherwise, a non-zero-dimensional array + containing the maximum values. The returned array must have the same data type + as ``x``. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.max.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -451,55 +526,41 @@ def prod( With :class:`ivy.Array` input: >>> x = ivy.array([1, 2, 3]) - >>> z = ivy.prod(x) + >>> z = ivy.max(x) >>> print(z) - ivy.array(6) + ivy.array(3) - >>> x = ivy.array([1, 0, 3]) - >>> z = ivy.prod(x) + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array(0) + >>> y = ivy.max(x, out=z) >>> print(z) - ivy.array(0) - - >>> x = ivy.array([[3., 4., 5.]]) - >>> y = ivy.prod(x, keepdims=True) - >>> print(y) - ivy.array([60.]) - - >>> x = ivy.array([2., 1.]) - >>> y = ivy.array(0.) - >>> ivy.prod(x, out=y) - >>> print(y) - ivy.array(2.) + ivy.array(2) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.prod(x, axis=1) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.max(x, axis=0, keepdims=True) >>> print(y) - ivy.array([2., 9.]) + ivy.array([[4, 6, 10]]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = ivy.prod(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.max(x) >>> print(y) { - a: ivy.array(-0.), - b: ivy.array(0.30800003) + a: ivy.array(2.), + b: ivy.array(5.) } - >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), - ... b=ivy.array([[ 4., 5.], [5., 6.]])) - >>> y = ivy.prod(x, axis=1, keepdims=True) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([[1, 2, 3],[-1,0,2]]), + ... b=ivy.array([[2, 3, 4], [0, 1, 2]])) + >>> z = ivy.max(x, axis=1) + >>> print(z) { - a: ivy.array([[2.], - [12.]]), - b: ivy.array([[20.], - [30.]]) + a: ivy.array([3, 2]), + b: ivy.array([4, 2]) } """ - return current_backend(x).prod( - x, axis=axis, dtype=dtype, keepdims=keepdims, out=out - ) + return current_backend(x).max(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -510,138 +571,114 @@ def prod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def std( +def mean( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the standard deviation of the input array ``x``. + Calculate the arithmetic mean of the input array ``x``. **Special Cases** - Let ``N`` equal the number of elements over which to compute the standard deviation. - - - If ``N - correction`` is less than or equal to ``0``, - the standard deviation is ``NaN``. - - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` - (i.e., ``NaN`` values propagate). + Let ``N`` equal the number of elements over which to compute the arithmetic mean. + - If ``N`` is ``0``, the arithmetic mean is ``NaN``. + - If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values + propagate). Parameters ---------- x - input array. + input array. Should have a floating-point data type. axis - axis or axes along which standard deviations must be computed. By default, the - standard deviation must be computed over the entire array. If a tuple of - integers, standard deviations must be computed over multiple axes. - Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other - than ``0`` has the effect of adjusting the divisor during the calculation of the - standard deviation according to ``N-c`` where ``N`` corresponds to the total - number of elements over which the standard deviation is computed and ``c`` - corresponds to the provided degrees of freedom adjustment. When computing the - standard deviation of a population, setting this parameter to ``0`` is the - standard choice (i.e., the provided array contains data constituting an - entire population). When computing the corrected sample standard deviation, - setting this parameter to ``1`` is the standard choice (i.e., the provided array - contains data sampled from a larger population; this is commonly referred to as - Bessel's correction). - Default: ``0``. + axis or axes along which arithmetic means must be computed. By default, the mean + must be computed over the entire array. If a Sequence of integers, arithmetic + means must be computed over multiple axes. Default: ``None``. keepdims - if ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + bool, if ``True``, the reduced axes (dimensions) must be included in the result + as singleton dimensions, and, accordingly, the result must be compatible with + the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced + axes (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - if the standard deviation was computed over the entire array, a zero-dimensional - array containing the standard deviation; otherwise, a non-zero-dimensional array - containing the standard deviations. The returned array must have the same data - type as ``x``. - + array, if the arithmetic mean was computed over the entire array, a + zero-dimensional array containing the arithmetic mean; otherwise, a + non-zero-dimensional array containing the arithmetic means. The returned + array must have the same data type as ``x``. .. note:: While this specification recommends that this function only accept input - arrays having a real-valued floating-point data type, specification-compliant - array libraries may choose to accept input arrays having an integer data - type. While mixed data type promotion is implementation-defined, if the input - array ``x`` has an integer data type, the returned array must have - the default real-valued floating-point data type. + arrays having a floating-point data type, specification-compliant array + libraries may choose to accept input arrays having an integer data type. + While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have the + default floating-point data type. + This function conforms to the `Array API Standard `_. This docstring is an extension of the - `docstring `_ - in the standard. + `docstring `_ in the standard. - Both the description and the type hints above assumes an array input for simplicity, - but this function is *nestable*, and therefore also accepts :class:`ivy.Container` - instances in place of any of the arguments. + Both the description and the type hints above assumes an array input for + simplicity, but this function is *nestable*, and therefore also accepts + :class:`ivy.Container` instances in place of any of the arguments. Examples -------- - >>> x = ivy.array([-1., 0., 1.]) - >>> y = ivy.std(x) - >>> print(y) - ivy.array(0.81649661) - - >>> x = ivy.array([-1., 0., 1.]) - >>> z = ivy.std(x, correction=1) - >>> print(z) - ivy.array(1.) + With :class:`ivy.Array` input: - >>> x = ivy.array([[0., 4.]]) - >>> y = ivy.std(x, keepdims=True) + >>> x = ivy.array([3., 4., 5.]) + >>> y = ivy.mean(x) >>> print(y) - ivy.array([[2.]]) + ivy.array(4.) - >>> x = ivy.array([2., 1.]) + >>> x = ivy.array([0., 1., 2.]) >>> y = ivy.array(0.) - >>> ivy.std(x, out=y) + >>> ivy.mean(x, out=y) >>> print(y) - ivy.array(0.5) + ivy.array(1.) - >>> x = ivy.array([[-1., -2.], [3., 3.]]) - >>> y = ivy.std(x, axis=1) + >>> x = ivy.array([[-1., -2., -3., 0., -1.], [1., 2., 3., 0., 1.]]) + >>> y = ivy.array([0., 0.]) + >>> ivy.mean(x, axis=1, out=y) >>> print(y) - ivy.array([0.5, 0. ]) + ivy.array([-1.4, 1.4]) + With :class:`ivy.Container` input: >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) - >>> y = x.std() + >>> y = ivy.mean(x) >>> print(y) { - a: ivy.array(0.81649661), - b: ivy.array(0.509902) + a: ivy.array(0.), + b: ivy.array(0.90000004) } - >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), - ... b=ivy.array([[ 4., 2.], [2., 1.]])) - >>> y = x.std(axis=1, keepdims=True) + >>> x = ivy.Container(a=ivy.array([[0., 1., 2.], [3., 4., 5.]]), + ... b=ivy.array([[3., 4., 5.], [6., 7., 8.]])) + >>> y = ivy.Container(a = ivy.zeros(3), b = ivy.zeros(3)) + >>> ivy.mean(x, axis=0, out=y) >>> print(y) { - a: ivy.array([[1.], - [1.5]]), - b: ivy.array([[1.], - [0.5]]) + a: ivy.array([1.5, 2.5, 3.5]), + b: ivy.array([4.5, 5.5, 6.5]) } """ - return current_backend(x).std( - x, axis=axis, correction=correction, keepdims=keepdims, out=out - ) + return current_backend(x).mean(x, axis=axis, keepdims=keepdims, out=out) + + +# Array API Standard # +# -------------------# -@handle_exceptions @handle_backend_invalid @handle_nestable @handle_array_like_without_promotion @@ -649,81 +686,63 @@ def std( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def sum( +def min( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, - keepdims: Optional[bool] = False, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the sum of the input array x. + Calculate the minimum value of the input array ``x``. + + .. note:: + When the number of elements over which to compute the minimum value is zero, the + minimum value is implementation-defined. Specification-compliant libraries may + choose to raise an error, return a sentinel value (e.g., if ``x`` is a + floating-point input array, return ``NaN``), or return the maximum possible value + for the input array ``x`` data type (e.g., if ``x`` is a floating-point array, + return ``+infinity``). **Special Cases** - Let ``N`` equal the number of elements over which to compute the sum. - - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). - For floating-point operands, - - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). - For both real-valued and complex floating-point operands, special cases must - be handled as if the operation is implemented by successive application of - :func:`ivy.add`: + - If ``x_i`` is ``NaN``, the minimum value is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - Input array. Should have a numeric data type. + Input array. Should have a real-valued data type. axis - Axis or axes along which sums must be computed. By default, the sum must be - computed over the entire array. If a tuple of integers, sums must be computed - over multiple axes. Default: ``None``. - dtype - Data type of the returned array. If ``None``, - If the default data type corresponding to the data type "kind" (integer or - floating-point) of ``x`` has a smaller range of values than the data type of - ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is - ``int32``, or ``x`` has data type ``uint64`` and the default data type is - ``int64``), the returned array must have the same data type as ``x``. - If ``x`` has a floating-point data type, the returned array must have the - default floating-point data type. - If ``x`` has a signed integer data type (e.g., ``int16``), the returned - array must have the default integer data type. - If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned - array must have an unsigned integer data type having the same number of bits - as the default integer data type (e.g., if the default integer data type is - ``int32``, the returned array must have a ``uint32`` data type). - - If the data type (either specified or resolved) differs from the data type of - ``x``, the input array should be cast to the specified data type before - computing the sum. Default: ``None``. - - .. note:: - keyword argument is intended to help prevent data type overflows. + axis or axes along which minimum values must be computed. By default, the + minimum value must be computed over the entire array. If a tuple of integers, + minimum values must be computed over multiple axes. Default: ``None``. keepdims - If ``True``, the reduced axes (dimensions) must be included in the result as - singleton dimensions, and, accordingly, the result must be compatible with the - input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes - (dimensions) must not be included in the result. Default: ``False``. + optional boolean, if ``True``, the reduced axes (dimensions) must be included + in the result as singleton dimensions, and, accordingly, the result must be + compatible with the input array (see :ref:`broadcasting`). Otherwise, + if ``False``, the reduced axes (dimensions) must not be included in the result. + Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - If the sum was computed over the entire array, a zero-dimensional array - containing the sum; otherwise, an array containing the sums. The returned array - must have a data type as described by the ``dtype`` parameter above. + if the minimum value was computed over the entire array, a zero-dimensional + array containing the minimum value; otherwise, a non-zero-dimensional array + containing the minimum values. The returned array must have the same data type + as ``x``. This function conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.min.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -734,52 +753,38 @@ def sum( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.41, 0.89]) - >>> y = ivy.sum(x) - >>> print(y) - ivy.array(1.3) - - >>> x = ivy.array([0.5, 0.7, 2.4]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) - >>> print(y) - ivy.array(3.6) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.min(x) + >>> print(z) + ivy.array(1) - >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.sum(x, axis = 1, keepdims = False) - >>> print(y) - ivy.array([3, 20]) + >>> x = ivy.array([0, 1, 2]) + >>> z = ivy.array([0, 0, 0]) + >>> y = ivy.min(x, out=z) + >>> print(z) + ivy.array(0) >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) - >>> y = ivy.array([0,0,0]) - >>> ivy.sum(x, axis = 0, keepdims = False, out = y) - >>> print(y) - ivy.array([4, 7, 12]) - - With :class:`ivy.NativeArray` input: - - >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.sum(x) + >>> y = ivy.min(x, axis=0, keepdims=True) >>> print(y) - ivy.array(1.9) + ivy.array([[0, 1, 2]]) - >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) - >>> y = ivy.array(0.0) - >>> ivy.sum(x, out=y) + >>> x = ivy.native_array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.min(x) >>> print(y) - ivy.array(8.) + ivy.array(0) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) - >>> y = ivy.sum(x) - >>> print(y) + >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([2, 3, 4])) + >>> z = ivy.min(x) + >>> print(z) { - a: ivy.array(3.), - b: ivy.array(12.) + a: ivy.array(1), + b: ivy.array(2) } """ - return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) + return current_backend(x).min(x, axis=axis, keepdims=keepdims, out=out) @handle_exceptions @@ -790,65 +795,72 @@ def sum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def var( +def prod( x: Union[ivy.Array, ivy.NativeArray], /, *, axis: Optional[Union[int, Sequence[int]]] = None, - correction: Union[int, float] = 0.0, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Calculate the variance of the input array x. + Calculate the product of input array x elements. **Special Cases** - Let N equal the number of elements over which to compute the variance. + Let ``N`` equal the number of elements over which to compute the product. - If N - correction is less than or equal to 0, the variance is NaN. + - If ``N`` is ``0``, the product is ``1`` (i.e., the empty product). - If x_i is NaN, the variance is NaN (i.e., NaN values propagate). + For both both real-valued and complex floating-point operands, special + cases must be handled as the operation is implemented by successive application + of :func:`ivy.multiply`: Parameters ---------- x - input array. Should have a floating-point data type. + input array. Should have a numeric data type. axis - axis or axes along which variances must be computed. By default, the variance - must be computed over the entire array. If a tuple of integers, variances must - be computed over multiple axes. Default: ``None``. - correction - degrees of freedom adjustment. Setting this parameter to a value other than 0 - has the effect of adjusting the divisor during the calculation of the variance - according to N-c where N corresponds to the total number of elements over which - the variance is computed and c corresponds to the provided degrees of freedom - adjustment. When computing the variance of a population, setting this parameter - to 0 is the standard choice (i.e., the provided array contains data constituting - an entire population). When computing the unbiased sample variance, setting this - parameter to 1 is the standard choice (i.e., the provided array contains data - sampled from a larger population; this is commonly referred to as Bessel's - correction). Default: ``0``. + axis or axes along which products must be computed. By default, the product must + be computed over the entire array. If a tuple of integers, products must be + computed over multiple axes. Default: ``None``. keepdims - if True, the reduced axes (dimensions) must be included in the result as + bool, if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see Broadcasting). Otherwise, if False, the reduced axes (dimensions) must not be included in the result. Default: ``False``. + dtype + data type of the returned array. If None, + if the default data type corresponding to the data type “kind” (integer or + floating-point) of x has a smaller range of values than the data type of x + (e.g., x has data type int64 and the default data type is int32, or x has data + type uint64 and the default data type is int64), the returned array must have + the same data type as x. if x has a floating-point data type, the returned array + must have the default floating-point data type. if x has a signed integer data + type (e.g., int16), the returned array must have the default integer data type. + if x has an unsigned integer data type (e.g., uint16), the returned array must + have an unsigned integer data type having the same number of bits as the default + integer data type (e.g., if the default integer data type is int32, the returned + array must have a uint32 data type). If the data type (either specified or + resolved) differs from the data type of x, the input array should be cast to the + specified data type before computing the product. Default: ``None``. out optional output array, for writing the result to. Returns ------- ret - if the variance was computed over the entire array, a zero-dimensional array - containing the variance; otherwise, a non-zero-dimensional array containing the - variances. The returned array must have the same data type as x. + array, if the product was computed over the entire array, a zero-dimensional + array containing the product; otherwise, a non-zero-dimensional array containing + the products. The returned array must have a data type as described by the dtype + parameter above. This method conforms to the `Array API Standard `_. This docstring is an extension of the `docstring `_ + API_specification/generated/array_api.prod.html>`_ in the standard. Both the description and the type hints above assumes an array input for simplicity, @@ -859,46 +871,58 @@ def var( -------- With :class:`ivy.Array` input: - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.var(x) - >>> print(y) - ivy.array(0.07472222) + >>> x = ivy.array([1, 2, 3]) + >>> z = ivy.prod(x) + >>> print(z) + ivy.array(6) - >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) - >>> y = ivy.array(0.0) - >>> ivy.var(x, out=y) + >>> x = ivy.array([1, 0, 3]) + >>> z = ivy.prod(x) + >>> print(z) + ivy.array(0) + + >>> x = ivy.array([[3., 4., 5.]]) + >>> y = ivy.prod(x, keepdims=True) >>> print(y) - ivy.array(0.07472222) + ivy.array([60.]) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> print(ivy.var(x, axis=1, keepdims=True)) - ivy.array([[0.00666667], - [0.11555555]]) + >>> x = ivy.array([2., 1.]) + >>> y = ivy.array(0.) + >>> ivy.prod(x, out=y) + >>> print(y) + ivy.array(2.) - >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) - >>> y = ivy.var(x, correction=1) + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.prod(x, axis=1) >>> print(y) - ivy.array(0.08966666) + ivy.array([2., 9.]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), - ... b=ivy.array([0.7, 0.1, 0.9])) - >>> y = ivy.var(x) + + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = ivy.prod(x) >>> print(y) { - a: ivy.array(0.12666667), - b: ivy.array(0.11555555) + a: ivy.array(-0.), + b: ivy.array(0.30800003) + } + + >>> x = ivy.Container(a=ivy.array([[1., 2.], [3., 4.]]), + ... b=ivy.array([[ 4., 5.], [5., 6.]])) + >>> y = ivy.prod(x, axis=1, keepdims=True) + >>> print(y) + { + a: ivy.array([[2.], + [12.]]), + b: ivy.array([[20.], + [30.]]) } """ - return current_backend(x).var( - x, axis=axis, correction=correction, keepdims=keepdims, out=out + return current_backend(x).prod( + x, axis=axis, dtype=dtype, keepdims=keepdims, out=out ) -# Extra # -# ------# - - @handle_exceptions @handle_backend_invalid @handle_nestable @@ -907,141 +931,135 @@ def var( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumsum( +def std( x: Union[ivy.Array, ivy.NativeArray], - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, + /, *, - dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + axis: Optional[Union[int, Sequence[int]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative sum of the elements along a given axis. + Calculate the standard deviation of the input array ``x``. + + **Special Cases** + + Let ``N`` equal the number of elements over which to compute the standard deviation. + + - If ``N - correction`` is less than or equal to ``0``, + the standard deviation is ``NaN``. + - If ``x_i`` is ``NaN``, the standard deviation is ``NaN`` + (i.e., ``NaN`` values propagate). Parameters ---------- x - Input array. + input array. axis - Axis along which the cumulative sum is computed. Default is ``0``. - exclusive - Whether to perform cumsum exclusively. Default is ``False``. - reverse - Whether to perform the cumsum from last to first element in the selected - axis. Default is ``False`` (from first to last element) - dtype - Data type of the returned array. Default is ``None``. - If None, if the default data type corresponding to the data type “kind” - (integer or floating-point) of x has a smaller range of values than the - data type of x (e.g., x has data type int64 and the default data type - is int32, or x has data type uint64 and the default data type is int64), - the returned array must have the same data type as x. - If x has a floating-point data type, the returned array must have the - default floating-point data type. - If x has a signed integer data type (e.g., int16), the returned array - must have the default integer data type. - If x has an unsigned integer data type (e.g., uint16), the returned - array must have an unsigned integer data type having the same number of - bits as the default integer data type (e.g., if the default integer data - type is int32, the returned array must have a uint32 data type). - If the data type (either specified or resolved) differs from the data type - of x, the input array should be cast to the specified data type before - computing the product. + axis or axes along which standard deviations must be computed. By default, the + standard deviation must be computed over the entire array. If a tuple of + integers, standard deviations must be computed over multiple axes. + Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other + than ``0`` has the effect of adjusting the divisor during the calculation of the + standard deviation according to ``N-c`` where ``N`` corresponds to the total + number of elements over which the standard deviation is computed and ``c`` + corresponds to the provided degrees of freedom adjustment. When computing the + standard deviation of a population, setting this parameter to ``0`` is the + standard choice (i.e., the provided array contains data constituting an + entire population). When computing the corrected sample standard deviation, + setting this parameter to ``1`` is the standard choice (i.e., the provided array + contains data sampled from a larger population; this is commonly referred to as + Bessel's correction). + Default: ``0``. + keepdims + if ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out - Optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Array which holds the result of applying cumsum at each - original array elements along the specified axis. + if the standard deviation was computed over the entire array, a zero-dimensional + array containing the standard deviation; otherwise, a non-zero-dimensional array + containing the standard deviations. The returned array must have the same data + type as ``x``. + + .. note:: + While this specification recommends that this function only accept input + arrays having a real-valued floating-point data type, specification-compliant + array libraries may choose to accept input arrays having an integer data + type. While mixed data type promotion is implementation-defined, if the input + array ``x`` has an integer data type, the returned array must have + the default real-valued floating-point data type. + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- - With :class:`ivy.Array` input: - - >>> x = ivy.array([1, 5, 2, 0]) - >>> y = ivy.cumsum(x, exclusive= True, reverse=False) + >>> x = ivy.array([-1., 0., 1.]) + >>> y = ivy.std(x) >>> print(y) - ivy.array([0, 1, 6, 8]) + ivy.array(0.81649661) - >>> x = ivy.array([[6, 4, 2], - ... [1, 3, 0]]) - >>> y = ivy.zeros((2,3)) - >>> ivy.cumsum(x, axis=0, exclusive=False, reverse=True, out=y) + >>> x = ivy.array([-1., 0., 1.]) + >>> z = ivy.std(x, correction=1) + >>> print(z) + ivy.array(1.) + + >>> x = ivy.array([[0., 4.]]) + >>> y = ivy.std(x, keepdims=True) >>> print(y) - ivy.array([[7, 7, 2], - [1, 3, 0]]) + ivy.array([[2.]]) - >>> x = ivy.array([[1, 5, 2], - ... [4, 3, 0]]) - >>> y = ivy.cumsum(x, axis=0, exclusive=True, reverse=True) + >>> x = ivy.array([2., 1.]) + >>> y = ivy.array(0.) + >>> ivy.std(x, out=y) >>> print(y) - ivy.array([[4, 3, 0], - [0, 0, 0]]) + ivy.array(0.5) - >>> x = ivy.array([[2, 4, 5], - ... [3, 6, 5], - ... [1, 3, 10]]) - >>> ivy.cumsum(x,axis=1,reverse=True, dtype='int64', out=x) - >>> print(x) - ivy.array([[11, 9, 5], - [14, 11, 5], - [14, 13, 10]]) + >>> x = ivy.array([[-1., -2.], [3., 3.]]) + >>> y = ivy.std(x, axis=1) + >>> print(y) + ivy.array([0.5, 0. ]) With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[1, 3, 5]]), - ... b=ivy.array([[3, 5, 7]])) - >>> y = ivy.cumsum(x, axis= 0) + >>> x = ivy.Container(a=ivy.array([-1., 0., 1.]), b=ivy.array([1.1, 0.2, 1.4])) + >>> y = x.std() >>> print(y) { - a: ivy.array([[1, 3, 5]]), - b: ivy.array([[3, 5, 7]]) + a: ivy.array(0.81649661), + b: ivy.array(0.509902) } - >>> x = ivy.Container(a=ivy.array([[1, 3, 4]]), - ... b=ivy.array([[3, 5, 8], - ... [5, 6, 5]]), - ... c=ivy.array([[2, 4, 1], - ... [3, 6, 9], - ... [0, 2, 3]])) - >>> y = ivy.Container(a = ivy.zeros((1, 3)), - ... b = ivy.zeros((2, 3)), - ... c = ivy.zeros((3,3))) - >>> ivy.cumsum(x,axis=1,reverse=True, out=y) + >>> x = ivy.Container(a=ivy.array([[1., 3.], [3., 6.]]), + ... b=ivy.array([[ 4., 2.], [2., 1.]])) + >>> y = x.std(axis=1, keepdims=True) >>> print(y) { - a: ivy.array([[8, 7, 4]]), - b: ivy.array([[16, 13, 8], - [16, 11, 5]]), - c: ivy.array([[7, 5, 1], - [18, 15, 9], - [5, 5, 3]]) - } - - >>> x = ivy.Container(a=ivy.array([[0], - ... [5]]), - ... b=ivy.array([[6, 8, 7], - ... [4, 2, 3]]), - ... c=ivy.array([[1, 2], - ... [3, 4], - ... [6, 4]])) - >>> ivy.cumsum(x,axis=0,out=x) - >>> print(x) - { - a: ivy.array([[0], - [5]]), - b: ivy.array([[6, 8, 7], - [10, 10, 10]]), - c: ivy.array([[1, 2], - [4, 6], - [10, 10]]) + a: ivy.array([[1.], + [1.5]]), + b: ivy.array([[1.], + [0.5]]) } """ - return current_backend(x).cumsum(x, axis, exclusive, reverse, dtype=dtype, out=out) + return current_backend(x).std( + x, axis=axis, correction=correction, keepdims=keepdims, out=out + ) @handle_exceptions @@ -1052,134 +1070,137 @@ def cumsum( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def cumprod( +def sum( x: Union[ivy.Array, ivy.NativeArray], /, *, - axis: int = 0, - exclusive: bool = False, - reverse: bool = False, + axis: Optional[Union[int, Sequence[int]]] = None, dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + keepdims: Optional[bool] = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Return the cumulative product of the elements along a given axis. + Calculate the sum of the input array x. + + **Special Cases** + + Let ``N`` equal the number of elements over which to compute the sum. + - If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum). + + For floating-point operands, + - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate). + + For both real-valued and complex floating-point operands, special cases must + be handled as if the operation is implemented by successive application of + :func:`ivy.add`: Parameters ---------- x - Input array. + Input array. Should have a numeric data type. axis - int , axis along which the cumulative product is computed. By default 0. - exclusive - optional bool, Whether to perform the cumprod exclusively. Defaults is False. - reverse - Whether to perform the cumprod from last to first element in the selected - axis. Default is ``False`` (from first to last element) + Axis or axes along which sums must be computed. By default, the sum must be + computed over the entire array. If a tuple of integers, sums must be computed + over multiple axes. Default: ``None``. + dtype + Data type of the returned array. If ``None``, + If the default data type corresponding to the data type "kind" (integer or + floating-point) of ``x`` has a smaller range of values than the data type of + ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is + ``int32``, or ``x`` has data type ``uint64`` and the default data type is + ``int64``), the returned array must have the same data type as ``x``. + If ``x`` has a floating-point data type, the returned array must have the + default floating-point data type. + If ``x`` has a signed integer data type (e.g., ``int16``), the returned + array must have the default integer data type. + If ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned + array must have an unsigned integer data type having the same number of bits + as the default integer data type (e.g., if the default integer data type is + ``int32``, the returned array must have a ``uint32`` data type). + + If the data type (either specified or resolved) differs from the data type of + ``x``, the input array should be cast to the specified data type before + computing the sum. Default: ``None``. + + .. note:: + keyword argument is intended to help prevent data type overflows. + + keepdims + If ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out - optional output array, for writing the result to. It must have a shape that the - inputs broadcast to. + optional output array, for writing the result to. Returns ------- ret - Input array with cumulatively multiplied elements along axis. + If the sum was computed over the entire array, a zero-dimensional array + containing the sum; otherwise, an array containing the sums. The returned array + must have a data type as described by the ``dtype`` parameter above. + + + This function conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. + + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. Examples -------- With :class:`ivy.Array` input: - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x) + >>> x = ivy.array([0.41, 0.89]) + >>> y = ivy.sum(x) >>> print(y) - ivy.array([2, 6, 24]) + ivy.array(1.3) - >>> x = ivy.array([2, 3, 4]) - >>> y = ivy.cumprod(x, exclusive=True) + >>> x = ivy.array([0.5, 0.7, 2.4]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) >>> print(y) - ivy.array([1, 2, 6]) + ivy.array(3.6) - >>> x = ivy.array([[2, 3], - [5, 7], - [11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.sum(x, axis = 1, keepdims = False) >>> print(y) - ivy.array([[ 1., 2.], - [ 1., 5.], - [ 1., 11.]]) - - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> ivy.cumprod(x, axis=0, exclusive=True, out=x) - >>> print(x) - ivy.array([[1, 1], - [2, 3], - [10, 21]]) + ivy.array([3, 20]) - >>> x = ivy.array([[2, 3],[5, 7],[11, 13]]) - >>> y = ivy.zeros((3, 2)) - >>> x.cumprod(axis=0, exclusive=True, out=y) - >>> print(x) - ivy.array([[1., 1.], - [2., 3.], - [10., 21.]]) + >>> x = ivy.array([[0, 1, 2], [4, 6, 10]]) + >>> y = ivy.array([0,0,0]) + >>> ivy.sum(x, axis = 0, keepdims = False, out = y) + >>> print(y) + ivy.array([4, 7, 12]) - With :class:`ivy.Container` input: + With :class:`ivy.NativeArray` input: - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x) + >>> x = ivy.native_array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.sum(x) >>> print(y) - { - a: ivy.array([2, 6, 24]), - b: ivy.array([3, 12, 60]) - } + ivy.array(1.9) - >>> x = ivy.Container(a=ivy.array([2, 3, 4]), b=ivy.array([3, 4, 5])) - >>> y = ivy.cumprod(x, exclusive=True) + >>> x = ivy.native_array([1.0, 2.0, 2.0, 3.0]) + >>> y = ivy.array(0.0) + >>> ivy.sum(x, out=y) >>> print(y) - { - a: ivy.array([1, 2, 6]), - b: ivy.array([1, 3, 12]) - } + ivy.array(8.) - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> y = ivy.Container(a = ivy.zeros((3, 2)), b = ivy.zeros((3, 2))) - >>> ivy.cumprod(x, axis=1, exclusive=True, out=y) - >>> print(y) - { - a: ivy.array([[1, 2], - [1, 5], - [1, 11]]), - b: ivy.array([[1, 3], - [1, 4], - [1, 5]]) - } + With :class:`ivy.Container` input: - >>> x = ivy.Container(a=ivy.array([[2, 3], - [5, 7], - [11, 13]]), - b=ivy.array([[3, 4], - [4, 5], - [5, 6]])) - >>> x.cumprod(axis=0, exclusive=True, out=x) - >>> print(x) + >>> x = ivy.Container(a=ivy.array([0., 1., 2.]), b=ivy.array([3., 4., 5.])) + >>> y = ivy.sum(x) + >>> print(y) { - a: ivy.array([[1, 1], - [2, 3], - [10, 21]]), - b: ivy.array([[1, 1], - [3, 4], - [15, 42]]) + a: ivy.array(3.), + b: ivy.array(12.) } """ - return current_backend(x).cumprod( - x, axis=axis, exclusive=exclusive, reverse=reverse, dtype=dtype, out=out - ) + return current_backend(x).sum(x, axis=axis, dtype=dtype, keepdims=keepdims, out=out) @handle_exceptions @@ -1190,123 +1211,106 @@ def cumprod( @to_native_arrays_and_back @handle_array_function @handle_device_shifting -def einsum( - equation: str, - *operands: Union[ivy.Array, ivy.NativeArray], +def var( + x: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[int, Sequence[int]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, out: Optional[ivy.Array] = None, ) -> ivy.Array: """ - Sum the product of the elements of the input operands along dimensions specified - using a notation based on the Einstein summation convention. + Calculate the variance of the input array x. + + **Special Cases** + + Let N equal the number of elements over which to compute the variance. + + If N - correction is less than or equal to 0, the variance is NaN. + + If x_i is NaN, the variance is NaN (i.e., NaN values propagate). Parameters ---------- - equation - A str describing the contraction, in the same format as numpy.einsum. - operands - seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes - should be consistent with equation. + x + input array. Should have a floating-point data type. + axis + axis or axes along which variances must be computed. By default, the variance + must be computed over the entire array. If a tuple of integers, variances must + be computed over multiple axes. Default: ``None``. + correction + degrees of freedom adjustment. Setting this parameter to a value other than 0 + has the effect of adjusting the divisor during the calculation of the variance + according to N-c where N corresponds to the total number of elements over which + the variance is computed and c corresponds to the provided degrees of freedom + adjustment. When computing the variance of a population, setting this parameter + to 0 is the standard choice (i.e., the provided array contains data constituting + an entire population). When computing the unbiased sample variance, setting this + parameter to 1 is the standard choice (i.e., the provided array contains data + sampled from a larger population; this is commonly referred to as Bessel's + correction). Default: ``0``. + keepdims + if True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see Broadcasting). Otherwise, if False, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. out optional output array, for writing the result to. Returns ------- ret - The array with sums computed. - - Functional Examples - ------------------- - - With :class:`ivy.Array` input: + if the variance was computed over the entire array, a zero-dimensional array + containing the variance; otherwise, a non-zero-dimensional array containing the + variances. The returned array must have the same data type as x. - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> y = ivy.einsum('ii', x) - >>> print(y) - ivy.array(12) - >>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> z = ivy.einsum('ij -> j', x) - >>> print(z) - ivy.array([ 9, 12, 15]) - - >>> A = ivy.array([0, 1, 2]) - >>> B = ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]) - >>> C = ivy.einsum('i,ij->i', A, B) - >>> print(C) - ivy.array([ 0, 22, 76]) - - >>> A = ivy.array([[1, 1, 1], - ... [2, 2, 2], - ... [5, 5, 5]]) - >>> B = ivy.array([[0, 1, 0], - ... [1, 1, 0], - ... [1, 1, 1]]) - >>> C = ivy.einsum('ij,jk->ik', A, B) - >>> print(C) - ivy.array([[ 2, 3, 1], - [ 4, 6, 2], - [10, 15, 5]]) + This method conforms to the `Array API Standard + `_. This docstring is an extension of the + `docstring `_ + in the standard. - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i->', A) - >>> print(C) - ivy.array(45) + Both the description and the type hints above assumes an array input for simplicity, + but this function is *nestable*, and therefore also accepts :class:`ivy.Container` + instances in place of any of the arguments. - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->i', A, B) - >>> print(C) - ivy.array([ 0, 6, 14, 24, 36, 50, 66, 84, 104, 126]) + Examples + -------- + With :class:`ivy.Array` input: - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i' - >>> print(C) - ivy.array(510) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.var(x) + >>> print(y) + ivy.array(0.07472222) - >>> A = ivy.arange(10) - >>> B = ivy.arange(5, 15) - >>> C = ivy.einsum('i,j->ij', A, B) - >>> print(C) - ivy.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [ 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], - [ 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], - [ 15, 18, 21, 24, 27, 30, 33, 36, 39, 42], - [ 20, 24, 28, 32, 36, 40, 44, 48, 52, 56], - [ 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], - [ 30, 36, 42, 48, 54, 60, 66, 72, 78, 84], - [ 35, 42, 49, 56, 63, 70, 77, 84, 91, 98], - [ 40, 48, 56, 64, 72, 80, 88, 96, 104, 112], - [ 45, 54, 63, 72, 81, 90, 99, 108, 117, 126]]) + >>> x = ivy.array([0.1, 0.2, 0.3, 0.3, 0.9, 0.10]) + >>> y = ivy.array(0.0) + >>> ivy.var(x, out=y) + >>> print(y) + ivy.array(0.07472222) - With a mix of :class:`ivy.Array` and :class:`ivy.Container` inputs: + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> print(ivy.var(x, axis=1, keepdims=True)) + ivy.array([[0.00666667], + [0.11555555]]) - >>> x = ivy.array([0, 1, 2]) - >>> y = ivy.Container(a=ivy.array([[ 0, 1, 2, 3], - ... [ 4, 5, 6, 7], - ... [ 8, 9, 10, 11]]), - ... b=ivy.array([[ 0, 1, 2], - ... [ 4, 5, 6], - ... [ 8, 9, 10]])) - >>> z = ivy.einsum('i,ij->i', x, y) - >>> print(z) - { - a: ivy.array([0, 22, 76]), - b: ivy.array([0, 15, 54]) - } + >>> x = ivy.array([[0.1, 0.2, 0.3], [0.3, 0.9, 0.10]]) + >>> y = ivy.var(x, correction=1) + >>> print(y) + ivy.array(0.08966666) With :class:`ivy.Container` input: - - >>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]), - ... b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]])) - >>> y = ivy.einsum('ii', x) + >>> x = ivy.Container(a=ivy.array([0.1, 0.2, 0.9]), + ... b=ivy.array([0.7, 0.1, 0.9])) + >>> y = ivy.var(x) >>> print(y) { - a: ivy.array(2), - b: ivy.array(15) + a: ivy.array(0.12666667), + b: ivy.array(0.11555555) } """ - return current_backend(operands[0]).einsum(equation, *operands, out=out) + return current_backend(x).var( + x, axis=axis, correction=correction, keepdims=keepdims, out=out + ) diff --git a/ivy/functional/ivy/utility.py b/ivy/functional/ivy/utility.py index 09bd432b60291..cefd9d446c711 100644 --- a/ivy/functional/ivy/utility.py +++ b/ivy/functional/ivy/utility.py @@ -237,6 +237,19 @@ def any( return ivy.current_backend(x).any(x, axis=axis, keepdims=keepdims, out=out) +@staticmethod +def load(filepath, format=None, type="module"): + if type == "module": + return ivy.Module.load(filepath) + elif type == "container": + if format is not None: + return ivy.Container.cont_load(filepath, format=format) + else: + return ivy.Container.cont_load(filepath) + else: + raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.") + + # Extra # # ----- # @@ -251,16 +264,3 @@ def save(item, filepath, format=None): item.save(filepath) else: raise ivy.utils.exceptions.IvyException("Unsupported item type for saving.") - - -@staticmethod -def load(filepath, format=None, type="module"): - if type == "module": - return ivy.Module.load(filepath) - elif type == "container": - if format is not None: - return ivy.Container.cont_load(filepath, format=format) - else: - return ivy.Container.cont_load(filepath) - else: - raise ivy.utils.exceptions.IvyException("Unsupported item type for loading.")