mlir::TF::CopyDeviceAndUnderscoredAttributes(call_op, cluster);
  cluster->setAttr(mlir::TF::kClusterOutlinedFunctionNameAttr, callee_name);
  cluster->setAttr(mlir::TF::kAllowSoftPlacementAttr,
                   builder.getBoolAttr(true));
}

// Encapsulate the first partitioned call that can be reached from
// `func` and is with compilation markers in a device cluster. For nested calls,

    case AttrValue::kS:
      return builder->getStringAttr(value.s());
    case AttrValue::kF:
      return builder->getFloatAttr(builder->getF32Type(), value.f());
    case AttrValue::kB:
      return builder->getBoolAttr(value.b());
    case AttrValue::kType: {
      mlir::Type type;
      TF_RETURN_IF_ERROR(ConvertDataType(value.type(), *builder, &type));
      return mlir::TypeAttr::get(type);
    }
    case AttrValue::kShape:

                builder.getBoolAttr(true));
  ConstOp shape = builder.create<ConstOp>(loc, shape_attr);
  shape->setAttr(kICIWeightDistributionMlirBridgeMarker,
                 builder.getBoolAttr(true));
  FillOp fill = builder.create<FillOp>(loc, shape, zero);
  fill->setAttr(kICIWeightDistributionMlirBridgeMarker,
                builder.getBoolAttr(true));
  return fill;
}

      auto tf_reduce_op = rewriter.create<BooleanReduce>(
          reduce_op.getLoc(), reduce_op->getResult(0).getType(), operand,
          reduction_indices,
          /*keep_dim=*/rewriter.getBoolAttr(false));
      auto tf_argreduce_op = rewriter.create<ArgReduce>(
          reduce_op.getLoc(), reduce_op->getResult(1).getType(), operand,
          reduction_indices);

    auto num_parallel_calls_op = rewriter.create<CastOp>(
        op.getLoc(), UnrankedTensorType::get(rewriter.getIntegerType(64)),
        batchInputOp.getNumParallelCalls(), rewriter.getBoolAttr(false));

    if (op.getMetadata() != batchInputOp.getMetadata()) {
      return failure();
    }

    auto fused_op = rewriter.create<MapAndBatchDatasetOp>(

    Type expected = fn_type.getInput(i);
    if (val.getType() != expected) {
      val =
          builder->create<TF::CastOp>(loc, expected, val,
                                      /*Truncate=*/builder->getBoolAttr(false));
    }
    operands.push_back(val);
  }
  return builder->create<func::CallOp>(loc, fn, operands).getOperation();
}

// Prepares for jump to the given block by introducing necessary tensor_cast

                                               rewriter.getF32Type());
    ::mlir::Value mean = rewriter.create<TF::MeanOp>(
        bn_op.getLoc(), mean_var_type, inputs, reduce_dim_op,
        /*keep_dims=*/rewriter.getBoolAttr(false));

    // Compute variance
    Value shape_value =
        getShapeValue(bn_op.getLoc(), bn_op.getOperand(), rewriter);
    auto broadcast_mean = broadcastToFeatureDim(

    return rewriter.createOrFold<TF::CastOp>(loc, new_type, val,
                                             rewriter.getBoolAttr(false));
  }
  return rewriter.createOrFold<TF::CastOp>(
      loc, UnrankedTensorType::get(new_ele_type), val,
      rewriter.getBoolAttr(false));
}

// Utility function to-
// 1. Create a tfl.const op with an int32_t values, from an MLIR Value, if the

        /*fused_activation_function=*/rewriter.getStringAttr("NONE"),
        /*weights_format=*/rewriter.getStringAttr("DEFAULT"),
        /*keep_num_dims=*/rewriter.getBoolAttr(true),
        /*asymmetric_quantize_inputs=*/mlir::BoolAttr());
    rewriter.replaceOp(bmm_op, {fc_op.getResult(0)});

    return success();
  };
};

  }

  if constexpr (is_input) {
    auto pi = builder.create<TF::TPUPartitionedInputV2Op>(
        op.getLoc(), op.getType(), op.getOperands(),
        builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),