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

#include "tensorflow/core/util/mirror_pad_mode.h"

#include "tensorflow/core/util/padding.h"

namespace tensorflow {

using shape_inference::Dimension;

using shape_inference::InferenceContext;

using shape_inference::Shape;

namespace {

Status GetAxisForPackAndUnpack(InferenceContext* c, int32 rank_after_pack,

int32* axis) {

TF_RETURN_IF_ERROR(c->GetAttr("axis", axis));

if (*axis < -1 * rank_after_pack || *axis >= rank_after_pack) {

return errors::InvalidArgument("Invalid axis: ", *axis, "; must be in [",

-1 * rank_after_pack, ",", rank_after_pack,

")");

}

if (*axis < 0) *axis = (rank_after_pack + *axis);

return Status::OK();

}

Status PadShapeFn(InferenceContext* c) {

// Paddings is a matrix of [input_rank, 2].

const Shape* paddings;

TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &paddings));

const Dimension* unused;

TF_RETURN_IF_ERROR(c->WithValue(c->Dim(paddings, 1), 2, &unused));

// n_dim and input.rank are equivalent.

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

const Dimension* n_dim = c->Dim(paddings, 0);

if (c->ValueKnown(n_dim)) {

TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(n_dim), &input));

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

TF_RETURN_IF_ERROR(c->WithValue(n_dim, c->Rank(input), &n_dim));

}

const Tensor* paddings_t = c->input_tensor(1);

// paddings_t is unknown

if (paddings_t == nullptr) {

if (c->ValueKnown(n_dim)) {

// Make output with n_dim unknown dims.

std::vector<const Dimension*> dims;

const auto value = c->Value(n_dim);

for (int i = 0; i < value; ++i) dims.push_back(c->UnknownDim());

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

} else {

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

}

return Status::OK();

}

// tensor value was provided for paddings_t; doublecheck n_dim value is the

// same.

const auto num_dims = c->Value(n_dim);

DCHECK_EQ(num_dims, paddings_t->shape().dim_size(0));

// paddings_t is known.

auto paddings_data = paddings_t->matrix<int32>();

std::vector<const Dimension*> dims(num_dims);

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

const int32 pad0 = paddings_data(i, 0);

const int32 pad1 = paddings_data(i, 1);

if (pad0 < 0 || pad1 < 0) {

return errors::InvalidArgument("Paddings must be non-negative");

}

TF_RETURN_IF_ERROR(c->Add(c->Dim(input, i), pad0 + pad1, &dims[i]));

}

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

return Status::OK();

}

} // namespace

REGISTER_OP("Pack")

.Input("values: N * T")

.Output("output: T")

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

.Attr("T: type")

.Attr("axis: int = 0")

.SetShapeFn([](InferenceContext* c) {

// Validate shapes of all inputs are compatible

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.");

}

if (!c->RankKnown(cur)) {

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

return Status::OK();

}

// Determine the axis that will be added, converting from negative

// axes to a positive point per negative indexing rules.

int32 rank = c->Rank(cur);

int32 axis;

TF_RETURN_IF_ERROR(GetAxisForPackAndUnpack(c, rank + 1, &axis));

// Copy all dimensions over, inserting a dimension of value #inputs

// at <axis>.

std::vector<const Dimension*> dims;

int index = 0;

while (index < axis) dims.push_back(c->Dim(cur, index++));

dims.push_back(c->MakeDim(c->num_inputs()));

while (index < rank) dims.push_back(c->Dim(cur, index++));

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

return Status::OK();

})

