Replaces "Arguments" with "Args" to be consistent with TF PR #45420.

PiperOrigin-RevId: 352505520
Change-Id: I3810dde4fb012701febe61ec1b27e8198916cba6

Author: Johannes Ballé

models/hific/archs.py

@@ -654,7 +654,7 @@ def call(self, latents, image_shape, mode: ModelMode) -> FactorizedPriorInfo:
 def estimate_entropy(entropy_model, inputs, spatial_shape=None) -> EntropyInfo:
   """Compresses `inputs` with the given entropy model and estimates entropy.
 
-  Arguments:
+  Args:
     entropy_model: An `EntropyModel` instance.
     inputs: The input tensor to be fed to the entropy model.
     spatial_shape: Shape of the input image (HxW). Must be provided for

models/toy_sources/compression_model.py

@@ -46,7 +46,7 @@ def ndim_source(self):
   def quantize(self, x):
     """Determines an equivalent vector quantizer for `x`.
 
-    Arguments:
+    Args:
       x: A batch of source vectors.
 
     Returns:
@@ -61,7 +61,7 @@ def quantize(self, x):
   def train_losses(self, x):
     """Computes the training losses for `x`.
 
-    Arguments:
+    Args:
       x: A batch of source vectors.
 
     Returns:
@@ -74,7 +74,7 @@ def train_losses(self, x):
   def test_losses(self, x):
     """Computes the rate and distortion for each element of `x`.
 
-    Arguments:
+    Args:
       x: A batch of source vectors.
 
     Returns:

models/toy_sources/sawbridge.py

@@ -12,7 +12,7 @@ def __init__(self, index_points, stationary=True, order=1,
                name="sawbridge"):
     """Initializer.
 
-    Arguments:
+    Args:
       index_points: 1-D `Tensor` representing the locations at which to evaluate
         the process. The intent is that all locations are in [0,1], but the
         process has a natural extrapolation outside this range so no error is

tensorflow_compression/python/distributions/deep_factorized.py

@@ -69,7 +69,7 @@ def __init__(self,
                allow_nan_stats=False, dtype=tf.float32, name="DeepFactorized"):
     """Initializer.
 
-    Arguments:
+    Args:
       batch_shape: Iterable of integers. The desired batch shape for the
         `Distribution` (rightmost dimensions which are assumed independent, but
         not identically distributed).
@@ -164,7 +164,7 @@ def _broadcast_inputs(self, inputs):
   def _logits_cumulative(self, inputs):
     """Evaluate logits of the cumulative densities.
 
-    Arguments:
+    Args:
       inputs: The values at which to evaluate the cumulative densities.
 
     Returns:

tensorflow_compression/python/distributions/helpers.py

@@ -42,7 +42,7 @@ def estimate_tails(func, target, shape, dtype):
   This operation is vectorized. The tensor shape of `x` is given by `shape`, and
   `target` must have a shape that is broadcastable to the output of `func(x)`.
 
-  Arguments:
+  Args:
     func: A callable that computes cumulative distribution function, survival
       function, or similar.
     target: The desired target value.
@@ -107,7 +107,7 @@ def quantization_offset(distribution):
   Note the offset is always in the range [-.5, .5] as it is assumed to be
   combined with a round quantizer.
 
-  Arguments:
+  Args:
     distribution: A `tfp.distributions.Distribution` object.
 
   Returns:
@@ -142,7 +142,7 @@ def lower_tail(distribution, tail_mass):
   functionality of the range coder implementation (using a Golomb-like
   universal code).
 
-  Arguments:
+  Args:
     distribution: A `tfp.distributions.Distribution` object.
     tail_mass: Float between 0 and 1. Desired probability mass for the tails.
 
@@ -178,7 +178,7 @@ def upper_tail(distribution, tail_mass):
   functionality of the range coder implementation (using a Golomb-like
   universal code).
 
-  Arguments:
+  Args:
     distribution: A `tfp.distributions.Distribution` object.
     tail_mass: Float between 0 and 1. Desired probability mass for the tails.
 

tensorflow_compression/python/distributions/round_adapters.py

@@ -44,7 +44,7 @@ class MonotonicAdapter(tfp.distributions.Distribution):
   def __init__(self, base, name="MonotonicAdapter"):
     """Initializer.
 
