Skip to content

Commit 98a54cd

Browse files
manuelcandalesfacebook-github-bot
authored andcommitted
SpecDB: Add specs: replication & reflection pad
Reviewed By: SS-JIA Differential Revision: D53543716 fbshipit-source-id: cf823d400336987c5da2f14189354571e40734f7
1 parent 8896942 commit 98a54cd

File tree

1 file changed

+201
-0
lines changed

1 file changed

+201
-0
lines changed

specdb/db.py

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2475,6 +2475,114 @@
24752475
],
24762476
outspec=[OutArg(ArgType.Tensor)],
24772477
),
2478+
Spec(
2479+
op="reflection_pad1d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2480+
inspec=[
2481+
InPosArg(
2482+
ArgType.Tensor,
2483+
name="self",
2484+
constraints=[
2485+
cp.Rank.In(lambda deps: [2, 3]),
2486+
],
2487+
),
2488+
InPosArg(
2489+
ArgType.LengthList,
2490+
name="padding",
2491+
deps=[0],
2492+
constraints=[
2493+
cp.Length.Eq(lambda deps: 2),
2494+
cp.Value.Lt(
2495+
lambda deps, length, ix: deps[0].size(-ix - 1)
2496+
if ix < 1
2497+
else None
2498+
),
2499+
# TODO(mcandales): Calibrate.
2500+
# self.size(-1) + padding[0] + padding[1] >= 0
2501+
],
2502+
),
2503+
],
2504+
outspec=[
2505+
OutArg(
2506+
ArgType.Tensor,
2507+
constraints=[
2508+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2509+
],
2510+
)
2511+
],
2512+
),
2513+
Spec(
2514+
op="reflection_pad2d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2515+
inspec=[
2516+
InPosArg(
2517+
ArgType.Tensor,
2518+
name="self",
2519+
constraints=[
2520+
cp.Rank.In(lambda deps: [3, 4]),
2521+
],
2522+
),
2523+
InPosArg(
2524+
ArgType.LengthList,
2525+
name="padding",
2526+
deps=[0],
2527+
constraints=[
2528+
cp.Length.Eq(lambda deps: 4),
2529+
cp.Value.Lt(
2530+
lambda deps, length, ix: deps[0].size(-ix - 1)
2531+
if ix < 2
2532+
else None
2533+
),
2534+
# TODO(mcandales): Calibrate.
2535+
# self.size(-1) + padding[0] + padding[1] >= 0
2536+
# self.size(-2) + padding[2] + padding[3] >= 0
2537+
],
2538+
),
2539+
],
2540+
outspec=[
2541+
OutArg(
2542+
ArgType.Tensor,
2543+
constraints=[
2544+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2545+
],
2546+
)
2547+
],
2548+
),
2549+
Spec(
2550+
op="reflection_pad3d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2551+
inspec=[
2552+
InPosArg(
2553+
ArgType.Tensor,
2554+
name="self",
2555+
constraints=[
2556+
cp.Rank.In(lambda deps: [4, 5]),
2557+
],
2558+
),
2559+
InPosArg(
2560+
ArgType.LengthList,
2561+
name="padding",
2562+
deps=[0],
2563+
constraints=[
2564+
cp.Length.Eq(lambda deps: 6),
2565+
cp.Value.Lt(
2566+
lambda deps, length, ix: deps[0].size(-ix - 1)
2567+
if ix < 3
2568+
else None
2569+
),
2570+
# TODO(mcandales): Calibrate.
2571+
# self.size(-1) + padding[0] + padding[1] >= 0
2572+
# self.size(-2) + padding[2] + padding[3] >= 0
2573+
# self.size(-3) + padding[4] + padding[5] >= 0
2574+
],
2575+
),
2576+
],
2577+
outspec=[
2578+
OutArg(
2579+
ArgType.Tensor,
2580+
constraints=[
2581+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2582+
],
2583+
)
2584+
],
2585+
),
24782586
Spec(
24792587
op="relu.default", # (Tensor self) -> Tensor
24802588
inspec=[
@@ -2565,6 +2673,99 @@
25652673
OutArg(ArgType.Tensor),
25662674
],
25672675
),
2676+
Spec(
2677+
op="replication_pad1d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2678+
inspec=[
2679+
InPosArg(
2680+
ArgType.Tensor,
2681+
name="self",
2682+
constraints=[
2683+
cp.Rank.In(lambda deps: [2, 3]),
2684+
],
2685+
),
2686+
InPosArg(
2687+
ArgType.LengthList,
2688+
name="padding",
2689+
deps=[0],
2690+
constraints=[
2691+
cp.Length.Eq(lambda deps: 2),
2692+
# TODO(mcandales): Calibrate.
2693+
# self.size(-1) + padding[0] + padding[1] >= 0
2694+
],
2695+
),
2696+
],
2697+
outspec=[
2698+
OutArg(
2699+
ArgType.Tensor,
2700+
constraints=[
2701+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2702+
],
2703+
)
2704+
],
2705+
),
2706+
Spec(
2707+
op="replication_pad2d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2708+
inspec=[
2709+
InPosArg(
2710+
ArgType.Tensor,
2711+
name="self",
2712+
constraints=[
2713+
cp.Rank.In(lambda deps: [3, 4]),
2714+
],
2715+
),
2716+
InPosArg(
2717+
ArgType.LengthList,
2718+
name="padding",
2719+
deps=[0],
2720+
constraints=[
2721+
cp.Length.Eq(lambda deps: 4),
2722+
# TODO(mcandales): Calibrate.
2723+
# self.size(-1) + padding[0] + padding[1] >= 0
2724+
# self.size(-2) + padding[2] + padding[3] >= 0
2725+
],
2726+
),
2727+
],
2728+
outspec=[
2729+
OutArg(
2730+
ArgType.Tensor,
2731+
constraints=[
2732+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2733+
],
2734+
)
2735+
],
2736+
),
2737+
Spec(
2738+
op="replication_pad3d.default", # (Tensor self, SymInt[2] padding) -> Tensor
2739+
inspec=[
2740+
InPosArg(
2741+
ArgType.Tensor,
2742+
name="self",
2743+
constraints=[
2744+
cp.Rank.In(lambda deps: [4, 5]),
2745+
],
2746+
),
2747+
InPosArg(
2748+
ArgType.LengthList,
2749+
name="padding",
2750+
deps=[0],
2751+
constraints=[
2752+
cp.Length.Eq(lambda deps: 6),
2753+
# TODO(mcandales): Calibrate.
2754+
# self.size(-1) + padding[0] + padding[1] >= 0
2755+
# self.size(-2) + padding[2] + padding[3] >= 0
2756+
# self.size(-3) + padding[4] + padding[5] >= 0
2757+
],
2758+
),
2759+
],
2760+
outspec=[
2761+
OutArg(
2762+
ArgType.Tensor,
2763+
constraints=[
2764+
cp.Dtype.Eq(lambda deps: deps[0].dtype),
2765+
],
2766+
)
2767+
],
2768+
),
25682769
Spec(
25692770
op="round.default", # (Tensor self) -> Tensor
25702771
inspec=[

0 commit comments

Comments
 (0)