.Doc(R"doc(

Packs a list of `N` rank-`R` tensors into one rank-`(R+1)` tensor.

Packs the `N` tensors in `values` into a tensor with rank one higher than each

tensor in `values`, by packing them along the `axis` dimension.

Given a list of tensors of shape `(A, B, C)`;

if `axis == 0` then the `output` tensor will have the shape `(N, A, B, C)`.

if `axis == 1` then the `output` tensor will have the shape `(A, N, B, C)`.

Etc.

For example:

```prettyprint

# 'x' is [1, 4]

# 'y' is [2, 5]

# 'z' is [3, 6]

pack([x, y, z]) => [[1, 4], [2, 5], [3, 6]] # Pack along first dim.

pack([x, y, z], axis=1) => [[1, 2, 3], [4, 5, 6]]

```

This is the opposite of `unpack`.

values: Must be of same shape and type.

axis: Dimension along which to pack. Negative values wrap around, so the

valid range is `[-(R+1), R+1)`.

output: The packed tensor.

)doc");

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

REGISTER_OP("Unpack")

.Input("value: T")

.Output("output: num * T")

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

.Attr("T: type")

.Attr("axis: int = 0")

.SetShapeFn([](InferenceContext* c) {

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

const Shape* out;

if (c->RankKnown(s)) {

// Determine the axis that will be removed, converting from negative

// axes to a positive point per negative indexing rules.

int32 rank = c->Rank(s);

int32 axis;

TF_RETURN_IF_ERROR(GetAxisForPackAndUnpack(c, rank, &axis));

// The axis dim matches the number of outputs.

const Dimension* unused;

TF_RETURN_IF_ERROR(

c->WithValue(c->Dim(s, axis), c->num_outputs(), &unused));

// Copy all dimensions, removing the <axis> dimension.

std::vector<const Dimension*> dims;

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

if (i != axis) dims.push_back(c->Dim(s, i));

}

out = c->MakeShape(dims);

} else {

// All outputs are the same shape, but it's not known.

out = c->UnknownShape();

}

for (int i = 0; i < c->num_outputs(); ++i) c->set_output(i, out);

return Status::OK();

})

.Doc(R"doc(

Unpacks a given dimension of a rank-`R` tensor into `num` rank-`(R-1)` tensors.

Unpacks `num` tensors from `value` by chipping it along the `axis` dimension.

For example, given a tensor of shape `(A, B, C, D)`;

If `axis == 0` then the i'th tensor in `output` is the slice `value[i, :, :, :]`

and each tensor in `output` will have shape `(B, C, D)`. (Note that the

dimension unpacked along is gone, unlike `split`).

If `axis == 1` then the i'th tensor in `output` is the slice `value[:, i, :, :]`

and each tensor in `output` will have shape `(A, C, D)`.

Etc.

This is the opposite of `pack`.

value: 1-D or higher, with `axis` dimension size equal to `num`.

axis: Dimension along which to unpack. Negative values wrap around, so the

valid range is `[-R, R)`.

output: The list of tensors unpacked from `value`.

)doc");

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

// TODO(josh11b): Remove the >= 2 constraint, once we can rewrite the graph

// in the N == 1 case to remove the node.

REGISTER_OP("Concat")

.Input("concat_dim: int32")

.Input("values: N * T")

.Output("output: T")

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

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* unused;

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

const Tensor* concat_dim_t = c->input_tensor(0);

if (concat_dim_t == nullptr) {

// Return an unknown shape with same rank as inputs, or an unknown rank

// if no input's rank is known.

// Find rank.

int32 rank = InferenceContext::kUnknownRank;

for (int i = 1; i < c->num_inputs(); ++i) {

if (rank == InferenceContext::kUnknownRank)

rank = c->Rank(c->input(i));

if (rank != InferenceContext::kUnknownRank) {

TF_RETURN_IF_ERROR(c->WithRank(c->input(i), rank, &unused));

}

}

if (rank == InferenceContext::kUnknownRank) {

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

return Status::OK();

}

if (rank == 0) {

return errors::InvalidArgument(

"Can't concatenate scalars (use tf.pack instead)");

}

// Build result of <rank> different unknown dims.

std::vector<const Dimension*> dims;

for (int i = 0; i < rank; ++i) dims.push_back(c->UnknownDim());

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

return Status::OK();

}

