docstr: Distribution object for Categorical data.
  Optionally works with batches and time steps, but be consistent in usage: i.e. if using time_step,
  mean, sigma and data for log_prob must all include a time_step dimension.
  
  :param n_classes: number of classes in distribution
  :param probs: probabilities for each class,
      of shape (n_classes),
      of shape (batch_size, n_classes) or
      of shape (batch_size, time_step, n_classes)
  :param (mx.nd or mx.sym) F: backend api (mx.sym if block has been hybridized).
docstr_tok: ["Distribution","object","for","Categorical","data",".","Optionally","works","with","batches","and","time","steps",",","but","be","consistent","in","usage",":","i.e",".","if","using","time_step",",","mean",",","sigma","and","data","for","log_prob","must","all","include","a","time_step","dimension",".",":","param","n_classes",":","number","of","classes","in","distribution",":","param","probs",":","probabilities","for","each","class",",","of","shape","(","n_classes",")",",","of","shape","(","batch_size",",","n_classes",")","or","of","shape","(","batch_size",",","time_step",",","n_classes",")",":","param","(","mx.nd","or","mx.sym",")","F",":","backend","api","(","mx.sym","if","block","has","been","hybridized",")","."]
code: def __init__(
      self, n_classes: int, probs: nd_sym_type, F: ModuleType = mx.nd
  ) -> None:
      """
      Distribution object for Categorical data.
      Optionally works with batches and time steps, but be consistent in usage: i.e. if using time_step,
      mean, sigma and data for log_prob must all include a time_step dimension.
  
      :param n_classes: number of classes in distribution
      :param probs: probabilities for each class,
          of shape (n_classes),
          of shape (batch_size, n_classes) or
          of shape (batch_size, time_step, n_classes)
      :param (mx.nd or mx.sym) F: backend api (mx.sym if block has been hybridized).
      """
      self.n_classes = n_classes
      self.probs = probs
      self.F = F
  
imports: from typing import List, Tuple, Union
  from types import ModuleType
  import math
  import mxnet
  from mxnet.gluon import nn
  from rl_coach.base_parameters import AgentParameters
  from rl_coach.core_types import ActionProbabilities
  from rl_coach.spaces import SpacesDefinition, BoxActionSpace, DiscreteActionSpace
  from rl_coach.utils import eps
  from rl_coach.architectures.mxnet_components.heads.head import Head, HeadLoss, LossInputSchema, NormalizedRSSInitializer
  from rl_coach.architectures.mxnet_components.heads.head import LOSS_OUT_TYPE_LOSS, LOSS_OUT_TYPE_REGULARIZATION
  from rl_coach.architectures.mxnet_components.utils import hybrid_clip, broadcast_like
label_desc: Use image node_id 1 to create a new CategoricalDist object  with example: obj = CategoricalDist(n_classes, probs, F)