/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");

you may not use this file except in compliance with the License.

You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software

distributed under the License is distributed on an "AS IS" BASIS,

WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

See the License for the specific language governing permissions and

limitations under the License.

==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"

#include "tensorflow/core/framework/op.h"

#include "tensorflow/core/framework/op_kernel.h"

#include "tensorflow/core/framework/resource_mgr.h"

#include "tensorflow/core/lib/core/status.h"

#include "tensorflow/core/public/version.h"

namespace tensorflow {

REGISTER_OP("KernelLabel")

.Output("result: string")

.SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("KernelLabelRequired")

.Input("input: int32")

.Output("result: string")

.SetShapeFn([](shape_inference::InferenceContext* c) {

shape_inference::ShapeHandle out;

TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &out));

c->set_output(0, c->Scalar());

return Status::OK();

});

REGISTER_OP("GraphDefVersion")

.Output("version: int32")

.SetIsStateful()

.SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("RequiresOlderGraphVersion")

.Output("version: int32")

.SetIsStateful()

.SetShapeFn([](shape_inference::InferenceContext* c) {

if (c->graph_def_version() != TF_GRAPH_DEF_VERSION - 1) {

return errors::InvalidArgument("Wrong graph version for shape");

}

return shape_inference::ScalarShape(c);

});

REGISTER_OP("Old")

.SetShapeFn(shape_inference::UnknownShape)

.Deprecated(8, "For reasons");

REGISTER_RESOURCE_HANDLE_OP(StubResource);

REGISTER_OP("ResourceInitializedOp")

.Input("resource: resource")

.Output("initialized: bool")

.SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ResourceCreateOp")

.Input("resource: resource")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("ResourceUsingOp")

.Input("resource: resource")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TestStringOutput")

.Input("input: float")

.Output("output1: float")

.Output("output2: string")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TestAttr")

.Output("out: T")

.Attr("T: {float, double}")

.SetShapeFn(shape_inference::UnknownShape);

namespace {

enum KernelLabel { DEFAULT_LABEL, OVERLOAD_1_LABEL, OVERLOAD_2_LABEL };

} // namespace

template <KernelLabel KL>

class KernelLabelOp : public OpKernel {

public:

using OpKernel::OpKernel;

void Compute(OpKernelContext* ctx) override {

Tensor* output;

OP_REQUIRES_OK(ctx,

ctx->allocate_output("result", TensorShape({}), &output));

switch (KL) {

case DEFAULT_LABEL:

output->scalar<string>()() = "My label is: default";

break;

case OVERLOAD_1_LABEL:

output->scalar<string>()() = "My label is: overload_1";

break;

case OVERLOAD_2_LABEL:

output->scalar<string>()() = "My label is: overload_2";

break;

}

}

};

REGISTER_KERNEL_BUILDER(Name("KernelLabel").Device(DEVICE_CPU),

KernelLabelOp<DEFAULT_LABEL>);

REGISTER_KERNEL_BUILDER(Name("KernelLabel")

.Device(DEVICE_CPU)

.Label("overload_1"),

KernelLabelOp<OVERLOAD_1_LABEL>);

REGISTER_KERNEL_BUILDER(Name("KernelLabel")

.Device(DEVICE_CPU)

.Label("overload_2"),

KernelLabelOp<OVERLOAD_2_LABEL>);

// All "KernelLabelRequired" kernels have labels

REGISTER_KERNEL_BUILDER(

Name("KernelLabelRequired").Device(DEVICE_CPU).Label("overload_1"),

KernelLabelOp<OVERLOAD_1_LABEL>);

REGISTER_KERNEL_BUILDER(

Name("KernelLabelRequired").Device(DEVICE_CPU).Label("overload_2"),

KernelLabelOp<OVERLOAD_2_LABEL>);

class GraphDefVersionOp : public OpKernel {

public:

explicit GraphDefVersionOp(OpKernelConstruction* ctx)

: OpKernel(ctx), graph_def_version_(ctx->graph_def_version()) {}

void Compute(OpKernelContext* ctx) override {

Tensor* output = nullptr;

OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output));

output->scalar<int>()() = graph_def_version_;

}

private:

const int graph_def_version_;

};

REGISTER_KERNEL_BUILDER(Name("GraphDefVersion").Device(DEVICE_CPU),

GraphDefVersionOp);

class OldOp : public OpKernel {

public:

explicit OldOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

void Compute(OpKernelContext* ctx) override {}

};

REGISTER_KERNEL_BUILDER(Name("Old").Device(DEVICE_CPU), OldOp);

// Stubbed-out resource to test resource handle ops.

class StubResource : public ResourceBase {

public:

string DebugString() override { return ""; }

};

REGISTER_RESOURCE_HANDLE_KERNEL(StubResource);