// Merge all the non-concat dims, and sum the concat dim to make an output

// shape.

const int32 concat_dim = concat_dim_t->scalar<int32>()();

if (concat_dim < 0) {

return errors::InvalidArgument("Expected concat_dim >= 0, but got ",

concat_dim);

}

const Shape* output_before;

const Shape* output_after;

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

TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, concat_dim + 1, &input));

TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &output_before));

const Dimension* output_middle = c->Dim(input, concat_dim);

TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &output_after));

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

const Shape* before;

const Shape* after;

input = c->input(i);

TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, concat_dim + 1, &input));

TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &before));

const Dimension* middle = c->Dim(input, concat_dim);

TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &after));

TF_RETURN_IF_ERROR(c->Merge(before, output_before, &output_before));

TF_RETURN_IF_ERROR(c->Add(output_middle, middle, &output_middle));

TF_RETURN_IF_ERROR(c->Merge(after, output_after, &output_after));

}

const Shape* s;

TF_RETURN_IF_ERROR(

c->Concatenate(output_before, c->Vector(output_middle), &s));

TF_RETURN_IF_ERROR(c->Concatenate(s, output_after, &s));

c->set_output(0, s);

return Status::OK();

})

.Doc(R"doc(

Concatenates tensors along one dimension.

concat_dim: 0-D. The dimension along which to concatenate. Must be in the

range [0, rank(values)).

values: The `N` Tensors to concatenate. Their ranks and types must match,

and their sizes must match in all dimensions except `concat_dim`.

output: A `Tensor` with the concatenation of values stacked along the

`concat_dim` dimension. This tensor's shape matches that of `values` except

in `concat_dim` where it has the sum of the sizes.

)doc");

REGISTER_OP("ConcatOffset")

.Input("concat_dim: int32")

.Input("shape: N * int32")

.Output("offset: N * int32")

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

.SetShapeFn([](InferenceContext* c) {

for (int i = 1; i < c->num_inputs(); ++i) {

c->set_output(i - 1, c->input(i));

}

return Status::OK();

})

.Doc(R"doc(

Computes offsets of concat inputs within its output.

For example:

```prettyprint

# 'x' is [2, 2, 7]

# 'y' is [2, 3, 7]

# 'z' is [2, 5, 7]

concat_offset(2, [x, y, z]) => [0, 0, 0], [0, 2, 0], [0, 5, 0]

```

concat_dim: The dimension along which to concatenate.

shape: The `N` int32 vectors representing shape of tensors being concatenated.

offset: The `N` int32 vectors representing the starting offset

of input tensors within the concatenated output.

This is typically used by gradient computations for a concat operation.

)doc");

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

REGISTER_OP("Split")

.Input("split_dim: int32")

.Input("value: T")

.Output("output: num_split * T")

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

.Attr("T: type")

.Doc(R"doc(

Splits a tensor into `num_split` tensors along one dimension.

split_dim: 0-D. The dimension along which to split. Must be in the range

`[0, rank(value))`.

num_split: The number of ways to split. Must evenly divide

`value.shape[split_dim]`.

value: The tensor to split.

output: They are identically shaped tensors, whose shape matches that of `value`

except along `split_dim`, where their sizes are

`values.shape[split_dim] / num_split`.

)doc");

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

REGISTER_OP("Const")

.Output("output: dtype")

.Attr("value: tensor")

.Attr("dtype: type")

.SetShapeFn([](InferenceContext* c) {

const TensorProto* proto = nullptr;

TF_RETURN_IF_ERROR(c->GetAttr("value", &proto));

TF_RETURN_IF_ERROR(TensorShape::IsValidShape(proto->tensor_shape()));

TensorShape shape(proto->tensor_shape());

std::vector<const Dimension*> dims;

for (int i = 0; i < shape.dims(); ++i) {

dims.push_back(c->MakeDim(shape.dim_size(i)));

}

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

return Status::OK();

})

