class_docstr: Numerically stable normalization layer.
class_docstr_tok: ["Numerically","stable","normalization","layer","."]
code: def forward(self, inputs):
      state = self.state
      observations = inputs
      if self._mode == "collect":
          # Accumulate statistics only in the collect mode, i.e. when collecting
          # data using the agent.
          for observation in observations[
              :, -1
          ]:  # (batch_size, time, ...)
              # Update statistics for each observation separately for simplicity.
              # Currently during data collection the batch size is 1 anyway.
              count = running_mean_and_variance_get_count(state)
              state = fastmath.cond(
                  count < self._sample_limit,
                  true_operand=(observation, state),
                  true_fun=lambda args: running_mean_and_variance_update(
                      *args
                  ),
                  false_operand=None,
                  false_fun=lambda _: state,
              )
  
      mean = running_mean_and_variance_get_mean(state)
      var = running_mean_and_variance_get_variance(state)
      norm_observations = (observations - mean) / (
          var**0.5 + self._epsilon
      )
      self.state = state
      return norm_observations
  
file_path: trax/trax/rl/normalization.py
imports: import gin
  import numpy
  from trax import fastmath
  from trax import layers
label_desc: Use image node_id 3 for calling the Normalize obj's underlying member method code with example usage: obj.forward(inputs) and returns: norm_observations