Home
last modified time | relevance | path

Searched defs:key_shape (Results 1 – 9 of 9) sorted by relevance

/aosp_15_r20/external/pytorch/torch/utils/
H A Dflop_counter.py241 def sdpa_flop_count(query_shape, key_shape, value_shape): argument
260 def sdpa_flop(query_shape, key_shape, value_shape, *args, out_shape=None, **kwargs) -> int: argument
436 def sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape): argument
464 def sdpa_backward_flop(grad_out_shape, query_shape, key_shape, value_shape, *args, out_shape=None, … argument
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dinitializable_lookup_table.h72 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function
225 TensorShape key_shape = keys_->shape(); in KeyValueTensorIterator() local
H A Dlookup_table_op.cc142 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function in tensorflow::lookup::MutableHashTableOfScalars
330 TensorShape key_shape() const final { return TensorShape(); } in key_shape() function in tensorflow::lookup::MutableHashTableOfTensors
644 TensorShape key_shape = MaybeVectorizeShape(key_shape_); in CheckKeyAndValueTensorsForImport() local
665 TensorShape key_shape() const override { return key_shape_; } in key_shape() function in tensorflow::lookup::MutableDenseHashTable
H A Dstateless_random_ops_v2.h25 TensorShape const& key_shape, in CheckKeyCounterShape()
H A Dmap_kernels.h205 TensorShape key_shape = it->first.shape(); in Compute() local
/aosp_15_r20/external/tensorflow/tensorflow/lite/kernels/
H A Dhashtable_lookup_test.cc42 std::initializer_list<int> key_shape, in HashtableLookupOpModel()
/aosp_15_r20/packages/modules/NeuralNetworks/common/cpu_operations/
DHashtableLookupTest.cpp62 std::initializer_list<uint32_t> key_shape, in HashtableLookupOpModel()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_evaluator.cc3471 Shape key_shape = sort->operand(0)->shape(); in HandleSort() local
/aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/nnapi/
H A Dnnapi_delegate_test.cc4889 std::initializer_list<int> key_shape, in HashtableLookupOpModel()