def assign_buckets(x: tf.Tensor,
                   bucket_boundaries: tf.Tensor,
                   side: Side = Side.LEFT) -> tf.Tensor:
  """Assigns every value in x to a bucket index defined by bucket_boundaries.

  Note that `x` and `bucket_boundaries` will be cast to a common type that can
  hold the largest of values.

  Args:
    x: a `Tensor` of values to be bucketized.
    bucket_boundaries:  The bucket boundaries `Tensor`. Note that the boundaries
      are going to be flattened.
    side: Controlls index of a bucket that is being assigned: LEFT means that
      a value is going to be assigned index of the rightmost boundary such that
      boundary <= value; RIGHT means that a value is assigned index of the
      leftmost boundary such that value < boundary.

  Returns:
    A `Tensor` of dtype int64 with the same shape as `x`, and 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(None, 'assign_buckets'):
    flat_x = tf.reshape(x, [-1])
    flat_boundaries = tf.reshape(bucket_boundaries, [-1])

    # Cast values or boundaries to the "largest" dtype to avoid truncating
    # larger values and avoid casting if dtypes are the same.
    if flat_x.dtype.max > flat_boundaries.dtype.max:
      flat_boundaries = tf.cast(flat_boundaries, flat_x.dtype)
    else:
      flat_x = tf.cast(flat_x, flat_boundaries.dtype)

    if side == Side.LEFT:
      # Ignore the last boundary to replicate behavior of the previously used
      # `BoostedTreesBucketize` for backwards compatibility.
      flat_boundaries = flat_boundaries[:-1]

    buckets = tf.searchsorted(
        flat_boundaries, flat_x, side=side.value, out_type=tf.int64)
    return tf.reshape(buckets, tf.shape(x))
Assigns every value in x to a bucket index defined by bucket_boundaries.

  Note that `x` and `bucket_boundaries` will be cast to a common type that can
  hold the largest of values.

  Args:
    x: a `Tensor` of values to be bucketized.
    bucket_boundaries:  The bucket boundaries `Tensor`. Note that the boundaries
      are going to be flattened.
    side: Controlls index of a bucket that is being assigned: LEFT means that
      a value is going to be assigned index of the rightmost boundary such that
      boundary <= value; RIGHT means that a value is assigned index of the
      leftmost boundary such that value < boundary.

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