code: def test_gradient_optim(self, input_dim, output_dim, batch_size):
      m = model_helper.ModelHelper()
      with core.NameScope("name_x"):
          fc1 = brew.fc(
              m, "data", "fc1", dim_in=input_dim, dim_out=output_dim
          )
          fc2 = brew.fc(
              m, fc1, "fc2", dim_in=output_dim, dim_out=output_dim
          )
          fc3 = brew.fc(
              m, fc2, "fc3", dim_in=output_dim, dim_out=output_dim
          )
          fc4 = brew.fc(
              m, fc3, "fc4", dim_in=output_dim, dim_out=output_dim
          )
          fc5 = brew.fc(
              m, fc4, "fc5", dim_in=output_dim, dim_out=output_dim
          )
          fc5.Relu([], fc5).Softmax([], "pred").LabelCrossEntropy(
              ["label"], ["xent"]
          ).AveragedLoss([], "loss")
      input_to_grad = m.AddGradientOperators(["name_x/loss"])
  
      blobs_before = count_blobs(m.net.Proto())
      optim_proto = memonger.share_grad_blobs(
          m.net,
          ["name_x/loss"],
          set(m.param_to_grad.values()),
          "name_x/",
          share_activations=False,
      )
      blobs_after = count_blobs(optim_proto)
      self.assertLess(blobs_after, blobs_before)
  
      optim_proto_wacts = memonger.share_grad_blobs(
          m.net,
          ["name_x/loss"],
          set(m.param_to_grad.values()),
          "name_x/",
          share_activations=True,
          dont_share_blobs=set([str(input_to_grad["name_x/fc1_w"])]),
      )
      blobs_wact_optim = count_blobs(optim_proto_wacts)
      self.assertLessEqual(blobs_wact_optim, blobs_after)
  
      # Check that the last activations are not shared
      self.assertTrue(has_blob(optim_proto, "name_x/fc5"))
      self.assertTrue(
          has_blob(optim_proto_wacts, "name_x/fc5"),
          "Dont remap final activation",
      )
  
      # Test networks produce exactly same gradients
      data = np.random.randn(batch_size, input_dim).astype(np.float32)
      label = np.random.randint(
          low=0, high=output_dim, size=(batch_size,)
      ).astype(np.int32)
      workspace.RunNetOnce(m.param_init_net)
      workspace.FeedBlob("name_x/data", data)
      workspace.FeedBlob("name_x/label", label)
      workspace.RunNetOnce(m.net)
      loss = workspace.FetchBlob("name_x/loss")
      grad = workspace.FetchBlob(str(input_to_grad["name_x/fc1_w"]))
      workspace.RunNetOnce(optim_proto)
      optimized_loss = workspace.FetchBlob("name_x/loss")
      optimized_grad = workspace.FetchBlob(
          str(input_to_grad["name_x/fc1_w"])
      )
      np.testing.assert_almost_equal(loss, optimized_loss)
      np.testing.assert_almost_equal(grad, optimized_grad)
  
      workspace.FeedBlob(
          str(input_to_grad["name_x/fc1_w"]), np.array([0.0])
      )
  
      # Run with the forward optimization
      workspace.RunNetOnce(optim_proto_wacts)
      optimized_loss = workspace.FetchBlob("name_x/loss")
      optimized_grad = workspace.FetchBlob(
          str(input_to_grad["name_x/fc1_w"])
      )
      np.testing.assert_almost_equal(loss, optimized_loss)
      np.testing.assert_almost_equal(grad, optimized_grad)
  
imports: import numpy
  from caffe2.python import workspace, memonger, core, model_helper, brew
  from caffe2.proto import caffe2_pb2
  import caffe2.python.hypothesis_test_util
  import hypothesis.strategies
  from hypothesis import given, settings
  import unittest
label_desc: Use image node_id 4 for calling the MemongerTest obj's underlying member method code with example usage: obj.test_gradient_optim(input_dim, output_dim, batch_size) without return types