Create a FunctionDef that takes one resource and one regular param
| 38 | |
| 39 | // Create a FunctionDef that takes one resource and one regular param |
| 40 | FunctionDef XTimesY() { |
| 41 | return FunctionDefHelper::Define( |
| 42 | // Name |
| 43 | "XTimesY", |
| 44 | // Args |
| 45 | {"x: float", "y: resource"}, |
| 46 | // Return values |
| 47 | {"z: float"}, |
| 48 | // Attr def |
| 49 | {}, |
| 50 | // Nodes |
| 51 | { |
| 52 | {{"y0"}, "ReadVariableOp", {"y"}, {{"dtype", DT_FLOAT}}}, |
| 53 | {{"z"}, "Mul", {"x", "y0"}, {{"T", DT_FLOAT}}}, |
| 54 | }); |
| 55 | } |
| 56 | |
| 57 | class XlaKernelCreatorTest : public ::testing::Test { |
| 58 | protected: |