def apply_buckets(
    x: common_types.ConsistentTensorType,
    bucket_boundaries: common_types.BucketBoundariesType,
    name: Optional[str] = None) -> common_types.ConsistentTensorType:
  """Returns a bucketized column, with a bucket index assigned to each input.

  Each element `e` in `x` is mapped to a positive index `i` for which
  `bucket_boundaries[i-1] <= e < bucket_boundaries[i]`, if it exists.
  If `e < bucket_boundaries[0]`, then `e` is mapped to `0`. If
  `e >= bucket_boundaries[-1]`, then `e` is mapped to `len(bucket_boundaries)`.
  NaNs are mapped to `len(bucket_boundaries)`.

  Example:

  >>> x = tf.constant([[4.0, float('nan'), 1.0], [float('-inf'), 7.5, 10.0]])
  >>> bucket_boundaries = tf.constant([[2.0, 5.0, 10.0]])
  >>> tft.apply_buckets(x, bucket_boundaries)
  <tf.Tensor: shape=(2, 3), dtype=int64, numpy=
  array([[1, 3, 0],
         [0, 2, 3]])>

  Args:
    x: A numeric input `Tensor`, `SparseTensor`, or `RaggedTensor` whose values
      should be mapped to buckets.  For `CompositeTensor`s, the non-missing
      values will be mapped to buckets and missing value left missing.
    bucket_boundaries: A rank 2 `Tensor` or list representing the bucket
      boundaries sorted in ascending order.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor`, `SparseTensor`, or `RaggedTensor` of the same shape as `x`, with
    each element in the returned tensor representing the bucketized value.
    Bucketized value is in the range [0, len(bucket_boundaries)].
  """
  with tf.compat.v1.name_scope(name, 'apply_buckets'):
    bucket_boundaries = tf.convert_to_tensor(bucket_boundaries)
    tf.compat.v1.assert_rank(bucket_boundaries, 2)

    bucketized_values = tf_utils.assign_buckets(
        tf_utils.get_values(x), bucket_boundaries, side=tf_utils.Side.RIGHT)

    # Attach the relevant metadata to result, so that the corresponding
    # output feature will have this metadata set.
    min_value = tf.constant(0, tf.int64)
    max_value = tf.shape(input=bucket_boundaries)[1]
    schema_inference.set_tensor_schema_override(
        bucketized_values, min_value, max_value)
    _annotate_buckets(bucketized_values, bucket_boundaries)
    compose_result_fn = _make_composite_tensor_wrapper_if_composite(x)
    return compose_result_fn(bucketized_values)
Returns a bucketized column, with a bucket index assigned to each input.

  Each element `e` in `x` is mapped to a positive index `i` for which
  `bucket_boundaries[i-1] <= e < bucket_boundaries[i]`, if it exists.
  If `e < bucket_boundaries[0]`, then `e` is mapped to `0`. If
  `e >= bucket_boundaries[-1]`, then `e` is mapped to `len(bucket_boundaries)`.
  NaNs are mapped to `len(bucket_boundaries)`.

  Example:

  >>> x = tf.constant([[4.0, float('nan'), 1.0], [float('-inf'), 7.5, 10.0]])
  >>> bucket_boundaries = tf.constant([[2.0, 5.0, 10.0]])
  >>> tft.apply_buckets(x, bucket_boundaries)
  <tf.Tensor: shape=(2, 3), dtype=int64, numpy=
  array([[1, 3, 0],
         [0, 2, 3]])>

  Args:
    x: A numeric input `Tensor`, `SparseTensor`, or `RaggedTensor` whose values
      should be mapped to buckets.  For `CompositeTensor`s, the non-missing
      values will be mapped to buckets and missing value left missing.
    bucket_boundaries: A rank 2 `Tensor` or list representing the bucket
      boundaries sorted in ascending order.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor`, `SparseTensor`, or `RaggedTensor` of the same shape as `x`, with
    each element in the returned tensor representing the bucketized value.
    Bucketized value is in the range [0, len(bucket_boundaries)].
  
tensorflow_transform/mappers.py