/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/experimental/rpc/ |
H A D | rpc_ops.py | 38 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 D | rpc_ops_test.py | 73 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.py | 156 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.py | 82 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 D | scatter_spmd_expander.cc | 41 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 D | gather_spmd_expander.cc | 314 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 D | lowered_backend_module.py | 285 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 D | graph_signature.py | 206 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 D | exported_program.py | 497 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.py | 52 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 D | utils.py | 267 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.py | 294 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 D | PTMCoreMLBackend.mm | 98 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 D | saved_model.cc | 114 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 D | quantization_context.cc | 143 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 D | preprocess.py | 90 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 D | insert_write_back_for_buffers_pass.py | 114 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 D | weights_to_outputs_pass.py | 38 for spec in gs.output_specs 80 gs.output_specs.append(
|
/aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/nnc/ |
H A D | context.h | 155 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 D | context.cpp | 179 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 D | rpc_ops.cc | 164 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 D | test_context.cpp | 106 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.py | 236 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 D | saved_model.cc | 173 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 D | saved_model_utils.cc | 173 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()
|