Avoid retaining two copies of each constant in `ConstantOp`.

Presently, the kernel keeps two copies of the constant tensor value, which can be large:
1. In the `ConstantOp::tensor_` field.
2. In the `OpKernel::def_` field (as an attr of the `NodeDef`).

Since we can be sure that `ConstantOp` will never need to access the
tensor value from `OpKernel::def_`, this change introduces a mechanism
for `OpKernel` implementations to store a stripped `NodeDef` in the
base class, and uses it in `ConstantOp` to avoid storing the tensor
value attr.

PiperOrigin-RevId: 184455793

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/core/framework/op_kernel.h b/tensorflow/core/framework/op_kernel.h
index b72f1405cf..a3dc96b3de 100644
--- a/tensorflow/core/framework/op_kernel.h
+++ b/tensorflow/core/framework/op_kernel.h
@@ -75,6 +75,14 @@ class OpKernel {
   // OpKernel won't be instantiated by the scheduler, so you may perform
   // expensive initialization in the descendant's constructor.
   explicit OpKernel(OpKernelConstruction* context);
+
+  // Specialized constructor that enables the descendant to provide a different
+  // `NodeDef` value. For example, this constructor can be used to provide a
+  // stripped-down `NodeDef` that does not contain the full set of attrs (such
+  // as tensor values) if the descendant stores them in a different form.
+  explicit OpKernel(OpKernelConstruction* context,
+                    std::unique_ptr<const NodeDef> node_def);
+
   virtual ~OpKernel();
 
   // An OpKernel's computation can be either synchronous or