code: def dist(self, params, nn_out):
      # TODO(colcarroll): Add a prior to constrain the scale (`nn_out[..., [1]]`)
      # separately before it goes into the likelihood.
      return normal_lib.Normal(
          loc=nn_out[..., [0]], scale=jax.nn.softplus(nn_out[..., [1]])
      )
  
imports: import dataclasses
  from typing import Any
  import jax
  from tensorflow_probability.substrates.jax.bijectors import softplus
  from tensorflow_probability.substrates.jax.distributions import distribution
  from tensorflow_probability.substrates.jax.distributions import inflated
  from tensorflow_probability.substrates.jax.distributions import logistic
  from tensorflow_probability.substrates.jax.distributions import lognormal
  from tensorflow_probability.substrates.jax.distributions import negative_binomial
  from tensorflow_probability.substrates.jax.distributions import normal
  from tensorflow_probability.substrates.jax.distributions import transformed_distribution
label_desc: Use image node_id 2 for calling the NormalLikelihoodVaryingNoise obj's underlying member method code with example usage: obj.dist(params, nn_out) and returns: normal_lib