Searched refs:devices_attr (Results 1 – 4 of 4) sorted by relevance
147 auto devices_attr = op->getAttr(kDevicesAttr); in GetDevicesFromOp() local148 if (!devices_attr) return mlir::success(); in GetDevicesFromOp()150 if (auto array_attr = devices_attr.dyn_cast<mlir::ArrayAttr>()) { in GetDevicesFromOp()153 } else if (auto dict_attr = devices_attr.dyn_cast<mlir::DictionaryAttr>()) { in GetDevicesFromOp()
79 auto devices_attr = in TEST() local81 ASSERT_NE(devices_attr, nullptr); in TEST()82 ASSERT_EQ(devices_attr.size(), 3); in TEST()85 auto device_meta_0 = devices_attr.get(cpu0).dyn_cast<mlir::DictionaryAttr>(); in TEST()91 devices_attr.get(gpu0).dyn_cast<mlir::TF::GpuDeviceMetadata>(); in TEST()97 auto device_meta_2 = devices_attr.get(gpu1).dyn_cast<mlir::DictionaryAttr>(); in TEST()
420 auto devices_attr = replicate.devices(); in HandleReplicateOp() local421 if (!devices_attr) return; in HandleReplicateOp()423 auto device_map = devices_attr.getValue(); in HandleReplicateOp()
563 llvm::Optional<DictionaryAttr> devices_attr; in build() local573 devices_attr.emplace(builder.getDictionaryAttr(device_list)); in build()576 BuildReplicateOp(&builder, &state, n, devices_attr, replicated_inputs, in build()