def bucketize_per_key(
    x: common_types.ConsistentTensorType,
    key: common_types.ConsistentTensorType,
    num_buckets: int,
    epsilon: Optional[float] = None,
    weights: Optional[common_types.ConsistentTensorType] = None,
    name: Optional[str] = None) -> common_types.ConsistentTensorType:
  """Returns a bucketized column, with a bucket index assigned to each input.

  Args:
    x: A numeric input `Tensor`, `SparseTensor`, or `RaggedTensor` with rank 1,
      whose values should be mapped to buckets.  `CompositeTensor`s will have
      their non-missing values mapped and missing values left as missing.
    key: A `Tensor`, `SparseTensor`, or `RaggedTensor` with the same shape as
      `x` and dtype tf.string.  If `x` is a `CompositeTensor`, `key` must
      exactly match `x` in everything except values, i.e. indices and
      dense_shape or nested row splits must be identical.
    num_buckets: Values in the input `x` are divided into approximately
      equal-sized buckets, where the number of buckets is num_buckets.
    epsilon: (Optional) see `bucketize`.
    weights: (Optional) A `Tensor`, `SparseTensor`, or `RaggedTensor` with the
      same shape as `x` and dtype tf.float32. Used as weights for quantiles
      calculation. If `x` is a `CompositeTensor`, `weights` must exactly match
      `x` in everything except values.
    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, actual_num_buckets). If the computed
    key vocabulary doesn't have an entry for `key` then the resulting bucket is
    -1.

  Raises:
    ValueError: If value of num_buckets is not > 1.
  """
  with tf.compat.v1.name_scope(name, 'bucketize_per_key'):
    if not isinstance(num_buckets, int):
      raise TypeError(
          'num_buckets must be an int, got {}'.format(type(num_buckets)))

    if num_buckets < 1:
      raise ValueError('Invalid num_buckets {}'.format(num_buckets))

    if epsilon is None:
      # See explanation in args documentation for epsilon.
      epsilon = min(1.0 / num_buckets, 0.01)

    (key_vocab, bucket_boundaries, scale_factor_per_key, shift_per_key,
     actual_num_buckets) = (
         analyzers._quantiles_per_key(  # pylint: disable=protected-access
             tf_utils.get_values(x),
             tf_utils.get_values(key),
             num_buckets,
             epsilon,
             weights=tf_utils.get_values(weights)))
    return _apply_buckets_with_keys(x, key, key_vocab, bucket_boundaries,
                                    scale_factor_per_key, shift_per_key,
                                    actual_num_buckets)
Returns a bucketized column, with a bucket index assigned to each input.

  Args:
    x: A numeric input `Tensor`, `SparseTensor`, or `RaggedTensor` with rank 1,
      whose values should be mapped to buckets.  `CompositeTensor`s will have
      their non-missing values mapped and missing values left as missing.
    key: A `Tensor`, `SparseTensor`, or `RaggedTensor` with the same shape as
      `x` and dtype tf.string.  If `x` is a `CompositeTensor`, `key` must
      exactly match `x` in everything except values, i.e. indices and
      dense_shape or nested row splits must be identical.
    num_buckets: Values in the input `x` are divided into approximately
      equal-sized buckets, where the number of buckets is num_buckets.
    epsilon: (Optional) see `bucketize`.
    weights: (Optional) A `Tensor`, `SparseTensor`, or `RaggedTensor` with the
      same shape as `x` and dtype tf.float32. Used as weights for quantiles
      calculation. If `x` is a `CompositeTensor`, `weights` must exactly match
      `x` in everything except values.
    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, actual_num_buckets). If the computed
    key vocabulary doesn't have an entry for `key` then the resulting bucket is
    -1.

  Raises:
    ValueError: If value of num_buckets is not > 1.
  
tensorflow_transform/mappers.py