code: def GaussianLogDensity(x, mu, log_var, name='GaussianLogDensity', EPSILON = 1e-6):
      '''GaussianLogDensity loss calculation for layer wise loss
      '''
      c = mx.sym.ones_like(log_var)*2.0 * 3.1416
      c = mx.symbol.log(c)
      var = mx.sym.exp(log_var)
      x_mu2 = mx.symbol.square(x - mu)   # [Issue] not sure the dim works or not?
      x_mu2_over_var = mx.symbol.broadcast_div(x_mu2, var + EPSILON)
      log_prob = -0.5 * (c + log_var + x_mu2_over_var)
      log_prob = mx.symbol.sum(log_prob, axis=1, name=name)   # keep_dims=True,
      return log_prof GaussianLogDensity(x, mu, log_var, name='GaussianLogDensity', EPSILON = 1e-6):
      '''GaussianLogDensity loss calculation for layer wise loss
      '''
      c = mx.sym.ones_like(log_var)*2.0 * 3.1416
      c = mx.symbol.log(c)
      var = mx.sym.exp(log_var)
      x_mu2 = mx.symbol.square(x - mu)   # [Issue] not sure the dim works or not?
      x_mu2_over_var = mx.symbol.broadcast_div(x_mu2, var + EPSILON)
      log_prob = -0.5 * (c + log_var + x_mu2_over_var)
      log_prob = mx.symbol.sum(log_prob, axis=1, name=name)   # keep_dims=True,
      return log_prob
docs: GaussianLogDensity loss calculation for layer wise loss