REGISTER_KERNEL_BUILDER(Name("ResourceInitializedOp").Device(DEVICE_CPU),

IsResourceInitialized<StubResource>);

class ResourceCreateOp : public OpKernel {

public:

explicit ResourceCreateOp(OpKernelConstruction* c) : OpKernel(c) {}

void Compute(OpKernelContext* c) override {

OP_REQUIRES_OK(c,

CreateResource(c, HandleFromInput(c, 0), new StubResource));

}

};

REGISTER_KERNEL_BUILDER(Name("ResourceCreateOp").Device(DEVICE_CPU),

ResourceCreateOp);

// Uses a ResourceHandle to check its validity.

class ResourceUsingOp : public OpKernel {

public:

explicit ResourceUsingOp(OpKernelConstruction* context) : OpKernel(context) {}

void Compute(OpKernelContext* ctx) override {

StubResource* unused;

OP_REQUIRES_OK(ctx, LookupResource<StubResource>(

ctx, HandleFromInput(ctx, 0), &unused));

}

};

REGISTER_KERNEL_BUILDER(Name("ResourceUsingOp").Device(DEVICE_CPU),

ResourceUsingOp);

class TestAttrOp : public OpKernel {

public:

explicit TestAttrOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

void Compute(OpKernelContext* ctx) override {

Tensor* output;

OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output));

output->scalar<float>()() = 1.0;

}

};

REGISTER_KERNEL_BUILDER(

Name("TestAttr").Device(DEVICE_CPU).TypeConstraint<float>("T"), TestAttrOp);

// Various test ops without kernels. These are used to test graph construction.

REGISTER_OP("A")