-    Arguments:
+    Args:
       base: A `tfp.distributions.Distribution` object representing a
         continuous-valued random variable.
       name: String. A name for this distribution.
@@ -196,7 +196,7 @@ class NoisyRoundAdapter(uniform_noise.UniformNoiseAdapter):
   def __init__(self, base, name="NoisyRoundAdapter"):
     """Initializer.
 
-    Arguments:
+    Args:
       base: A `tfp.distributions.Distribution` object representing a
         continuous-valued random variable.
       name: String. A name for this distribution.
@@ -225,7 +225,7 @@ class SoftRoundAdapter(MonotonicAdapter):
   def __init__(self, base, alpha, name="SoftRoundAdapter"):
     """Initializer.
 
-    Arguments:
+    Args:
       base: A `tfp.distributions.Distribution` object representing a
         continuous-valued random variable.
       alpha: Float or tf.Tensor. Controls smoothness of the approximation.
@@ -247,7 +247,7 @@ class NoisySoftRoundAdapter(uniform_noise.UniformNoiseAdapter):
   def __init__(self, base, alpha, name="NoisySoftRoundAdapter"):
     """Initializer.
 
-    Arguments:
+    Args:
       base: A `tfp.distributions.Distribution` object representing a
         continuous-valued random variable.
       alpha: Float or tf.Tensor. Controls smoothness of soft round.

tensorflow_compression/python/distributions/uniform_noise.py

@@ -36,7 +36,7 @@ def _logsum_expbig_minus_expsmall(big, small):
   This assumes `small <= big` and arguments that can be broadcast against each
   other.
 
-  Arguments:
+  Args:
     big: Floating-point `tf.Tensor`.
     small: Floating-point `tf.Tensor`.
 
@@ -73,7 +73,7 @@ class UniformNoiseAdapter(tfp.distributions.Distribution):
   def __init__(self, base, name="UniformNoiseAdapter"):
     """Initializer.
 
-    Arguments:
+    Args:
       base: A `tfp.distributions.Distribution` object representing a
         continuous-valued random variable.
       name: String. A name for this distribution.

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -51,7 +51,7 @@ def __init__(
   ):
     """Initializes the instance.
 
-    Arguments:
+    Args:
       prior: A `tfp.distributions.Distribution` object. 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
@@ -226,7 +226,7 @@ def _build_tables(self, prior):
     `compression=True`, and then distribute the model to a sender and a
     receiver.
 
-    Arguments:
+    Args:
       prior: The `tfp.distributions.Distribution` object (see initializer).
     """
     # TODO(jonycgn, relational): Consider not using offset when soft quantization
@@ -343,7 +343,7 @@ def get_config(self):
   def from_config(cls, config):
     """Instantiates an entropy model from a configuration dictionary.
 
-    Arguments:
+    Args:
       config: A `dict`, typically the output of `get_config`.
 
     Returns:

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -83,7 +83,7 @@ def __init__(self,
                no_variables=False):
     """Initializes the instance.
 
-    Arguments:
+    Args:
       prior: A `tfp.distributions.Distribution` object. 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
