code: 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)
  
imports: import os
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import tensorflow
  from tensorflow_transform import analyzers
  from tensorflow_transform import common
  from tensorflow_transform import common_types
  from tensorflow_transform import gaussianization
  from tensorflow_transform import schema_inference
  from tensorflow_transform import tf_utils
  from tfx_bsl.types import tfx_namedtuple
label_desc: Use image node_id 34 for calling a global function with example usage: _apply_buckets_with_keys(x, key, key_vocab, bucket_boundaries, scale_factor_per_key, shift_per_key, num_buckets, name) and returns: compose_result_fn