Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 28 additions & 3 deletions utils/model_ops/utils/torchop_yaml.py
Original file line number Diff line number Diff line change
Expand Up @@ -332,13 +332,18 @@ def add_test_case_yaml(

if len(kwmap) > 0:
if not old_format:
test_case_yaml["kwargs"] = kwmap
test_case_yaml[inputs]["kwargs"] = kwmap
else:
test_case_yaml["kwmap"] = kwmap

return test_case_yaml


class YamlFmtDumper(yaml.Dumper):
def increase_indent(self, flow=False, indentless=False):
return super().increase_indent(flow, indentless=False)


@dataclass
class _ArgResult:
skip_this_node: bool = False
Expand Down Expand Up @@ -1669,7 +1674,17 @@ def _filter_cases(cases: list[dict[str, Any]]):
print(f"Total no. of test cases: {len(config['cases'])}")

with open(os.path.join(output_dir, model_name + ".yaml"), "w") as f:
yaml.dump(config, f, sort_keys=False)
yaml.dump(
config,
f,
Dumper=YamlFmtDumper,
default_flow_style=False,
sort_keys=False,
allow_unicode=True,
explicit_start=False,
indent=2,
width=4096,
)

if not supress_spyre:
if not USE_OLDFORMAT:
Expand All @@ -1682,7 +1697,17 @@ def _filter_cases(cases: list[dict[str, Any]]):
else:
config["cases"] = _filter_cases(self.test_cases_norm_yaml)
with open(os.path.join(output_dir, model_name + "_spyre.yaml"), "w") as f:
yaml.dump(config, f, sort_keys=False)
yaml.dump(
config,
f,
Dumper=YamlFmtDumper,
default_flow_style=False,
sort_keys=False,
allow_unicode=True,
explicit_start=False,
indent=2,
width=4096,
)


# for debug purpose
Expand Down
Loading