@@ -167,7 +167,7 @@ def __call__(self, bottleneck, training=True):
     """Perturbs a tensor with (quantization) noise and estimates bitcost.
 
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
         least `self.coding_rank` dimensions, and the innermost dimensions must
         be broadcastable to `self.prior_shape`.
@@ -209,7 +209,7 @@ def quantize(self, bottleneck):
     The gradient of this rounding operation is overridden with the identity
     (straight-through gradient estimator).
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be quantized. The innermost
         dimensions must be broadcastable to `self.prior_shape`.
 
@@ -231,7 +231,7 @@ def compress(self, bottleneck):
     i.e. are compressed into one string each. Any additional dimensions to the
     left are treated as batch dimensions.
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
         least `self.coding_rank` dimensions, and the innermost dimensions must
         be broadcastable to `self.prior_shape`.
@@ -279,7 +279,7 @@ def decompress(self, strings, broadcast_shape):
     Reconstructs the quantized tensor from bit strings produced by `compress()`.
     It is necessary to provide a part of the output shape in `broadcast_shape`.
 
-    Arguments:
+    Args:
       strings: `tf.Tensor` containing the compressed bit strings.
       broadcast_shape: Iterable of ints. The part of the output tensor shape
         between the shape of `strings` on the left and
@@ -338,7 +338,7 @@ def get_config(self):
   def from_config(cls, config):
     """Instantiates an entropy model from a configuration dictionary.
 
-    Arguments:
+    Args:
       config: A `dict`, typically the output of `get_config`.
 
     Returns:

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -140,7 +140,7 @@ def __init__(self,
                no_variables=False):
     """Initializes the instance.
 
-    Arguments:
+    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
@@ -285,7 +285,7 @@ def __call__(self, bottleneck, indexes, training=True):
     """Perturbs a tensor with (quantization) noise and estimates bitcost.
 
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed.
       indexes: `tf.Tensor` specifying the scalar distribution for each element
         in `bottleneck`. See class docstring for examples.
@@ -340,7 +340,7 @@ def quantize(self, bottleneck, indexes):
     The gradient of this rounding operation is overridden with the identity
     (straight-through gradient estimator).
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be quantized.
       indexes: `tf.Tensor` specifying the scalar distribution for each element
         in `bottleneck`. See class docstring for examples.
@@ -365,7 +365,7 @@ def compress(self, bottleneck, indexes):
     i.e. are compressed into one string each. Any additional dimensions to the
     left are treated as batch dimensions.
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed.
       indexes: `tf.Tensor` specifying the scalar distribution for each element
         in `bottleneck`. See class docstring for examples.
@@ -412,7 +412,7 @@ def decompress(self, strings, indexes):
 
     Reconstructs the quantized tensor from bit strings produced by `compress()`.
 
-    Arguments:
+    Args:
       strings: `tf.Tensor` containing the compressed bit strings.
       indexes: `tf.Tensor` specifying the scalar distribution for each output
         element. See class docstring for examples.
@@ -478,7 +478,7 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
                range_coder_precision=12, no_variables=False):
     """Initializes the instance.
 
-    Arguments:
+    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
@@ -527,7 +527,7 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
   def __call__(self, bottleneck, scale_indexes, loc=None, training=True):
     """Perturbs a tensor with (quantization) noise and estimates bitcost.
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed.
       scale_indexes: `tf.Tensor` indexing the scale parameter for each element
         in `bottleneck`. Must have the same shape as `bottleneck`.
@@ -568,7 +568,7 @@ def quantize(self, bottleneck, scale_indexes, loc=None):
     The gradient of this rounding operation is overridden with the identity
     (straight-through gradient estimator).
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be quantized.
       scale_indexes: `tf.Tensor` indexing the scale parameter for each element
         in `bottleneck`. Must have the same shape as `bottleneck`.
@@ -598,7 +598,7 @@ def compress(self, bottleneck, scale_indexes, loc=None):
     i.e. are compressed into one string each. Any additional dimensions to the
     left are treated as batch dimensions.
 
-    Arguments:
+    Args:
       bottleneck: `tf.Tensor` containing the data to be compressed.
       scale_indexes: `tf.Tensor` indexing the scale parameter for each element
         in `bottleneck`. Must have the same shape as `bottleneck`.
@@ -622,7 +622,7 @@ def decompress(self, strings, scale_indexes, loc=None):
 
     Reconstructs the quantized tensor from bit strings produced by `compress()`.
 
-    Arguments:
+    Args:
       strings: `tf.Tensor` containing the compressed bit strings.
       scale_indexes: `tf.Tensor` indexing the scale parameter for each output
         element.