.Output("out: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("B")

.Output("out: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Foo1")

.Input("a: float32")

.Input("b: int32")

.Input("c: int32")

.Output("d: float32")

.Output("e: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Foo2")

.Input("a: float32")

.Input("b: string")

.Input("c: string")

.Output("d: float32")

.Output("e: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Foo3")

.Input("a: float32")

.Input("b: string")

.Input("c: float32")

.Output("d: float32")

.Output("e: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("CopyOp").Input("a: T").Output("b: T").Attr("T: type").SetShapeFn(

shape_inference::UnknownShape);

REGISTER_OP("None").SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntOutput")

.Output("a: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Int64Output")

.Output("out: int64")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefOutput")

.Output("a: Ref(int32)")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("FloatOutput")

.Output("a: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoFloatOutputs")

.Output("a: float32")

.Output("b: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("FiveFloatOutputs")

.Output("a: float32")

.Output("b: float32")

.Output("c: float32")

.Output("d: float32")

.Output("e: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefOutputFloatOutput")

.Output("a: Ref(float32)")

.Output("b: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefInputFloatInput")

.Input("a: Ref(float)")

.Input("b: float")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntInput")

.Input("a: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntInputIntOutput")

.Input("a: int32")

.Output("b: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("FloatInput")

.Input("a: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoIntOutputs")

.Output("a: int32")

.Output("b: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntOutputFloatOutput")

.Output("a: int32")

.Output("b: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("FloatOutputStringOutput")

.Output("a: float32")

.Output("b: string")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoIntInputs")

.Input("a: int32")

.Input("b: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoFloatInputs")

.Input("a: float32")

.Input("b: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntInputFloatInput")

.Input("a: int32")

.Input("b: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefInputIntInput")

.Input("a: Ref(int32)")

.Input("b: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoFloatInputsFloatOutput")

.Input("a: float32")

.Input("b: float32")

.Output("c: float32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoFloatInputsIntOutput")

.Input("a: float32")

.Input("b: float32")

.Output("c: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefInputFloatInputIntOutput")

.Input("a: Ref(float32)")

.Input("b: float32")

.Output("c: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("ListInput")

.Input("a: N * T")

.Attr("N: int >= 1")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("ListOutput")

.Output("a: T")

.Attr("T: list(type) >= 1")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Unary").Input("a: T").Output("b: T").Attr("T: type").SetShapeFn(

shape_inference::UnknownShape);

REGISTER_OP("OpWithDefaultAttr")

.Output("a: int32")

.Attr("default_float: float = 123.0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("OpWithFutureDefaultAttr")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("IntAttr")

.Output("out: int64")

.Attr("foo: int = 1")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("StringListAttr")

.Attr("a: list(string)")

.Attr("b: string")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("DefaultAttrs")

.Attr("string_val: string = 'abc'")

.Attr("string_list_val: list(string) = ['abc', '']")

.Attr("int_val: int = 123")

.Attr("int_list_val: list(int) = [1, 2, 3]")

.Attr("float_val: float = 10.0")

.Attr("float_list_val: list(float) = [10.0]")

.Attr("bool_val: bool = true")

.Attr("bool_list_val: list(bool) = [true, false]")

.Attr("type_val: type = DT_INT32")

.Attr("type_list_val: list(type) = [DT_INT32, DT_FLOAT]")

.Attr("shape_val: shape = { dim { size: 2 } dim { size: 1 } }")

.Attr("shape_list_val: list(shape) = [{}, { dim { size: 1} }]")

.Attr("tensor_val: tensor = { dtype: DT_INT32 tensor_shape: {} int_val: 1}")

.Attr(

"tensor_list_val: list(tensor) = "

"[{ dtype: DT_INT32 tensor_shape: {} int_val: 1}]")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("FuncAttr")

.Attr("f: func")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Simple")

.Input("a: int32")

.Output("out: float")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("OutT").Output("a: T").Attr("T: type").SetShapeFn(

shape_inference::UnknownShape);

REGISTER_OP("ReservedInput")

.Input("input: int32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Polymorphic")

.Input("a: T")

.Output("out: T")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("PolymorphicOut")

.Output("out: T")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("PolymorphicDefaultOut")

.Output("out: T")

.Attr("T: type = DT_STRING")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Binary")

.Input("a: T")

.Input("b: T")

.Output("out: T")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Restrict")

.Input("a: T")

.Output("out: T")

.Attr("T: {string, bool}")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TypeList")

.Input("a: T")

.Attr("T: list(type) >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TypeListTwice")

.Input("a: T")

.Input("b: T")

.Attr("T: list(type) >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("OutTypeList")

.Output("out: T")

.Attr("T: list(type) >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TypeListRestrict")

.Input("a: T")

.Attr("T: list({string, bool})")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("OutTypeListRestrict")

.Output("out: t")

.Attr("t: list({string, bool})")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("Attr").Attr("a: int").SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrFloat")

.Attr("a: float")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrBool")

.Attr("a: bool")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrBoolList")

.Attr("a: list(bool)")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrMin")

.Attr("a: int >= 5")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrListMin")

.Attr("a: list(int) >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrEnum")

.Attr("a: {'apples', 'oranges'}")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrEnumList")

.Attr("a: list({'apples', 'oranges'})")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrShape")

.Attr("a: shape")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrShapeList")

.Attr("a: list(shape)")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrPartialShape")

.Attr("a: shape")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrPartialShapeList")

.Attr("a: list(shape)")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrDefault")

.Attr("a: string = 'banana'")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrListDefault")

.Attr("a: list(int) = [5, 15]")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrEmptyListDefault")

.Attr("a: list(float) = []")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("ReservedAttr")

.Attr("range: int")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrTypeDefault")

.Input("a: T")

.Attr("T: type = DT_INT32")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("AttrListTypeDefault")

.Input("a: N * T")

.Input("b: N * T")

.Attr("T: type = DT_INT32")

.Attr("N: int")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NIntsIn")

.Input("a: N * int32")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NPolymorphicIn")

.Input("a: N * T")

.Attr("T: type")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NPolymorphicRestrictIn")

.Input("a: N * T")

.Attr("T: {string, bool}")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NInTwice")

.Input("a: N * int32")

.Input("b: N * string")

.Attr("N: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NInPolymorphicTwice")

.Input("a: N * T")

.Input("b: N * T")

.Attr("T: type")

.Attr("N: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NInTwoTypeVariables")

.Input("a: N * S")

.Input("b: N * T")

.Attr("S: type")

.Attr("T: type")

.Attr("N: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("InPolymorphicTwice")

.Input("a: N * T")

.Input("b: M * T")

.Attr("T: type")

.Attr("N: int >= 0")

.Attr("M: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NIntsOut")

.Output("a: N * int32")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NIntsOutDefault")

.Output("a: N * int32")

.Attr("N: int >= 2 = 3")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NPolymorphicOut")

.Output("a: N * T")

.Attr("T: type")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NPolymorphicOutDefault")

.Output("a: N * T")

.Attr("T: type = DT_BOOL")

.Attr("N: int >= 2 = 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("NPolymorphicRestrictOut")

.Output("a: N * T")

.Attr("T: {string, bool}")

.Attr("N: int >= 2")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefIn")

.Input("a: Ref(T)")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("TwoRefsIn")

.Input("a: Ref(T)")

.Input("b: Ref(T)")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("RefOut")

.Output("a: Ref(T)")

.Attr("T: type")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("SimpleStruct")

.Output("a: n_a * int32")

.Attr("n_a: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("MixedStruct")

.Output("a: n_a * int32")

.Output("b: float")

.Attr("n_a: int >= 0")

.SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("ComplexStruct")

.Output("a: n_a * int32")

.Output("b: n_b * int64")

.Output("c: t_c")

.Attr("n_a: int >= 0")

.Attr("n_b: int >= 0")

.Attr("t_c: list(type) >= 0")

.SetShapeFn(shape_inference::UnknownShape);

} // end namespace tensorflow