class_docstr: Tests for ops in byteps.mxnet.
docstr: Test that the byteps_push_pull correctly sums 1D, 2D, 3D tensors.
code: def test_byteps_push_pull(self):
      """Test that the byteps_push_pull correctly sums 1D, 2D, 3D tensors."""
      dtypes = ["float16", "float32", "float64"]
      dims = [1, 2, 3]
      count = 0
      ctx = self._current_context()
      shapes = [(), (17), (17, 17), (17, 17, 17)]
      for dtype, dim in itertools.product(dtypes, dims):
          # MXNet uses gpu_id as part of the seed, so to get identical seeds
          # we must set a context.
          mx.random.seed(10 + 10 * bps.rank(), ctx=ctx)
          tensor = mx.nd.random.uniform(
              -100, 100, shape=shapes[dim], ctx=ctx
          )
          tensor = tensor.astype(dtype)
          input = tensor.asnumpy()
  
          bps.byteps_declare_tensor("tensor_" + str(count))
          bps.byteps_push_pull(tensor, name="tensor_" + str(count))
          tensor.wait_to_read()
          output = tensor.asnumpy()
          assert np.allclose(input, output)
          count += 1
  
      print("test_byteps_push_pull passed")
  
file_path: byteps/tests/test_mxnet.py
imports: import itertools
  import unittest
  import byteps.mxnet
  import mxnet
  import numpy
  from meta_test import MetaTest
label_desc: Use image node_id 3 for calling the MXTest obj's underlying member method code with example usage: obj.test_byteps_push_pull() without return types