| /external/tensorflow/tensorflow/core/kernels/ |
| D | stateless_random_ops_v2.h | 25 TensorShape const& key_shape, in CheckKeyCounterShape() argument
27 if (!(key_shape.dims() == 1 && key_shape.dim_size(0) == RNG_KEY_SIZE)) { in CheckKeyCounterShape()
30 key_shape.DebugString(), in CheckKeyCounterShape()
|
| D | initializable_lookup_table.h | 71 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function
179 TensorShape key_shape = keys_->shape(); in KeyValueTensorIterator() local
180 if (!key_shape.IsSameSize(values_->shape())) { in KeyValueTensorIterator()
184 key_shape.DebugString(), " vs ", values_->shape().DebugString()); in KeyValueTensorIterator()
186 if (key_shape.num_elements() == 0) { in KeyValueTensorIterator()
|
| D | map_kernels.h | 205 TensorShape key_shape = it->first.shape(); in Compute() local
211 ctx, it->first.shape() == key_shape, in Compute()
|
| D | lookup_table_op.cc | 143 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function in tensorflow::lookup::MutableHashTableOfScalars
292 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function in tensorflow::lookup::MutableHashTableOfTensors
557 TensorShape key_shape = MaybeVectorizeShape(key_shape_); in CheckKeyAndValueTensorsForImport() local
564 expected_value_shape.RemoveLastDims(key_shape.dims()); in CheckKeyAndValueTensorsForImport()
578 TensorShape key_shape() const override { return key_shape_; } in key_shape() function in tensorflow::lookup::MutableDenseHashTable
844 output_shape.RemoveLastDims(table->key_shape().dims()); in Compute()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | lookup_interface.cc | 25 if (!TensorShapeUtils::EndsWith(shape, key_shape())) { in CheckKeyShape()
28 key_shape().DebugString()); in CheckKeyShape()
52 for (int i = 0; i < key_shape().dims(); ++i) { in CheckKeyAndValueTensorsHelper()
87 for (int i = 0; i < key_shape().dims(); ++i) { in CheckFindArguments()
|
| D | lookup_interface.h | 101 virtual TensorShape key_shape() const = 0;
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | hashtable_lookup_test.cc | 42 std::initializer_list<int> key_shape, in HashtableLookupOpModel() argument
51 BuildInterpreter({lookup_shape, key_shape, value_shape}); in HashtableLookupOpModel()
|
| /external/tensorflow/tensorflow/python/keras/layers/ |
| D | multi_head_attention.py | 296 key_shape = config.pop("key_shape")
299 if None in [query_shape, key_shape, value_shape]:
306 …layer._build_from_signature(query_shape, value_shape, key_shape) # pylint: disable=protected-acce…
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | random_ops.cc | 132 const xla::Shape key_shape = in Compile() local
143 xla::XlaOp keys = xla::RngUniform(zero, max_value, key_shape); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_evaluator.cc | 2171 Shape key_shape = sort->operand(0)->shape(); in HandleSort() local
2172 auto rank = key_shape.rank(); in HandleSort()
2181 int64 sort_dim_elements = key_shape.dimensions(sort_dim); in HandleSort()
2186 key_shape, zero_base, AsInt64Slice(key_shape.dimensions()), increment, in HandleSort()
|
| /external/tensorflow/tensorflow/lite/delegates/nnapi/ |
| D | nnapi_delegate_test.cc | 4812 std::initializer_list<int> key_shape, in HashtableLookupOpModel() argument
4821 BuildInterpreterWithNNAPI({lookup_shape, key_shape, value_shape}); in HashtableLookupOpModel()
|