.Doc(R"doc(

Returns a constant tensor.

value: Attr `value` is the tensor to return.

)doc");

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

// TODO(mgubin): Update the doc when the freeze_graph script supports converting

// into memmapped format.

REGISTER_OP("ImmutableConst")

.Attr("dtype: type")

.Attr("shape: shape")

.Attr("memory_region_name: string")

.Output("tensor: dtype")

.SetShapeFn([](InferenceContext* c) {

TensorShape shape_from_attr;

TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape_from_attr));

TensorShapeProto shape_proto;

shape_from_attr.AsProto(&shape_proto);

const Shape* output_shape;

TF_RETURN_IF_ERROR(

c->MakeShapeFromShapeProto(shape_proto, &output_shape));

c->set_output(0, output_shape);

return Status::OK();

})

.Doc(R"doc(

Returns immutable tensor from memory region.

The current implementation memmaps the tensor from a file.

dtype: Type of the returned tensor.

shape: Shape of the returned tensor.

memory_region_name: Name of readonly memory region used by the tensor, see

NewReadOnlyMemoryRegionFromFile in tensorflow::Env.

)doc");

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

REGISTER_OP("ZerosLike")

.Input("x: T")

.Output("y: T")

.Attr("T: type")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Returns a tensor of zeros with the same shape and type as x.

x: a tensor of type T.

y: a tensor of the same shape and type as x but filled with zeros.

)doc");

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

REGISTER_OP("Diag")

.Input("diagonal: T")

.Output("output: T")

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

.SetShapeFn([](InferenceContext* c) {

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

TF_RETURN_IF_ERROR(c->WithRankAtMost(in, 3, &in));

// Output shape is original concatenated with itself.

const Shape* out;

TF_RETURN_IF_ERROR(c->Concatenate(in, in, &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Returns a diagonal tensor with a given diagonal values.

Given a `diagonal`, this operation returns a tensor with the `diagonal` and

everything else padded with zeros. The diagonal is computed as follows:

Assume `diagonal` has dimensions [D1,..., Dk], then the output is a tensor of

rank 2k with dimensions [D1,..., Dk, D1,..., Dk] where:

`output[i1,..., ik, i1,..., ik] = diagonal[i1, ..., ik]` and 0 everywhere else.

For example:

```prettyprint

# 'diagonal' is [1, 2, 3, 4]

tf.diag(diagonal) ==> [[1, 0, 0, 0]

[0, 2, 0, 0]

[0, 0, 3, 0]

[0, 0, 0, 4]]

```

diagonal: Rank k tensor where k is at most 3.

)doc");

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

REGISTER_OP("DiagPart")

.Input("input: T")

.Output("diagonal: T")

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

.SetShapeFn([](InferenceContext* c) {

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

if (!c->RankKnown(in)) {

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

return Status::OK();

}

// Rank must be even, and result will have rank <rank/2>.

const int32 rank = c->Rank(in);

if ((rank % 2) != 0 || rank > 6) {

return errors::InvalidArgument(

"Input must have even rank <= 6, input rank is ", rank);

}

const int32 mid = rank / 2;

// output dim[i] is the merge of in.dim[i] and in.dim[i+mid].

std::vector<const Dimension*> dims(mid);

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

TF_RETURN_IF_ERROR(

c->Merge(c->Dim(in, i), c->Dim(in, i + mid), &dims[i]));

}

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

return Status::OK();

})

.Doc(R"doc(

Returns the diagonal part of the tensor.

This operation returns a tensor with the `diagonal` part

of the `input`. The `diagonal` part is computed as follows:

Assume `input` has dimensions `[D1,..., Dk, D1,..., Dk]`, then the output is a

tensor of rank `k` with dimensions `[D1,..., Dk]` where:

`diagonal[i1,..., ik] = input[i1, ..., ik, i1,..., ik]`.

For example:

```prettyprint

# 'input' is [[1, 0, 0, 0]

[0, 2, 0, 0]

[0, 0, 3, 0]

[0, 0, 0, 4]]

tf.diag_part(input) ==> [1, 2, 3, 4]

```

input: Rank k tensor where k is 2, 4, or 6.

diagonal: The extracted diagonal.

)doc");

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

REGISTER_OP("BatchMatrixDiag")

.Input("diagonal: T")

.Output("output: T")

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* in;

TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &in));

