Tensorflow embedded column throwing exception based on the size of vocabulary

Question

I'm stuck with a very weird behavior in tensorflow while trying to use tf.feature_column.embedding_column API. While I don't think is relevant I'm generating the input data via tf.data.Dataset.from_generator.

Now quickly jumping to the issue. When I'm trying to use the feature_column it works or fails based on size of vocabulary size. Code snippet for creating feature column

    for colmn_name in indicator_colms:
        feature_col = feature_column.categorical_column_with_vocabulary_list(
            colmn_name, unique_values(colmn_name))
        indicator_column = feature_column.indicator_column(feature_col)
        feature_columns.append(indicator_column)

For small vocabulary sizes everything works fine but the moment I try to add more elements to my vocabulary I weirdly get the following exception:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/keras/feature_column/dense_features.py in call(self, features, cols_to_output_tensors, training)
    165           tensor = column.get_dense_tensor(
--> 166               transformation_cache, self._state_manager, training=training)
    167         except TypeError:

TypeError: get_dense_tensor() got an unexpected keyword argument 'training'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get(self, key, state_manager, training)
   2457       transformed = column.transform_feature(
-> 2458           self, state_manager, training=training)
   2459     except TypeError:

TypeError: transform_feature() got an unexpected keyword argument 'training'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get(self, key, state_manager, training)
   2457       transformed = column.transform_feature(
-> 2458           self, state_manager, training=training)
   2459     except TypeError:

TypeError: transform_feature() got an unexpected keyword argument 'training'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-69-8136fbe34af6> in <module>
      2     feature_columns = _create_feature_columns()
      3     feature_layer = tf.keras.layers.DenseFeatures(feature_columns)
----> 4     demo(feature_columns)

<ipython-input-10-65d9d3c8db7a> in demo(feature_column)
      1 def demo(feature_column):
      2     demo_feature_layer = layers.DenseFeatures(feature_column)
----> 3     print(demo_feature_layer(example_batch).numpy())

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/keras/feature_column/dense_features.py in call(self, features, cols_to_output_tensors, training)
    167         except TypeError:
    168           tensor = column.get_dense_tensor(transformation_cache,
--> 169                                            self._state_manager)
    170         processed_tensors = self._process_dense_tensor(column, tensor)
    171         if cols_to_output_tensors is not None:

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get_dense_tensor(self, transformation_cache, state_manager)
   4301     # Feature has been already transformed. Return the intermediate
   4302     # representation created by transform_feature.
-> 4303     return transformation_cache.get(self, state_manager)
   4304 
   4305   @deprecation.deprecated(_FEATURE_COLUMN_DEPRECATION_DATE,

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get(self, key, state_manager, training)
   2458           self, state_manager, training=training)
   2459     except TypeError:
-> 2460       transformed = column.transform_feature(self, state_manager)
   2461     if transformed is None:
   2462       raise ValueError('Column {} is not supported.'.format(column.name))

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in transform_feature(self, transformation_cache, state_manager)
   4238     """
   4239     id_weight_pair = self.categorical_column.get_sparse_tensors(
-> 4240         transformation_cache, state_manager)
   4241     return self._transform_id_weight_pair(id_weight_pair,
   4242                                           self.variable_shape[-1])

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get_sparse_tensors(self, transformation_cache, state_manager)
   3725     """See `CategoricalColumn` base class."""
   3726     return CategoricalColumn.IdWeightPair(
-> 3727         transformation_cache.get(self, state_manager), None)
   3728 
   3729   @deprecation.deprecated(_FEATURE_COLUMN_DEPRECATION_DATE,

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in get(self, key, state_manager, training)
   2458           self, state_manager, training=training)
   2459     except TypeError:
-> 2460       transformed = column.transform_feature(self, state_manager)
   2461     if transformed is None:
   2462       raise ValueError('Column {} is not supported.'.format(column.name))

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in transform_feature(self, transformation_cache, state_manager)
   3703     input_tensor = _to_sparse_input_and_drop_ignore_values(
   3704         transformation_cache.get(self.key, state_manager))
-> 3705     return self._transform_input_tensor(input_tensor, state_manager)
   3706 
   3707   @deprecation.deprecated(_FEATURE_COLUMN_DEPRECATION_DATE,

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/feature_column/feature_column_v2.py in _transform_input_tensor(self, input_tensor, state_manager)
   3691             num_oov_buckets=self.num_oov_buckets,
   3692             dtype=key_dtype,
-> 3693             name=name)
   3694       if state_manager is not None:
   3695         state_manager.add_resource(self, name, table)

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/ops/lookup_ops.py in index_table_from_tensor(vocabulary_list, num_oov_buckets, default_value, hasher_spec, dtype, name)
   1507 
   1508   with ops.name_scope(name, "string_to_index"):
-> 1509     keys = ops.convert_to_tensor(vocabulary_list)
   1510     if keys.dtype.is_integer != dtype.is_integer:
   1511       raise ValueError(

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/profiler/trace.py in wrapped(*args, **kwargs)
    161         with Trace(trace_name, **trace_kwargs):
    162           return func(*args, **kwargs)
--> 163       return func(*args, **kwargs)
    164 
    165     return wrapped

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in convert_to_tensor(value, dtype, name, as_ref, preferred_dtype, dtype_hint, ctx, accepted_result_types)
   1564 
   1565     if ret is None:
-> 1566       ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
   1567 
   1568     if ret is NotImplemented:

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py in _constant_tensor_conversion_function(v, dtype, name, as_ref)
    344                                          as_ref=False):
    345   _ = as_ref
--> 346   return constant(v, dtype=dtype, name=name)
    347 
    348 

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py in constant(value, dtype, shape, name)
    270   """
    271   return _constant_impl(value, dtype, shape, name, verify_shape=False,
--> 272                         allow_broadcast=True)
    273 
    274 

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py in _constant_impl(value, dtype, shape, name, verify_shape, allow_broadcast)
    281       with trace.Trace("tf.constant"):
    282         return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
--> 283     return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    284 
    285   g = ops.get_default_graph()

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py in _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    306 def _constant_eager_impl(ctx, value, dtype, shape, verify_shape):
    307   """Creates a constant on the current device."""
--> 308   t = convert_to_eager_tensor(value, ctx, dtype)
    309   if shape is None:
    310     return t

~/anaconda3/envs/amazonei_tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py in convert_to_eager_tensor(value, ctx, dtype)
    104       dtype = dtypes.as_dtype(dtype).as_datatype_enum
    105   ctx.ensure_initialized()
--> 106   return ops.EagerTensor(value, ctx.device_name, dtype)
    107 
    108 

ValueError: Can't convert Python sequence with mixed types to Tensor.

I'd absolutely appreciate any help on this. P.S. I'm running this on sagemaker I've tried training it as a training job as well as on the notebook. The behavior is consistent.