Method Compile

tensorflow/compiler/tf2xla/kernels/shape_op.cc:110–128

  explicit SizeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}

  void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape input_shape = ctx->InputShape(0);
    OP_REQUIRES(ctx,
                FastBoundsCheck(input_shape.num_elements(),
                                std::numeric_limits<int32>::max()),
                errors::InvalidArgument("Size does not work for tensors > "
                                        "int32 max."));
    Tensor size_constant(DT_INT32, TensorShape({}));
    const int rank = input_shape.dims();
    xla::XlaBuilder* builder = ctx->builder();
    auto size = xla::One(builder, xla::U32);
    for (int64 i = 0; i < rank; ++i) {
      size = xla::Mul(
          size, xla::ConvertElementType(xla::GetDimensionSize(ctx->Input(0), i),
                                        xla::U32));
    }
    size = xla::ConvertElementType(size, ctx->output_xla_type(0));
    ctx->SetOutput(0, size);
  }
};

REGISTER_XLA_OP(Name("Size").CompilationOnly().IsMetadataOp(), SizeOp);