| 41 | } |
| 42 | |
| 43 | void Compile(XlaOpKernelContext* context) override { |
| 44 | const TensorShape lhs_shape = context->InputShape(0); |
| 45 | const TensorShape rhs_shape = context->InputShape(1); |
| 46 | const TensorShape padding_shape = context->InputShape("padding"); |
| 47 | std::vector<int64> window_strides; |
| 48 | std::vector<int64> lhs_dilation; |
| 49 | std::vector<int64> rhs_dilation; |
| 50 | int64 feature_group_count; |
| 51 | OP_REQUIRES_OK(context, context->ConstantInputAsIntVector("window_strides", |
| 52 | &window_strides)); |
| 53 | OP_REQUIRES_OK(context, context->ConstantInputAsIntVector("lhs_dilation", |
| 54 | &lhs_dilation)); |
| 55 | OP_REQUIRES_OK(context, context->ConstantInputAsIntVector("rhs_dilation", |
| 56 | &rhs_dilation)); |
| 57 | OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar( |
| 58 | "feature_group_count", &feature_group_count)); |
| 59 | |
| 60 | OP_REQUIRES(context, |
| 61 | TensorShapeUtils::IsMatrix(padding_shape) && |
| 62 | padding_shape.dim_size(1) == 2, |
| 63 | errors::InvalidArgument( |
| 64 | "padding must be a matrix with minor dimension 2, got ", |
| 65 | padding_shape.DebugString())); |
| 66 | xla::Literal padding_literal; |
| 67 | OP_REQUIRES_OK(context, context->ConstantInputAsInt64Literal( |
| 68 | "padding", &padding_literal)); |
| 69 | std::vector<std::pair<int64, int64>> padding(padding_shape.dim_size(0)); |
| 70 | for (int i = 0; i < padding.size(); ++i) { |
| 71 | padding[i] = {padding_literal.Get<int64>({i, 0}), |
| 72 | padding_literal.Get<int64>({i, 1})}; |
| 73 | } |
| 74 | |
| 75 | // We do only minimal checking, relying on XLA to check the shape |
| 76 | // invariants. |
| 77 | xla::XlaOp output = xla::ConvGeneralDilated( |
| 78 | context->Input(0), context->Input(1), window_strides, padding, |
| 79 | lhs_dilation, rhs_dilation, dnums_, feature_group_count, |
| 80 | /*batch_group_count=*/1, &precision_config_); |
| 81 | context->SetOutput(0, output); |
| 82 | } |
| 83 | |
| 84 | private: |
| 85 | xla::ConvolutionDimensionNumbers dnums_; |
nothing calls this directly
no test coverage detected