code: def _create_high_dimensional_ragged_dataset(
      self, strategy, include_weights=False, weight=0.5
  ):
      ragged_features = (
          ragged_tensor.RaggedTensor.from_row_lengths(
              row_lengths=self.feature_watched_row_lengths_high_dimensional,
              values=self.feature_watched_values_high_dimensional,
          ),
          ragged_tensor.RaggedTensor.from_row_lengths(
              row_lengths=self.feature_favorited_row_lengths_high_dimensional,
              values=self.feature_favorited_values_high_dimensional,
          ),
          ragged_tensor.RaggedTensor.from_row_lengths(
              row_lengths=self.feature_friends_row_lengths_high_dimensional,
              values=self.feature_friends_values_high_dimensional,
          ),
      )
      if include_weights:
          weights = []
          for ragged in ragged_features:
              values = (
                  array_ops.ones_like(
                      ragged.values, dtype=dtypes.float32
                  )
                  * weight
              )
              weights.append(
                  ragged_tensor.RaggedTensor(
                      row_lengths=ragged.row_lengths(), values=values
                  )
              )
          ragged_features = (ragged_features, tuple(weights))
  
      dataset = dataset_ops.DatasetV2.from_tensors(ragged_features)
      # Data is batched to self.data_batch_size, rebatch to global batch size.
      return (
          dataset.unbatch()
          .repeat()
          .batch(
              self.batch_size * strategy.num_replicas_in_sync,
              drop_remainder=True,
          )
      )
  
imports: import os
  from typing import Tuple
  from absl import flags
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import remote
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import init_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import tpu_embedding_v2
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.util import nest
label_desc: Use image node_id 10 for calling the TPUEmbeddingBaseTest obj's underlying member method code with example usage: obj._create_high_dimensional_ragged_dataset(strategy, include_weights, weight) and returns: dataset