if (!c->RankKnown(in)) {

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

return Status::OK();

}

const int32 rank = c->Rank(in);

const Shape* out;

TF_RETURN_IF_ERROR(

c->Concatenate(in, c->Vector(c->Dim(in, rank - 1)), &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Returns a batched diagonal tensor with a given batched diagonal values.

Given a `diagonal`, this operation returns a tensor with the `diagonal` and

everything else padded with zeros. The diagonal is computed as follows:

Assume `diagonal` has `k` dimensions `[I, J, K, ..., N]`, then the output is a

tensor of rank `k+1` with dimensions [I, J, K, ..., N, N]` where:

`output[i, j, k, ..., m, n] = 1{m=n} * diagonal[i, j, k, ..., n]`.

For example:

```prettyprint

# 'diagonal' is [[1, 2, 3, 4], [5, 6, 7, 8]]

and diagonal.shape = (2, 4)

tf.batch_matrix_diag(diagonal) ==> [[[1, 0, 0, 0]

[0, 2, 0, 0]

[0, 0, 3, 0]

[0, 0, 0, 4]],

[[5, 0, 0, 0]

[0, 6, 0, 0]

[0, 0, 7, 0]

[0, 0, 0, 8]]]

which has shape (2, 4, 4)

```

diagonal: Rank `k`, where `k >= 1`.

output: Rank `k+1`, with `output.shape = diagonal.shape + [diagonal.shape[-1]]`.

)doc");

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

REGISTER_OP("BatchMatrixSetDiag")

.Input("input: T")

.Input("diagonal: T")

.Output("output: T")

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* input;

const Shape* diag;

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

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

const Dimension* square_dim;

TF_RETURN_IF_ERROR(

c->Merge(c->Dim(input, -2), c->Dim(input, -1), &square_dim));

TF_RETURN_IF_ERROR(c->Merge(square_dim, c->Dim(diag, -1), &square_dim));

const Shape* output;

TF_RETURN_IF_ERROR(c->Concatenate(diag, c->Vector(square_dim), &output));

TF_RETURN_IF_ERROR(c->Merge(input, output, &output));

c->set_output(0, output);

return Status::OK();

})

.Doc(R"doc(

Returns a batched matrix tensor with new batched diagonal values.

Given `input` and `diagonal`, this operation returns a tensor with the

same shape and values as `input`, except for the diagonals of the innermost

matrices. These will be overwritten by the values in `diagonal`.

The batched matrices must be square.

The output is computed as follows:

Assume `input` has `k+1` dimensions `[I, J, K, ..., N, N]` and `diagonal` has

`k` dimensions `[I, J, K, ..., N]`. Then the output is a

tensor of rank `k+1` with dimensions [I, J, K, ..., N, N]` where:

* `output[i, j, k, ..., m, n] = diagonal[i, j, k, ..., n]` for `m == n`.

* `output[i, j, k, ..., m, n] = input[i, j, k, ..., m, n]` for `m != n`.

input: Rank `k+1`, where `k >= 1`.

diagonal: Rank `k`, where `k >= 1`.

output: Rank `k+1`, with `output.shape = input.shape`.

)doc");

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

REGISTER_OP("BatchMatrixDiagPart")

.Input("input: T")

.Output("diagonal: T")

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* in;

TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &in));

if (!c->RankKnown(in)) {

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

return Status::OK();

}

const int32 rank = c->Rank(in);

// Last two dims must match.

const Dimension* unused;

TF_RETURN_IF_ERROR(

c->Merge(c->Dim(in, rank - 1), c->Dim(in, rank - 2), &unused));

