Method Compile

tensorflow/compiler/tf2xla/kernels/bcast_ops.cc:38–64

  }

  void Compile(XlaOpKernelContext* ctx) override {
    OP_REQUIRES(
        ctx, ctx->num_inputs() == 2,
        errors::Unimplemented("Broadcast for n-ary operations (n > 2)"));
    absl::InlinedVector<BCast::Vec, 2> shapes;
    for (int i = 0; i < ctx->num_inputs(); ++i) {
      const TensorShape in_shape = ctx->InputShape(i);
      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(in_shape),
                  errors::InvalidArgument("In[", i, "] must be a vector.",
                                          in_shape.DebugString()));
      std::vector<int64> shape;
      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(i, &shape));
      shapes.push_back(BCast::Vec(shape.begin(), shape.end()));
    }
    BCast bcast(shapes[0], shapes[1]);
    OP_REQUIRES(ctx, bcast.IsValid(),
                errors::InvalidArgument(
                    "Incompatible shapes: [", absl::StrJoin(shapes[0], ","),
                    "] vs. [", absl::StrJoin(shapes[1], ","), "]"));

    const int64 len = bcast.output_shape().size();
    Tensor output(DT_INT32, TensorShape({len}));
    for (int64 i = 0; i < len; ++i) {
      output.flat<int32>()(i) = static_cast<int32>(bcast.output_shape()[i]);
    }
    ctx->SetConstantOutput(0, output);
  }

 private:
  TF_DISALLOW_COPY_AND_ASSIGN(BCastArgsOp);