Internal change

PiperOrigin-RevId: 256796501
Change-Id: Ide683b4eb0bf82e73b9fd2decf89f90993b59c5c

Author: Googler

tensorflow_compression/cc/kernels/range_coding_helper_kernels.cc

@@ -232,37 +232,5 @@ class ArrayFingerprintOp : public tensorflow::OpKernel {
 
 REGISTER_KERNEL_BUILDER(Name("ArrayFingerprint").Device(tensorflow::DEVICE_CPU),
                         ArrayFingerprintOp);
-
-class CheckArrayFingerprintOp : public tensorflow::OpKernel {
- public:
-  using OpKernel::OpKernel;
-
-  void Compute(tensorflow::OpKernelContext* context) override {
-    const Tensor& input = context->input(0);
-    const Tensor& fingerprint = context->input(1);
-    OP_REQUIRES(context, tensorflow::DataTypeCanUseMemcpy(input.dtype()),
-                InvalidArgument("Data type not supported: ",
-                                tensorflow::DataTypeString(input.dtype())));
-    OP_REQUIRES(context, TensorShapeUtils::IsScalar(fingerprint.shape()),
-                InvalidArgument("`fingerprint` should be a scalar"));
-
-    const int64 size =
-        input.shape().num_elements() * tensorflow::DataTypeSize(input.dtype());
-    auto data = input.bit_casted_shaped<char, 1>({size});
-
-    OP_REQUIRES(context,
-                static_cast<uint64>(fingerprint.scalar<int64>()()) ==
-                    ::util::Fingerprint64(data.data(),
-                                            static_cast<size_t>(data.size())),
-                tensorflow::errors::DataLoss("Fingerprint mismatch"));
-
-    context->set_output(0, input);
-  }
-};
-
-REGISTER_KERNEL_BUILDER(
-    Name("CheckArrayFingerprint").Device(tensorflow::DEVICE_CPU),
-    CheckArrayFingerprintOp);
-
 }  // namespace
 }  // namespace tensorflow_compression

tensorflow_compression/cc/kernels/range_coding_helper_kernels_test.cc

@@ -161,19 +161,6 @@ class FingerprintOpTest : public tensorflow::OpsTestBase {
     inputs_.clear();
     inputs_.emplace_back(tensor);
   }
-
-  void MakeCheckFingerprintOp(Tensor* tensor, Tensor* fingerprint) {
-    TF_ASSERT_OK(
-        tensorflow::NodeDefBuilder("check_fingerprint", "CheckArrayFingerprint")
-            .Input(tensorflow::FakeInput(tensor->dtype()))
-            .Input(tensorflow::FakeInput(fingerprint->dtype()))
-            .Finalize(node_def()));
-    TF_ASSERT_OK(InitOp());
-
-    inputs_.clear();
-    inputs_.emplace_back(tensor);
-    inputs_.emplace_back(fingerprint);
-  }
 };
 
 TEST_F(FingerprintOpTest, Verify) {
@@ -196,18 +183,17 @@ TEST_F(FingerprintOpTest, Verify) {
 
     MakeFingerprintOp(&tensor);
     TF_ASSERT_OK(RunOpKernel());
-
-    Tensor fingerprint = *GetOutput(0);
-
-    MakeCheckFingerprintOp(&tensor, &fingerprint);
-    TF_ASSERT_OK(RunOpKernel());
+    Tensor fingerprint0 = *GetOutput(0);
 
     // Change one byte in the buffer.
     const int64 pos = rand.Uniform(length);
     buffer(pos) = ~buffer(pos);
 
-    MakeCheckFingerprintOp(&tensor, &fingerprint);
-    ASSERT_FALSE(RunOpKernel().ok());
+    MakeFingerprintOp(&tensor);
+    TF_ASSERT_OK(RunOpKernel());
+    Tensor fingerprint1 = *GetOutput(0);
+
+    EXPECT_NE(fingerprint0.tensor_data(), fingerprint1.tensor_data());
   }
 }
 
@@ -220,17 +206,6 @@ TEST_F(FingerprintOpTest, FingerprintShapeFn) {
   INFER_OK(op, "[1,2]", "[]");
   INFER_OK(op, "[1,2,3]", "[]");
 }
