Home
last modified time | relevance | path

Searched refs:output_specs (Results 1 – 25 of 65) sorted by relevance

123

/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/experimental/rpc/
H A Drpc_ops.py38 output_specs = nest.map_structure(type_spec.type_spec_from_value,
40 output_specs_proto = nested_structure_coder.encode_structure(output_specs)
220 output_specs=None, argument
299 output_specs=get_output_specs_from_function(concrete_fn),
307 output_specs=get_output_specs_from_function(func),
358 output_specs = nested_structure_coder.decode_proto(m.output_specs)
360 self._method_registry[m.method] = output_specs
364 self._add_method(m.method, output_specs, input_specs, self._client_handle,
367 def _add_method(self, method_name, output_specs, input_specs, client_handle, argument
385 return StatusOrResult(status_or, deleter, output_specs)
[all …]
H A Drpc_ops_test.py73 output_specs=rpc_ops.get_output_specs_from_function(concrete_remote_fn),
119 output_specs=tensor_spec.TensorSpec((), dtypes.int32))
129 output_specs=tensor_spec.TensorSpec((), dtypes.int32))
176 output_specs=tensor_spec.TensorSpec((), dtypes.int32))
184 output_specs=tensor_spec.TensorSpec((), dtypes.int32))
194 output_specs=tensor_spec.TensorSpec((), dtypes.int32))
526 "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
557 output_specs=[tensor_spec.TensorSpec([], dtypes.int32)])
593 "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
599 "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
[all …]
/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A D_sharding_prop.py156 output_specs: OutputSpecType,
163 if isinstance(output_specs, DTensorSpec):
174 output_specs.tensor_meta = output_tensor_meta
175 elif isinstance(output_specs, (tuple, list)):
177 output_specs
183 for i, spec in enumerate(output_specs):
264 assert isinstance(output_strategy.output_specs, DTensorSpec)
306 output_specs: OutputSpecType = output_strategy.output_specs
307 if isinstance(output_specs, DTensorSpec):
308 output_specs = tuple(
[all …]
H A D_op_schema.py82 output_specs: Union[DTensorSpec, Tuple[Optional[DTensorSpec], ...]]
97 if isinstance(self.output_specs, DTensorSpec):
98 return self.output_specs
117 output_spec_str = _pretty_print_spec(self.output_specs)
150 output_spec = self.strategies[0].output_specs
281 args_schema.append(_pretty_print_spec(arg.strategies[0].output_specs))
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dscatter_spmd_expander.cc41 std::vector<ShardingSpec> output_specs(tensor_rank); in GetOutputLayout() local
48 output_specs[i].set_sharding_spec(Layout::kUnshardedDim); in GetOutputLayout()
54 output_specs[i] = tensor_layout->dim(i); in GetOutputLayout()
55 if (Layout::IsShardedSpec(output_specs[i])) in GetOutputLayout()
56 used_mesh_dims.emplace(output_specs[i].sharding_spec()); in GetOutputLayout()
65 if (Layout::IsUnshardedSpec(output_specs[i]) && in GetOutputLayout()
68 output_specs[i] = update_spec; in GetOutputLayout()
72 return Layout::GetLayout(output_specs, mesh); in GetOutputLayout()
H A Dgather_spmd_expander.cc314 std::vector<ShardingSpec> output_specs(params_rank - index_dimensions + in GatherNdGetOutputLayoutFromInput() local
322 output_specs[i + params_offset] = params_spec; in GatherNdGetOutputLayoutFromInput()
325 output_specs[i + params_offset].set_sharding_spec(Layout::kUnshardedDim); in GatherNdGetOutputLayoutFromInput()
331 output_specs[i] = indices_layout->dim(i); in GatherNdGetOutputLayoutFromInput()
333 output_specs[i].set_sharding_spec(Layout::kUnshardedDim); in GatherNdGetOutputLayoutFromInput()
335 return Layout::GetLayout(output_specs, mesh); in GatherNdGetOutputLayoutFromInput()
/aosp_15_r20/external/executorch/exir/
H A Dlowered_backend_module.py285 output_specs = [
315 input_specs=input_specs, output_specs=output_specs
479 output_specs = []
495 for output_spec in old_signature.output_specs:
606 output_specs.append(
617 output_specs.append(
626 output_specs.append(
635 output_specs.append(
644 input_specs=input_specs, output_specs=output_specs
866 original_program._graph_signature.output_specs = [
[all …]
/aosp_15_r20/external/pytorch/torch/export/
H A Dgraph_signature.py206 output_specs: List[OutputSpec]
277 for s in self.output_specs:
321 for s in self.output_specs
331 for s in self.output_specs
363 for spec in self.output_specs:
405 for s in self.output_specs:
429 for o in self.output_specs:
592 output_specs = [to_output_spec(idx, o) for idx, o in enumerate(outputs)]
593 return ExportGraphSignature(input_specs=input_specs, output_specs=output_specs)
H A Dexported_program.py497 output_specs = [
503 for i, spec in enumerate(ep.graph_signature.output_specs)
515 for i, node in enumerate(new_outputs[len(output_specs) :]):
526 output_specs.append(
537 input_specs=input_specs, output_specs=output_specs
928 output_spec = self.graph_signature.output_specs[i]
1063 old_signature.output_specs
1065 old_output_spec = old_signature.output_specs[i]
1078 input_specs=new_input_specs, output_specs=new_output_specs
/aosp_15_r20/external/pytorch/torch/distributed/tensor/_ops/
H A D_tensor_ops.py52 output_specs=DTensorSpec(
112 PlacementStrategy(output_specs=output_spec)
163 PlacementStrategy(output_specs=output_spec, input_specs=(arg_spec,))
198 output_specs=replica_spec,
215 output_specs=input_spec,
236 output_specs=arg_spec, input_specs=(arg_spec, replica_spec)
276 slice_strategy.strategies.append(PlacementStrategy(output_specs=out_spec))
286 PlacementStrategy(output_specs=unshard_spec)
341 PlacementStrategy(output_specs=arg_spec)
353 PlacementStrategy(output_specs=replicate_spec)
[all …]
H A Dutils.py267 output_specs = tuple(spec_list[:input_index])
270 output_specs = spec_list[0] # type: ignore[assignment]
274 output_specs=output_specs,
H A D_math_ops.py294 output_specs=DTensorSpec(
457 output_specs=replicate_spec,
494 output_specs=output_target_spec,
539 output_specs=tgt_spec,
659 output_specs=(output_expected_spec, total_weight_expected_spec),
771 output_specs=grad_in_expected_spec,
866 output_specs=output_target_spec,
1020 output_specs=tuple(output_specs_list),
/aosp_15_r20/external/pytorch/torch/csrc/jit/backends/coreml/objc/
H A DPTMCoreMLBackend.mm98 GenericList pack_outputs(const std::vector<TensorSpec>& output_specs, id<MLFeatureProvider> outputP…
100 for (const TensorSpec& spec : output_specs) {
120 if(output_specs.size() > 1){
139 std::vector<TensorSpec> output_specs;
145 output_specs = extra_json["outputs"].get<std::vector<TensorSpec>>();
150 if (!type_validity(input_specs) || !type_validity(output_specs)) {
183 model_wrapper.outputs = output_specs;
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/saved_model/
H A Dsaved_model.cc114 output_specs; in MapFunctionSignaturesFromTFSavedModelMLIR() local
125 output_specs.push_back(std::move(statusor_spec).ValueOrDie()); in MapFunctionSignaturesFromTFSavedModelMLIR()
136 sig_info.output_specs = output_specs; in MapFunctionSignaturesFromTFSavedModelMLIR()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/lite/quantization/
H A Dquantization_context.cc143 llvm::SmallVector<Attribute, 4> output_specs; in Finalize() local
149 output_specs.push_back(original_output_specs[res]); in Finalize()
151 output_specs.push_back(TypeAttr::get(requantize.params)); in Finalize()
153 output_specs.push_back(TypeAttr::get(state.params)); in Finalize()
156 op->setAttr("output_specs", ArrayAttr::get(context, output_specs)); in Finalize()
/aosp_15_r20/external/pytorch/torch/backends/_coreml/
H A Dpreprocess.py90 output_specs,
114 assert len(spec.description.output) == len(output_specs) # type: ignore[attr-defined]
116 for index, output in enumerate(output_specs):
/aosp_15_r20/external/executorch/exir/passes/
H A Dinsert_write_back_for_buffers_pass.py114 for out_spec in ep.graph_signature.output_specs
126 for output_spec in ep.graph_signature.output_specs:
137 output_specs=new_output_specs,
H A Dweights_to_outputs_pass.py38 for spec in gs.output_specs
80 gs.output_specs.append(
/aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/nnc/
H A Dcontext.h155 const std::vector<OutputSpec>& output_specs() const { in output_specs() function
159 void set_output_specs(const std::vector<OutputSpec>& output_specs) { in set_output_specs() argument
160 output_specs_ = output_specs; in set_output_specs()
H A Dcontext.cpp179 std::vector<c10::IValue> output_specs; in serialize() local
180 output_specs.reserve(output_specs_.size()); in serialize()
182 output_specs.emplace_back(output_spec.serialize()); in serialize()
184 dict.insert("output_specs", Tup(std::move(output_specs))); in serialize()
/aosp_15_r20/external/tensorflow/tensorflow/distribute/experimental/rpc/kernels/
H A Drpc_ops.cc164 StructuredValue output_specs; member
185 const StructuredValue& output_specs) { in Register() argument
192 fn_metadata.output_specs = output_specs; in Register()
290 *registered_method->mutable_output_specs() = it.second.output_specs; in List()
325 const StructuredValue& output_specs) { in Register() argument
333 input_specs, output_specs); in Register()
/aosp_15_r20/external/pytorch/test/mobile/nnc/
H A Dtest_context.cpp106 EXPECT_EQ(f2.output_specs().size(), 1); in TEST()
107 EXPECT_EQ(f2.output_specs()[0].sizes_, std::vector<int64_t>({1000})); in TEST()
108 EXPECT_EQ(f2.output_specs()[0].dtype_, at::kFloat); in TEST()
/aosp_15_r20/external/pytorch/torch/distributed/tensor/experimental/
H A D_tp_transform.py236 output_specs=_get_output_spec_from_output_sharding(output_sharding),
276 output_specs=DTensorSpec(
281 _populate_tensor_meta(node, placement.output_specs)
496 output_spec = placement_strategies[input_arg].output_specs
511 Node, lambda arg: placement_strategies[arg].output_specs, node.args
/aosp_15_r20/external/tensorflow/tensorflow/core/tfrt/saved_model/
H A Dsaved_model.cc173 DCHECK(signature.output_specs.empty()); in GetFunctionSignaturesFromTFSavedModelMLIR()
174 signature.output_specs.reserve(sig_info.output_specs.size()); in GetFunctionSignaturesFromTFSavedModelMLIR()
175 for (auto& spec : sig_info.output_specs) { in GetFunctionSignaturesFromTFSavedModelMLIR()
176 signature.output_specs.push_back(TensorSpec(spec.first, spec.second)); in GetFunctionSignaturesFromTFSavedModelMLIR()
413 signature.output_specs.reserve(signature_def.outputs().size()); in GetSignaturesFromSignatureDef()
418 signature.output_specs.push_back( in GetSignaturesFromSignatureDef()
/aosp_15_r20/external/tensorflow/tensorflow/c/experimental/saved_model/core/
H A Dsaved_model_utils.cc173 std::vector<const TensorSpecProto*> output_specs; in ValidateSavedFunctionCompatibleWithFunctionDef() local
174 TF_RETURN_IF_ERROR(FlattenSignature(output_signature, &output_specs)); in ValidateSavedFunctionCompatibleWithFunctionDef()
175 if (output_specs.size() != function_def->signature().output_arg_size()) { in ValidateSavedFunctionCompatibleWithFunctionDef()
179 " outputs, but the SavedConcreteFunction has ", output_specs.size(), in ValidateSavedFunctionCompatibleWithFunctionDef()

123