Python tensorflow.TypeSpec() Examples

def __init__(self, type_specs: Dict[Text, tf.TypeSpec]):
    """Initializer.

    Args:
      type_specs: a mapping from names of tensors to their TypeSpecs. When
        calling convert(), the dict of tensors passed in must contain the
        same names, and each TensorAlike must be compatible to their TypeSpecs.
    """
    self._handlers = _make_handlers(type_specs)
    all_fields = []
    seen_column_names = set()
    for tensor_name, handler in self._handlers:
      for f in handler.arrow_fields():
        if f.name in seen_column_names:
          raise ValueError("Handler for tensor {} produces a column of a "
                           "conflicting name: {}".format(tensor_name, f.name))
        seen_column_names.add(f.name)
        all_fields.append(f)
    self._arrow_schema = pa.schema(all_fields) 

def convert(self, tensors: Dict[Text, TensorAlike]) -> pa.RecordBatch:
    """Converts a dict of tensors to a RecordBatch.

    Args:
      tensors: must contain the same keys as the dict passed to the initialier.
        and each TensorAlike must be compatible with the corresponding TypeSpec.

    Returns:
      a RecordBatch, whose schema equals to self.arrow_schema().
    """
    assert len(self._handlers) == len(tensors)
    arrays = []
    for tensor_name, handler in self._handlers:
      arrays.extend(handler.convert(tensors[tensor_name]))

    return pa.record_batch(arrays, schema=self._arrow_schema) 

def GetTensor(self, record_batch: pa.RecordBatch,
                produce_eager_tensors: bool) -> Any:
    """Converts the RecordBatch to Tensor or CompositeTensor.

    The result must be of the same (not only compatible) TypeSpec as
    self.type_spec.

    Args:
      record_batch: a RecordBatch that is of the same Schema as what was
        passed at initialization time.
      produce_eager_tensors: if True, returns Eager Tensors, otherwise returns
        ndarrays or Tensor value objects.

    Returns:
      A Tensor or a CompositeTensor. Note that their types may vary depending
      on whether the TF eager mode is on.
    """ 

def _make_handlers(
    type_specs: Dict[Text, tf.TypeSpec]) -> List[Tuple[Text, _TypeHandler]]:
  return [
      (tensor_name, _get_handler(tensor_name, type_spec))
      for tensor_name, type_spec in sorted(type_specs.items())
  ] 

def output_type_specs(self) -> Dict[Text, tf.TypeSpec]:
    """Returns the tf.TypeSpecs of the decoded tensors.

    Returns:
      A dict whose keys are the same as keys of the dict returned by
      `decode_record()` and values are the tf.TypeSpec of the corresponding
      (composite) tensor.
    """
    return {
        k: v._type_spec for k, v in  # pylint: disable=protected-access
        self.decode_record.get_concrete_function().structured_outputs.items()
    } 

def type_spec(self) -> tf.TypeSpec:
    return typing.cast(tf.TypeSpec,
                       tf.SparseTensorSpec(tf.TensorShape([None] + self._shape),
                                           self._dtype)) 

def type_spec(self) -> tf.TypeSpec:
    return typing.cast(tf.TypeSpec, tf.SparseTensorSpec(
        tf.TensorShape([None, None]), self._dtype)) 

def type_spec(self) -> tf.TypeSpec:
    # TF's type stub is not correct about TypeSpec and its sub-classes.
    return typing.cast(tf.TypeSpec, tf.TensorSpec(self._shape, self._dtype)) 

def type_spec(self) -> tf.TypeSpec:
    """Returns the TypeSpec of the converted Tensor or CompositeTensor.""" 

def OriginalTypeSpecs(self) -> Dict[Text, tf.TypeSpec]:
    """Returns the origin's type specs.

    A TFXIO 'Y' may be a result of projection of another TFXIO 'X', in which
    case then 'X' is the origin of 'Y'. And this method returns what
    X.TensorAdapter().TypeSpecs() would return.

    May equal to `self.TypeSpecs()`.

    Returns: a mapping from tensor names to `tf.TypeSpec`s.
    """
    return self._original_type_specs 

def __new__(cls,
              arrow_schema: pa.Schema,
              tensor_representations: TensorRepresentations,
              original_type_specs: Optional[Dict[Text, tf.TypeSpec]] = None):
    return super(TensorAdapterConfig, cls).__new__(
        cls, arrow_schema, tensor_representations, original_type_specs) 

def _get_handler(
    tensor_name: Text, type_spec: tf.TypeSpec) -> _TypeHandler:
  """Returns a TypeHandler that can handle `type_spec`."""
  for handler_cls in _ALL_HANDLERS_CLS:
    if handler_cls.can_handle(type_spec):
      return handler_cls(tensor_name, type_spec)
  raise ValueError(
      "No handler found for tensor {} of spec {}. "
      "Note that tensors with dtype == tf.bool cannot be handled in general -- "
      "consider casting them to tf.uint8."
      .format(tensor_name, type_spec)) 

