| 125 | } |
| 126 | |
| 127 | const xla::XlaComputation* XlaContext::GetOrCreateMul(const DataType type) { |
| 128 | return LookupOrCreate(type, &mul_func_, [type] { |
| 129 | const string type_string = DataTypeString(type); |
| 130 | VLOG(1) << "Building Mul() for " << type_string; |
| 131 | xla::XlaBuilder b("mul<" + type_string + ">"); |
| 132 | xla::PrimitiveType xla_type; |
| 133 | TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type)); |
| 134 | auto x = |
| 135 | xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla_type, {}), "x"); |
| 136 | auto y = |
| 137 | xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla_type, {}), "y"); |
| 138 | xla::Mul(x, y); |
| 139 | return b.Build().ConsumeValueOrDie(); |
| 140 | }); |
| 141 | } |
| 142 | |
| 143 | const xla::XlaComputation* XlaContext::LookupOrCreate( |
| 144 | DataType type, ComputationMap* out, |
no test coverage detected