docstr: Runs training for one step and verifies gradient-based updates.
  
  This uses a linear regressor as the base model.
  
  Args:
    example: An instance of `tf.train.Example`.
    model_fn: A function that builds a linear regression model.
    dense_layer_index: The index of the dense layer in the linear regressor.
    max_neighbors: The maximum number of neighbors for graph regularization.
    weight: Initial value for the weights variable in the linear regressor.
    expected_grad_from_weight: The expected gradient of the loss with respect
      to the weights variable.
    distributed_strategy: An instance of `tf.distribute.Strategy` specifying
      the distributed strategy to use for training.
    neighbor_cache_client: Client to lookup and update cached neighbor state.
docstr_tok: ["Runs","training","for","one","step","and","verifies","gradient-based","updates",".","This","uses","a","linear","regressor","as","the","base","model",".","Args",":","example",":","An","instance","of","`","tf.train.Example","`",".","model_fn",":","A","function","that","builds","a","linear","regression","model",".","dense_layer_index",":","The","index","of","the","dense","layer","in","the","linear","regressor",".","max_neighbors",":","The","maximum","number","of","neighbors","for","graph","regularization",".","weight",":","Initial","value","for","the","weights","variable","in","the","linear","regressor",".","expected_grad_from_weight",":","The","expected","gradient","of","the","loss","with","respect","to","the","weights","variable",".","distributed_strategy",":","An","instance","of","`","tf.distribute.Strategy","`","specifying","the","distributed","strategy","to","use","for","training",".","neighbor_cache_client",":","Client","to","lookup","and","update","cached","neighbor","state","."]
code: def _train_and_check_params(
      self,
      example,
      model_fn,
      dense_layer_index,
      max_neighbors,
      weight,
      expected_grad_from_weight,
      distributed_strategy=None,
      neighbor_cache_client=None,
  ):
      """Runs training for one step and verifies gradient-based updates.
  
      This uses a linear regressor as the base model.
  
      Args:
        example: An instance of `tf.train.Example`.
        model_fn: A function that builds a linear regression model.
        dense_layer_index: The index of the dense layer in the linear regressor.
        max_neighbors: The maximum number of neighbors for graph regularization.
        weight: Initial value for the weights variable in the linear regressor.
        expected_grad_from_weight: The expected gradient of the loss with respect
          to the weights variable.
        distributed_strategy: An instance of `tf.distribute.Strategy` specifying
          the distributed strategy to use for training.
        neighbor_cache_client: Client to lookup and update cached neighbor state.
      """
  
      dataset = make_dataset(
          example,
          input_shape=[2],
          training=True,
          max_neighbors=max_neighbors,
          include_key=neighbor_cache_client is not None,
      )
  
      def _create_and_compile_graph_reg_model(
          model_fn, weight, max_neighbors
      ):
          """Creates and compiles a graph regularized model.
  
          Args:
            model_fn: A function that builds a linear regression model.
            weight: Initial value for the weights variable in the linear regressor.
            max_neighbors: The maximum number of neighbors for graph regularization.
  
          Returns:
            A pair containing the unregularized model and the graph regularized
            model as `tf.keras.Model` instances.
          """
          model = model_fn((2,), weight)
          graph_reg_config = configs.make_graph_reg_config(
              max_neighbors=max_neighbors, multiplier=1
          )
          graph_reg_model = (
              graph_regularization.GraphRegularizationWithCaching(
                  model, graph_reg_config, neighbor_cache_client
              )
          )
          graph_reg_model.compile(
              optimizer=tf.keras.optimizers.SGD(LEARNING_RATE),
              loss="MSE",
          )
          return model, graph_reg_model
  
      if distributed_strategy:
          with distributed_strategy.scope():
              (
                  model,
                  graph_reg_model,
              ) = _create_and_compile_graph_reg_model(
                  model_fn, weight, max_neighbors
              )
      else:
          model, graph_reg_model = _create_and_compile_graph_reg_model(
              model_fn, weight, max_neighbors
          )
  
      graph_reg_model.fit(x=dataset, epochs=1, steps_per_epoch=1)
  
      # Compute the new weight value based on the gradient.
      expected_weight = weight - LEARNING_RATE * (
          expected_grad_from_weight
      )
  
      # Check that the weight parameter of the linear regressor has the correct
      # value.
      self.assertAllClose(
          expected_weight,
          tf.keras.backend.get_value(
              model.layers[dense_layer_index].weights[0]
          ),
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from neural_structured_learning import configs
  from research.carls import graph_regularization
  from research.carls import neighbor_cache_client
  import numpy
  import tensorflow
  from google.protobuf import text_format
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 3 for calling the GraphRegularizationWithCachingTest obj's underlying member method code with example usage: obj._train_and_check_params(example, model_fn, dense_layer_index, max_neighbors, weight, expected_grad_from_weight, distributed_strategy, neighbor_cache_client) and returns: model,  graph_reg_model