def batched_placeholders_from_specs(specs):
  """Returns placeholders for the given tf.TypeSpecs or feature specs.

  Args:
    specs: a Dict[Text, Union[tf.TypeSpec, FeatureSpec]]. Note that the values
      in this dict must be of the same type. Mixing is not allowed.

  Returns:
    A dictionary from strings to `Tensor` or `SparseTensor`s.

  Raises:
    ValueError: when the TypeSpec or feature spec has an unsupported dtype.
  """
  if not (all([_is_feature_spec(s) for s in six.itervalues(specs)]) or
          all([isinstance(s, tf.TypeSpec) for s in six.itervalues(specs)])):
    raise TypeError('Specs must be all tf.TypeSpecs or feature specs. '
                    'Mixing is not allowed. Got: {}'.format(specs))

  result = {}
  for name, spec in six.iteritems(specs):
    if spec.dtype not in (tf.int64, tf.float32, tf.string):
      raise ValueError('Feature {} ({}, {}) had invalid dtype'
                       .format(name, spec, type(spec)))
    if isinstance(spec, tf.TypeSpec):
      result[name] = _batched_placeholder_from_typespec(name, spec)
    else:
      result[name] = _batched_placeholder_from_feature_spec(name, spec)

  return result


# Older TFX versions may refer to this function instead.
# TODO(b/150721482): remove once TFX 0.21.1 is released. 

def _get_type_spec(tensor_alike: TensorAlike):
  """Returns the TypeSpec of a TensorAlike."""
  if isinstance(tensor_alike, tf.Tensor):
    return tf.TensorSpec.from_tensor(tensor_alike)
  elif isinstance(tensor_alike, composite_tensor.CompositeTensor):
    return tensor_alike._type_spec  # pylint:disable=protected-access
  raise TypeError("Not a Tensor or CompositeTensor: {}".format(
      type(tensor_alike))) 

def __init__(self, tensor_name: Text, type_spec: tf.TypeSpec):
    super(_VarLenSparseTensorHandler, self).__init__(tensor_name, type_spec)
    self._values_arrow_type = _tf_dtype_to_arrow_type(type_spec.dtype) 

def can_handle(type_spec: tf.TypeSpec) -> bool:
    """Returns `True` if the handler can handle the given `tf.TypeSpec`.""" 

def __init__(self, tensor_name: Text, type_spec: tf.TypeSpec):
    self._tensor_name = tensor_name
    self._type_spec = type_spec 

def _batched_placeholder_from_typespec(name, typespec):
  """Creates a batched placeholder from a tf.TypeSpec."""
  scope_name = _sanitize_scope_name(name)
  if isinstance(typespec, tf.TensorSpec):
    return tf.compat.v1.placeholder(
        typespec.dtype, typespec.shape, name=scope_name)
  if isinstance(typespec, tf.SparseTensorSpec):
    return tf.compat.v1.sparse_placeholder(
        typespec.dtype,
        # BUGGY TF
        typespec.shape.as_list(),
        name=scope_name)

  raise ValueError('Unsupported typespec: {}({}) for feature {}'.format(
      typespec, type(typespec), name)) 

def _RunInPlaceImpl(
      self, preprocessing_fn: Any,
      metadata: dataset_metadata.DatasetMetadata,
      feature_spec_or_typespecs: Dict[Text, Any],
      transform_output_path: Text) -> _Status:
    """Runs a transformation iteration in-place without looking at the data.

    Args:
      preprocessing_fn: The tf.Transform preprocessing_fn.
      metadata: A DatasetMetadata object for the input data.
      feature_spec_or_typespecs: a Dict[Text, Union[FeatureSpec, tf.TypeSpec]]
      transform_output_path: An absolute path to write the output to.

    Returns:
      Status of the execution.
    """

    absl.logging.debug('Processing an in-place transform')

    raw_metadata_dir = os.path.join(transform_output_path,
                                    tft.TFTransformOutput.RAW_METADATA_DIR)
    metadata_io.write_metadata(metadata, raw_metadata_dir)

    with tf.compat.v1.Graph().as_default() as graph:
      with tf.compat.v1.Session(graph=graph) as sess:

        input_signature = impl_helper.batched_placeholders_from_specs(
            schema_utils.schema_as_feature_spec(
                _GetSchemaProto(metadata)).feature_spec)

        # In order to avoid a bug where import_graph_def fails when the
        # input_map and return_elements of an imported graph are the same
        # (b/34288791), we avoid using the placeholder of an input column as an
        # output of a graph. We do this by applying tf.identity to all inputs of
        # the preprocessing_fn.  Note this applies at the level of raw tensors.
        # TODO(b/34288791): Remove this workaround and use a shallow copy of
        # inputs instead.  A shallow copy is needed in case
        # self._preprocessing_fn mutates its input.
        copied_inputs = impl_helper.copy_tensors(input_signature)

        output_signature = preprocessing_fn(copied_inputs)
        sess.run(tf.compat.v1.global_variables_initializer())
        sess.run(tf.compat.v1.tables_initializer())
        transform_fn_path = os.path.join(transform_output_path,
                                         tft.TFTransformOutput.TRANSFORM_FN_DIR)
        saved_transform_io.write_saved_transform_from_session(
            sess, input_signature, output_signature, transform_fn_path)

        transformed_metadata = dataset_metadata.DatasetMetadata(
            schema=tft.schema_inference.infer_feature_schema(
                output_signature, graph, sess))

    transformed_metadata_dir = os.path.join(
        transform_output_path, tft.TFTransformOutput.TRANSFORMED_METADATA_DIR)
    metadata_io.write_metadata(transformed_metadata, transformed_metadata_dir)

    return _Status.OK()