Repository: ivy-llc/ivy

Issue:
weighted_moments



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/tensorflow/nn.py b/ivy/functional/frontends/tensorflow/nn.py
--- a/ivy/functional/frontends/tensorflow/nn.py
+++ b/ivy/functional/frontends/tensorflow/nn.py
@@ -2,7 +2,7 @@
 import ivy
 from ivy.functional.frontends.tensorflow.func_wrapper import to_ivy_arrays_and_back
 from ivy.func_wrapper import with_unsupported_dtypes
-
+from ivy.functional.frontends.tensorflow import check_tensorflow_casting
 
 def _reduce_strides_dilations(dim, stride, dilations):
     if not isinstance(stride, int):
@@ -487,4 +487,38 @@ def pool(
 # log_poisson_loss
 @to_ivy_arrays_and_back
 def log_poisson_loss(targets, log_input, compute_full_loss=False, name=None):
-    return ivy.log_poisson_loss(targets, log_input, compute_full_loss, name)
+    return ivy.log_poisson_loss(targets, log_input, compute_full_loss, name)  
+ 
+
+#weighted_moments
+@to_ivy_arrays_and_back
+def weighted_moments(x, axes, frequency_weights, keepdims=False, name=None):
+    fw_x_prod=frequency_weights * x
+    fw_x_prod=ivy.array(fw_x_prod)
+    weighted_input_sum = ivy.sum(fw_x_prod, axis=axes, keepdims=True).astype(fw_x_prod.dtype)
+
+    broadcasted_weights = frequency_weights + ivy.zeros_like(x)
+    broadcasted_weights=ivy.array(broadcasted_weights)
+    sum_of_weights = ivy.sum(broadcasted_weights, axis=axes, keepdims=True).astype(broadcasted_weights.dtype)
+
+    divisor = ivy.reciprocal(sum_of_weights)
+
+    weighted_input_sum, divisor = check_tensorflow_casting(weighted_input_sum, divisor)
+    weighted_mean = ivy.multiply(weighted_input_sum, divisor)
+
+    x, weighted_mean = check_tensorflow_casting(x, weighted_mean)
+    squared_difference = ivy.square(ivy.subtract(x, weighted_mean))
+    if isinstance(squared_difference, complex):
+        squared_difference = squared_difference.real - squared_difference.imag*1j
+
+    fw_sq_diff_prod=frequency_weights * squared_difference
+    fw_sq_diff_prod=ivy.array(fw_sq_diff_prod)
+    weighted_distsq = ivy.sum(fw_sq_diff_prod, axis=axes, keepdims=True).astype(fw_sq_diff_prod.dtype)
+
+    weighted_distsq, divisor=check_tensorflow_casting(weighted_distsq, divisor)
+    weighted_variance = ivy.multiply(weighted_distsq, divisor)
+
+    if not keepdims:
+        weighted_mean = ivy.squeeze(weighted_mean, axis=axes)
+        weighted_variance = ivy.squeeze(weighted_variance, axis=axes)
+    return weighted_mean, weighted_variance
\ No newline at end of file


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_nn.py b/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_nn.py
--- a/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_nn.py
+++ b/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_nn.py
@@ -1603,3 +1603,42 @@ def test_log_poisson_loss(
         compute_full_loss=compute_full_loss,
         atol=1e-2,
     )
+    
+    
+# weighted moments
+@handle_frontend_test(
+    fn_tree="tensorflow.nn.weighted_moments",
+    dtype_and_x_and_axis=_statistical_dtype_values(function="mean"),
+    dtype_and_fw=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        num_arrays=1,
+        min_value=0.00001,
+    ),
+    keepdims=st.booleans(),
+    test_with_out=st.just(False),
+)
+def test_tensorflow_weighted_moments(
+    *,
+    dtype_and_x_and_axis,
+    dtype_and_fw,
+    keepdims,
+    frontend,
+    test_flags,
+    fn_tree,
+    on_device,
+):
+    input_dtype, x, axis = dtype_and_x_and_axis
+    fw_dtype, fw = dtype_and_fw
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        rtol=1e-1,
+        atol=1e-1,
+        x=x[0],
+        axes=axis,
+        frequency_weights=fw[0],
+        keepdims=keepdims,
+    )    
\ No newline at end of file