code: def bucketize(
      x: common_types.ConsistentTensorType,
      num_buckets: int,
      epsilon: Optional[float] = None,
      weights: Optional[tf.Tensor] = None,
      elementwise: bool = False,
      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` whose values
          should be mapped to buckets.  For a `CompositeTensor` only non-missing
          values will be included in the quantiles computation, and the result of
          `bucketize` will be a `CompositeTensor` with non-missing values mapped to
          buckets. If elementwise=True then `x` must be dense.
        num_buckets: Values in the input `x` are divided into approximately
          equal-sized buckets, where the number of buckets is `num_buckets`.
        epsilon: (Optional) Error tolerance, typically a small fraction close to
          zero. If a value is not specified by the caller, a suitable value is
          computed based on experimental results.  For `num_buckets` less than 100,
          the value of 0.01 is chosen to handle a dataset of up to ~1 trillion input
          data values.  If `num_buckets` is larger, then epsilon is set to
          (1/`num_buckets`) to enforce a stricter error tolerance, because more
          buckets will result in smaller range for each bucket, and so we want the
          boundaries to be less fuzzy. See analyzers.quantiles() for details.
        weights: (Optional) Weights tensor for the quantiles. Tensor must have the
          same shape as x.
        elementwise: (Optional) If true, bucketize each element of the tensor
          independently.
        name: (Optional) A name for this operation.
  
      Returns:
        A `Tensor` 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). Sometimes the actual number of buckets
        can be different than num_buckets hint, for example in case the number of
        distinct values is smaller than num_buckets, or in cases where the
        input values are not uniformly distributed.
        NaN values are mapped to the last bucket. Values with NaN weights are
        ignored in bucket boundaries calculation.
  
      Raises:
        TypeError: If num_buckets is not an int.
        ValueError: If value of num_buckets is not > 1.
        ValueError: If elementwise=True and x is a `CompositeTensor`.
      """
      with tf.compat.v1.name_scope(name, "bucketize"):
          if not isinstance(num_buckets, int):
              raise TypeError(
                  "num_buckets must be an int, got %s"
                  % type(num_buckets)
              )
  
          if num_buckets < 1:
              raise ValueError("Invalid num_buckets %d" % num_buckets)
  
          if (
              isinstance(x, (tf.SparseTensor, tf.RaggedTensor))
              and elementwise
          ):
              raise ValueError(
                  "bucketize requires `x` to be dense if `elementwise=True`"
              )
  
          if epsilon is None:
              # See explanation in args documentation for epsilon.
              epsilon = min(1.0 / num_buckets, 0.01)
  
          x_values = tf_utils.get_values(x)
          bucket_boundaries = analyzers.quantiles(
              x_values,
              num_buckets,
              epsilon,
              weights,
              reduce_instance_dims=not elementwise,
          )
  
          if not elementwise:
              return apply_buckets(x, bucket_boundaries)
  
          num_features = tf.math.reduce_prod(x.get_shape()[1:])
          bucket_boundaries = tf.reshape(
              bucket_boundaries, [num_features, -1]
          )
          x_reshaped = tf.reshape(x, [-1, num_features])
          bucketized = []
          for idx, boundaries in enumerate(
              tf.unstack(bucket_boundaries, axis=0)
          ):
              bucketized.append(
                  apply_buckets(
                      x_reshaped[:, idx],
                      tf.expand_dims(boundaries, axis=0),
                  )
              )
          return tf.reshape(
              tf.stack(bucketized, axis=1),
              [-1] + x.get_shape().as_list()[1:],
          )
  
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 30 for calling a global function with example usage: bucketize(x, num_buckets, epsilon, weights, elementwise, name) and returns: tf, apply_buckets