def _apply_buckets_with_keys(
    x: common_types.ConsistentTensorType,
    key: common_types.ConsistentTensorType,
    key_vocab: tf.Tensor,
    bucket_boundaries: tf.Tensor,
    scale_factor_per_key: tf.Tensor,
    shift_per_key: tf.Tensor,
    num_buckets: int,
    name: Optional[int] = None) -> common_types.ConsistentTensorType:
  """Bucketize an input where boundaries depend on the index.

  Args:
    x: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor`.
    key: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor` with the same size as
      `x`.
    key_vocab: A vocab containing all keys.  Must be exhaustive, an out-of-vocab
      entry in `key` will cause a crash.
    bucket_boundaries: A rank-1 Tensor.
    scale_factor_per_key: A rank-1 Tensor of shape (key_size,).
    shift_per_key: A rank-1 Tensor of shape (key_size,).
    num_buckets: A scalar.
    name: (Optional) A name for this operation.

  Returns:
    A tensor with the same shape as `x` and dtype tf.int64. If any value in
    `key` is not present in `key_vocab` then the resulting bucket will be -1.
  """
  with tf.compat.v1.name_scope(name, 'apply_buckets_with_keys'):
    x_values = tf.cast(tf_utils.get_values(x), tf.float32)
    compose_result_fn = _make_composite_tensor_wrapper_if_composite(x)
    key_values = tf_utils.get_values(key)

    # Convert `key_values` to indices in key_vocab.
    key_indices = tf_utils.lookup_key(key_values, key_vocab)

    adjusted_key_indices = tf.where(
        key_indices < 0, _fill_shape(0, tf.shape(key_indices), tf.int64),
        key_indices)

    # Apply the per-key offsets to x, which produces offset buckets (where the
    # bucket offset is an integer offset).  Then remove this offset to get the
    # actual per-key buckets for x.
    scale_factors = tf.gather(scale_factor_per_key, adjusted_key_indices)
    shifts = tf.gather(shift_per_key, adjusted_key_indices)

    transformed_x = x_values * scale_factors + shifts

    offset_buckets = tf_utils.assign_buckets(
        transformed_x, bucket_boundaries, side=tf_utils.Side.RIGHT)

    max_bucket = num_buckets - 1

    # Shift the bucket numbers back to the correct range [0, num_buckets].
    # We use max_bucket-1 due to different keys sharing 1 boundary.
    corrected_buckets = offset_buckets - (
        (max_bucket - 1) * adjusted_key_indices)
    bucketized_values = tf.clip_by_value(corrected_buckets, 0, max_bucket)

    # Set values with missing keys as -1.
    bucketized_values = tf.where(key_indices < 0, key_indices,
                                 bucketized_values)

    # Attach the relevant metadata to result, so that the corresponding
    # output feature will have this metadata set.
    min_value = tf.constant(0, tf.int64)
    schema_inference.set_tensor_schema_override(
        bucketized_values, min_value, max_bucket)

    return compose_result_fn(bucketized_values)
Bucketize an input where boundaries depend on the index.

  Args:
    x: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor`.
    key: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor` with the same size as
      `x`.
    key_vocab: A vocab containing all keys.  Must be exhaustive, an out-of-vocab
      entry in `key` will cause a crash.
    bucket_boundaries: A rank-1 Tensor.
    scale_factor_per_key: A rank-1 Tensor of shape (key_size,).
    shift_per_key: A rank-1 Tensor of shape (key_size,).
    num_buckets: A scalar.
    name: (Optional) A name for this operation.

  Returns:
    A tensor with the same shape as `x` and dtype tf.int64. If any value in
    `key` is not present in `key_vocab` then the resulting bucket will be -1.
  
tensorflow_transform/mappers.py