class_docstr: Indexed entropy model for location-scale family of random variables.
  
  This class is a common special case of `ContinuousIndexedEntropyModel`. The
  specified distribution is parameterized with `num_scales` values of scale
  parameters. An element-wise location parameter is handled by shifting the
  distributions to zero.
  
  This method is illustrated in Figure 10 of:
  > "Nonlinear Transform Coding"<br />
  > J. Ballé, P.A. Chou, D. Minnen, S. Singh, N. Johnston, E. Agustsson,
  > S.J. Hwang, G. Toderici<br />
  > https://doi.org/10.1109/JSTSP.2020.3034501
  
  Note: this only works for shift-invariant `tfpd.Distribution` objects, where
  the `loc` parameter really denotes a translation (i.e., not for the log-normal
  distribution).
docstr: Initializes the instance.
  
  Args:
    prior_fn: A callable returning a `tfp.distributions.Distribution` object,
      typically a `Distribution` class or factory function. This is a density
      model fitting the marginal distribution of the bottleneck data with
      additive uniform noise, which is shared a priori between the sender and
      the receiver. For best results, the distributions should be flexible
      enough to have a unit-width uniform distribution as a special case,
      since this is the marginal distribution for bottleneck dimensions that
      are constant. The callable will receive keyword arguments as determined
      by `parameter_fns`.
    num_scales: Integer. Values in `indexes` must be in the range
      `[0, num_scales)`.
    scale_fn: Callable. `indexes` is passed to the callable, and the return
      value is given as `scale` keyword argument to `prior_fn`.
    coding_rank: Integer. Number of innermost dimensions considered a coding
      unit. Each coding unit is compressed to its own bit string, and the
      bits in the `__call__` method are summed over each coding unit.
    compression: Boolean. If set to `True`, the range coding tables used by
      `compress()` and `decompress()` will be built on instantiation. If set
      to `False`, these two methods will not be accessible.
    stateless: Boolean. If `False`, range coding tables are created as
      `Variable`s. This allows the entropy model to be serialized using the
      `SavedModel` protocol, so that both the encoder and the decoder use
      identical tables when loading the stored model. If `True`, creates range
      coding tables as `Tensor`s. This makes the entropy model stateless and
      allows it to be constructed within a `tf.function` body, for when the
      range coding tables are provided manually. If `compression=False`, then
      `stateless=True` is implied and the provided value is ignored.
    expected_grads: If True, will use analytical expected gradients during
      backpropagation w.r.t. additive uniform noise.
    tail_mass: Float. Approximate probability mass which is encoded using an
      Elias gamma code embedded into the range coder.
    range_coder_precision: Integer. Precision passed to the range coding op.
    bottleneck_dtype: `tf.dtypes.DType`. Data type of bottleneck tensor.
      Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
    prior_dtype: `tf.dtypes.DType`. Data type of prior and probability
      computations. Defaults to `tf.float32`.
    laplace_tail_mass: Float, or a float-valued tf.Tensor. If positive,
      will augment the prior with a `NoisyLaplace` mixture component for
      training stability. (experimental)
