(methods: Dict[str, ExportedProgram])
| 90 | # and if it is return a map of the indices in the model output that the |
| 91 | # gradient outputs start at and that the parameter outputs start at. |
| 92 | def _get_training_metadata(methods: Dict[str, ExportedProgram]) -> Dict[str, int]: |
| 93 | gradients_method_prefix = "__et_training_gradients_index_" |
| 94 | parameters_method_prefix = "__et_training_parameters_index_" |
| 95 | fqn_method_prefix = "__et_training_fqn_" |
| 96 | training_metadata = {} |
| 97 | for name, method in methods.items(): |
| 98 | found_grad = False |
| 99 | found_param = False |
| 100 | fqns = [] |
| 101 | i = 0 |
| 102 | for output_spec in method.graph_signature.output_specs: |
| 103 | if output_spec.kind == OutputKind.GRADIENT_TO_PARAMETER: |
| 104 | if not found_grad: |
| 105 | training_metadata[gradients_method_prefix + name] = i |
| 106 | found_grad = True |
| 107 | fqns.append(output_spec.target) |
| 108 | elif output_spec.kind == OutputKind.TOKEN and not found_param: |
| 109 | assert found_grad # Params must come after gradients |
| 110 | training_metadata[parameters_method_prefix + name] = i |
| 111 | found_param = True |
| 112 | i += 1 |
| 113 | if len(fqns) > 0: |
| 114 | training_metadata[fqn_method_prefix + name] = fqns |
| 115 | return training_metadata |
| 116 | |
| 117 | |
| 118 | def emit_program( |
no test coverage detected