def infer_feature_schema(
    features: Mapping[str, common_types.TensorType],
    graph: tf.compat.v1.Graph,
    session: Optional[tf.compat.v1.Session] = None) -> schema_pb2.Schema:
  """Given a dict of tensors, creates a `Schema`.

  Infers a schema, in the format of a tf.Transform `Schema`, for the given
  dictionary of tensors.

  If there is an override specified, we override the inferred schema for the
  given feature's tensor.  An override has the meaning that we should set
  is_categorical=True.  If session is not provided then we just set
  is_categorical=True, and if the session is provided then was also compute
  values of the tensors representing the min and max values and set them in the
  schema.

  If annotations have been specified, they are added to the output schema.

  Args:
    features: A dict mapping column names to `Tensor`, `SparseTensor` or
      `RaggedTensor`. The `Tensor`, `SparseTensor` or `RaggedTensor` should have
      a 0'th dimension which is interpreted as the batch dimension.
    graph: A `tf.Graph` used to determine schema overrides.
    session: (optional) A `tf.Session` used to compute schema overrides.  If
      None, schema overrides will not be computed. It is assumed that if all
      analyzers have been evaluated, `session` is passed to this API.

  Returns:
    A `Schema` proto.
  """
  tensor_ranges = _get_tensor_ranges(graph)
  if session is None:
    tensor_ranges = {hashable: (None, None) for hashable in tensor_ranges}
    tensor_annotations = {}
    global_annotations = []
  else:
    tensor_ranges = session.run(tensor_ranges)
    tensor_annotations, global_annotations = _get_schema_annotations(
        graph, session)
  sparse_output_annotations = _get_sparse_output_annotations_v1(graph, session)
  modified_sparse_output_annotations = {}
  modified_tensor_ranges = {}
  feature_annotations = {}
  for name, tensor in features.items():
    hashable_values = tf_utils.hashable_tensor_or_op(
        values := _get_tensor_values(tensor)
    )
    if hashable_values in tensor_ranges:
      assert values.dtype == tf.int64
      modified_tensor_ranges[name] = tensor_ranges[hashable_values]
    if hashable_values in sparse_output_annotations:
      modified_sparse_output_annotations[name] = sparse_output_annotations[
          hashable_values
      ]
    feature_annotations[name] = tensor_annotations.get(hashable_values, [])

  return _infer_feature_schema_common(
      features,
      modified_tensor_ranges,
      feature_annotations,
      global_annotations,
      modified_sparse_output_annotations,
      # A session is passed to compute schema overrides which exist only if all
      # analyzers have been evaluated. Hence, if `session` was provided to this
      # API assume TFT's evaluation is complete.
      is_evaluation_complete=session is not None,
  )
Given a dict of tensors, creates a `Schema`.

  Infers a schema, in the format of a tf.Transform `Schema`, for the given
  dictionary of tensors.

  If there is an override specified, we override the inferred schema for the
  given feature's tensor.  An override has the meaning that we should set
  is_categorical=True.  If session is not provided then we just set
  is_categorical=True, and if the session is provided then was also compute
  values of the tensors representing the min and max values and set them in the
  schema.

  If annotations have been specified, they are added to the output schema.

  Args:
    features: A dict mapping column names to `Tensor`, `SparseTensor` or
      `RaggedTensor`. The `Tensor`, `SparseTensor` or `RaggedTensor` should have
      a 0'th dimension which is interpreted as the batch dimension.
    graph: A `tf.Graph` used to determine schema overrides.
    session: (optional) A `tf.Session` used to compute schema overrides.  If
      None, schema overrides will not be computed. It is assumed that if all
      analyzers have been evaluated, `session` is passed to this API.

  Returns:
    A `Schema` proto.
  
tensorflow_transform/schema_inference.py