Fix Ivy Failing Test: jax - norms.batch_norm

ivy/data_classes/array/layers.py

@@ -529,22 +529,22 @@ def conv1d_transpose(
         bias: Optional[ivy.Array] = None,
         out: Optional[ivy.Array] = None,
     ) -> ivy.Array:
-        """ivy.Array instance method variant of ivy.conv1d_transpose. This
-        method simply wraps the function, and so the docstring for
-        ivy.conv1d_transpose also applies to this method with minimal changes.
+        """ivy.Container static method variant of ivy.conv1d_transpose. This method
+        simply wraps the function, and so the docstring for ivy.conv1d_transpose also
+        applies to this method with minimal changes.
 
         Parameters
         ----------
-        self
+        x
             Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*.
         filters
             Convolution filters *[fw,d_out,d_in]*.
         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 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.
         output_shape
             Shape of the output (Default value = None)
         filter_format
@@ -556,11 +556,22 @@ def conv1d_transpose(
             corresponds to input with shape (batch_size, channels, width).
         dilations
             The dilation factor for each dimension of input. (Default value = 1)
+        key_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 key_chains, otherwise key_chains
+            will be skipped. Default is ``True``.
+        prune_unapplied
+            Whether to prune key_chains for which the function was not applied.
+            Default is ``False``.
+        map_sequences
+            Whether to also map method to sequences (lists, tuples).
+            Default is ``False``.
         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.
+            optional output container, for writing the result to. It must have a shape that the
+            inputs broadcast to.
 
         Returns
         -------
@@ -569,14 +580,15 @@ def conv1d_transpose(
 
         Examples
         --------
-        >>> x = ivy.array([[[1., 2.], [3., 4.], [6., 7.], [9., 11.]]])  # NWC
-        >>> filters = ivy.array([[[0., 1.], [1., 1.]]])  # WIO (I == C)
-        >>> result = x.conv1d_transpose(filters, (1,), 'VALID')
-        >>> print(result)
-        ivy.array([[[ 2.,  3.],
-        ...         [ 4.,  7.],
-        ...         [ 7., 13.],
-        ...         [11., 20.]]])
+        >>> x = ivy.Container(a=ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]),
+        ...                   b=ivy.random_normal(mean=0, std=1, shape=[1, 56, 3]))
+        >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 6, 3])
+        >>> y = ivy.Container.static_conv1d_transpose(x, filters, [2], 'SAME')
+        >>> print(y.shape)
+        {
+            a: ivy.Shape(1, 56, 6),
+            b: ivy.Shape(1, 112, 6)
+        }
         """
         return ivy.conv1d_transpose(
             self._data,

ivy/data_classes/container/layers.py

@@ -1547,7 +1547,7 @@ def conv2d(
     def _static_conv1d_transpose(
         x: Union[ivy.Array, ivy.NativeArray, ivy.Container],
         filters: Union[ivy.Array, ivy.NativeArray, ivy.Container],
-        strides: Union[int, Tuple[int], ivy.Container],
+        strides: Union[int, Tuple[int]],
         padding: Union[str, ivy.Container],
         /,
         *,
@@ -1562,9 +1562,9 @@ def _static_conv1d_transpose(
         bias: Optional[ivy.Container] = None,
         out: Optional[ivy.Container] = None,
     ) -> ivy.Container:
-        """ivy.Container static method variant of ivy.conv1d_transpose. This
-        method simply wraps the function, and so the docstring for
-        ivy.conv1d_transpose also applies to this method with minimal changes.
+        """ivy.Container static method variant of ivy.conv1d_transpose. This method
+        simply wraps the function, and so the docstring for ivy.conv1d_transpose also
+        applies to this method with minimal changes.
 
         Parameters
         ----------
@@ -1575,9 +1575,9 @@ def _static_conv1d_transpose(
         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 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.
         output_shape
             Shape of the output (Default value = None)
         filter_format
@@ -1603,8 +1603,8 @@ def _static_conv1d_transpose(
         bias
             Bias array of shape *[d_out]*.
         out
-            optional output container, for writing the result to. It must have a shape
-            that the inputs broadcast to.
+            optional output container, for writing the result to. It must have a shape that the
+            inputs broadcast to.
 
         Returns
         -------
@@ -1616,11 +1616,11 @@ def _static_conv1d_transpose(
         >>> x = ivy.Container(a=ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]),
         ...                   b=ivy.random_normal(mean=0, std=1, shape=[1, 56, 3]))
         >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 6, 3])
-        >>> y = ivy.Container.static_conv1d_transpose(x, filters, 2, 'SAME')
+        >>> y = ivy.Container.static_conv1d_transpose(x, filters, [2], 'SAME')
         >>> print(y.shape)
         {
-            a: [1,56,6],
-            b: [1,112,6]
+            a: ivy.Shape(1, 56, 6),
+            b: ivy.Shape(1, 112, 6)
         }
         """
         return ContainerBase.cont_multi_map_in_function(
@@ -1672,9 +1672,9 @@ def conv1d_transpose(
         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 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.
         output_shape
             Shape of the output (Default value = None)
         filter_format
@@ -1700,8 +1700,8 @@ def conv1d_transpose(
         bias
             Bias array of shape *[d_out]*.
         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. It must have a shape that the
+            inputs broadcast to.
 
         Returns
         -------
@@ -1713,7 +1713,7 @@ def conv1d_transpose(
         >>> x = ivy.Container(a=ivy.random_normal(mean=0, std=1, shape=[1, 28, 3]),
         ...                   b=ivy.random_normal(mean=0, std=1, shape=[1, 56, 3]))
         >>> filters = ivy.random_normal(mean=0, std=1, shape=[3, 6, 3])
-        >>> y = x.conv1d_transpose(filters, 2, 'SAME')
+        >>> y = x.conv1d_transpose(filters, [2], 'SAME')
         >>> print(y.shape)
         {
             a: ivy.Shape(1, 56, 6),

ivy/functional/backends/numpy/layers.py

@@ -181,6 +181,41 @@ def conv1d_transpose(
     bias: Optional[np.ndarray] = None,
     out: Optional[np.ndarray] = None,
 ) -> np.ndarray:
+    """Compute a 1-D transpose convolution given 3-D input x and filters arrays.
+
+    Parameters
+    ----------
+    x
+        Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*.
+    filters
+        Convolution filters *[fw,d_out,d_in]*.
+    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)
+    filter_format
+        Either "channel_first" or "channel_last". "channel_first" corresponds
+        to "IOW",input data formats, while "channel_last" corresponds to "WOI".
+    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.
+
+    Returns
+    -------
+    ret
+        The result of the transpose convolution operation.
+    """
     if data_format == "NCW":
         x = np.transpose(x, (0, 2, 1))
     if filter_format == "channel_last":

ivy/functional/backends/tensorflow/experimental/norms.py

@@ -94,12 +94,12 @@ def batch_norm(
     xdims = len(x.shape)
     if data_format == "NCS":
         x = tf.transpose(x, perm=(0, *range(2, xdims), 1))
-
+    x_dtype = x.dtype
     runningmean = mean
     runningvariance = variance
     if training:
         n = tf.size(x) if xdims == 1 else tf.divide(tf.size(x), tf.shape(x)[-1])
-        n = tf.cast(n, x.dtype) if n.dtype != x.dtype else n
+        n = tf.cast(n, x_dtype) if n.dtype != x_dtype else n
         dims = (0, *range(1, xdims - 1))
         mean = tf.math.reduce_mean(x, axis=dims)
         variance = tf.math.reduce_variance(x, axis=dims)
@@ -114,9 +114,18 @@ def batch_norm(
             else runningvariance
         )
 
-    inv = 1.0 / tf.math.sqrt(variance + eps)
-    offset = 0 if offset is None else offset
+    one = tf.constant(1.0, dtype=x_dtype)
+    eps_tensor = tf.constant(eps, dtype=x_dtype)
+    mean = tf.cast(mean, x_dtype)
+    variance = tf.cast(variance, x_dtype)
+
+    inv = one / tf.math.sqrt(variance + eps_tensor)
+    if offset is None:
+        offset = tf.constant(0, dtype=x_dtype)
+    else:
+        offset = tf.cast(offset, x_dtype)
     if scale is not None:
+        scale = tf.cast(scale, x_dtype)
         inv = tf.math.multiply(inv, scale)
     xnormalized = tf.math.add(tf.math.multiply(x, inv), offset)
     xnormalized = tf.math.subtract(xnormalized, tf.math.multiply(mean, inv))

ivy/functional/backends/tensorflow/layers.py

@@ -160,44 +160,86 @@ def conv1d(
 
 @with_unsupported_dtypes({"2.15.0 and below": ("bfloat16", "complex")}, backend_version)
 def conv1d_transpose(
-    x: Union[tf.Tensor, tf.Variable],
-    filters: Union[tf.Tensor, tf.Variable],
+    x: Union[ivy.Array, ivy.NativeArray],
+    filters: Union[ivy.Array, ivy.NativeArray],
     strides: Union[int, Tuple[int]],
     padding: str,
     /,
     *,
-    output_shape: Optional[Union[ivy.NativeShape, Sequence[int]]] = None,
+    output_shape: Optional[Union[ivy.Shape, ivy.NativeShape]] = None,
     filter_format: str = "channel_last",
     data_format: str = "NWC",
     dilations: Union[int, Tuple[int]] = 1,
-    bias: Optional[Union[tf.Tensor, tf.Variable]] = None,
-    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
-):
+    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.
+
+    Parameters
+    ----------
+    x
+        Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*.
+    filters
+        Convolution filters *[fw,d_out,d_in]*.
+    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)
+    filter_format
+        Either "channel_first" or "channel_last". "channel_first" corresponds
+        to "IOW",input data formats, while "channel_last" corresponds to "WOI".
+    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.
+
+    Returns
+    -------
+    ret
+        The result of the transpose convolution operation.
+    """
     if ivy.dev(x) == "cpu" and (
         (dilations > 1) if isinstance(dilations, int) else any(d > 1 for d in dilations)
     ):
         raise ivy.utils.exceptions.IvyException(
             "Tensorflow does not support dilations greater than 1 when device is cpu"
         )
+
     permuted_x = False
     if data_format == "NCW" and ivy.dev(x) == "cpu":
         x = tf.transpose(x, (0, 2, 1))
         data_format = "NWC"
         permuted_x = True
+
     if filter_format == "channel_first":
         filters = tf.transpose(filters, (2, 1, 0))
+
     output_shape, padding = _transpose_out_pad(
         x.shape, filters.shape, strides, padding, 1, dilations, data_format
     )
+
     res = tf.nn.conv1d_transpose(
         x, filters, output_shape, strides, padding, data_format, dilations
     )
+
     if bias is not None:
         if data_format[1] == "C":
             bias = tf.reshape(bias, [1, -1, 1])
         res = tf.math.add(res, bias)
+
     if permuted_x:
         res = tf.transpose(res, (0, 2, 1))
+
     return res
 
 

ivy/functional/backends/torch/layers.py

@@ -333,7 +333,6 @@ def conv1d_v_1p9p0_and_above(
     },
     backend_version,
 )
-# noinspection PyUnresolvedReferences
 def conv1d_transpose(
     x: torch.Tensor,
     filters: torch.Tensor,
@@ -347,7 +346,42 @@ def conv1d_transpose(
     dilations: Union[int, Tuple[int]] = 1,
     bias: Optional[torch.Tensor] = None,
     out: Optional[torch.Tensor] = None,
-):
+) -> torch.Tensor:
+    """Compute a 1-D transpose convolution given 3-D input x and filters arrays.
+
+    Parameters
+    ----------
+    x
+        Input image *[batch_size,w,d_in]* or *[batch_size,d_in,w]*.
+    filters
+        Convolution filters *[fw,d_out,d_in]*.
+    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)
+    filter_format
+        Either "channel_first" or "channel_last". "channel_first" corresponds
+        to "IOW",input data formats, while "channel_last" corresponds to "WOI".
+    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.
+
+    Returns
+    -------
+    ret
+        The result of the transpose convolution operation.
+    """
     if data_format == "NWC":
         x = x.permute(0, 2, 1)
     if filter_format == "channel_last":