code: def train_step(params, inputs, targets, learning_rate=0.1):
      grad_fn = extensions.grad(loss_fn)
      grads = grad_fn(params, inputs, targets)
      new_w = params[0] - (grads[0] * learning_rate)
      new_b = params[1] - (grads[1] * learning_rate)
  
      return new_w, new_b
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: train_step(params, inputs, targets, learning_rate) and returns: new_w,  new_b