// Output shape has all dims but last of input.

std::vector<const Dimension*> dims;

for (int i = 0; i < rank - 1; ++i) dims.push_back(c->Dim(in, i));

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

return Status::OK();

})

.Doc(R"doc(

Returns the batched diagonal part of a batched tensor.

This operation returns a tensor with the `diagonal` part

of the batched `input`. The `diagonal` part is computed as follows:

Assume `input` has `k` dimensions `[I, J, K, ..., N, N]`, then the output is a

tensor of rank `k - 1` with dimensions `[I, J, K, ..., N]` where:

`diagonal[i, j, k, ..., n] = input[i, j, k, ..., n, n]`.

The input must be at least a matrix.

For example:

```prettyprint

# 'input' is [[[1, 0, 0, 0]

[0, 2, 0, 0]

[0, 0, 3, 0]

[0, 0, 0, 4]],

[[5, 0, 0, 0]

[0, 6, 0, 0]

[0, 0, 7, 0]

[0, 0, 0, 8]]]

and input.shape = (2, 4, 4)

tf.batch_matrix_diag_part(input) ==> [[1, 2, 3, 4], [5, 6, 7, 8]]

which has shape (2, 4)

```

input: Rank `k` tensor where `k >= 2` and the last two dimensions are equal.

diagonal: The extracted diagonal(s) having shape

`diagonal.shape = input.shape[:-1]`.

)doc");

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

REGISTER_OP("BatchMatrixBandPart")

.Input("input: T")

.Input("num_lower: int64")

.Input("num_upper: int64")

.Output("band: T")

.Attr("T: type")

.SetShapeFn(shape_inference::UnchangedShape)

.Doc(R"doc(

Copy a tensor setting everything outside a central band in each innermost matrix

to zero.

The `band` part is computed as follows:

Assume `input` has `k` dimensions `[I, J, K, ..., M, N]`, then the output is a

tensor with the same shape where

`band[i, j, k, ..., m, n] = in_band(m, n) * input[i, j, k, ..., m, n]`.

The indicator function 'in_band(m, n)` is one if

`(num_lower < 0 || (m-n) <= num_lower)) &&

(num_upper < 0 || (n-m) <= num_upper)`, and zero otherwise.

For example:

```prettyprint

# if 'input' is [[ 0, 1, 2, 3]

[-1, 0, 1, 2]

[-2, -1, 0, 1]

[-3, -2, -1, 0]],

tf.batch_matrix_band_part(input, 1, -1) ==> [[ 0, 1, 2, 3]

[-1, 0, 1, 2]

[ 0, -1, 0, 1]

[ 0, 0, -1, 0]],

tf.batch_matrix_band_part(input, 2, 1) ==> [[ 0, 1, 0, 0]

[-1, 0, 1, 0]

[-2, -1, 0, 1]

[ 0, -2, -1, 0]]

```

Useful special cases:

```prettyprint

tf.batch_matrix_band_part(input, 0, -1) ==> Upper triangular part.

tf.batch_matrix_band_part(input, -1, 0) ==> Lower triangular part.

tf.batch_matrix_band_part(input, 0, 0) ==> Diagonal.

```

input: Rank `k` tensor.

num_lower: 0-D tensor. Number of subdiagonals to keep. If negative, keep entire

lower triangle.

num_upper: 0-D tensor. Number of superdiagonals to keep. If negative, keep

entire upper triangle.

band: Rank `k` tensor of the same shape as input. The extracted banded tensor.

)doc");

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

REGISTER_OP("Reverse")

.Input("tensor: T")

.Input("dims: bool")

.Output("output: T")

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

.SetShapeFn([](InferenceContext* c) {

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

const Shape* dims;

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

const Dimension* dims_dim = c->Dim(dims, 0);

if (c->ValueKnown(dims_dim)) {

TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(dims_dim), &input));

}

