QUESTION: It'd be nice to support DT_INT16, DT_UINT8, etc. here for the type of the second input argument. Should we have the framework do some sort of integer promotion automatically, or should that be something that users have to do explicitly with a conversion operator in the graph?
| 43 | // automatically, or should that be something that users have to |
| 44 | // do explicitly with a conversion operator in the graph? |
| 45 | explicit GatherOp(OpKernelConstruction* c) : OpKernel(c) { |
| 46 | // Set batch_dims_ to 0 if the attribute does not exist. |
| 47 | if (c->HasAttr("batch_dims")) { |
| 48 | OP_REQUIRES_OK(c, c->GetAttr("batch_dims", &batch_dims_)); |
| 49 | } else { |
| 50 | batch_dims_ = 0; |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | void Compute(OpKernelContext* c) override { |
| 55 | const Tensor& params = c->input(0); |