code: def update_state(self, state):
      for key in state:
          self._states[key]["sum"].assign_add(
              tf.math.reduce_sum(state[key], axis=0)
          )
          self._states[key]["count"].assign_add(
              tf.cast(tf.shape(state[key])[0], tf.float32)
          )
  
imports: import tensorflow
  import numpy
label_desc: Use image node_id 3 for calling the DictMean obj's underlying member method code with example usage: obj.update_state(state) without return types