if (c->Rank(input) > 8) {

return errors::InvalidArgument(

"reverse does not work on tensors with more than 8 dimensions");

}

c->set_output(0, input);

return Status::OK();

})

.Doc(R"Doc(

Reverses specific dimensions of a tensor.

Given a `tensor`, and a `bool` tensor `dims` representing the dimensions

of `tensor`, this operation reverses each dimension i of `tensor` where

`dims[i]` is `True`.

`tensor` can have up to 8 dimensions. The number of dimensions

of `tensor` must equal the number of elements in `dims`. In other words:

`rank(tensor) = size(dims)`

For example:

```prettyprint

# tensor 't' is [[[[ 0, 1, 2, 3],

# [ 4, 5, 6, 7],

# [ 8, 9, 10, 11]],

# [[12, 13, 14, 15],

# [16, 17, 18, 19],

# [20, 21, 22, 23]]]]

# tensor 't' shape is [1, 2, 3, 4]

# 'dims' is [False, False, False, True]

reverse(t, dims) ==> [[[[ 3, 2, 1, 0],

[ 7, 6, 5, 4],

[ 11, 10, 9, 8]],

[[15, 14, 13, 12],

[19, 18, 17, 16],

[23, 22, 21, 20]]]]

# 'dims' is [False, True, False, False]

reverse(t, dims) ==> [[[[12, 13, 14, 15],

[16, 17, 18, 19],

[20, 21, 22, 23]

[[ 0, 1, 2, 3],

[ 4, 5, 6, 7],

[ 8, 9, 10, 11]]]]

# 'dims' is [False, False, True, False]

reverse(t, dims) ==> [[[[8, 9, 10, 11],

[4, 5, 6, 7],

[0, 1, 2, 3]]

[[20, 21, 22, 23],

[16, 17, 18, 19],

[12, 13, 14, 15]]]]

```

tensor: Up to 8-D.

dims: 1-D. The dimensions to reverse.

output: The same shape as `tensor`.

)Doc");

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

REGISTER_OP("EditDistance")

.Input("hypothesis_indices: int64")

.Input("hypothesis_values: T")

.Input("hypothesis_shape: int64")

.Input("truth_indices: int64")

.Input("truth_values: T")

.Input("truth_shape: int64")

.Attr("normalize: bool = true")

.Attr("T: type")

.Output("output: float")

.Doc(R"doc(

Computes the (possibly normalized) Levenshtein Edit Distance.

The inputs are variable-length sequences provided by SparseTensors

(hypothesis_indices, hypothesis_values, hypothesis_shape)

and

(truth_indices, truth_values, truth_shape).

The inputs are:

hypothesis_indices: The indices of the hypothesis list SparseTensor.

This is an N x R int64 matrix.

hypothesis_values: The values of the hypothesis list SparseTensor.

This is an N-length vector.

hypothesis_shape: The shape of the hypothesis list SparseTensor.

This is an R-length vector.

truth_indices: The indices of the truth list SparseTensor.

This is an M x R int64 matrix.

truth_values: The values of the truth list SparseTensor.

This is an M-length vector.

truth_shape: The shape of the truth list SparseTensor.

This is an R-length vector.

truth_shape: truth indices, vector.

normalize: boolean (if true, edit distances are normalized by length of truth).

The output is:

output: A dense float tensor with rank R - 1.

For the example input:

// hypothesis represents a 2x1 matrix with variable-length values:

// (0,0) = ["a"]

// (1,0) = ["b"]

hypothesis_indices = [[0, 0, 0],

[1, 0, 0]]

hypothesis_values = ["a", "b"]

hypothesis_shape = [2, 1, 1]

// truth represents a 2x2 matrix with variable-length values:

// (0,0) = []

// (0,1) = ["a"]

// (1,0) = ["b", "c"]

// (1,1) = ["a"]

truth_indices = [[0, 1, 0],

[1, 0, 0],

[1, 0, 1],

[1, 1, 0]]

truth_values = ["a", "b", "c", "a"]

truth_shape = [2, 2, 2]

normalize = true

The output will be:

// output is a 2x2 matrix with edit distances normalized by truth lengths.

output = [[inf, 1.0], // (0,0): no truth, (0,1): no hypothesis

[0.5, 1.0]] // (1,0): addition, (1,1): no hypothesis

)doc");

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

