Home
last modified time | relevance | path

Searched defs:all_dims (Results 1 – 11 of 11) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/
H A DBatchRulesReduceOps.cpp137 auto all_dims = range(0, std::max((int64_t)1, logical_dim)); in boxed_reduction_batch_rule() local
156 auto all_dims = range(0, self.dim() - 1); in boxed_reduction_batch_rule() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DSortingCommon.cuh144 int all_dims = self_info.dims; in run_launcher() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/quantized/cpu/
H A DReduceOps.cpp182 auto all_dims = std::vector<int64_t>(self.dim()); in is_std_inner_dim_fast_path() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dgather_scatter_handler.cc463 std::vector<int64_t> all_dims(operand.base_shape().rank()); in PartitionGatherTrivialIndexedOperandDimension() local
933 std::vector<int64_t> all_dims(indices.base_shape().rank()); in PartitionScatterIndexPassthroughPartition() local
H A Dspmd_partitioner.cc1171 std::vector<int64_t> all_dims(shape.rank()); in Replicate() local
/aosp_15_r20/external/pytorch/torch/csrc/jit/runtime/
H A Dargument_spec.h244 int32_t all_dims = 0; in CompleteArgumentSpec() local
/aosp_15_r20/external/tensorflow/tensorflow/core/grappler/clusters/
H A Dsingle_machine_test.cc184 auto all_dims = ops::Const(root.WithOpName("all_dims"), {0}, {1}); in TEST_F() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_creation_utils.cc426 std::vector<int64_t> all_dims(operand->shape().rank()); in MakeReduceHlo() local
/aosp_15_r20/external/tensorflow/tensorflow/cc/gradients/
H A Darray_grad.cc157 auto all_dims = Range(scope, Const(scope, 0), Rank(scope, grad_inputs[0]), in FillGrad() local
/aosp_15_r20/external/pytorch/functorch/csrc/dim/
H A Ddim.cpp1229 Slice<mpy::hdl<Dim>> all_dims; in run_torch_function() local
2765 bool all_dims = true; in py_split() local
/aosp_15_r20/external/pytorch/aten/src/ATen/test/
H A Dvulkan_api_test.cpp6721 std::vector<std::vector<int64_t>> all_dims; in TEST_F() local
6745 std::vector<std::vector<int64_t>> all_dims; in TEST_F() local
6786 std::vector<std::vector<int64_t>> all_dims; in TEST_F() local