auto permuted = rewriter.create<TransposeOp>(loc, reshaped, permutation);
    auto output_batch = padded_shape_splits[0];
    for (int64_t i = 0; i < block_rank; ++i) {
      output_batch =
          rewriter.create<MulOp>(loc, output_batch, block_shape_splits[i]);
    }
    SmallVector<Value, 4> output_shape_vals{output_batch};
    for (int64_t i = 0; i < block_rank; ++i) {

  // CHECK-DAG: [[OUTPUT_BATCH_PART:%.+]] = "tf.Mul"([[PADDED_SHAPE_SPLITS]]#0, [[BLOCK_SHAPE_SPLITS]]#0)
  // CHECK-DAG: [[OUTPUT_BATCH:%.+]] = "tf.Mul"([[OUTPUT_BATCH_PART]], [[BLOCK_SHAPE_SPLITS]]#1)
  // CHECK-DAG: [[OUTPUT_SHAPE:%.+]] = "tf.ConcatV2"([[OUTPUT_BATCH]], [[OUTER_SHAPE_0]], [[OUTER_SHAPE_1]], [[PADDED_SHAPE_SPLITS]]#3, [[ZERO_I64]])
  // CHECK-DAG: [[RESHAPED:%.+]] = "tf.Reshape"([[PADDED]], [[RESHAPED_SHAPE]])

  def _create_and_save_vocab_table_lookup_model_tf1(
      self,
      output_path: str,
      tags: Collection[str],
      signature_def_key: str,
  ) -> Tuple[Mapping[str, core.Tensor], Mapping[str, core.Tensor]]:
    """Creates and saves a simple model that uses a vocab table.

    Args:
      output_path: Path to the directory to save the created model.