REGISTER_OP("Fill")

.Input("dims: int32")

.Input("value: T")

.Output("output: T")

.Attr("T: type")

.SetShapeFn([](InferenceContext* c) {

const Shape* out;

TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Creates a tensor filled with a scalar value.

This operation creates a tensor of shape `dims` and fills it with `value`.

For example:

```prettyprint

# Output tensor has shape [2, 3].

fill([2, 3], 9) ==> [[9, 9, 9]

[9, 9, 9]]

```

dims: 1-D. Represents the shape of the output tensor.

value: 0-D (scalar). Value to fill the returned tensor.

)doc");

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

REGISTER_OP("Gather")

.Input("params: Tparams")

.Input("indices: Tindices")

.Attr("validate_indices: bool = true")

.Output("output: Tparams")

.Attr("Tparams: type")

.Attr("Tindices: {int32,int64}")

.SetShapeFn([](InferenceContext* c) {

const Shape* unused;

TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &unused));

const Shape* params_subshape;

TF_RETURN_IF_ERROR(c->Subshape(c->input(0), 1, &params_subshape));

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

const Shape* out;

TF_RETURN_IF_ERROR(c->Concatenate(indices_shape, params_subshape, &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Gather slices from `params` according to `indices`.

`indices` must be an integer tensor of any dimension (usually 0-D or 1-D).

Produces an output tensor with shape `indices.shape + params.shape[1:]` where:

# Scalar indices

output[:, ..., :] = params[indices, :, ... :]

# Vector indices

output[i, :, ..., :] = params[indices[i], :, ... :]

# Higher rank indices

output[i, ..., j, :, ... :] = params[indices[i, ..., j], :, ..., :]

If `indices` is a permutation and `len(indices) == params.shape[0]` then

this operation will permute `params` accordingly.

<div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">

<img style="width:100%" src="../../images/Gather.png" alt>

</div>

)doc");

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

REGISTER_OP("GatherNd")

.Input("params: Tparams")

.Input("indices: Tindices")

.Output("output: Tparams")

.Attr("Tparams: type")

.Attr("Tindices: {int32,int64}")

.SetShapeFn([](InferenceContext* c) {

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

const Shape* indices;

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

const Dimension* r_dim = c->Dim(indices, -1);

if (!c->RankKnown(params) || !c->ValueKnown(r_dim)) {

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

return Status::OK();

}

if (c->Value(r_dim) > c->Rank(params)) {

return errors::InvalidArgument(

"indices.shape[-1] must be <= params.rank, but saw indices shape: ",

c->DebugString(indices), " and params shape: ",

c->DebugString(params));

}

// Remove r_dim from indices to get output.

const Shape* indices_slice;

const Shape* params_slice;

TF_RETURN_IF_ERROR(c->Subshape(indices, 0, -1, &indices_slice));

TF_RETURN_IF_ERROR(c->Subshape(params, c->Value(r_dim), &params_slice));

const Shape* out;

TF_RETURN_IF_ERROR(c->Concatenate(indices_slice, params_slice, &out));

c->set_output(0, out);

return Status::OK();

})

.Doc(R"doc(

Gather values or slices from `params` according to `indices`.

`params` is a Tensor of rank `R` and `indices` is a Tensor of rank `M`.

`indices` must be integer tensor, containing indices into `params`.

It must be shape `[d_0, ..., d_N, R]` where `0 < R <= M`.

The innermost dimension of `indices` (with length `R`) corresponds to

indices into elements (if `R = M`) or slices (if `R < M`) along the `N`th

dimension of `params`.

Produces an output tensor with shape