// supported.
  const auto& device_list =
      devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))->getSecond();

  // Create the state variable for each device.
  for (llvm::StringRef device : device_list) {
    state_vars.push_back(builder->create<TF::VarHandleOp>(
        loc,
        llvm::ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(llvm::ArrayRef<TensorType>{key_type},

    auto static_shape_attr =
        mlir::DenseIntElementsAttr::get(static_shape_type, static_shape);
    return rewriter.create<TF::ConstOp>(loc, static_shape_attr).getOutput();
  }

  // If the shape is not static, create a new ShapeOp.
  BoolAttr false_attr = rewriter.getBoolAttr(false);
  return rewriter
      .create<TF::ShapeOp>(loc, input,
                           /*use_32bit=*/false_attr)
      .getOutput();
}

        return false;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(op->getLoc(), cast_type,
                                                        op.getResult());
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),
                                                           expressed_type, q);

    TF_RETURN_IF_ERROR(c_event->Create());
    return std::move(c_event);
  }

  absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
      std::optional<std::variant<StreamPriority, int>> priority =
          std::nullopt) override {
    auto stream = std::make_unique<CStream>(&device_, stream_executor_, this);
    TF_RETURN_IF_ERROR(stream->Create());
    return std::move(stream);
  }

 private:

    IntegerAttr zp_attr;
    if (!matchPattern(zp, m_Constant(&zp_attr))) {
      return failure();
    }
    rewriter.setInsertionPoint(zp.getDefiningOp());
    auto zp_tensor = rewriter.create<TF::ConstOp>(
        loc, RankedTensorType::get({}, zp.getType()), zp_attr);
    auto zp_cast = rewriter.create<CastOp>(

  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module_ref->getBodyRegion());

  llvm::SmallVector<mlir::Type, 8> result_types;
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kNumCoresPerReplicaAttr,

            /*quant_dim=*/-1, num_bits, narrow_range, is_signed);
        builder.setInsertionPoint(block, insertion_point);
        auto q_op = builder.create<quantfork::QuantizeCastOp>(
            loc, params.getValue(), arg);
        auto dq_op = builder.create<quantfork::DequantizeCastOp>(
            loc, input_type, q_op.getResult());
        arg.replaceAllUsesWith(dq_op.getResult());
        q_op.setOperand(arg);
      }

  if (!dq_op) {
    auto mul_op = builder.create<TF::MulOp>(loc, value, multiplier);
    return mul_op.getResult();
  }
  auto q_op = dq_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
  if (!q_op) return {};

  Value float_value = q_op.getArg();
  Value new_value = builder.create<TF::MulOp>(loc, float_value, multiplier);

}

// Convenience function to create a `TransposeOp` with a given `permutation`.
// The Location is set as `input`'s loc.
TransposeOp CreateTransposeOp(Value input, const ArrayRef<int64_t> permutation,
                              PatternRewriter& rewriter) {
  return rewriter.create<TransposeOp>(
      input.getLoc(), input, rewriter.getDenseI64ArrayAttr(permutation));
}

    if (!status.ok()) {
      LOG(WARNING) << "cannot create directory '" << path
                   << "': " << status.message();
      return;
    }

    path += "/mlir_reproducer_";

    if (!tensorflow::Env::Default()->CreateUniqueFileName(&path, ".mlir")) {
      LOG(WARNING) << "cannot create unique filename, won't enable MLIR crash "
                      "reproducer.";
      return;