def _join_boundary_rows(
    boundary_matrix: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
  """Joins boundaries per key, by scaling and shifting them.

  This returns a new list of boundaries which is composed from the given 2-d
  array. For each row we compute a scale factor, and a shift value which are
  used to compute the transformed boundaries, and should be used to transform
  a value before its bucket is computed.

  Neighboring key bucket boundaries have their adjacent boundaries merged into
  one.

  Args:
    boundary_matrix: A 2-d np.array where each row is a list of boundaries for a
      certain key.

  Returns:
    A 4-tuple of (boundaries, scale, shift, num_buckets).
    The returned boundaries is a 1-d np.array of size:
    ((num_buckets - 2) * num_keys) + 1
  """
  boundary_matrix = _make_strictly_increasing_boundaries_rows(boundary_matrix)

  num_buckets = np.array(boundary_matrix.shape[1] + 1, dtype=np.int64)

  # Min boundary for each row.
  min_boundary = np.min(boundary_matrix, axis=1)

  # Max boundary for each row.
  max_boundary = np.max(boundary_matrix, axis=1)

  boundary_difference = max_boundary - min_boundary
  scale = np.divide(
      1.0,
      boundary_difference,
      out=np.ones_like(boundary_difference),
      where=boundary_difference != 0)

  # Shifts what would shift values so that when applied to min[key_id] we
  # get: min[key_id] * scale[key_id] + shift[key_id] = key_id
  # Therefore shift is defined as:
  # shift[key_id] = key_id -  min[key_id] * scale[key_id]
  shift = np.arange(scale.size, dtype=np.float32) - min_boundary * scale

  scaled_buckets = (
      boundary_matrix[:, 1:] * np.expand_dims(scale, axis=1) +
      np.expand_dims(shift, axis=1))
  boundaries = np.insert(scaled_buckets.flatten(), 0, 0.)

  return boundaries, scale, shift, num_buckets
Joins boundaries per key, by scaling and shifting them.

  This returns a new list of boundaries which is composed from the given 2-d
  array. For each row we compute a scale factor, and a shift value which are
  used to compute the transformed boundaries, and should be used to transform
  a value before its bucket is computed.

  Neighboring key bucket boundaries have their adjacent boundaries merged into
  one.

  Args:
    boundary_matrix: A 2-d np.array where each row is a list of boundaries for a
      certain key.

  Returns:
    A 4-tuple of (boundaries, scale, shift, num_buckets).
    The returned boundaries is a 1-d np.array of size:
    ((num_buckets - 2) * num_keys) + 1
  
tensorflow_transform/beam/analyzer_impls.py