Home
last modified time | relevance | path

Searched defs:split_dim (Results 1 – 25 of 36) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dsplit_v_op.cc66 const int32_t split_dim = in ComputeEasyCases() local
198 const TensorShape& input_shape, int32_t split_dim, in SplitHasAlignedOutputsInFirstDimension()
221 const TensorShape& input_shape, int32_t split_dim, in operator ()()
307 const int32_t split_dim = in Compute() local
386 const int32_t split_dim = in Compute() local
H A Dsplit_v_op_test.cc60 static Graph* MakeGraph(int split_dim, const std::vector<int64_t>& size_splits, in MakeGraph()
93 #define BM_SPLITV_2D(split_dim, num_split, dim0, dim1) \ argument
110 #define BM_SPLITV_3D(split_dim, num_split, dim0, dim1, dim2) \ argument
H A Dsplit_op.cc56 const int32_t split_dim = in ComputeEasyCases() local
133 const TensorShape& input_shape, int32_t split_dim, in operator ()()
214 const int32_t split_dim = in Compute() local
284 const int32_t split_dim = in Compute() local
H A Dsplit_op_test.cc29 static Graph* MakeGraph(int split_dim, int num_split, in MakeGraph()
59 #define BM_SPLIT_2D(split_dim, num_split, chunk_size0, chunk_size1) \ argument
/aosp_15_r20/external/XNNPACK/src/subgraph/
H A Deven-split.c397 size_t split_dim, in check_output_value()
478 size_t split_dim, in xnn_define_even_split_n()
625 size_t split_dim, in xnn_define_even_split2()
638 size_t split_dim, in xnn_define_even_split3()
652 size_t split_dim, in xnn_define_even_split4()
/aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/xnnpack/
H A Dsplit_tester.h36 inline SplitTester& SplitDimension(int32_t split_dim) { in SplitDimension()
62 int32_t split_dim = SplitDimension(); in OutputShape() local
H A Dsplit_tester.cc163 std::array<int32_t, 1> split_dim = {SplitDimension()}; in CreateTfLiteModel() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dsplit_op.cc49 int32_t split_dim = split_dim_orig < 0 ? split_dim_orig + input_shape.dims() in Compile() local
115 int64_t split_dim = split_dim_orig < 0 ? split_dim_orig + input_shape.dims() in Compile() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/lite/experimental/tac/transforms/
H A Ddevice_transform_patterns.cc154 RankedTensorType input_type, int split_dim, in SliceOutputs()
304 int split_dim = -1; in matchAndRewrite() local
325 auto split_dim = splitv_op.split_dim().getDefiningOp(); in matchAndRewrite() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dreduce_scatter_utils.h26 int64_t split_dim; member
H A Dall_to_all_decomposer.cc58 int64_t split_dim = *all_to_all->split_dimension(); in ExpandInstruction() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_reduce_scatter_creator.cc58 const int64_t split_dim = ar_spec->split_dim; in Run() local
/aosp_15_r20/external/pytorch/torch/distributed/pipelining/
H A Dmicrobatch.py61 def __init__(self, split_dim): argument
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dsplit_spmd_expander.cc40 int64_t split_dim, const llvm::DenseMap<int, Layout>& layouts) { in MergeLayoutsForSplitOutput()
/aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/linalg/
H A Dlinear_operator_util_test.py350 def test_blockwise_input(self, op_dimension_values, split_dim): argument
/aosp_15_r20/external/tensorflow/tensorflow/cc/framework/
H A Dgradient_checker_test.cc145 auto split_dim = Const(scope, 1, {}); in TEST() local
/aosp_15_r20/external/tensorflow/tensorflow/core/util/sparse/
H A Dsparse_tensor.h497 const int split_dim, const int num_split, in Split()
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/
H A Darray_ops.cc588 int64_t split_dim = c->Value(split_dimension); in __anon38bbb0e80a02() local
649 int64_t split_dim = c->Value(split_dimension); in __anon38bbb0e80b02() local
H A Darray_ops_test.cc1204 Tensor split_dim = test::AsTensor<int32>({1, 2}); in TEST() local
/aosp_15_r20/external/tensorflow/tensorflow/python/ops/
H A Dsparse_ops.py984 split_dim=None): argument
/aosp_15_r20/external/ComputeLibrary/src/cpu/operators/internal/
H A DCpuGemmAssemblyDispatch.cpp537 const unsigned int split_dim = scheduling_hint.split_dimension(); in run() local
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_functional_ops.cc1456 std::vector<NodeDef>& ndefs, int split_dim, in InitializeShardedVarOnTPU()
1821 int split_dim = -1; in ReplaceAndPartitionXLAShardingVariable() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dprng.cc163 int64_t split_dim; member
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/
H A Dhlo_function_importer.cc1025 uint64_t split_dim = all_to_all->split_dimension().value(); in ImportInstructionImpl() local
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DTensorIterator.cpp1684 auto split_dim = iter.get_dim_to_split(); in operator ++() local

12