code: def train_op(x, xl):
      g, v = gv(x, xl)  # returns private gradients, loss
      opt(lr, g)
      return v
  
file_path: objax/examples/image_classification/mnist_dp.py
imports: import argparse
  import os
  from functools import partial
  import numpy
  import tensorflow_datasets
  from tqdm import trange
  import objax
  from objax.functional import one_hot, tanh, average_pool_2d
  from objax.jaxboard import SummaryWriter, Summary
  from objax.nn import Sequential, Conv2D
  from objax.util import EasyDict
label_desc: Use image node_id 2 for calling a global function with example usage: train_op(x, xl) and returns: v