def supply_missing_tensor(batch_size: int, tensor_shape: tf.TensorShape,
                          tensor_dtype: tf.DType) -> tf.Tensor:
  """Supplies a `tf.Tensor` compatible with `tensor`.

  Supports only string and numeric dtypes.
  Args:
    batch_size: an integer representing the size of the batch returned.
    tensor_shape: a `tf.TensorShape`. The returned tensor will have shape
      compatible with this.
    tensor_dtype: The dtype of the returned tensors.

  Returns:
    A batch of `tf.Tensor` tensors.
  """
  # If tensor rank is 0 or unknown, return a scalar.
  if tensor_shape.ndims is None or tensor_shape.ndims == 0:
    return tf.zeros([], dtype=tensor_dtype)

  input_shape = tensor_shape.as_list()
  result_shape = [input_shape[0] or batch_size]

  for shape in input_shape[1:]:
    if shape is None:
      result_shape = result_shape + [1]
    else:
      result_shape = result_shape + [shape]
  return tf.zeros(result_shape, dtype=tensor_dtype)
Supplies a `tf.Tensor` compatible with `tensor`.

  Supports only string and numeric dtypes.
  Args:
    batch_size: an integer representing the size of the batch returned.
    tensor_shape: a `tf.TensorShape`. The returned tensor will have shape
      compatible with this.
    tensor_dtype: The dtype of the returned tensors.

  Returns:
    A batch of `tf.Tensor` tensors.
  
tensorflow_transform/tf2_utils.py