| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cpu/ |
| H A D | SoftMaxKernel.cpp | 38 int64_t dim_size) { in _vec_log_softmax_lastdim()
118 int64_t dim_size) { in _vec_softmax_lastdim()
152 int64_t dim_size) { in _vec_softmax_lastdim()
218 int64_t dim_size) { in _vec_host_softmax_backward_lastdim()
270 int64_t dim_size) { in _vec_softmax_backward()
354 int64_t dim_size) { in _vec_softmax_backward()
481 int64_t dim_size) { in _vec_log_softmax_backward()
564 int64_t dim_size) { in _vec_log_softmax_backward()
674 int64_t dim_size = input.size(input.ndimension() - 1); in apply() local
696 int64_t dim_size) { in _vec_softmax()
[all …]
|
| H A D | SortingKernel.cpp | 44 auto dim_size = values.size(dim); in _dim_apply() local
147 int64_t dim_size, bool descending, bool stable) { in sort_kernel_impl()
196 ) { in sort_kernel()
|
| /aosp_15_r20/external/federated-compute/fcp/aggregation/core/ |
| H A D | tensor_shape.cc | 32 for (auto dim_size : dim_sizes_) { in NumElements() local
43 for (int64_t dim_size : shape_proto.dim_sizes()) { in FromProto() local
56 for (auto dim_size : dim_sizes()) { in ToProto() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/ |
| H A D | SoftMax.cu | 121 uint64_t dim_size, uint64_t inner_size) { in SpatialSoftMax_getBlockSize()
137 uint64_t outer_size, uint64_t dim_size, uint64_t inner_size, in SpatialSoftMax_getLaunchSizes()
157 inline dim3 SoftMax_getBlockSize(int ILP, uint64_t dim_size) { in SoftMax_getBlockSize()
176 inline dim3 SoftMaxForward_getBlockSize(uint64_t dim_size) { in SoftMaxForward_getBlockSize()
257 index_t outer_size, index_t dim_size, index_t inner_size) in cunn_SpatialSoftMaxForward()
315 uint32_t outer_size, uint32_t dim_size, uint32_t inner_size) in cunn_SpatialSoftMaxBackward()
830 int64_t dim_size = input.size(dim); in host_softmax() local
970 int64_t dim_size = output.size(dim); in host_softmax_backward() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | roll_op.cc | 81 gtl::InlinedVector<int32, 4> dim_size(num_dims); in Compute() local
122 const int num_dims, const gtl::ArraySlice<int32> dim_size, in DoRoll()
186 const gtl::ArraySlice<int32> dim_size, const T* input, in DoRollWithMemcpy()
|
| H A D | shape_ops.h | 54 int64_t dim_size = shape.dim_size(i); in Compute() local
83 int64_t dim_size = shape.dim_size(j); in Compute() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| H A D | shape_util.cc | 30 int64_t dim_size = input_shape.dim_size(i); in TensorShapeToConstant() local
41 int64_t dim_size = input_shape.dim_size(i); in TensorShapeToConstant() local
|
| H A D | mirror_pad_op.cc | 47 int64_t dim_size = original_shape.dimensions(dimno); in DoMirrorPad() local
133 int64_t dim_size = original_shape.dimensions(dimno); in DoMirrorPadGrad() local
|
| H A D | strided_slice_op.cc | 155 xla::XlaOp dim_size; in EmitDynamicSlice() local
489 xla::XlaOp dim_size = xla::Slice(ctx->Input(0), {i}, {i + 1}, {1}); in CompileAsDynamicUpdateSlice() local
616 auto dim_size = xla::Slice(dynamic_shape, {dim}, {dim + 1}, {1}); in Compile() local
|
| H A D | sharding_util_ops.cc | 209 const int64 dim_size = input_shape.dim_size(dim); in CompileInternal() local
350 const int64_t dim_size = output_shape[dim]; in CompileInternal() local
|
| H A D | while_op.cc | 540 int32_t dim_size = shape.dimensions(0); in Compile() local
549 int32_t dim_size = shape.dimensions(dim); in Compile() local
|
| H A D | cross_op.cc | 54 for (auto dim_size : in0_shape.dim_sizes()) { in Compile() local
|
| /aosp_15_r20/external/pytorch/torch/distributed/_shard/sharding_spec/ |
| H A D | _internals.py | 162 def get_split_size(dim_size, chunks): argument
176 def get_chunked_dim_size(dim_size, split_size, idx): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/distributed_runtime/rpc/ |
| H A D | grpc_tensor_coding.cc | 98 int64_t dim_size = val.shape().dim_size(d); in EncodeSkeleton() local
110 int64_t dim_size = val.shape().dim_size(d); in EncodeSkeleton() local
|
| /aosp_15_r20/external/ComputeLibrary/src/core/utils/helpers/ |
| H A D | tensor_transform.cpp | 60 const int dim_size = input_shape[index]; in calculate_start_on_index() local
108 const int dim_size = input_shape[index]; in calculate_end_on_index() local
|
| /aosp_15_r20/external/executorch/kernels/portable/cpu/util/ |
| H A D | copy_ops_util.cpp | 259 const ssize_t dim_size = in.size(dim); in check_unbind_copy_args() local
435 const ssize_t dim_size = in.size(dim); in check_split_with_sizes_copy_args() local
625 const ssize_t dim_size = input.size(dim); in check_split_copy_args() local
|
| H A D | activation_ops_util.cpp | 32 const size_t dim_size = in.size(non_negative_dim); in check_glu_args() local
|
| /aosp_15_r20/external/executorch/backends/vulkan/runtime/graph/ops/impl/ |
| H A D | Slice.cpp | 166 int64_t dim_size = graph.size_at<int64_t>(dim, in_ref); in get_slice_sizes() local
209 const int64_t dim_size = graph.size_at<int64_t>(dim, in_ref); in check_slice_view_args() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/ |
| H A D | FunctionalInverses.cpp | 240 auto dim_size = base.sym_size(dim); in split_Tensor_inverse() local
256 auto dim_size = base.sym_size(dim); in split_with_sizes_inverse() local
448 const auto dim_size = base.sym_size(dim); in chunk_inverse() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/ |
| H A D | tensor_layout.cc | 304 StatusOr<int64_t> Mesh::dim_size(absl::string_view name) const { in dim_size() function in tensorflow::dtensor::Mesh
831 StatusOr<int64> dim_size = mesh().dim_size(spec.sharding_spec()); in GetShardVector() local
1008 int64_t dim_size = global_shape.dim_size(spec_index); in LocalShapeFromGlobalShape() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | SoftMax.cpp | 168 int64_t dim_size = input.size(dim); local
275 int64_t dim_size = grad.size(dim); in host_softmax_backward() local
|
| H A D | TopKImpl.h | 20 const int64_t dim_size, in topk_impl_loop()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ |
| H A D | hlo_utils.h | 65 int64_t dim_size = xla_ty.dimensions(dim); in ConvertTensorShapeToType() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/ |
| H A D | xla_op_kernel.h | 276 bool is_dynamic_dimension(int64_t dim_size) { return dim_size == -1; } in is_dynamic_dimension()
|
| /aosp_15_r20/external/tensorflow/tensorflow/lite/toco/graph_transformations/ |
| H A D | resolve_constant_slice.cc | 54 int dim_size = size[i]; in Slice() local
|