Searched refs:out_dim_sizes (Results 1 – 3 of 3) sorted by relevance
86 std::vector<int64> out_dim_sizes; in SparseTensorReduceHelper() local88 out_dim_sizes.reserve(ndims); in SparseTensorReduceHelper()93 out_dim_sizes.push_back(1); // A reduced axis. in SparseTensorReduceHelper()95 out_dim_sizes.push_back(sp_shape[d]); in SparseTensorReduceHelper()99 out_dim_sizes = sp.PickDims(reduction.group_by_dims); in SparseTensorReduceHelper()102 reduction.reduced_shape = TensorShape(out_dim_sizes); in SparseTensorReduceHelper()322 auto out_dim_sizes = reduction.reduced_shape.dim_sizes(); in Compute() local323 std::copy(out_dim_sizes.begin(), out_dim_sizes.end(), &out_shape_flat(0)); in Compute()
476 const LocalOp& operand, absl::Span<const int64> out_dim_sizes, in BroadcastInDim() argument478 return xla::BroadcastInDim(operand.op(), out_dim_sizes, broadcast_dimensions); in BroadcastInDim()
244 absl::Span<const int64> out_dim_sizes,