Skip to content

Commit c8624c8

Browse files
Charlie Yanwangkuiyi
Charlie Yan
authored and
wangkuiyi
committed
Ruff migration M3: convert lambda assignments to def (E731)
GitOrigin-RevId: da046f5
1 parent 53e21cc commit c8624c8

19 files changed

Lines changed: 126 additions & 57 deletions

‎axlearn/common/array_serialization_test.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -326,13 +326,15 @@ def get_tensorstore_spec_for_deserialization(arrays: list[jax.Array]):
326326
"bucket": "fake-bucket",
327327
"path": f"fake-path/{time.time()}",
328328
}
329-
create_spec = lambda arr: {
330-
"driver": "zarr",
331-
"kvstore": kvstore_spec,
332-
"dtype": str(arr.dtype),
333-
"create": True,
334-
"delete_existing": True,
335-
}
329+
330+
def create_spec(arr):
331+
return {
332+
"driver": "zarr",
333+
"kvstore": kvstore_spec,
334+
"dtype": str(arr.dtype),
335+
"create": True,
336+
"delete_existing": True,
337+
}
336338

337339
try:
338340
# Yield the temp path so our mock can redirect GCS writes to a local file.

‎axlearn/common/attention_bias.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,10 @@ def _nonzero(self) -> Sequence[BaseAttentionBias]: # pytype: disable=invalid-an
255255
256256
Returned biases are not guaranteed to be nonzero, but are guaranteed to not return None.
257257
"""
258-
filt = lambda b: b.has_value()
258+
259+
def filt(b):
260+
return b.has_value()
261+
259262
return list(filter(filt, self.biases))
260263

261264
def bias_and_residual(self, cls: Type[B]) -> "BiasAndResidual[B]":
@@ -565,9 +568,11 @@ def from_sequence(
565568
pass
566569

567570
# Combine masks.
568-
mask = lambda query_position, key_position: jnp.all(
569-
jnp.stack([b.mask(query_position, key_position) for b in biases]), axis=0
570-
)
571+
def mask(query_position, key_position):
572+
return jnp.all(
573+
jnp.stack([b.mask(query_position, key_position) for b in biases]), axis=0
574+
)
575+
571576
return MaskFnAttentionBias(
572577
mask=mask,
573578
target_positions=biases[0].target_positions,

‎axlearn/common/config.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -753,9 +753,14 @@ def visit(
753753
exit_fn: Called after an enter-able object has been traversed.
754754
"""
755755
if not enter_fn:
756-
enter_fn = lambda key, val, items: items
756+
757+
def enter_fn(key, val, items):
758+
return items
759+
757760
if not exit_fn:
758-
exit_fn = lambda key, val: None
761+
762+
def exit_fn(key, val):
763+
return None
759764

760765
def _visit(key: str, val: Any):
761766
val_items = enter_fn(key, val, _default_enter_fn(key, val))

‎axlearn/common/config_test.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ class CConfig5(PConfig):
143143

144144
@config_class
145145
class CConfig6(PConfig):
146-
foo = lambda self: self
146+
foo = lambda self: self # noqa: E731 — lambda assignment is the pattern under test
147147

148148
del CConfig6
149149

@@ -599,7 +599,9 @@ def fn_with_kwargs(**var_kwargs):
599599
)
600600
def test_config_for_function(self, kwargs, fn=None, expected=None):
601601
if fn is None:
602-
fn = lambda x, y, z: (x, y, z)
602+
603+
def fn(x, y, z):
604+
return (x, y, z)
603605

604606
def build_and_invoke():
605607
cfg = config.config_for_function(fn)

‎axlearn/common/ein_ops.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -441,7 +441,9 @@ def find_ellipsis_index(axes: _Axes) -> tuple[int, bool, int]:
441441
return lhs, rhs
442442

443443
def count_explicit_axes(axes: _Axes) -> int:
444-
count = lambda ax: 0 if ax == ellipsis else 1
444+
def count(ax):
445+
return 0 if ax == ellipsis else 1
446+
445447
return sum(count(ax) for ax in axes)
446448

447449
ndim = len(in_shape)

‎axlearn/common/flash_attention/gpu_attention_benchmark.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -258,7 +258,10 @@ def bench_flash_attention(
258258
argnums=(0, 1, 2),
259259
)
260260
else:
261-
fn = lambda q, k, v, b: base_fn(dict(query=q, key=k, value=v, bias=b))
261+
262+
def fn(q, k, v, b):
263+
return base_fn(dict(query=q, key=k, value=v, bias=b))
264+
262265
return measure(fn, q, k, v, bias)
263266

264267

‎axlearn/common/flash_attention/tpu_attention.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1163,7 +1163,9 @@ def build(self, input_batch: Nested[Tensor | BaseAttentionBias]):
11631163
# Need key limit mask if padding is applied to keys
11641164
original_mask_len = unpadded_k_len if unpadded_k_len % block_size != 0 else None
11651165

1166-
mul_block_len = lambda seq_len: seq_len + (-seq_len % block_size)
1166+
def mul_block_len(seq_len):
1167+
return seq_len + (-seq_len % block_size)
1168+
11671169
mask_shape = (mul_block_len(query.shape[1]), mul_block_len(key.shape[1]))
11681170
splash_mask = _to_splash_mask(mask, mask_shape=mask_shape, unpadded_k_len=original_mask_len)
11691171
mesh = thread_resources.env.physical_mesh

‎axlearn/common/gradient_accumulation_test.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,8 @@ def check_sharding(path, value):
7979
jax.tree.map(check_sharding, tree_paths(input_batch), input_batch)
8080
return input_batch
8181

82-
callback = lambda path, sharding: self.assertEqual(expected[path], sharding.spec)
82+
def callback(path, sharding):
83+
return self.assertEqual(expected[path], sharding.spec)
8384

8485
callback_sharding(
8586
input_batch=input_batch,

‎axlearn/common/inference.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,11 @@ def __call__(self, input_batch: NestedTensor) -> Output:
104104
"""
105105
with self._mesh:
106106
# TODO(zhucheng_tu,tom_gunter): Handle a mixture of pre-sharded and host-local inputs.
107-
is_host_local_input_check = lambda x: (
108-
(isinstance(x, jax.Array) and len(x.devices()) == 1) or isinstance(x, np.ndarray)
109-
)
107+
def is_host_local_input_check(x):
108+
return (isinstance(x, jax.Array) and len(x.devices()) == 1) or isinstance(
109+
x, np.ndarray
110+
)
111+
110112
all_host_local_inputs = all(
111113
is_host_local_input_check(t) for t in jax.tree_util.tree_leaves(input_batch)
112114
)

‎axlearn/common/launch_trainer_test.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@ def test_get_trainer_config(
8181
) as mock_fn:
8282
trainer_config_fn = None
8383
if trainer_config is not None:
84-
trainer_config_fn = lambda: trainer_config
84+
85+
def trainer_config_fn():
86+
return trainer_config
8587

8688
cfg = launch_trainer.get_trainer_config(
8789
flag_values=fv, trainer_config_fn=trainer_config_fn
@@ -138,7 +140,8 @@ def test_crash_on_hang_timeout_seconds_flag(self):
138140
fv = _flag_values_from_dict(flag_values)
139141

140142
# Mock the trainer config function
141-
trainer_config_fn = lambda: mock_trainer_config
143+
def trainer_config_fn():
144+
return mock_trainer_config
142145

143146
# Mock get_named_trainer_config to avoid dependency issues
144147
with mock.patch(
@@ -182,7 +185,8 @@ def test_crash_on_hang_timeout_seconds_not_overridden(self):
182185
fv = _flag_values_from_dict(flag_values)
183186

184187
# Mock the trainer config function
185-
trainer_config_fn = lambda: mock_trainer_config
188+
def trainer_config_fn():
189+
return mock_trainer_config
186190

187191
# Mock get_named_trainer_config to avoid dependency issues
188192
with mock.patch(

0 commit comments

Comments
 (0)