code: def test_simple_regression_correctness(self):
      # Verify that optimizing a simple linear regression by gradient descent
      # recovers the known-correct weights.
      batch_shape = [4, 3]
      num_timesteps = 10
      num_features = 2
      design_matrix = self._build_placeholder(
          np.random.randn(
              *(batch_shape + [num_timesteps, num_features])
          )
      )
  
      true_weights = self._build_placeholder([4.0, -3.0])
      predicted_time_series = tf.linalg.matmul(
          design_matrix, true_weights[..., tf.newaxis]
      )
  
      linear_regression = LinearRegression(
          design_matrix=design_matrix,
          weights_prior=independent.Independent(
              cauchy.Cauchy(
                  loc=self._build_placeholder(np.zeros([num_features])),
                  scale=self._build_placeholder(
                      np.ones([num_features])
                  ),
              ),
              reinterpreted_batch_ndims=1,
          ),
      )
      observation_noise_scale_prior = lognormal.LogNormal(
          loc=self._build_placeholder(-2),
          scale=self._build_placeholder(0.1),
      )
      model = Sum(
          components=[linear_regression],
          observation_noise_scale_prior=observation_noise_scale_prior,
      )
  
      learnable_weights = tf.Variable(
          tf.zeros([num_features], dtype=true_weights.dtype)
      )
  
      def build_loss():
          learnable_ssm = model.make_state_space_model(
              num_timesteps=num_timesteps,
              param_vals={
                  "LinearRegression/_weights": learnable_weights,
                  "observation_noise_scale": observation_noise_scale_prior.mode(),
              },
          )
          return -learnable_ssm.log_prob(predicted_time_series)
  
      # We provide graph- and eager-mode optimization for TF 2.0 compatibility.
      num_train_steps = 80
      optimizer = tf1.train.AdamOptimizer(learning_rate=0.1)
      if tf.executing_eagerly():
          for _ in range(num_train_steps):
              optimizer.minimize(build_loss)
      else:
          train_op = optimizer.minimize(build_loss())
          self.evaluate(tf1.global_variables_initializer())
          for _ in range(num_train_steps):
              _ = self.evaluate(train_op)
      self.assertAllClose(
          *self.evaluate((true_weights, learnable_weights)), atol=0.2
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import cauchy
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.regression import LinearRegression
  from tensorflow_probability.python.sts.components.regression import SparseLinearRegression
  from tensorflow_probability.python.sts.components.sum import Sum
label_desc: Use image node_id 2 for calling the _LinearRegressionTest obj's underlying member method code with example usage: obj.test_simple_regression_correctness() and returns: unknown