code: def _run_step(inputs, w, b, k):
      with backprop.GradientTape() as g:
          g.watch([w, b])
          logits = nn_ops.conv2d_v2(
              inputs, k, strides=[1, 1, 1, 1], padding="SAME"
          )
          logits = array_ops.reshape(logits, [logits.shape[0], -1])
          logits = math_ops.matmul(logits, w)
          logits = logits + b
          loss = math_ops.reduce_sum(logits, axis=[0, 1])
      gw, gb = g.gradient(loss, [w, b])
      for v, v_grad in zip([w, b], [gw, gb]):
          v.assign_sub(_LR * v_grad)
      return gw, gb, loss
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling a global function with example usage: _run_step(inputs, w, b, k) and returns: gw,  gb,  loss