code: def _computeMnistMlpGrads(
      self,
      math_ops_lib,
      nn_ops_lib,
      backprop_lib,
      cast,
      num_iters,
      hidden_layers,
      hidden_size,
      batch_size,
  ):
      batch_size = 1
      image_size = 28 * 28
      num_classes = 10
  
      def model(x, hidden_weights, softmax_weight, labels):
          with backprop_lib.GradientTape() as tape:
              for weight in hidden_weights + [softmax_weight]:
                  tape.watch(weight)
              for hidden_weight in hidden_weights:
                  x = math_ops_lib.mat_mul(x, hidden_weight)
                  x = nn_ops_lib.relu(x)
              logits = math_ops_lib.mat_mul(x, softmax_weight)
              loss = (
                  nn_ops_lib.sparse_softmax_cross_entropy_with_logits(
                      logits=logits, labels=labels
                  )
              )
  
          grads = tape.gradient(loss, hidden_weights + [softmax_weight])
          return grads
  
      x = maybe_cast(array_ops.ones([batch_size, image_size]), cast)
      hidden_weights = []
      for i in range(hidden_layers):
          hidden_weights.append(
              maybe_cast(
                  random_ops.random_uniform(
                      [hidden_size if i else image_size, hidden_size]
                  ),
                  cast,
              )
          )
      softmax_weight = maybe_cast(
          random_ops.random_uniform([hidden_size, num_classes]), cast
      )
      labels = maybe_cast(
          array_ops.zeros([batch_size], dtype=dtypes.int32), cast
      )
  
      with context_lib.set_default(get_immediate_execution_context()):
          # Warm up.
          for _ in range(10):
              model(x, hidden_weights, softmax_weight, labels)
          runtimes = timeit.repeat(
              lambda: model(x, hidden_weights, softmax_weight, labels),
              repeat=num_iters,
              number=10,
          )
      return min(runtimes) / 10
  
imports: import timeit
  from absl.testing import parameterized
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework.experimental import _unified_api
  from tensorflow.python.framework.experimental import context_stack
  from tensorflow.python.framework.experimental import def_function
  from tensorflow.python.framework.experimental import math_ops
  from tensorflow.python.framework.experimental import nn_ops
  from tensorflow.python.framework.experimental import tape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_grad
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the UnifiedTapeBenchmark obj's underlying member method code with example usage: obj._computeMnistMlpGrads(math_ops_lib, nn_ops_lib, backprop_lib, cast, num_iters, hidden_layers, hidden_size, batch_size) and returns: unknown, grads