Repository: ivy-llc/ivy

Issue:
 sufficient_statistics



Discussion:
This issue has been labelled as stale. If you would like to continue working on this issue, then please add another comment or create a PR . If a PR has already been created which refers to this issue, then you should explicitly mention this issue in the relevant PR. For more information please refer to this page in our documentation: https://lets-unify.ai/ivy/contributing/1_the_basics.html#todo-list-issues

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
@@ -484,6 +484,41 @@ def pool(
     )
 
 
+
+# sufficient_statistics
+@to_ivy_arrays_and_back
+def sufficient_statistics(x, axes, shift=None, keepdims=False, name=None):
+    count = 1
+    shape = ivy.shape(x)
+    axes = list(set(axes))
+    for a in axes:
+        if ivy.to_scalar(a) < 0:
+            index = x.ndim + ivy.to_scalar(a)
+        else:
+            index = ivy.to_scalar(a)
+        count *= shape[index]
+    count = ivy.array(count, dtype=ivy.dtype(x))
+    if shift is None:
+        sum_of_elements = ivy.sum(x, axis=axes, keepdims=keepdims)
+        sum_of_squares = ivy.sum(ivy.square(x), axis=axes, keepdims=keepdims)
+    else:
+        sum_of_elements = ivy.sum(
+            (ivy.subtract(x, shift)), axis=axes, keepdims=keepdims
+        )
+        sum_of_squares = ivy.sum(
+            (ivy.square(ivy.subtract(x, shift))), axis=axes, keepdims=keepdims
+        )
+        if shift.ndim == 0:
+            ivy.reshape(shift, ())
+
+    if count.ndim == 0:
+        ivy.reshape(count, ())
+    if sum_of_elements.ndim == 0:
+        ivy.reshape(sum_of_elements, ())
+    if sum_of_squares.ndim == 0:
+        ivy.reshape(sum_of_squares, ())
+    return count, sum_of_elements, sum_of_squares, shift
+
 # log_poisson_loss
 @to_ivy_arrays_and_back
 def log_poisson_loss(targets, log_input, compute_full_loss=False, name=None):


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
@@ -1565,6 +1565,59 @@ def test_tensorflow_pool(
     )
 
 
+# sufficient_statistics
+@st.composite
+def _axes_value(draw):
+    s = draw(
+        helpers.get_shape(
+            min_num_dims=1,
+            max_num_dims=5,
+            min_dim_size=1,
+            max_dim_size=5,
+        )
+    )
+    dtype_and_x = draw(
+        helpers.dtype_values_axis(
+            available_dtypes=helpers.get_dtypes("float"),
+            shape=s,
+            valid_axis=True,
+            force_tuple_axis=True,
+        )
+    )
+    return dtype_and_x
+
+
+@handle_frontend_test(
+    fn_tree="tensorflow.nn.sufficient_statistics",
+    dtypes_x_axes_shift=_axes_value(),
+    sh=helpers.dtype_and_values(available_dtypes=helpers.get_dtypes("float"), shape=()),
+    keepdims=st.booleans(),
+)
+def test_tensorflow_sufficient_statistics(
+    *,
+    dtypes_x_axes_shift,
+    sh,
+    keepdims,
+    frontend,
+    test_flags,
+    fn_tree,
+    on_device,
+):
+    input_dtypes, x, a = dtypes_x_axes_shift
+    return helpers.test_frontend_function(
+        input_dtypes=input_dtypes,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        x=x[0],
+        axes=a,
+        shift=sh[1][0],
+        keepdims=keepdims,
+        name=None,
+    )
+  
+
 @handle_frontend_test(
     fn_tree="tensorflow.nn.log_poisson_loss",
     dtype_target_log_inputs=helpers.dtype_and_values(