| 46 | namespace { |
| 47 | |
| 48 | std::unique_ptr<const NodeDef> StripTensorDataFromNodeDef( |
| 49 | OpKernelConstruction* ctx) { |
| 50 | #ifndef __ANDROID__ |
| 51 | DCHECK_EQ(NodeDef::descriptor()->field_count(), 6) |
| 52 | << "The NodeDef format has changed, and the attr-stripping code may need " |
| 53 | << "to be updated."; |
| 54 | #endif |
| 55 | const NodeDef& original = ctx->def(); |
| 56 | NodeDef* ret = new NodeDef; |
| 57 | ret->set_name(original.name()); |
| 58 | ret->set_op(original.op()); |
| 59 | ret->set_device(original.device()); |
| 60 | // Strip the "value" attr from the returned NodeDef. |
| 61 | // NOTE(mrry): The present implementation of `OpKernel::OpKernel()` only uses |
| 62 | // attrs that affect the cardinality of list-typed inputs and outputs, so it |
| 63 | // is safe to drop other attrs from the NodeDef. |
| 64 | AddNodeAttr("dtype", ctx->output_type(0), ret); |
| 65 | MergeDebugInfo(original, ret); |
| 66 | return std::unique_ptr<const NodeDef>(ret); |
| 67 | } |
| 68 | |
| 69 | } // namespace |
| 70 |
no test coverage detected