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 …]
136 TensorShape ta_shape; in SetZeroValue() local137 ta_shape.AddDim(max_array_size_); in SetZeroValue()138 ta_shape.AppendShape(shape_); in SetZeroValue()140 ta_shape.dim_sizes()); in SetZeroValue()144 TensorShape ta_shape; in SetZeroValue() local145 ta_shape.AddDim(max_array_size_); in SetZeroValue()146 ta_shape.AppendShape(shape_); in SetZeroValue()149 ta_shape.dim_sizes()), in SetZeroValue()169 TensorShape ta_shape; in GetOrCreateTensorArrayGradient() local170 ta_shape.AddDim(max_array_size_); in GetOrCreateTensorArrayGradient()[all …]