1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/framework/common_shape_fns.h"
17#include "tensorflow/core/framework/op.h"
18
19namespace tensorflow {
20
21REGISTER_OP("_ScopedAllocator")
22 .Output("output: T")
23 .Attr("shapes: list(shape)")
24 .Attr("shape: shape")
25 .Attr("T: type")
26 .Attr("sa_name: string")
27 .Attr("id: int")
28 .Attr("expected_call_count: int")
29 .SetIsStateful()
30 .SetShapeFn(shape_inference::ExplicitShape)
31 .Doc(R"doc(
32Allocates a mutable tensor that becomes available to appropriately annotated
33downstream Ops as backing store for their output tensor allocations via the
34ScopedAllocatorMgr.
35Returns a reference to this value.
36
37This is an experimental op for internal use only. It is possible to use this
38op in unsafe ways.
39
40'shapes' is a list of the shapes of the tensors that are to be allocated
41by this ScopedAllocator.
42'shape' is the shape of the output of this Op, i.e. the 1D backing tensor
43from which the individual allocated tensors are aliased.
44'sa_name' is the name assigned to the Node, for connectivity specification
45and debugging.
46'id' is a non-negative integer 'scope_id' handled by the ScopedAllocatorMgr.
47'expected_call_count' is the number of individual tensors expected to
48be allocated from the backing tensor.
49)doc");
50
51REGISTER_OP("_ScopedAllocatorConcat")
52 .Output("output: T")
53 .Input("backing: T")
54 .Input("inputs: N * T")
55 .Attr("shape: shape")
56 .Attr("T: type")
57 .Attr("reshape: bool = false")
58 .Attr("sa_name: string")
59 .Attr("id: int")
60 .Attr("N: int >= 2")
61 .SetIsStateful()
62 .SetShapeFn(shape_inference::ExplicitShape)
63 .Doc(R"doc(
64Acts like a Concat Op that merges multiple tensors into one, however it must
65only be used in conjunction with a ScopedAllocator which is backing the memory
66of all of its input tensors so that actually it just outputs a read-only
67reference to that ScopedAllocator's backing tensor.
68
69This is an experimental op for internal use only. It is possible to use this
70op in unsafe ways.
71
72'backing' is the backing tensor, i.e. the output of an upstream ScopedAllocator.
73'inputs' is a list of nominal input tensors, all of which must be aliases
74to regions of the backing tensor. These will be outputs of upstream nodes
75that allocate their outputs from the same ScopedAllocator.
76'shape' is the shape of the output, which will usually be the same shape as
77the input backing tensor.
78'reshape' is true iff the output shape is to be different from that of
79the input backing tensor.
80'sa_name' is the Node name of the upstream ScopedAllocator.
81'id' is the scope_id identifying the upstream ScopedAllocator.
82'N' is the number of nominal inputs to be concatenated.
83)doc");
84
85REGISTER_OP("_ScopedAllocatorSplit")
86 .Output("output: N * T")
87 .Input("concat: T")
88 .Input("split: N * T")
89 .Attr("T: type")
90 .Attr("sa_name: string")
91 .Attr("id: int")
92 .Attr("N: int >= 2")
93 .Attr("shapes: list(shape)")
94 .SetIsStateful()
95 .SetShapeFn(shape_inference::ExplicitShapes)
96 .Doc(R"doc(
97Acts roughly like a SplitV Op that splits one tensor into multiple tensors
98but must only be used in conjunction with corresponding ScopedAllocator
99and ScopedAllocatorConcat instances. In practice it is provided as inputs
100the backing tensor as first input, which contains the concatenated values,
101and a list of alias tensors as its other input and it simply outputs that
102second list.
103
104This is an experimental op for internal use only. It is possible to use this
105op in unsafe ways.
106
107'concat' is the single output produced by an upstream ScopedAllocatorConcat
108node. This is actually the backing tensor from a ScopedAllocator node
109upstream of the ScopedAllocatorConcat.
110'split' is a list of tensors aliased from the backing tensor. It will
111become the output of this ScopedAllocatorSplit node.
112'type' is the common DataType of all of the input and output tensors.
113'sa_name' is the Node name of the upstream ScopedAllocator.
114'id' is the scope_id identifying the upstream ScopedAllocator.
115'N' is the number of split tensors.
116'shapes' is a list of the split tensor shapes.
117)doc");
118
119} // end namespace tensorflow
120