docstr_tok: ["Initializes","the","instance",".","Args",":","prior_fn",":","A","callable","returning","a","`","tfp.distributions.Distribution","`","object",",","typically","a","`","Distribution","`","class","or","factory","function",".","This","is","a","density","model","fitting","the","marginal","distribution","of","the","bottleneck","data","with","additive","uniform","noise",",","which","is","shared","a","priori","between","the","sender","and","the","receiver",".","For","best","results",",","the","distributions","should","be","flexible","enough","to","have","a","unit-width","uniform","distribution","as","a","special","case",",","since","this","is","the","marginal","distribution","for","bottleneck","dimensions","that","are","constant",".","The","callable","will","receive","keyword","arguments","as","determined","by","`","parameter_fns","`",".","num_scales",":","Integer",".","Values","in","`","indexes","`","must","be","in","the","range","`","[","0",",","num_scales",")","`",".","scale_fn",":","Callable",".","`","indexes","`","is","passed","to","the","callable",",","and","the","return","value","is","given","as","`","scale","`","keyword","argument","to","`","prior_fn","`",".","coding_rank",":","Integer",".","Number","of","innermost","dimensions","considered","a","coding","unit",".","Each","coding","unit","is","compressed","to","its","own","bit","string",",","and","the","bits","in","the","`","__call__","`","method","are","summed","over","each","coding","unit",".","compression",":","Boolean",".","If","set","to","`","True","`",",","the","range","coding","tables","used","by","`","compress","(",")","`","and","`","decompress","(",")","`","will","be","built","on","instantiation",".","If","set","to","`","False","`",",","these","two","methods","will","not","be","accessible",".","stateless",":","Boolean",".","If","`","False","`",",","range","coding","tables","are","created","as","`","Variable","`","s",".","This","allows","the","entropy","model","to","be","serialized","using","the","`","SavedModel","`","protocol",",","so","that","both","the","encoder","and","the","decoder","use","identical","tables","when","loading","the","stored","model",".","If","`","True","`",",","creates","range","coding","tables","as","`","Tensor","`","s",".","This","makes","the","entropy","model","stateless","and","allows","it","to","be","constructed","within","a","`","tf.function","`","body",",","for","when","the","range","coding","tables","are","provided","manually",".","If","`","compression=False","`",",","then","`","stateless=True","`","is","implied","and","the","provided","value","is","ignored",".","expected_grads",":","If","True",",","will","use","analytical","expected","gradients","during","backpropagation","w.r.t",".","additive","uniform","noise",".","tail_mass",":","Float",".","Approximate","probability","mass","which","is","encoded","using","an","Elias","gamma","code","embedded","into","the","range","coder",".","range_coder_precision",":","Integer",".","Precision","passed","to","the","range","coding","op",".","bottleneck_dtype",":","`","tf.dtypes.DType","`",".","Data","type","of","bottleneck","tensor",".","Defaults","to","`","tf.keras.mixed_precision.global_policy","(",")",".compute_dtype","`",".","prior_dtype",":","`","tf.dtypes.DType","`",".","Data","type","of","prior","and","probability","computations",".","Defaults","to","`","tf.float32","`",".","laplace_tail_mass",":","Float",",","or","a","float-valued","tf.Tensor",".","If","positive",",","will","augment","the","prior","with","a","`","NoisyLaplace","`","mixture","component","for","training","stability",".","(","experimental",")"]
code: def __init__(
      self,
      prior_fn,
      num_scales,
      scale_fn,
      coding_rank,
      compression=False,
      stateless=False,
      expected_grads=False,
      tail_mass=2**-8,
      range_coder_precision=12,
      bottleneck_dtype=None,
      prior_dtype=tf.float32,
      laplace_tail_mass=0,
  ):
      """Initializes the instance.
  
      Args:
        prior_fn: A callable returning a `tfp.distributions.Distribution` object,
          typically a `Distribution` class or factory function. This is a density
          model fitting the marginal distribution of the bottleneck data with
          additive uniform noise, which is shared a priori between the sender and
          the receiver. For best results, the distributions should be flexible
          enough to have a unit-width uniform distribution as a special case,
          since this is the marginal distribution for bottleneck dimensions that
          are constant. The callable will receive keyword arguments as determined
          by `parameter_fns`.
        num_scales: Integer. Values in `indexes` must be in the range
          `[0, num_scales)`.
        scale_fn: Callable. `indexes` is passed to the callable, and the return
          value is given as `scale` keyword argument to `prior_fn`.
        coding_rank: Integer. Number of innermost dimensions considered a coding
          unit. Each coding unit is compressed to its own bit string, and the
          bits in the `__call__` method are summed over each coding unit.
        compression: Boolean. If set to `True`, the range coding tables used by
          `compress()` and `decompress()` will be built on instantiation. If set
          to `False`, these two methods will not be accessible.
        stateless: Boolean. If `False`, range coding tables are created as
          `Variable`s. This allows the entropy model to be serialized using the
          `SavedModel` protocol, so that both the encoder and the decoder use
          identical tables when loading the stored model. If `True`, creates range
          coding tables as `Tensor`s. This makes the entropy model stateless and
          allows it to be constructed within a `tf.function` body, for when the
          range coding tables are provided manually. If `compression=False`, then
          `stateless=True` is implied and the provided value is ignored.
        expected_grads: If True, will use analytical expected gradients during
          backpropagation w.r.t. additive uniform noise.
        tail_mass: Float. Approximate probability mass which is encoded using an
          Elias gamma code embedded into the range coder.
        range_coder_precision: Integer. Precision passed to the range coding op.
        bottleneck_dtype: `tf.dtypes.DType`. Data type of bottleneck tensor.
          Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
        prior_dtype: `tf.dtypes.DType`. Data type of prior and probability
          computations. Defaults to `tf.float32`.
        laplace_tail_mass: Float, or a float-valued tf.Tensor. If positive,
          will augment the prior with a `NoisyLaplace` mixture component for
          training stability. (experimental)
      """
      num_scales = int(num_scales)
      super().__init__(
          prior_fn=prior_fn,
          index_ranges=(num_scales,),
          parameter_fns=dict(
              loc=lambda _: 0.0,
              scale=scale_fn,
          ),
          coding_rank=coding_rank,
          channel_axis=None,
          compression=compression,
          stateless=stateless,
          expected_grads=expected_grads,
          tail_mass=tail_mass,
          range_coder_precision=range_coder_precision,
          bottleneck_dtype=bottleneck_dtype,
          prior_dtype=prior_dtype,
          laplace_tail_mass=laplace_tail_mass,
      )
  
imports: import tensorflow
  from tensorflow_compression.python.entropy_models import continuous_base
  from tensorflow_compression.python.ops import gen_ops
  from tensorflow_compression.python.ops import math_ops
  from tensorflow_compression.python.ops import round_ops
label_desc: Use image node_id 1 to create a new LocationScaleIndexedEntropyModel object from inherited base classes: ContinuousIndexedEntropyModel with example: obj = LocationScaleIndexedEntropyModel(prior_fn, num_scales, scale_fn, coding_rank, compression, stateless, expected_grads, tail_mass, range_coder_precision, bottleneck_dtype, prior_dtype, laplace_tail_mass)