code: def test_sparse_nd_where():
      def get_forward_expected_output(condition, x, y):
          original_shape = x.shape
          out = np.zeros(original_shape)
          if condition.shape == x.shape:
              for index, c in np.ndenumerate(condition):
                  if c != 0:
                      out[index] = x[index]
                  else:
                      out[index] = y[index]
          else:
              raise RuntimeError("Invalid condition shape for where op")
  
          out = out.reshape(original_shape)
          return out
  
      def get_forward_inputs_same_shape(shape):
          condition_np = np.random.randint(
              0, 2, np.prod(shape)
          ).reshape(shape)
          x_np = np.random.randint(1, 6, np.prod(shape)).reshape(shape)
          y_np = np.random.randint(7, 11, np.prod(shape)).reshape(shape)
          return condition_np, x_np, y_np
  
      def get_backward_input(shape):
          return np.random.randint(20, 30, np.prod(shape)).reshape(
              shape
          )
  
      def get_backward_expected_outputs(grad_in, condition):
          shape = grad_in.shape
          grad_cond = np.zeros(condition.shape)
          grad_x = np.empty(shape)
          grad_y = np.empty(shape)
  
          for index, c in np.ndenumerate(condition):
              if 0 != c:
                  grad_x[index] = grad_in[index]
                  grad_y[index] = 0
              else:
                  grad_x[index] = 0
                  grad_y[index] = grad_in[index]
  
          return grad_cond, grad_x, grad_y
  
      def test_where_helper(shape):
          condition_np, x_np, y_np = get_forward_inputs_same_shape(
              shape
          )
  
          out_expected = get_forward_expected_output(
              condition_np, x_np, y_np
          )
  
          grad_in_np = get_backward_input(shape)
          (
              grad_expected_cond,
              grad_expected_x,
              grad_expected_y,
          ) = get_backward_expected_outputs(grad_in_np, condition_np)
  
          condition = mx.sym.Variable("condition", stype="csr")
          x = mx.sym.Variable("x")
          y = mx.sym.Variable("y")
          grad_in_mx = mx.nd.array(grad_in_np, dtype=np.int32)
          where_sym = mx.sym.where(condition, x, y)
  
          cond_nd = mx.nd.array(condition_np)
          args = {
              "condition": cond_nd.tostype("csr"),
              "x": mx.nd.array(x_np),
              "y": mx.nd.array(y_np),
          }
          args_grad = {
              "condition": mx.nd.zeros_like(cond_nd),
              "x": mx.nd.array(x_np).tostype("csr"),
              "y": mx.nd.array(y_np),
          }
          # test req='write'
          where_exe_write = where_sym._bind(
              ctx=default_device(),
              args=args,
              args_grad=args_grad,
              grad_req="write",
          )
  
          # test forward req='write'
          outputs = where_exe_write.forward(is_train=True)
          assert same(outputs[0].asnumpy(), out_expected)
          # test backward req='write'
          where_exe_write.backward(grad_in_mx.astype("float32"))
          assert same(
              where_exe_write.grad_dict["x"].asnumpy(), grad_expected_x
          )
          assert same(
              where_exe_write.grad_dict["y"].asnumpy(), grad_expected_y
          )
          assert same(
              where_exe_write.grad_dict["condition"].asnumpy(),
              grad_expected_cond,
          )
  
          # test req='add'
          x_grad_init = np.random.randint(
              30, 40, np.prod(shape)
          ).reshape(shape)
          y_grad_init = np.random.randint(
              40, 50, np.prod(shape)
          ).reshape(shape)
          where_exe_add = where_sym._bind(
              ctx=default_device(),
              args=args,
              args_grad=args_grad,
              grad_req="add",
          )
          where_exe_add.grad_dict["x"][:] = x_grad_init
          where_exe_add.grad_dict["y"][:] = y_grad_init
          # test forward req='add'
          outputs = where_exe_add.forward(is_train=True)
          assert same(outputs[0].asnumpy(), out_expected)
  
      def test_where_numeric_gradient(shape):
          condition = mx.sym.Variable("condition", stype="csr")
          x = mx.sym.Variable("x")
          y = mx.sym.Variable("y")
          where_sym = mx.sym.where(condition, x, y)
          condition_np, x_np, y_np = get_forward_inputs_same_shape(
              shape
          )
          check_numeric_gradient(
              where_sym,
              [condition_np, x_np, y_np],
              grad_nodes=["x", "y"],
          )
  
      test_where_helper((5, 9))
      test_where_numeric_gradient((5, 9))
  
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 35 for calling a global function with example usage: test_sparse_nd_where() and returns: out, condition_np,  x_np,  y_np, np, grad_cond,  grad_x,  grad_y