/* 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/numeric_op.h"

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

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

namespace tensorflow {

using shape_inference::Dimension;

using shape_inference::InferenceContext;

using shape_inference::Shape;

REGISTER_OP("AddN")

.Input("inputs: N * T")

.Output("sum: T")

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

.Attr("T: numbertype")

.SetIsCommutative()

.SetIsAggregate()

.SetShapeFn([](InferenceContext* c) {

const Shape* cur = c->input(c->num_inputs() - 1);

for (int i = c->num_inputs() - 2; i >= 0; --i) {

TF_RETURN_WITH_CONTEXT_IF_ERROR(c->Merge(c->input(i), cur, &cur),

"From merging shape ", i,

" with other shapes.");

}

c->set_output(0, cur);

return Status::OK();

})

.Doc(R"doc(

Add all input tensors element wise.

inputs: Must all be the same size and shape.

)doc");

namespace {

// Shape inference function for binary operators that broadcast their inputs.

Status BroadcastBinaryOpShapeFn(InferenceContext* c) {

const Shape* shape_x = c->input(0);

const Shape* shape_y = c->input(1);

if (!c->RankKnown(shape_x) || !c->RankKnown(shape_y)) {

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

return Status::OK();

}

const int32 rank_x = c->Rank(shape_x);

const int32 rank_y = c->Rank(shape_y);

const int32 rank_out = std::max(rank_x, rank_y);

// To compute the broadcast dimensions, we zip together shape_x and shape_y

// and

// pad with 1 to make them the same length.

std::vector<const Dimension*> dims;

const Dimension* dim_one = rank_x == rank_y ? nullptr : c->MakeDim(1);

for (int i = 0; i < rank_out; ++i) {

const auto* dim_x = i < (rank_out - rank_x)

? dim_one

: c->Dim(shape_x, i - (rank_out - rank_x));

const auto* dim_y = i < (rank_out - rank_y)

? dim_one

: c->Dim(shape_y, i - (rank_out - rank_y));

if (!c->ValueKnown(dim_x) || !c->ValueKnown(dim_y)) {

// One or both dimensions is unknown.

//

// - If either dimension is greater than 1, we assume that the program is

// correct, and the other dimension will be broadcast to match it.

// TODO(cwhipkey): For shape inference, if we eliminate the shape checks

// in C++ op code, we must still assert that the unknown dim is either 1

// or the same as the known dim.

// - If either dimension is 1, the other dimension is the output.

if (c->Value(dim_x) > 1) {

dims.push_back(dim_x);

} else if (c->Value(dim_y) > 1) {

dims.push_back(dim_y);

} else if (c->Value(dim_x) == 1) {

dims.push_back(dim_y);

} else if (c->Value(dim_y) == 1) {

dims.push_back(dim_x);

} else {

dims.push_back(c->UnknownDim());

}

} else if (c->Value(dim_x) == 1 || c->Value(dim_y) == 1) {

if (c->Value(dim_x) == 1 && dim_y != dim_one) {

// We will broadcast dim_x to dim_y.

dims.push_back(dim_y);

} else {

DCHECK_EQ(c->Value(dim_y), 1);

// We will broadcast dim_y to dim_x.

dims.push_back(dim_x);

}

} else {

const Dimension* dim;

TF_RETURN_IF_ERROR(c->Merge(dim_x, dim_y, &dim));

dims.push_back(dim);

}

}

c->set_output(0, c->MakeShape(dims));

return Status::OK();

}

} // namespace

// --------------------------------------------------------------------------

REGISTER_OP("BatchMatMul")

.Input("x: T")

.Input("y: T")

.Output("output: T")

.Attr("T: {half, float, double, int32, complex64, complex128}")

.Attr("adj_x: bool = false")

.Attr("adj_y: bool = false")

.SetShapeFn([](InferenceContext* c) {

const Shape* a_shape;

const Shape* b_shape;

TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &a_shape));

TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 3, &b_shape));

// Determine output rows and cols.

bool adj_x;

bool adj_y;

TF_RETURN_IF_ERROR(c->GetAttr("adj_x", &adj_x));

TF_RETURN_IF_ERROR(c->GetAttr("adj_y", &adj_y));

const Dimension* output_rows = c->Dim(a_shape, adj_x ? -1 : -2);

const Dimension* output_cols = c->Dim(b_shape, adj_y ? -2 : -1);

// Batch dims match between inputs.

const Shape* a_batch_dims;

const Shape* b_batch_dims;

const Shape* batch_dims;

TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_dims));

TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims));

TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims));

// Assert inner dims match.

const Dimension* unused;

TF_RETURN_IF_ERROR(c->Merge(c->Dim(a_shape, adj_x ? -2 : -1),

c->Dim(b_shape, adj_y ? -1 : -2), &unused));

const Shape* out;

TF_RETURN_IF_ERROR(c->Concatenate(

batch_dims, c->Matrix(output_rows, output_cols), &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Multiplies slices of two tensors in batches.

Multiplies all slices of `Tensor` `x` and `y` (each slice can be

viewed as an element of a batch), and arranges the individual results

in a single output tensor of the same batch size. Each of the

individual slices can optionally be adjointed (to adjoint a matrix

means to transpose and conjugate it) before multiplication by setting

the `adj_x` or `adj_y` flag to `True`, which are by default `False`.

The input tensors `x` and `y` are 3-D or higher with shape `[..., r_x, c_x]`

and `[..., r_y, c_y]`.

The output tensor is 3-D or higher with shape `[..., r_o, c_o]`, where:

r_o = c_x if adj_x else r_x

c_o = r_y if adj_y else c_y

It is computed as:

output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :])

x: 3-D or higher with shape `[..., r_x, c_x]`.

y: 3-D or higher with shape `[..., r_y, c_y]`.

output: 3-D or higher with shape `[..., r_o, c_o]`

adj_x: If `True`, adjoint the slices of `x`. Defaults to `False`.

adj_y: If `True`, adjoint the slices of `y`. Defaults to `False`.

)doc");

// --------------------------------------------------------------------------

// Casting Ops

//

// NOTE: Only a smaller number of types are supported by

// Cast. The exact casting rule is TBD. The current

// implementation uses C++ static cast rules for numeric

// types, which may be changed in the future.

REGISTER_OP("Cast")

.Input("x: SrcT")

.Output("y: DstT")

.Attr("SrcT: type")

.Attr("DstT: type")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Cast x of type SrcT to y of DstT.

)doc");

REGISTER_OP("_HostCast")

.Input("x: SrcT")

.Output("y: DstT")

.Attr("SrcT: type")

.Attr("DstT: type")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Cast x of type SrcT to y of DstT.

_HostCast requires its input and produces its output in host memory.

)doc");

// --------------------------------------------------------------------------

REGISTER_OP("Abs")

.Input("x: T")

.Output("y: T")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Computes the absolute value of a tensor.

Given a tensor `x`, this operation returns a tensor containing the absolute

value of each element in `x`. For example, if x is an input element and y is

an output element, this operation computes \\(y = |x|\\).

)doc");

REGISTER_OP("ComplexAbs")

.Input("x: T")

.Output("y: Tout")

.Attr("T: {complex64, complex128} = DT_COMPLEX64")

.Attr("Tout: {float, double} = DT_FLOAT")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Computes the complex absolute value of a tensor.

Given a tensor `x` of complex numbers, this operation returns a tensor of type

`float` or `double` that is the absolute value of each element in `x`. All

elements in `x` must be complex numbers of the form \\(a + bj\\). The absolute

value is computed as \\( \sqrt{a^2 + b^2}\\).

For example:

```

# tensor 'x' is [[-2.25 + 4.75j], [-3.25 + 5.75j]]

tf.complex_abs(x) ==> [5.25594902, 6.60492229]

```

)doc");

// Declares cwise unary operations signature: 't -> 't

#define UNARY() \

Input("x: T") \

.Output("y: T") \

.Attr("T: {half, float, double, int32, int64, complex64, complex128}") \

.SetShapeFn(shape_inference::UnchangedShape)

#define UNARY_REAL() \

Input("x: T") \

.Output("y: T") \

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

.SetShapeFn(shape_inference::UnchangedShape)

#define UNARY_COMPLEX() \

Input("x: T") \

.Output("y: T") \

.Attr("T: {half, float, double, complex64, complex128}") \

.SetShapeFn(shape_inference::UnchangedShape)

#define UNARY_GRADIENT_COMPLEX() \

Input("x: T") \

.Input("y: T") \

.Output("z: T") \

.Attr("T: {half, float, double, complex64, complex128}") \

.SetShapeFn(shape_inference::UnchangedShape)

REGISTER_OP("Neg")

.UNARY()

.Doc(R"doc(

Computes numerical negative value element-wise.

I.e., \\(y = -x\\).

)doc");

REGISTER_OP("Inv")

.UNARY()

.Doc(R"doc(

Computes the reciprocal of x element-wise.

I.e., \\(y = 1 / x\\).

)doc");

REGISTER_OP("Square")

.UNARY()

.Doc(R"doc(

Computes square of x element-wise.

I.e., \\(y = x * x = x^2\\).

)doc");

REGISTER_OP("Sqrt")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes square root of x element-wise.

I.e., \\(y = \sqrt{x} = x^{1/2}\\).

)doc");

REGISTER_OP("Rsqrt")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes reciprocal of square root of x element-wise.

I.e., \\(y = 1 / \sqrt{x}\\).

)doc");

REGISTER_OP("Exp")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes exponential of x element-wise. \\(y = e^x\\).

)doc");

REGISTER_OP("Log")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes natural logarithm of x element-wise.

I.e., \\(y = \log_e x\\).

)doc");

REGISTER_OP("Tanh")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes hyperbolic tangent of `x` element-wise.

)doc");

REGISTER_OP("TanhGrad").UNARY_GRADIENT_COMPLEX().Doc(R"doc(

Computes the gradient for the tanh of `x` wrt its input.

Specifically, `grad = dy * (1 - y*y)`, where `y = tanh(x)`, and `dy`

is the corresponding input gradient.

)doc");

REGISTER_OP("Lgamma")

.UNARY_REAL()

.Doc(R"doc(

Computes the log of the absolute value of `Gamma(x)` element-wise.

)doc");

REGISTER_OP("Digamma")

.UNARY_REAL()

.Doc(R"doc(

Computes Psi, the derivative of Lgamma (the log of the absolute value of

`Gamma(x)`), element-wise.

)doc");

REGISTER_OP("Erf")

.UNARY_REAL()

.Doc(R"doc(

Computes the Gauss error function of `x` element-wise.

)doc");

REGISTER_OP("Erfc")

.UNARY_REAL()

.Doc(R"doc(

Computes the complementary error function of `x` element-wise.

)doc");

REGISTER_OP("Sigmoid")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes sigmoid of `x` element-wise.

Specifically, `y = 1 / (1 + exp(-x))`.

)doc");

REGISTER_OP("SigmoidGrad").UNARY_GRADIENT_COMPLEX().Doc(R"doc(

Computes the gradient of the sigmoid of `x` wrt its input.

Specifically, `grad = dy * y * (1 - y)`, where `y = sigmoid(x)`, and

`dy` is the corresponding input gradient.

)doc");

REGISTER_OP("Sin")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes sin of x element-wise.

)doc");

REGISTER_OP("Cos")

.UNARY_COMPLEX()

.Doc(R"doc(

Computes cos of x element-wise.

)doc");

REGISTER_OP("Tan")

.UNARY()

.Doc(R"doc(

Computes tan of x element-wise.

)doc");

REGISTER_OP("Asin")

.UNARY()

.Doc(R"doc(

Computes asin of x element-wise.

)doc");

REGISTER_OP("Acos")

.UNARY()

.Doc(R"doc(

Computes acos of x element-wise.

)doc");

REGISTER_OP("Atan")

.UNARY()

.Doc(R"doc(

Computes atan of x element-wise.

)doc");

#undef UNARY

#undef UNARY_REAL

#undef UNARY_COMPLEX

REGISTER_OP("IsNan")

.Input("x: T")

.Output("y: bool")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns which elements of x are NaN.

)doc");

REGISTER_OP("IsInf")

.Input("x: T")

.Output("y: bool")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns which elements of x are Inf.

)doc");

REGISTER_OP("IsFinite")

.Input("x: T")

.Output("y: bool")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns which elements of x are finite.

)doc");

REGISTER_OP("Sign")

.Input("x: T")

.Output("y: T")

.Attr("T: {half, float, double, int32, int64, complex64, complex128}")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns an element-wise indication of the sign of a number.

`y = sign(x) = -1` if `x < 0`; 0 if `x == 0`; 1 if `x > 0`.

For complex numbers, `y = sign(x) = x / |x|` if `x != 0`, otherwise `y = 0`.

)doc");

REGISTER_OP("Floor")

.Input("x: T")

.Output("y: T")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns element-wise largest integer not greater than x.

)doc");

REGISTER_OP("Ceil")

.Input("x: T")

.Output("y: T")

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

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns element-wise smallest integer in not less than x.

)doc");

// Declares cwise binary operations signature: 't, 't -> 't.

#define BINARY_MORE() \

Input("x: T").Input("y: T").Output("z: T").Attr( \

"T: {half, float, double, uint8, int8, int16, int32, int64, complex64, complex128}")

#define BINARY_FEWER() \

Input("x: T").Input("y: T").Output("z: T").Attr( \

"T: {half, float, double, int32, int64, complex64, complex128}")

// TODO(mrry): Restore `SetIsCommutative()` for non-string types.

REGISTER_OP("Add")

.Input("x: T")

.Input("y: T")

.Output("z: T")

.Attr(

"T: {half, float, double, uint8, int8, int16, int32, int64, complex64, "

"complex128, string}")

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns x + y element-wise.

*NOTE*: Add supports broadcasting. AddN does not.

)doc");

REGISTER_OP("Sub")

.BINARY_FEWER()

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns x - y element-wise.

)doc");

REGISTER_OP("Mul")

.BINARY_MORE()

.SetIsCommutative()

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns x * y element-wise.

)doc");

REGISTER_OP("Div").BINARY_MORE().SetShapeFn(BroadcastBinaryOpShapeFn).Doc(R"doc(

Returns x / y element-wise.

)doc");

REGISTER_OP("SquaredDifference")

.BINARY_FEWER()

.SetIsCommutative()

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns (x - y)(x - y) element-wise.

)doc");

#undef BINARY_FEWER

#undef BINARY_MORE

REGISTER_OP("Maximum")

.Input("x: T")

.Input("y: T")

.Output("z: T")

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

.SetIsCommutative()

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns the max of x and y (i.e. x > y ? x : y) element-wise, broadcasts.

)doc");

REGISTER_OP("Minimum")

.Input("x: T")

.Input("y: T")

.Output("z: T")

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

.SetIsCommutative()

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns the min of x and y (i.e. x < y ? x : y) element-wise, broadcasts.

)doc");

REGISTER_OP("Mod")

.Input("x: T")

.Input("y: T")

.Output("z: T")

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

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Returns element-wise remainder of division.

)doc");

REGISTER_OP("Pow")

.Input("x: T")

.Input("y: T")

.Output("z: T")

.Attr("T: {half, float, double, int32, int64, complex64, complex128}")

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Computes the power of one value to another.

Given a tensor `x` and a tensor `y`, this operation computes \\(x^y\\) for

corresponding elements in `x` and `y`. For example:

```

# tensor 'x' is [[2, 2]], [3, 3]]

# tensor 'y' is [[8, 16], [2, 3]]

tf.pow(x, y) ==> [[256, 65536], [9, 27]]

```

)doc");

REGISTER_OP("Igammac")

.Input("a: T")

.Input("x: T")

.Output("z: T")

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

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Compute the upper regularized incomplete Gamma function `Q(a, x)`.

The upper regularized incomplete Gamma function is defined as:

```

Q(a, x) = Gamma(a, x) / Gamma(x) = 1 - P(a, x)

```

where

```

Gamma(a, x) = int_{x}^{\infty} t^{a-1} exp(-t) dt

```

is the upper incomplete Gama function.

Note, above `P(a, x)` (`Igamma`) is the lower regularized complete

Gamma function.

)doc");

REGISTER_OP("Igamma")

.Input("a: T")

.Input("x: T")

.Output("z: T")

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

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Compute the lower regularized incomplete Gamma function `Q(a, x)`.

The lower regularized incomplete Gamma function is defined as:

```

P(a, x) = gamma(a, x) / Gamma(x) = 1 - Q(a, x)

```

where

```

gamma(a, x) = int_{0}^{x} t^{a-1} exp(-t) dt

```

is the lower incomplete Gamma function.

Note, above `Q(a, x)` (`Igammac`) is the upper regularized complete

Gamma function.

)doc");

REGISTER_OP("Zeta")

.Input("x: T")

.Input("q: T")

.Output("z: T")

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

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Compute the Hurwitz zeta function \\(\zeta(x, q)\\).

The Hurwitz zeta function is defined as:

```

\zeta(x, q) = \sum_{n=0}^{\infty} (q + n)^{-x}

```

)doc");

REGISTER_OP("Polygamma")

.Input("a: T")

.Input("x: T")

.Output("z: T")

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

.SetShapeFn(BroadcastBinaryOpShapeFn)

.Doc(R"doc(

Compute the polygamma function \\(\psi^{(n)}(x)\\).

The polygamma function is defined as:

```

\psi^{(n)}(x) = \frac{d^n}{dx^n} \psi(x)

```

where \\(\psi(x)\\) is the digamma function.

)doc");

// --------------------------------------------------------------------------

// Declares cwise binary comparison operations signature: 't, 't -> bool,

// where 't has a natural total order.

#define COMPARISON() \

Input("x: T") \

.Input("y: T") \

.Output("z: bool") \

.Attr("T: realnumbertype") \

.SetShapeFn(BroadcastBinaryOpShapeFn)

REGISTER_OP("Less")

.COMPARISON()

.Doc(R"doc(

Returns the truth value of (x < y) element-wise.

)doc");

REGISTER_OP("LessEqual")

.COMPARISON()

.Doc(R"doc(

Returns the truth value of (x <= y) element-wise.

)doc");

REGISTER_OP("Greater")

.COMPARISON()

.Doc(R"doc(

Returns the truth value of (x > y) element-wise.

)doc");

REGISTER_OP("GreaterEqual")

.COMPARISON()

.Doc(R"doc(

Returns the truth value of (x >= y) element-wise.

)doc");

#undef COMPARISON

// --------------------------------------------------------------------------

#define EQUALITY_COMPARISON() \

Input("x: T") \

.Input("y: T") \

.Output("z: bool") \

.SetIsCommutative() \

.Attr( \

"T: {half, float, double, uint8, int8, int16, int32, int64, " \

"complex64, " \

"quint8, qint8, qint32, string, bool, complex128}") \

.SetShapeFn(BroadcastBinaryOpShapeFn)

REGISTER_OP("Equal")

.EQUALITY_COMPARISON()

.Doc(R"doc(

Returns the truth value of (x == y) element-wise.

)doc");

REGISTER_OP("NotEqual")

.EQUALITY_COMPARISON()

.Doc(R"doc(

Returns the truth value of (x != y) element-wise.

)doc");

#undef EQUALITY_COMPARISON

// --------------------------------------------------------------------------

REGISTER_OP("LogicalNot")

.Input("x: bool")

.Output("y: bool")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns the truth value of NOT x element-wise.

)doc");

#define BINARY_LOGICAL() \

Input("x: bool") \

.Input("y: bool") \

.Output("z: bool") \

.SetIsCommutative() \

.SetShapeFn(BroadcastBinaryOpShapeFn)

REGISTER_OP("LogicalAnd")

.BINARY_LOGICAL()

.Doc(R"doc(

Returns the truth value of x AND y element-wise.

)doc");

REGISTER_OP("LogicalOr")

.BINARY_LOGICAL()

.Doc(R"doc(

Returns the truth value of x OR y element-wise.

)doc");

#undef BINARY_LOGICAL

// --------------------------------------------------------------------------

REGISTER_OP("Select")

.Input("condition: bool")

.Input("t: T")

.Input("e: T")

.Output("output: T")

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* cond = c->input(0);

const Shape* data = c->input(1);

TF_RETURN_IF_ERROR(c->Merge(data, c->input(2), &data));

// Validate condition's shape if possible.

if (c->RankKnown(data)) {

const int32 data_rank = c->Rank(data);

if (data_rank == 0 || data_rank == 1) {

// Cond must match inputs since they are scalar or vector.

TF_RETURN_IF_ERROR(c->Merge(data, cond, &data));

} else {

// cond must be the shape [data.dim[0]] or data.

if (c->RankKnown(cond)) {

if (c->Rank(cond) == 1) {

// Must be a vector whose first dimension matches first dimension

// of the data vectors.

const Dimension* merged_dim;

TF_RETURN_IF_ERROR(

c->Merge(c->Dim(data, 0), c->Dim(cond, 0), &merged_dim));

if (merged_dim != c->Dim(data, 0)) {

// Merging used the cond dim. Update data to refer to it.

std::vector<const Dimension*> dims{merged_dim};

for (int i = 1; i < data_rank; ++i) {

dims.push_back(c->Dim(data, i));

}

data = c->MakeShape(dims);

}

} else {

// Must be the same as the data vectors.

TF_RETURN_IF_ERROR(c->Merge(data, cond, &data));

}

} else {

// We want to express that it's either [data.dim[0]] or data.

// - rank(cond) must be 1 or rank(data).

// - cond.dim[0] is same as data.dim[0]

// But neither of these are expressible with unknown rank cond.

// TODO(cwhipkey): improve this case.

}

}

} else if (c->RankKnown(cond)) {

if (c->Rank(cond) == 1) {

// cond is vector, so data is either the same shape, or a shape with

// higher rank or the same first dimension.

// TODO(cwhipkey): make the call to WithRankAtLeast do something when

// <data> is known. Then we could assert the first dimensions are the

// same.

// TF_RETURN_IF_ERROR(c->WithRankAtLeast(data, 1, &data));

} else {

// If cond is a non-vector, it must be the same shape as data.

TF_RETURN_IF_ERROR(c->Merge(data, cond, &data));

}

}

c->set_output(0, data);

return Status::OK();

})

.Doc(R"doc(

Selects elements from `t` or `e`, depending on `condition`.

The `t`, and `e` tensors must all have the same shape,

and the output will also have that shape. The `condition` tensor

must be a scalar if `t` and `e` are scalars. If `t` and `e` are vectors

or higher rank, then `condition` must be either a vector with size

matching the first dimension of `t`, or must have the same shape as `t`.

The `condition` tensor acts as a mask that chooses, based on the value at each

element, whether the corresponding element / row in the output should be

taken from `t` (if true) or `e` (if false).

If `condition` is a vector and `t` and `e` are higher rank matrices, then

it chooses which row (outer dimension) to copy from `t` and `e`.

If `condition` has the same shape as `t` and `e`, then it chooses which

element to copy from `t` and `e`.

For example:

```prettyprint

# 'condition' tensor is [[True, False]

# [False, True]]

# 't' is [[1, 2],

# [3, 4]]

# 'e' is [[5, 6],

# [7, 8]]

select(condition, t, e) ==> [[1, 6],

[7, 4]]

# 'condition' tensor is [True, False]

# 't' is [[1, 2],

# [3, 4]]

# 'e' is [[5, 6],

# [7, 8]]

select(condition, t, e) ==> [[1, 2],

[7, 8]]

```

t:= A `Tensor` which may have the same shape as `condition`.

If `condition` is rank 1, `t` may have higher rank,

but its first dimension must match the size of `condition`.

e:= A `Tensor` with the same type and shape as `t`.

output:= A `Tensor` with the same type and shape as `t` and `e`.

)doc");

// --------------------------------------------------------------------------

REGISTER_OP("MatMul")

.Input("a: T")

.Input("b: T")

.Output("product: T")

.Attr("transpose_a: bool = false")

.Attr("transpose_b: bool = false")

.Attr("T: {half, float, double, int32, complex64, complex128}")

.SetShapeFn(shape_inference::MatMulShape)

.Doc(R"doc(

Multiply the matrix "a" by the matrix "b".

The inputs must be two-dimensional matrices and the inner dimension of

"a" (after being transposed if transpose_a is true) must match the

outer dimension of "b" (after being transposed if transposed_b is

true).

*Note*: The default kernel implementation for MatMul on GPUs uses

cublas.

transpose_a: If true, "a" is transposed before multiplication.

transpose_b: If true, "b" is transposed before multiplication.

)doc");

REGISTER_OP("SparseMatMul")

.Input("a: Ta")

.Input("b: Tb")

.Output("product: float")

.Attr("transpose_a: bool = false")

.Attr("transpose_b: bool = false")

.Attr("a_is_sparse: bool = false")

.Attr("b_is_sparse: bool = false")

.Attr("Ta: {float, bfloat16} = DT_FLOAT")

.Attr("Tb: {float, bfloat16} = DT_FLOAT")

.SetShapeFn(shape_inference::MatMulShape)

.Doc(R"doc(

Multiply matrix "a" by matrix "b".

The inputs must be two-dimensional matrices and the inner dimension of "a" must

match the outer dimension of "b". This op is optimized for the case where at

least one of "a" or "b" is sparse. The breakeven for using this versus a dense

matrix multiply on one platform was 30% zero values in the sparse matrix.

)doc");

// --------------------------------------------------------------------------

// For operations where the output is a reduction function along some

// dimensions of the input.

REGISTER_OP("Sum")

.Input("input: T")

.Input("reduction_indices: int32")

.Output("output: T")

.Attr("keep_dims: bool = false")

.Attr("T: numbertype")

.Doc(R"doc(

Computes the sum of elements across dimensions of a tensor.

Reduces `input` along the dimensions given in `reduction_indices`. Unless

`keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in

`reduction_indices`. If `keep_dims` is true, the reduced dimensions are

retained with length 1.

input: The tensor to reduce.

reduction_indices: The dimensions to reduce.

keep_dims: If true, retain reduced dimensions with length 1.

output: The reduced tensor.

)doc");

REGISTER_OP("Mean")

.Input("input: T")

.Input("reduction_indices: int32")

.Output("output: T")

.Attr("keep_dims: bool = false")

.Attr("T: numbertype")

.Doc(R"doc(

Computes the mean of elements across dimensions of a tensor.

Reduces `input` along the dimensions given in `reduction_indices`. Unless

`keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in

`reduction_indices`. If `keep_dims` is true, the reduced dimensions are

retained with length 1.

input: The tensor to reduce.

reduction_indices: The dimensions to reduce.

keep_dims: If true, retain reduced dimensions with length 1.

output: The reduced tensor.

)doc");

REGISTER_OP("Prod")

.Input("input: T")

.Input("reduction_indices: int32")

.Output("output: T")

.Attr("keep_dims: bool = false")

.Attr("T: numbertype")

.Doc(R"doc(

Computes the product of elements across dimensions of a tensor.

Reduces `input` along the dimensions given in `reduction_indices`. Unless

`keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in

`reduction_indices`. If `keep_dims` is true, the reduced dimensions are

retained with length 1.

input: The tensor to reduce.

reduction_indices: The dimensions to reduce.

keep_dims: If true, retain reduced dimensions with length 1.

output: The reduced tensor.

)doc");

REGISTER_OP("Min")

.Input("input: T")

.Input("reduction_indices: int32")

.Output("output: T")

.Attr("keep_dims: bool = false")

.Attr("T: numbertype")

.Doc(R"doc(

Computes the minimum of elements across dimensions of a tensor.

Reduces `input` along the dimensions given in `reduction_indices`. Unless

`keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in

`reduction_indices`. If `keep_dims` is true, the reduced dimensions are

retained with length 1.

input: The tensor to reduce.

reduction_indices: The dimensions to reduce.

keep_dims: If true, retain reduced dimensions with length 1.