code: def update_a_and_b_with_forgetting(
      a_prev: types.Tensor,
      b_prev: types.Tensor,
      r: types.Tensor,
      x: types.Tensor,
      gamma: float,
  ) -> Tuple[types.Tensor, types.Tensor]:
      r"""Update the covariance matrix `a` and the weighted sum of rewards `b`.
  
      This function updates the covariance matrix `a` and the sum of weighted
      rewards `b` using a forgetting factor `gamma`.
  
      Args:
        a_prev: previous estimate of `a`.
        b_prev: previous estimate of `b`.
        r: a `Tensor` of shape [`batch_size`]. This is the rewards of the batched
          observations.
        x: a `Tensor` of shape [`batch_size`, `context_dim`]. This is the matrix
          with the (batched) observations.
        gamma: a float forgetting factor in [0.0, 1.0].
  
      Returns:
        The updated estimates of `a` and `b`.
      """
      a_new = gamma * a_prev + tf.matmul(x, x, transpose_a=True)
      b_new = (
          gamma * b_prev
          + bandit_utils.sum_reward_weighted_observations(r, x)
      )
  
      return a_new, b_new
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import enum
  from typing import Optional, Sequence, Text, Tuple
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.policies import linalg
  from tf_agents.bandits.policies import linear_bandit_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: update_a_and_b_with_forgetting(a_prev, b_prev, r, x, gamma) and returns: a_new,  b_new