Skip to content

Commit 09b1282

Browse files
ds-hwangchanglan
ds-hwang
authored and
changlan
committed
Relax dtype constraint for compute_conv(_transpose)_paddings.
GitOrigin-RevId: 2230d78
1 parent fc4d8ec commit 09b1282

2 files changed

Lines changed: 238 additions & 33 deletions

File tree

‎axlearn/common/convolution.py‎

Lines changed: 52 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -300,7 +300,6 @@ def compute_conv_paddings(
300300
ValueError: If anchor is not between left_time_padding and right_time_padding.
301301
"""
302302
chex.assert_rank(in_paddings, 2)
303-
chex.assert_type(in_paddings, jnp.bool)
304303
dilation = dilation or 1
305304
conv_padding = conv_explicit_padding(
306305
window=(window,), strides=(stride,), padding=conv_padding, dilation=(dilation,)
@@ -491,18 +490,23 @@ def forward(
491490
output = super().forward(x)
492491

493492
# Compute paddings conv output.
494-
output_paddings = compute_conv_paddings(
493+
output_paddings = self.conv_paddings(paddings)
494+
# Apply padding to the outputs.
495+
output_paddings = maybe_shard(output_paddings, cfg.output_partition_spec)
496+
output = output * safe_not(output_paddings)[..., None]
497+
return output, output_paddings
498+
499+
@nowrap
500+
def conv_paddings(self, paddings: Tensor) -> Tensor:
501+
cfg = self.config
502+
return compute_conv_paddings(
495503
paddings,
496504
window=cfg.window,
497505
stride=cfg.strides,
498506
conv_padding=cfg.padding,
499507
dilation=cfg.dilation,
500508
anchor=cfg.anchor,
501509
)
502-
# Apply padding to the outputs.
503-
output_paddings = maybe_shard(output_paddings, cfg.output_partition_spec)
504-
output = output * (1 - output_paddings[..., None])
505-
return output, output_paddings
506510

507511

508512
# The accuracy of the output of this layer currently doesn't match that of PyTorch
@@ -708,26 +712,30 @@ def forward(
708712
output: A Tensor of shape [batch_size, seq_len, frequency, output_dim].
709713
paddings: 0/1 boolean Tensor of shape [batch_size, seq_len].
710714
"""
711-
cfg = self.config
712715
# Apply padding to the input.
713716
assert len(x.shape) == len(paddings.shape) + 2
714717
x = x * safe_not(paddings)[..., None, None]
715718

716719
# Apply Conv2D.
717720
output = super().forward(x)
718721
# Compute paddings conv output.
722+
output_paddings = self.conv_paddings(paddings)
723+
# Apply padding to the outputs.
724+
output = output * safe_not(output_paddings)[..., None, None]
725+
return output, output_paddings
726+
727+
@nowrap
728+
def conv_paddings(self, paddings: Tensor) -> Tensor:
729+
cfg = self.config
719730
dilation = 1 if cfg.dilation is None else cfg.dilation[0]
720-
output_paddings = compute_conv_paddings(
731+
return compute_conv_paddings(
721732
paddings,
722733
window=cfg.window[0],
723734
stride=cfg.strides[0],
724735
conv_padding=cfg.padding,
725736
dilation=dilation,
726737
anchor=cfg.anchor,
727738
)
728-
# Apply padding to the outputs.
729-
output = output * (1 - output_paddings[..., None, None])
730-
return output, output_paddings
731739

732740

733741
class Conv3D(BaseConv):
@@ -1305,7 +1313,6 @@ def compute_conv_transpose_paddings(
13051313
"""
13061314

13071315
chex.assert_rank(in_paddings, 2)
1308-
chex.assert_type(in_paddings, jnp.bool)
13091316
conv_padding = conv_transpose_explicit_padding(
13101317
window=(window,), strides=(stride,), padding=conv_padding, dilation=(dilation,)
13111318
)
@@ -1436,15 +1443,8 @@ def forward(
14361443
output_paddings = None
14371444
else:
14381445
# Compute paddings conv output.
1439-
output_paddings = compute_conv_transpose_paddings(
1440-
paddings,
1441-
window=cfg.window,
1442-
stride=cfg.strides,
1443-
conv_padding=cfg.padding,
1444-
dilation=cfg.dilation,
1445-
anchor=cfg.anchor,
1446-
)
1447-
output = output * (1 - output_paddings[..., None])
1446+
output_paddings = self.conv_paddings(paddings)
1447+
output = output * safe_not(output_paddings)[..., None]
14481448
return output, output_paddings
14491449

14501450
def _conv(
@@ -1470,6 +1470,18 @@ def _conv(
14701470
output += self.parameters["bias"]
14711471
return output
14721472

1473+
@nowrap
1474+
def conv_paddings(self, paddings: Tensor) -> Tensor:
1475+
cfg = self.config
1476+
return compute_conv_transpose_paddings(
1477+
paddings,
1478+
window=cfg.window,
1479+
stride=cfg.strides,
1480+
conv_padding=cfg.padding,
1481+
dilation=cfg.dilation,
1482+
anchor=cfg.anchor,
1483+
)
1484+
14731485
def output_shape(self, *, input_shape: Sequence[Optional[int]]) -> Sequence[Optional[int]]:
14741486
cfg = self.config
14751487
if len(input_shape) != 3:
@@ -1645,25 +1657,29 @@ def forward(
16451657
output: A Tensor of shape [batch_size, seq_len, frequency, output_dim].
16461658
paddings: 0/1 boolean Tensor of shape [batch_size, seq_len].
16471659
"""
1648-
cfg = self.config
16491660
# Apply padding to the input.
16501661
assert len(x.shape) == len(paddings.shape) + 2
16511662
x = x * safe_not(paddings)[..., None, None]
16521663

16531664
# Apply Conv2D.
16541665
output = super().forward(x)
16551666
# Compute paddings conv output.
1656-
output_paddings = compute_conv_transpose_paddings(
1667+
output_paddings = self.conv_paddings(paddings)
1668+
# Apply padding to the outputs.
1669+
output = output * safe_not(output_paddings)[..., None, None]
1670+
return output, output_paddings
1671+
1672+
@nowrap
1673+
def conv_paddings(self, paddings: Tensor) -> Tensor:
1674+
cfg = self.config
1675+
return compute_conv_transpose_paddings(
16571676
paddings,
16581677
window=cfg.window[0],
16591678
stride=cfg.strides[0],
16601679
conv_padding=cfg.padding,
16611680
dilation=cfg.dilation[0],
16621681
anchor=cfg.anchor,
16631682
)
1664-
# Apply padding to the outputs.
1665-
output = output * (1 - output_paddings[..., None, None])
1666-
return output, output_paddings
16671683

16681684

16691685
class Conv3DTranspose(BaseConv):
@@ -1823,12 +1839,18 @@ def forward(self, inputs: Tensor, *, paddings: Tensor) -> tuple[Tensor, Tensor]:
18231839
# Stack inputs over the time dimension.
18241840
stacked_inputs = jnp.reshape(inputs[:, : output_length * cfg.stride, :], new_shape)
18251841
# An output frame is padding if at least one of the stacked input frames is padding.
1826-
stacked_paddings = compute_conv_paddings(
1827-
paddings, window=cfg.stride, stride=cfg.stride, conv_padding=(padding,)
1828-
)
1829-
stacked_inputs = stacked_inputs * (1 - stacked_paddings)[:, :, None]
1842+
stacked_paddings = self.conv_paddings(paddings)
1843+
stacked_inputs = stacked_inputs * safe_not(stacked_paddings)[:, :, None]
18301844
return stacked_inputs, stacked_paddings
18311845

1846+
@nowrap
1847+
def conv_paddings(self, paddings: Tensor) -> Tensor:
1848+
cfg = self.config
1849+
conv_padding = cfg.padding if isinstance(cfg.padding, str) else (cfg.padding,)
1850+
return compute_conv_paddings(
1851+
paddings, window=cfg.stride, stride=cfg.stride, conv_padding=conv_padding
1852+
)
1853+
18321854
@nowrap
18331855
def output_shape(self, *, input_shape: Sequence[Optional[int]]) -> Sequence[Optional[int]]:
18341856
"""Computes stacked output shape.

‎axlearn/common/convolution_test.py‎

Lines changed: 186 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,9 @@ def test_conv_output_shape(self, in_shape, window, strides, padding, dilation, e
8787
([0, 0, 0, 1, 1, 1], [0, 0], 2, "VALID"),
8888
([0, 0, 1, 1, 1, 1], [0, 1], 2, "VALID"),
8989
)
90-
def test_conv_padding(self, input_paddings, expected_paddings, stride: int, padding_cfg: str):
90+
def test_compute_conv_paddings(
91+
self, input_paddings, expected_paddings, stride: int, padding_cfg: str
92+
):
9193
"""Tests conv_output_shape() with SAME and VALID padding cfg."""
9294
# This test is from lingvo
9395
# https://github.com/tensorflow/lingvo/blob/master/lingvo/core/conv_layers_with_time_padding_test.py#L157.
@@ -100,6 +102,64 @@ def test_conv_padding(self, input_paddings, expected_paddings, stride: int, padd
100102
)
101103
assert_allclose(out_paddings[0], expected_paddings)
102104

105+
@parameterized.parameters(
106+
# window=5, stride=2, dilation=1
107+
# SAME padding=(2, 2)
108+
# pad | |pad
109+
# segment_ids: 0 0|1 1 1 1 2 2|0 0
110+
# |___^___|
111+
# |___^___|
112+
# |___^___|
113+
("SAME", 1, [1, 1, 1, 1, 2, 2], [1, 1, 2]),
114+
# VALID padding=(0, 0)
115+
# | |
116+
# segment_ids:|1 1 1 1 2 2|
117+
# |^_______|
118+
("VALID", 1, [1, 1, 1, 1, 2, 2], [1]),
119+
# CAUSAL padding=(3, 1)
120+
# pad | |pad
121+
# segment_ids: 0 0 0|1 1 1 1 2 2|0
122+
# |_____^_|
123+
# |_____^_|
124+
# |_____^_|
125+
("CAUSAL", 1, [1, 1, 1, 1, 2, 2], [1, 1, 2]),
126+
# window=5, stride=2, dilation=2
127+
# dilate_window = 9, pad_total = 8
128+
# SAME padding=(4, 4)
129+
# pad| |pad
130+
# segment_ids: 0 0 0 0|1 2 3 4 5 6 7 8 9|0 0 0 0
131+
# |_______^_______|
132+
# |_______^_______|
133+
# |_______^_______|
134+
# |_______^_______|
135+
# |_______^_______|
136+
("SAME", 2, [1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 3, 5, 7, 9]),
137+
# VALID padding=(0, 0)
138+
# | |pad
139+
# segment_ids:|1 2 3 4 5 6 7 8 9 0|
140+
# |^_______________|
141+
("VALID", 2, [1, 2, 3, 4, 5, 6, 7, 8, 9, 0], [1]),
142+
# CAUSAL padding=(7, 1)
143+
# pad | |pad
144+
# segment_ids: 0 0 0 0 0 0 0|1 2 3 4 5 6 7 8 9|0
145+
# |_____________^_|
146+
# |_____________^_|
147+
# |_____________^_|
148+
# |_____________^_|
149+
# |_____________^_|
150+
("CAUSAL", 2, [1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 3, 5, 7, 9]),
151+
)
152+
def test_compute_conv_segment_ids(self, padding, dilation, segment_ids, expected):
153+
window, stride = 5, 2
154+
out_segment_ids = compute_conv_paddings(
155+
jnp.array([segment_ids], jnp.int32),
156+
window=window,
157+
stride=stride,
158+
conv_padding=padding,
159+
dilation=dilation,
160+
)[0]
161+
assert_allclose(out_segment_ids, expected)
162+
103163
@parameterized.parameters(
104164
(5, 1, "SAME", 1, (2, 2)),
105165
(5, 2, "SAME", 1, (2, 2)),
@@ -1089,6 +1149,129 @@ def test_compute_conv_transpose_paddings(self, window, strides, padding, dilatio
10891149
expected = jnp.array(expected).astype(out_paddings.dtype)
10901150
self.assertNestedEqual(out_paddings[0], expected)
10911151

1152+
@parameterized.parameters(
1153+
# window=3, stride=1
1154+
# SAME padding=(1, 1)
1155+
# pad| |pad
1156+
# segment_ids: 0|1 1 2 2|0
1157+
# |_____|
1158+
# * 1 * -> 1
1159+
# * 1 * -> 1
1160+
# * 2 * -> 2
1161+
# * 2 * -> 2
1162+
("SAME", 1, 1, [1, 1, 2, 2], [1, 1, 2, 2]),
1163+
# VALID padding=(2, 2)
1164+
# pad | |pad
1165+
# segment_ids: 0 0|1 1 2 2|* *
1166+
# |_________|
1167+
# * * 1 -> 1
1168+
# * * 1 -> 1
1169+
# * * 2 -> 2
1170+
# * * 2 -> 2
1171+
# * * 2 -> 2
1172+
# * * 2 -> 2
1173+
("VALID", 1, 1, [1, 1, 2, 2], [1, 1, 2, 2, 2, 2]),
1174+
# CAUSAL padding=(2, 0)
1175+
# pad | |pad
1176+
# segment_ids: 0 0|1 1 2 2|
1177+
# |_____|
1178+
# * * 1 -> 1
1179+
# * * 1 -> 1
1180+
# * * 2 -> 2
1181+
# * * 2 -> 2
1182+
("CAUSAL", 1, 1, [1, 1, 2, 2], [1, 1, 2, 2]),
1183+
# window=3, stride=2
1184+
# SAME padding=(2, 1)
1185+
# pad | |pad
1186+
# segment_ids: 0 0|1 * 2 * 3 * 4|*
1187+
# |_____________|
1188+
# * * 1 -> 1
1189+
# * * 1 -> 1
1190+
# * * 2 -> 2
1191+
# * * 2 -> 2
1192+
# * * 3 -> 3
1193+
# * * 3 -> 3
1194+
# * * 4 -> 4
1195+
# * * 4 -> 4
1196+
("SAME", 2, 1, [1, 2, 3, 4], [1, 1, 2, 2, 3, 3, 4, 4]),
1197+
# VALID padding=(2, 2)
1198+
# pad | |pad
1199+
# segment_ids: 0 0|1 * 2 * 3 * 4|* *
1200+
# |_______________|
1201+
# * * 1 -> 1
1202+
# * * 1 -> 1
1203+
# * * 2 -> 2
1204+
# * * 2 -> 2
1205+
# * * 3 -> 3
1206+
# * * 3 -> 3
1207+
# * * 4 -> 4
1208+
# * * 4 -> 4
1209+
# * * 0 -> 0
1210+
("VALID", 2, 1, [1, 2, 3, 4], [1, 1, 2, 2, 3, 3, 4, 4, 4]),
1211+
# CAUSAL padding=(2, 1)
1212+
# pad | |pad
1213+
# segment_ids: 0 0|1 * 2 * 3 * 4|*
1214+
# |_____________|
1215+
# * * 1 -> 1
1216+
# * * 1 -> 1
1217+
# * * 2 -> 2
1218+
# * * 2 -> 2
1219+
# * * 3 -> 3
1220+
# * * 3 -> 3
1221+
# * * 4 -> 4
1222+
# * * 4 -> 4
1223+
("CAUSAL", 2, 1, [1, 2, 3, 4], [1, 1, 2, 2, 3, 3, 4, 4]),
1224+
# window=3, stride=3
1225+
# SAME, VALID, CAUSAL padding=(2, 2)
1226+
# pad | |pad
1227+
# segment_ids: 0 0|1 * * 2 * * 3 * * 4|* *
1228+
# |_____________________|
1229+
# * * 1 -> 1
1230+
# * * 1 -> 1
1231+
# * * 1 -> 1
1232+
# * * 2 -> 2
1233+
# * * 2 -> 2
1234+
# * * 2 -> 2
1235+
# * * 3 -> 3
1236+
# * * 3 -> 3
1237+
# * * 3 -> 3
1238+
# * * 4 -> 4
1239+
# * * 4 -> 4
1240+
# * * 4 -> 4
1241+
("CAUSAL", 3, 1, [1, 2, 3, 4], [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]),
1242+
# window=3, stride=2, dilation=2
1243+
# dilate_window = 5
1244+
# SAME padding=(3, 2)
1245+
# pad | |pad
1246+
# segment_ids: 0 0 0|1 * * 2 * * 3 * * 4|* *
1247+
# |_____________|
1248+
# * * * 1 * -> 1
1249+
# * * * 1 * -> 1
1250+
# * * * 1 * -> 2
1251+
# * * * 2 * -> 2
1252+
# * * * 2 * -> 2
1253+
# * * * 2 * -> 2
1254+
# * * * 3 * -> 3
1255+
# * * * 3 * -> 3
1256+
# * * * 3 * -> 3
1257+
# * * * 4 * -> 4
1258+
# * * * 4 * -> 4
1259+
# * * * 4 * -> 4
1260+
("SAME", 3, 2, [1, 2, 3, 4], [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]),
1261+
)
1262+
def test_compute_conv_transpose_segment_ids(
1263+
self, padding, stride, dilation, segment_ids, expected
1264+
):
1265+
window = 3
1266+
out_segment_ids = convolution.compute_conv_transpose_paddings(
1267+
jnp.array([segment_ids], jnp.int32),
1268+
window=window,
1269+
stride=stride,
1270+
conv_padding=padding,
1271+
dilation=dilation,
1272+
)[0]
1273+
assert_allclose(out_segment_ids, expected)
1274+
10921275
@parameterized.product(
10931276
window=[1, 3],
10941277
strides=[1, 2, 3],
@@ -1131,7 +1314,7 @@ def test_compute_conv_transpose_paddings_with_conv_paddings(
11311314
in_paddings, window=window, stride=strides, conv_padding=padding, dilation=dilation
11321315
)
11331316

1134-
recon_paddings = convolution.compute_conv_paddings(
1317+
recon_paddings = compute_conv_paddings(
11351318
out_paddings, window=window, stride=strides, conv_padding=padding, dilation=dilation
11361319
)
11371320
self.assertNestedEqual(recon_paddings[0], in_paddings[0])
@@ -1161,7 +1344,7 @@ def test_compute_conv_transpose_paddings_against_conv_paddings(
11611344
anchor = None
11621345

11631346
in_paddings = jnp.array(in_paddings, dtype=jnp.bool)[None, :]
1164-
ref_paddings = convolution.compute_conv_paddings(
1347+
ref_paddings = compute_conv_paddings(
11651348
in_paddings,
11661349
window=window,
11671350
stride=strides,

0 commit comments

Comments
 (0)