Searched refs:ta_shape (Results 1 – 2 of 2) sorted by relevance
79 TensorShape ta_shape; in MaybeInitializeTensorArray() local80 ta_shape.AddDim(resource->max_array_size()); in MaybeInitializeTensorArray()81 ta_shape.AppendShape(elem_shape); in MaybeInitializeTensorArray()82 if (ta_shape != shape) { in MaybeInitializeTensorArray()84 "Mismatched TensorArray sizes: ", ta_shape.DebugString(), " vs ", in MaybeInitializeTensorArray()162 TensorShape ta_shape; in Compile() local163 ta_shape.AddDim(size); in Compile()164 ta_shape.AppendShape(shape); in Compile()166 value = xla::Broadcast(zero, ta_shape.dim_sizes()); in Compile()259 TensorShape ta_shape; in Compile() local[all …]
137 TensorShape ta_shape; in SetZeroValue() local138 ta_shape.AddDim(max_array_size_); in SetZeroValue()139 ta_shape.AppendShape(shape_); in SetZeroValue()141 ta_shape.dim_sizes()); in SetZeroValue()145 TensorShape ta_shape; in SetZeroValue() local146 ta_shape.AddDim(max_array_size_); in SetZeroValue()147 ta_shape.AppendShape(shape_); in SetZeroValue()150 ta_shape.dim_sizes()), in SetZeroValue()170 TensorShape ta_shape; in GetOrCreateTensorArrayGradient() local171 ta_shape.AddDim(max_array_size_); in GetOrCreateTensorArrayGradient()[all …]