code: def step_tf(optim, steps=1, kwargs={}):
      net = TinyNetTF()
      optim = optim(**kwargs)
      for _ in range(steps):
          with tf.GradientTape() as tape:
              out = net.forward()
          grads = tape.gradient(out, [net.x, net.W])
          optim.apply_gradients(zip(grads, [net.x, net.W]))
      return net.x.numpy(), net.W.numpy()
  
imports: import unittest
  import numpy
  import tensorflow
  import tensorflow_addons
  from tinygrad.tensor import Tensor
  from tinygrad.nn.optim import LAMB
label_desc: Use image node_id 2 for calling a global function with example usage: step_tf(optim, steps, kwargs) and returns: net,  net