code: def init_var(self, mesh):
      # Initialize TF randon normal variables(without using DTensor).
      w_initializer = stateless_random_ops.stateless_random_normal(
          shape=[28 * 28, 10], seed=[0, 1]
      )
      b_initializer = stateless_random_ops.stateless_random_normal(
          shape=[10], seed=[1, 2]
      )
      # A filter with 3x3 shape, 1 input channel and 1 output channel.
      k_initializer = stateless_random_ops.stateless_random_normal(
          [3, 3, 1, 1], seed=[2, 3]
      )
  
      n_w = variables.Variable(w_initializer)
      n_b = variables.Variable(b_initializer)
      n_k = variables.Variable(k_initializer)
  
      # Initialize DTensor variables.
      w_initializer_on_mesh = api.copy_to_mesh(
          w_initializer, Layout.replicated(mesh, 2)
      )
      b_initializer_on_mesh = api.copy_to_mesh(
          b_initializer, Layout.replicated(mesh, rank=1)
      )
      k_initializer_on_mesh = api.copy_to_mesh(
          k_initializer, Layout.replicated(mesh, rank=4)
      )
  
      w = d_variable.DVariable(w_initializer_on_mesh)
      b = d_variable.DVariable(b_initializer_on_mesh)
      k = d_variable.DVariable(k_initializer_on_mesh)
  
      return (n_w, n_b, n_k), (w, b, k)
  
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 the DTensorMNISTTest obj's underlying member method code with example usage: obj.init_var(mesh) and returns: n_w,  n_b,  n_k,  w,  b,  k