Fixing Fatal Process Aborts in TensorFlow Lookup Tables

Few things in Python development are more disruptive than a sudden SIGABRT or core dumped error that crashes the entire Python interpreter. In managed environments, invalid inputs should raise catchable exceptions like ValueError or tf.errors.InvalidArgumentError. When a C++ kernel assertion fails instead, your process dies instantly without giving your application a chance to recover.

While investigating open issues in the tensorflow/tensorflow repository, I found an issue (#125503) where calling lookup table export operations with mismatched data types triggered an immediate crash.

Here is the breakdown of why it happened, how to trace it in TensorFlow's C++ kernels, and how we resolved it.

Consider a standard tf.lookup.StaticHashTable initialized with int64 keys and values:

import tensorflow as tf

table = tf.lookup.StaticHashTable(
    tf.lookup.KeyValueTensorInitializer(
        tf.constant([0], dtype=tf.int64),
        tf.constant([1], dtype=tf.int64),
    ),
    default_value=-1,
)

tf.raw_ops.LookupTableExportV2(
    table_handle=table.resource_handle,
    Tkeys=tf.string,
    Tvalues=tf.string,
)

Running this script immediately terminates the Python process:

Check failed: dtype() == expected_dtype (7 vs. 9)
*** Check failure stack trace: ***
Aborted (core dumped)

The error 7 vs. 9 refers to TensorFlow's internal enum data types: DT_INT64 (7) versus DT_STRING (9).

In TensorFlow, lookup tables are managed in tensorflow/core/kernels/lookup_table_op.cc.

When you interact with a lookup table via operations like LookupTableFind, LookupTableInsert, or LookupTableImport, the C++ kernel first performs signature verification:

// LookupTableFindOp validates both inputs and outputs
DataTypeVector expected_inputs = {expected_input_0_, table->key_dtype(),
                                  table->value_dtype()};
DataTypeVector expected_outputs = {table->value_dtype()};
OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));

ctx->MatchSignature(...) checks whether the data types requested by the Op match the actual runtime types of the underlying table. If there is a mismatch, it returns an errors::InvalidArgument status, which safely bubbles up to Python as an InvalidArgumentError.

However, looking at LookupTableExportOp:

class LookupTableExportOp : public LookupTableOpKernel {
 public:
  using LookupTableOpKernel::LookupTableOpKernel;

  void Compute(OpKernelContext* ctx) override {
    lookup::LookupInterface* table;
    OP_REQUIRES_OK(ctx, GetTable(ctx, &table));
    core::ScopedUnref unref_me(table);

    // Bypassed signature matching entirely!
    OP_REQUIRES_OK(ctx, table->ExportValues(ctx));
  }
};

LookupTableExportOp never verified signatures. It immediately called table->ExportValues(ctx), which attempts to allocate the output tensors using the table's internal key_dtype() and value_dtype().

During allocation, OpKernelContext::allocate_output asserts that the allocated tensor's data type matches the Op's output specification using a hard CHECK_EQ assertion:

CHECK_EQ(dtype(), expected_dtype);

Because the assertion is hardcoded in the allocator, any type mismatch triggered a process-level abort (SIGABRT) rather than returning an error status.

The fix was straightforward and matches the pattern established across all other lookup table kernels in lookup_table_op.cc:

class LookupTableExportOp : public LookupTableOpKernel {
 public:
  using LookupTableOpKernel::LookupTableOpKernel;

  void Compute(OpKernelContext* ctx) override {
    lookup::LookupInterface* table;
    OP_REQUIRES_OK(ctx, GetTable(ctx, &table));
    core::ScopedUnref unref_me(table);

    DataTypeVector expected_inputs = {expected_input_0_};
    DataTypeVector expected_outputs = {table->key_dtype(),
                                       table->value_dtype()};
    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));

    OP_REQUIRES_OK(ctx, table->ExportValues(ctx));
  }
};

With MatchSignature in place, any attempt to export a table with invalid Tkeys or Tvalues is intercepted before tensor allocation.

We added a unit test in tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py:

def testExportSignatureMismatch(self, is_anonymous):
  if is_anonymous and not tf2.enabled():
    self.skipTest(SKIP_ANONYMOUS_IN_TF1_REASON)
  table = self.getHashTable()(
      lookup_ops.KeyValueTensorInitializer(
          constant_op.constant([0], dtype=dtypes.int64),
          constant_op.constant([1], dtype=dtypes.int64)),
      -1,
      experimental_is_anonymous=is_anonymous)
  self.initialize_table(table)
  with self.assertRaises((errors_impl.InvalidArgumentError, ValueError)):
    self.evaluate(
        gen_lookup_ops.lookup_table_export_v2(
            table.resource_handle,
            Tkeys=dtypes.string,
            Tvalues=dtypes.string))

Instead of crashing, the test cleanly verifies that Python catches tf.errors.InvalidArgumentError.

The change is submitted in Pull Request #125856.

CHECK macros protect memory integrity, but they should be impossible for user inputs to trigger.MatchSignature, any operation omitting it is likely an oversight that risks hard crashes.