def testFullTypesForFlatTensors(self):
spec = TwoTensorsSpec([5], dtypes.int32, [5, 8], dtypes.float32, "red")
full_type_list = fulltypes_for_flat_tensors(spec)
expect = [
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET),
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET)
]
self.assertEqual(len(spec._flat_tensor_specs), len(full_type_list))
self.assertEqual(expect, full_type_list)
def test_fn():
ta = tensor_array_ops.TensorArray(dtypes.int32, size=1)
h = math_ops.cast(ta.flow, dtypes.variant)
t = full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_PRODUCT,
args=[full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_ARRAY)])
h.op.experimental_set_type(t)
ta = tensor_array_ops.TensorArray(dtypes.int32, flow=h)
ta = ta.write(0, constant_op.constant(1))
return ta.stack()
def full_type_from_spec(element_spec):
"""Returns a FullTypeDef for the element tensor representation.
Args:
element_spec: A nested structure of `tf.TypeSpec` objects representing the
element type specification.
Returns:
A FullTypeDef for the element tensor representation.
"""
args = []
for ts in nest.flatten(element_spec):
if isinstance(ts, NoneTensorSpec):
# NoneTensorSpec does not correspond to an output
continue
if isinstance(ts, sparse_tensor.SparseTensorSpec):
# Currently, this represents a SparseTensor spec as a single ouput (that
# is a variant) as a TFT_TENSOR. When shape information is added to
# fulltype, either the shape needs to reflect this
# (e.g. TFT_TENSOR[..., shape=COOSparseShape]) or a new TFT_SPARSE data
# type should be created.
fts = get_flat_tensor_specs(ts)
if (len(fts) != 1) or (fts[0].dtype != dtypes.variant):
raise TypeError("Only sparse tensors as variants is supported")
type_id = full_type_pb2.TFT_TENSOR
elif isinstance(ts, ragged_tensor.RaggedTensorSpec):
type_id = full_type_pb2.TFT_RAGGED
elif isinstance(ts, tensor_spec.TensorSpec):
type_id = full_type_pb2.TFT_TENSOR
else:
# The intent of defaulting to TFT_UNSET is so other cases can fallback
# to the behavior prior to full type or a reasonable default. Users
# can define their own type specs, so it is important to have a reasonable
# default.
type_id = full_type_pb2.TFT_UNSET
if type_id != full_type_pb2.TFT_UNSET:
args.append(
full_type_pb2.FullTypeDef(
type_id=type_id,
args=[
full_type_pb2.FullTypeDef(type_id=DT_TO_FT.get(
ts.dtype, full_type_pb2.TFT_ANY))
]))
else:
args.append(
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET))
element_type = full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_PRODUCT,
args=args)
return element_type
def testFullTypesForFlatTensors(self):
a = TwoComposites(
ragged_factory_ops.constant([[1, 2], [3]]),
ragged_factory_ops.constant([[5], [6, 7, 8]]))
a_spec = type_spec.type_spec_from_value(a)
full_type_list = fulltypes_for_flat_tensors(a_spec)
expect = [
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)]),
]
self.assertEqual(len(a_spec._flat_tensor_specs), len(full_type_list))
self.assertEqual(expect, full_type_list)
def iterator_full_type_from_spec(element_spec):
"""Returns a FullTypeDef for an iterator for the elements.
Args:
element_spec: A nested structure of `tf.TypeSpec` objects representing the
element type specification.
Returns:
A FullTypeDef for an iterator for the element tensor representation.
"""
args = fulltypes_for_flat_tensors(element_spec)
return full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_PRODUCT,
args=[
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_ITERATOR,
args=[
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_PRODUCT, args=args)
])
])
def _translate_to_fulltype_for_flat_tensors(
spec: type_spec.TypeSpec) -> List[full_type_pb2.FullTypeDef]:
"""Convert a TypeSec to a list of FullTypeDef.
The FullTypeDef created corresponds to the encoding used with datasets
(and map_fn) that uses variants (and not FullTypeDef corresponding to the
default "component" encoding).
Currently, the only use of this is for information about the contents of
ragged tensors, so only ragged tensors return useful full type information
and other types return TFT_UNSET. While this could be improved in the future,
this function is intended for temporary use and expected to be removed
when type inference support is sufficient.
Args:
spec: A TypeSpec for one element of a dataset or map_fn.
Returns:
A list of FullTypeDef corresponding to SPEC. The length of this list
is always the same as the length of spec._flat_tensor_specs.
"""
if isinstance(spec, RaggedTensorSpec):
dt = spec.dtype
elem_t = _DT_TO_FT.get(dt)
if elem_t is None:
logging.vlog(1, "dtype %s that has no conversion to fulltype.", dt)
elif elem_t == full_type_pb2.TFT_LEGACY_VARIANT:
logging.vlog(
1, "Ragged tensors containing variants are not supported.", dt)
else:
assert len(spec._flat_tensor_specs) == 1 # pylint: disable=protected-access
return [
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[full_type_pb2.FullTypeDef(type_id=elem_t)])
]
return [
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET)
for t in spec._flat_tensor_specs # pylint: disable=protected-access
]
def _set_handle_data(list_handle, element_shape, element_dtype):
"""Sets type information on `list_handle` for consistency with graphs."""
# TODO(b/169968286): It would be better if we had a consistent story for
# creating handle data from eager operations (shared with VarHandleOp).
if isinstance(list_handle, ops.EagerTensor):
if tensor_util.is_tf_type(element_shape):
element_shape = tensor_shape.TensorShape(None)
elif not isinstance(element_shape, tensor_shape.TensorShape):
element_shape = tensor_shape.TensorShape(element_shape)
handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData(
)
handle_data.is_set = True
# TODO(b/191472076): This duplicates type inference. Clean up.
handle_data.shape_and_type.append(
cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
shape=element_shape.as_proto(),
dtype=element_dtype.as_datatype_enum,
type=full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_ARRAY)))
list_handle._handle_data = handle_data # pylint: disable=protected-access
def testFullTypesForFlatTensors(self):
spec = tensor_spec.TensorSpec([1], np.float32)
full_type_list = fulltypes_for_flat_tensors(spec)
expect = [full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET)]
self.assertEqual(len(spec._flat_tensor_specs), len(full_type_list))
self.assertEqual(expect, full_type_list)
def testFulTypesForFlatTensors(self):
# Note that the batchable tensor list encoding for a StructuredTensor
# contains a separate tensor for each leaf field.
# In this example, _flat_tensor_specs in class StructuredTensorSpec is
# called three times and it returns results with length 2, 3 and 11
# for "g", "h" and `struct` respectively.
fields = self._lambda_for_fields()
rank = 4
if callable(fields):
fields = fields(
) # deferred construction: fields may include tensors.
struct = StructuredTensor.from_fields_and_rank(fields, rank)
spec = type_spec.type_spec_from_value(struct)
flat_specs = spec._flat_tensor_specs
fulltype = fulltypes_for_flat_tensors(spec)
expected_ft_list = [
# a, b
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET),
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET),
# c, d, e, f
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UINT8)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT)
]),
# g
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)
]),
# h
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)
]),
full_type_pb2.FullTypeDef(
type_id=full_type_pb2.TFT_RAGGED,
args=[
full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32)
]),
]
self.assertEqual(len(expected_ft_list), len(flat_specs))
self.assertEqual(fulltype, expected_ft_list)
def testFullTypesForFlatTensors(self, dt):
st_spec = sparse_tensor.SparseTensorSpec(dtype=dt)
full_type_list = fulltypes_for_flat_tensors(st_spec)
expect = [full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET)]
self.assertEqual(len(st_spec._flat_tensor_specs), len(full_type_list))
self.assertEqual(expect, full_type_list)
def fulltype_list_to_product(fulltype_list):
"""Convert a list of FullType Def into a single FullType Def."""
return full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_PRODUCT,
args=fulltype_list)