code: def test_sparse_square_sum():
      dim0 = 30
      dim1 = 30
      axes = [0, 1]
      keepdims = [False, True]
      densities = [0, 0.01, 0.2, 0.5, 1.0]
      for density in densities:
          shape = rand_shape_2d(dim0, dim1)
          rsp = rand_ndarray(shape, "row_sparse", density)
          dns = rsp.tostype("default")
          for axis in axes:
              for keepdim in keepdims:
                  ret = mx.nd._internal._square_sum(
                      rsp, axis=axis, keepdims=keepdim
                  )
                  if axis == 1 and keepdim:
                      assert ret.stype == "row_sparse"
                  else:
                      assert ret.stype == "default"
                  ret_expected = mx.nd.sum(
                      dns * dns, axis=axis, keepdims=keepdim
                  )
                  # check forward result
                  assert_almost_equal(
                      ret.asnumpy(), ret_expected.asnumpy()
                  )
  
                  rsp_data = mx.sym.Variable("data", stype="row_sparse")
                  test = mx.symbol._internal._square_sum(
                      rsp_data, axis=axis, keepdims=keepdim
                  )
  
                  # check symbolic backward since ograd can be an rsp
                  # and cannot be checked through check_numeric_gradient
                  # because it will add a loss layer as the output layer
                  # which makes ograd of the square_sum dense
                  if axis == 1 and keepdim:
                      dns_data = mx.sym.Variable("data")
                      baseline = mx.sym.sum(
                          mx.sym.square(dns_data),
                          axis=axis,
                          keepdims=keepdim,
                      )
                      igrad_expected = mx.nd.empty(dns.shape)
                      baseline_exec = baseline._bind(
                          default_device(),
                          args=[dns],
                          args_grad=[igrad_expected],
                      )
                      baseline_exec.forward(is_train=True)
                      baseline_exec.backward([ret_expected])
                      # check backward when ograd is row sparse
                      check_symbolic_backward(
                          test,
                          [rsp],
                          [ret_expected.tostype("row_sparse")],
                          [igrad_expected.asnumpy()],
                          grad_stypes={"data": "row_sparse"},
                      )
  
                      # check backward when ograd is dense
                      # the stype of output of the square_sum is deteremined in symbol binding stage.
                      # The ograd stype of the last layer is the same as the output stype of the last layer.
                      # Need to add one more layer after square_sum to trigger the kernel for ograd
                      # with default stype in square_sum op.
                      baseline1 = baseline + 1
                      baseline_exec1 = baseline1._bind(
                          default_device(),
                          args=[dns],
                          args_grad=[igrad_expected],
                      )
                      baseline_exec1.forward(is_train=True)
                      baseline_exec1.backward([ret_expected])
                      test1 = test + 1
                      check_symbolic_backward(
                          test1,
                          [rsp],
                          [ret_expected],
                          [igrad_expected.asnumpy()],
                          grad_stypes={"data": "row_sparse"},
                      )
  
                  # check numeric gradient
                  check_numeric_gradient(
                      test,
                      [rsp],
                      grad_stype_dict={"data": "row_sparse"},
                      atol=1e-2,
                      rtol=0.1,
                  )
  
imports: from mxnet.test_utils import *
  from mxnet.base import MXNetError
  import pytest
  from common import assertRaises
  import random
  import warnings
label_desc: Use image node_id 27 for calling a global function with example usage: test_sparse_square_sum() without return types