-
-TEST_F(FingerprintOpTest, CheckFingerprintShapeFn) {
-  tensorflow::ShapeInferenceTestOp op("CheckArrayFingerprint");
-
-  INFER_OK(op, "?;?", "in0");
-  INFER_OK(op, "[];?", "in0");
-  INFER_OK(op, "[1,2];?", "in0");
-  INFER_OK(op, "[1,2,3];?", "in0");
-  INFER_ERROR("rank 0", op, "?;[1]");
-}
-
 }  // namespace
 }  // namespace tensorflow_compression
 

tensorflow_compression/cc/ops/range_coding_ops.cc

@@ -26,7 +26,6 @@ using tensorflow::shape_inference::DimensionHandle;
 using tensorflow::shape_inference::InferenceContext;
 using tensorflow::shape_inference::ShapeHandle;
 
-// clang-format off
 REGISTER_OP("RangeEncode")
     .Input("data: int16")
     .Input("cdf: int32")
@@ -96,7 +95,7 @@ REGISTER_OP("RangeDecode")
     .Output("decoded: int16")
     .Attr("precision: int >= 1")
     .Attr("debug_level: int = 1")
-    .SetShapeFn([] (InferenceContext* c) {
+    .SetShapeFn([](InferenceContext* c) {
       ShapeHandle out;
       TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &out));
       c->set_output(0, out);
@@ -133,7 +132,7 @@ REGISTER_OP("UnboundedIndexRangeEncode")
     .Attr("precision: int >= 1")
     .Attr("overflow_width: int >= 1")
     .Attr("debug_level: int = 1")
-    .SetShapeFn(tensorflow::shape_inference::ScalarShape)
+    .SetShapeFn(shape_inference::ScalarShape)
     .Doc(R"doc(
 Range encodes unbounded integer `data` using an indexed probability table.
 
@@ -211,7 +210,7 @@ REGISTER_OP("UnboundedIndexRangeDecode")
     .Attr("precision: int >= 1")
     .Attr("overflow_width: int >= 1")
     .Attr("debug_level: int = 1")
-    .SetShapeFn([] (InferenceContext* c) {
+    .SetShapeFn([](InferenceContext* c) {
       c->set_output(0, c->input(1));
       return Status::OK();
     })
@@ -246,7 +245,7 @@ REGISTER_OP("PmfToQuantizedCdf")
     .Input("pmf: float")
     .Output("cdf: int32")
     .Attr("precision: int >= 1")
-    .SetShapeFn([] (InferenceContext* c) {
+    .SetShapeFn([](InferenceContext* c) {
       ShapeHandle in;
       TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &in));
       DimensionHandle last;
@@ -275,34 +274,12 @@ REGISTER_OP("ArrayFingerprint")
     .Input("input: T")
     .Output("fingerprint: int64")
     .Attr("T: realnumbertype")
-    .SetShapeFn(tensorflow::shape_inference::ScalarShape)
+    .SetShapeFn(shape_inference::ScalarShape)
     .Doc(R"doc(
 Produces fingerprint of the input data.
 
 input: Tensor to be fingerprinted.
 fingerprint: Fingerprint value of input.
 )doc");
-
-REGISTER_OP("CheckArrayFingerprint")
-    .Input("input: T")
-    .Input("fingerprint: int64")
-    .Output("output: T")
-    .Attr("T: realnumbertype")
-    .SetShapeFn([](InferenceContext* c) {
-      ShapeHandle unused;
-      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
-      c->set_output(0, c->input(0));
-      return tensorflow::Status::OK();
-    })
-    .Doc(R"doc(
-Computes the fingerprint of `input` and checks the computed value against
-`fingerprint`. If the check fails, then this op returns an error status.
-
-input: Tensor to be fingerprinted and checked.
-fingerprint: Fingerprint value to be checked against.
-output: The same as input.
-)doc");
-// clang-format on
-
 }  // namespace
 }  // namespace tensorflow_compression

tensorflow_compression/python/layers/entropy_models.py

@@ -633,7 +633,10 @@ def cdf_initializer(shape, dtype=None, partition_info=None):
     # here, or the variable will return the wrong dynamic shape later. A
     # placeholder with default gets the trick done (see initializer above).
     quantized_cdf = self.add_variable(
-        "quantized_cdf", shape=None, dtype=tf.int32, trainable=False,
+        "quantized_cdf",
+        shape=(channels, None),
+        dtype=tf.int32,
+        trainable=False,
         initializer=cdf_initializer)
     cdf_length = self.add_variable(
         "cdf_length", shape=(channels,), dtype=tf.int32, trainable=False,

tensorflow_compression/python/layers/parameterizers.py

@@ -49,21 +49,30 @@ class Parameterizer(object):
 
 
 class StaticParameterizer(Parameterizer):
-  """A parameterizer that always returns a constant tensor.
+  """A parameterizer that returns a non-variable.
 
-  No variables are created, hence the parameter never changes.
+  No variables are created, and `getter` is ignored. If `value` is a `Tensor`,
+  the parameter can depend on some other computation. Otherwise, it never
+  changes.
 
   Args:
-    initializer: An initializer object which will be called to produce the
-      static parameter.
+    value: Either a constant or `Tensor` value, or a callable which returns such
+      a thing given a shape and dtype argument (for example, an initializer).
   """
 
-  def __init__(self, initializer):
-    self.initializer = initializer
+  def __init__(self, value):
+    self.value = value
 
   def __call__(self, getter, name, shape, dtype, initializer, regularizer=None):
-    del getter, name, initializer, regularizer  # unused
-    return self.initializer(shape, dtype)
+    del getter, name, initializer  # unused
+    if regularizer is not None:
+      raise NotImplementedError("Regularizers are not currently supported for "
+                                "static parameterizers.")
+    if callable(self.value):
+      # Treat value as initializer.
+      return self.value(shape, dtype=dtype)
+    else:
+      return self.value
 
 
 class RDFTParameterizer(Parameterizer):
@@ -92,35 +101,41 @@ def __call__(self, getter, name, shape, dtype, initializer, regularizer=None):
     size = var_shape[0]
     for s in var_shape[1:-2]:
       size *= s
-    irdft_matrix = spectral_ops.irdft_matrix(var_shape[:-2], dtype=var_dtype)
     if self.dc:
       rdft_shape = (size, var_shape[-2] * var_shape[-1])
     else:
-      irdft_matrix = irdft_matrix[:, 1:]
       rdft_shape = (size - 1, var_shape[-2] * var_shape[-1])
     rdft_dtype = var_dtype
     rdft_name = name + "_rdft"
 
     def rdft_initializer(shape, dtype=None, partition_info=None):
+      """Initializer wrapper."""
       assert tuple(shape) == rdft_shape, shape
       assert dtype == rdft_dtype, dtype
       init = initializer(
           var_shape, dtype=var_dtype, partition_info=partition_info)
       init = tf.reshape(init, (-1, rdft_shape[-1]))
+      irdft_matrix = spectral_ops.irdft_matrix(var_shape[:-2], dtype=var_dtype)
+      if not self.dc:
+        irdft_matrix = irdft_matrix[:, 1:]
       init = tf.linalg.matmul(irdft_matrix, init, transpose_a=True)
       return init
 
     def reparam(rdft):
+      irdft_matrix = spectral_ops.irdft_matrix(var_shape[:-2], dtype=var_dtype)
+      if not self.dc:
+        irdft_matrix = irdft_matrix[:, 1:]
       var = tf.linalg.matmul(irdft_matrix, rdft)
       var = tf.reshape(var, var_shape)
       return var
 
+    reparam_regularizer = None
     if regularizer is not None:
-      regularizer = lambda rdft: regularizer(reparam(rdft))
+      reparam_regularizer = lambda rdft: regularizer(reparam(rdft))
 
     rdft = getter(
         name=rdft_name, shape=rdft_shape, dtype=rdft_dtype,
-        initializer=rdft_initializer, regularizer=regularizer)
+        initializer=rdft_initializer, regularizer=reparam_regularizer)
     return reparam(rdft)
 
 
@@ -165,10 +180,11 @@ def reparam(var):
       var = tf.math.square(var) - pedestal
       return var
 
+    reparam_regularizer = None
     if regularizer is not None:
-      regularizer = lambda var: regularizer(reparam(var))
+      reparam_regularizer = lambda var: regularizer(reparam(var))
 
     var = getter(
         name=reparam_name, shape=shape, dtype=dtype,
-        initializer=reparam_initializer, regularizer=regularizer)
+        initializer=reparam_initializer, regularizer=reparam_regularizer)
     return reparam(var)