# Path: tensorflow_transform/common_types.py
# def is_ragged_feature_available() -> bool:
# def is_ragged_feature(spec: FeatureSpecType) -> bool:


import copy
import itertools
import tensorflow as tf
from typing import Collection, Iterable, Mapping, Optional, Tuple
from tensorflow_transform import common_types
from tensorflow.python import tf2
from tensorflow.python.framework import ops
from tensorflow.python.framework.func_graph import FuncGraph

  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)


def supply_missing_inputs(
    structured_inputs: Mapping[str, common_types.TensorType],
    batch_size: int,
    missing_keys: Optional[Collection[str]] = None

# Path: tensorflow_transform/common_types.py
# def is_ragged_feature_available() -> bool:
# def is_ragged_feature(spec: FeatureSpecType) -> bool:


# Path: tensorflow_transform/tf2_utils.py
import copy
import itertools
import tensorflow as tf
from typing import Collection, Iterable, Mapping, Optional, Tuple
from tensorflow_transform import common_types
from tensorflow.python import tf2
from tensorflow.python.framework import ops
from tensorflow.python.framework.func_graph import FuncGraph


  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)


def supply_missing_inputs(
    structured_inputs: Mapping[str, common_types.TensorType],
    batch_size: int,
    missing_keys: Optional[Collection[str]] = None

) -> Mapping[str, common_types.TensorType]: