Fix black

.github/workflows/style_check.yml

@@ -66,4 +66,4 @@ jobs:
  shell: bash
  working-directory: ${{github.workspace}}
  run: |
- black --check --diff .
+ black -l 80 --check --diff .

fast_rnnt/python/fast_rnnt/__init__.py

@@ -13,5 +13,3 @@
 from .rnnt_loss import rnnt_loss_pruned
 from .rnnt_loss import rnnt_loss_simple
 from .rnnt_loss import rnnt_loss_smoothed
-
-

fast_rnnt/python/fast_rnnt/mutual_information.py

@@ -160,7 +160,8 @@ def forward(
  if return_grad or px.requires_grad or py.requires_grad:
  ans_grad = torch.ones(B, device=px.device, dtype=px.dtype)
  (px_grad, py_grad) = _fast_rnnt.mutual_information_backward(
- px, py, boundary, p, ans_grad)
+ px, py, boundary, p, ans_grad
+ )
  ctx.save_for_backward(px_grad, py_grad)
  assert len(pxy_grads) == 2
  pxy_grads[0] = px_grad
@@ -290,8 +291,9 @@ def mutual_information_recursion(
  px, py = px.contiguous(), py.contiguous()
 
  pxy_grads = [None, None]
- scores = MutualInformationRecursionFunction.apply(px, py, pxy_grads,
- boundary, return_grad)
+ scores = MutualInformationRecursionFunction.apply(
+ px, py, pxy_grads, boundary, return_grad
+ )
  px_grad, py_grad = pxy_grads
  return (scores, (px_grad, py_grad)) if return_grad else scores
 
@@ -388,16 +390,18 @@ def joint_mutual_information_recursion(
  p = torch.empty(B, S + 1, T + 1, device=px_tot.device, dtype=px_tot.dtype)
 
  # note, tot_probs is without grad.
- tot_probs = _fast_rnnt.mutual_information_forward(px_tot, py_tot, boundary, p)
+ tot_probs = _fast_rnnt.mutual_information_forward(
+ px_tot, py_tot, boundary, p
+ )
 
  # this is a kind of "fake gradient" that we use, in effect to compute
  # occupation probabilities. The backprop will work regardless of the
  # actual derivative w.r.t. the total probs.
  ans_grad = torch.ones(B, device=px_tot.device, dtype=px_tot.dtype)
 
- (px_grad,
- py_grad) = _fast_rnnt.mutual_information_backward(px_tot, py_tot, boundary, p,
-  ans_grad)
+ (px_grad, py_grad) = _fast_rnnt.mutual_information_backward(
+  px_tot, py_tot, boundary, p, ans_grad
+ )
 
  px_grad = px_grad.reshape(1, B, -1)
  py_grad = py_grad.reshape(1, B, -1)

fast_rnnt/python/fast_rnnt/rnnt_loss.py

@@ -170,7 +170,7 @@ def get_rnnt_logprobs(
  am.transpose(1, 2), # (B, C, T)
  dim=1,
  index=symbols.unsqueeze(2).expand(B, S, T),
- ) # (B, S, T)
+ )  # (B, S, T)
 
  if rnnt_type == "regular":
  px_am = torch.cat(
@@ -291,7 +291,9 @@ def rnnt_loss_simple(
  T = T0 if rnnt_type != "regular" else T0 - 1
  if boundary is None:
  offset = torch.tensor(
- (T - 1) / 2, dtype=px.dtype, device=px.device,
+ (T - 1) / 2,
+ dtype=px.dtype,
+ device=px.device,
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
@@ -495,7 +497,9 @@ def rnnt_loss(
  T = T0 if rnnt_type != "regular" else T0 - 1
  if boundary is None:
  offset = torch.tensor(
- (T - 1) / 2, dtype=px.dtype, device=px.device,
+ (T - 1) / 2,
+ dtype=px.dtype,
+ device=px.device,
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
@@ -770,9 +774,7 @@ def do_rnnt_pruning(
  lm_pruning = torch.gather(
  lm,
  dim=1,
- index=ranges.reshape(B, T * s_range, 1).expand(
- (B, T * s_range, C)
- ),
+ index=ranges.reshape(B, T * s_range, 1).expand((B, T * s_range, C)),
  ).reshape(B, T, s_range, C)
  return am_pruning, lm_pruning
 
@@ -1057,7 +1059,9 @@ def rnnt_loss_pruned(
  T = T0 if rnnt_type != "regular" else T0 - 1
  if boundary is None:
  offset = torch.tensor(
- (T - 1) / 2, dtype=px.dtype, device=px.device,
+ (T - 1) / 2,
+ dtype=px.dtype,
+ device=px.device,
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
@@ -1248,7 +1252,7 @@ def get_rnnt_logprobs_smoothed(
  am.transpose(1, 2), # (B, C, T)
  dim=1,
  index=symbols.unsqueeze(2).expand(B, S, T),
- ) # (B, S, T)
+ )  # (B, S, T)
 
  if rnnt_type == "regular":
  px_am = torch.cat(
@@ -1413,7 +1417,9 @@ def rnnt_loss_smoothed(
  T = T0 if rnnt_type != "regular" else T0 - 1
  if boundary is None:
  offset = torch.tensor(
- (T - 1) / 2, dtype=px.dtype, device=px.device,
+ (T - 1) / 2,
+ dtype=px.dtype,
+ device=px.device,
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2