black reformat

fast_rnnt/python/fast_rnnt/mutual_information.py

@@ -390,9 +390,7 @@ 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

fast_rnnt/python/fast_rnnt/rnnt_loss.py

@@ -203,9 +203,7 @@ def get_rnnt_logprobs(
  dim=2,
  ) # now: [B][S][T+1], index [:,:,T] has -inf..
 
- px_lm = torch.gather(
- lm[:, :S], dim=2, index=symbols.unsqueeze(-1)
- ) # [B][S][1]
+ px_lm = torch.gather(lm[:, :S], dim=2, index=symbols.unsqueeze(-1)) # [B][S][1]
 
  px = px_am + px_lm # [B][S][T+1], last slice with indexes out of
  # boundary is -inf
@@ -314,9 +312,9 @@ def rnnt_loss_simple(
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
- penalty = offset.reshape(B, 1, 1) - torch.arange(
- T0, device=px.device
- ).reshape(1, 1, T0)
+ penalty = offset.reshape(B, 1, 1) - torch.arange(T0, device=px.device).reshape(
+ 1, 1, T0
+ )
  penalty = penalty * delay_penalty
  px += penalty.to(px.dtype)
 
@@ -424,18 +422,14 @@ def get_rnnt_logprobs_joint(
  px = torch.cat(
  (
  px,
- torch.full(
- (B, S, 1), float("-inf"), device=px.device, dtype=px.dtype
- ),
+ torch.full((B, S, 1), float("-inf"), device=px.device, dtype=px.dtype),
  ),
  dim=2,
  ) # now: [B][S][T+1], index [:,:,T] has -inf..
 
  px[:, :, :T] -= normalizers[:, :S, :]
 
- py = (
- logits[:, :, :, termination_symbol].permute((0, 2, 1)).clone()
- ) # [B][S+1][T]
+ py = logits[:, :, :, termination_symbol].permute((0, 2, 1)).clone() # [B][S+1][T]
  py -= normalizers
 
  if rnnt_type == "regular":
@@ -521,9 +515,9 @@ def rnnt_loss(
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
- penalty = offset.reshape(B, 1, 1) - torch.arange(
- T0, device=px.device
- ).reshape(1, 1, T0)
+ penalty = offset.reshape(B, 1, 1) - torch.arange(T0, device=px.device).reshape(
+ 1, 1, T0
+ )
  penalty = penalty * delay_penalty
  px += penalty.to(px.dtype)
 
@@ -578,9 +572,7 @@ def _monotonic_lower_bound(x: Tensor) -> Tensor:
  return x
 
 
-def _adjust_pruning_lower_bound(
- s_begin: Tensor, s_range: int
-) -> Tensor:
+def _adjust_pruning_lower_bound(s_begin: Tensor, s_range: int) -> Tensor:
  """Adjust s_begin (pruning lower bounds) to make it satisfy the following
  constraints
 
@@ -617,17 +609,13 @@ def _adjust_pruning_lower_bound(
  (B, T) = s_begin.shape
  s_begin = _monotonic_lower_bound(s_begin)
  # do the magic transformation
- s_begin = -(
- s_begin - (s_range - 1) * torch.arange(0, T, device=s_begin.device)
- )
+ s_begin = -(s_begin - (s_range - 1) * torch.arange(0, T, device=s_begin.device))
  # make the transformed tensor to be non-decreasing
  s_begin = _monotonic_lower_bound(s_begin)
  # make start symbol to be zero.
  s_begin = torch.clamp(s_begin, min=0)
  # do the magic transformation again to recover s_begin
- s_begin = -(
- s_begin - (s_range - 1) * torch.arange(0, T, device=s_begin.device)
- )
+ s_begin = -(s_begin - (s_range - 1) * torch.arange(0, T, device=s_begin.device))
  return s_begin
 
 
@@ -773,9 +761,7 @@ def get_rnnt_prune_ranges(
  return ranges
 
 
-def do_rnnt_pruning(
- am: Tensor, lm: Tensor, ranges: Tensor
-) -> Tuple[Tensor, Tensor]:
+def do_rnnt_pruning(am: Tensor, lm: Tensor, ranges: Tensor) -> Tuple[Tensor, Tensor]:
  """Prune the output of encoder(am) and prediction network(lm) with ranges
  generated by `get_rnnt_prune_ranges`.
 
@@ -840,10 +826,7 @@ def _roll_by_shifts(src: Tensor, shifts: torch.LongTensor):
  assert shifts.shape == (B, T), shifts.shape
 
  index = (
- torch.arange(S, device=src.device)
- .view((1, S))
- .repeat((T, 1))
- .repeat((B, 1, 1))
+ torch.arange(S, device=src.device).view((1, S)).repeat((T, 1)).repeat((B, 1, 1))
  )
  index = (index - shifts.reshape(B, T, 1)) % S
  return torch.gather(src, 2, index)
@@ -979,9 +962,7 @@ def get_rnnt_logprobs_pruned(
  px = torch.cat(
  (
  px,
- torch.full(
- (B, S, 1), float("-inf"), device=px.device, dtype=px.dtype
- ),
+ torch.full((B, S, 1), float("-inf"), device=px.device, dtype=px.dtype),
  ),
  dim=2,
  ) # now: [B][S][T+1], index [:,:,T] has -inf..
@@ -1101,9 +1082,9 @@ def rnnt_loss_pruned(
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
- penalty = offset.reshape(B, 1, 1) - torch.arange(
- T0, device=px.device
- ).reshape(1, 1, T0)
+ penalty = offset.reshape(B, 1, 1) - torch.arange(T0, device=px.device).reshape(
+ 1, 1, T0
+ )
  penalty = penalty * delay_penalty
  px += penalty.to(px.dtype)
 
@@ -1272,9 +1253,7 @@ def get_rnnt_logprobs_smoothed(
  + torch.finfo(lm_probs.dtype).tiny
  ) # [1][1][C]
  amonly_normalizers = (
- torch.mv(am_probs.reshape(-1, C), unigram_lm.reshape(C))
- .reshape(B, T, 1)
- .log()
+ torch.mv(am_probs.reshape(-1, C), unigram_lm.reshape(C)).reshape(B, T, 1).log()
  + am_max
  ) # [B][T][1]
  amonly_normalizers = amonly_normalizers.transpose(1, 2) # [B][1][T]
@@ -1308,9 +1287,7 @@ def get_rnnt_logprobs_smoothed(
  dim=2,
  ) # now: [B][S][T+1], index [:,:,T] has -inf..
 
- px_lm = torch.gather(
- lm[:, :S], dim=2, index=symbols.unsqueeze(-1)
- ) # [B][S][1]
+ px_lm = torch.gather(lm[:, :S], dim=2, index=symbols.unsqueeze(-1)) # [B][S][1]
  px_lm_unigram = torch.gather(
  unigram_lm.expand(B, S, C), dim=2, index=symbols.unsqueeze(-1)
  ) # [B][S][1]
@@ -1343,14 +1320,10 @@ def get_rnnt_logprobs_smoothed(
  am_only_scale = 1.0e-20
 
  px_interp = (
- px * combined_scale
- + px_lmonly * lm_only_scale
- + px_amonly * am_only_scale
+ px * combined_scale + px_lmonly * lm_only_scale + px_amonly * am_only_scale
  )
  py_interp = (
- py * combined_scale
- + py_lmonly * lm_only_scale
- + py_amonly * am_only_scale
+ py * combined_scale + py_lmonly * lm_only_scale + py_amonly * am_only_scale
  )
 
  if rnnt_type == "regular":
@@ -1463,9 +1436,9 @@ def rnnt_loss_smoothed(
  ).expand(B, 1, 1)
  else:
  offset = (boundary[:, 3] - 1) / 2
- penalty = offset.reshape(B, 1, 1) - torch.arange(
- T0, device=px.device
- ).reshape(1, 1, T0)
+ penalty = offset.reshape(B, 1, 1) - torch.arange(T0, device=px.device).reshape(
+ 1, 1, T0
+ )
  penalty = penalty * delay_penalty
  px += penalty.to(px.dtype)
 

fast_rnnt/python/tests/mutual_information_test.py

@@ -64,12 +64,8 @@ def test_mutual_information_basic(self):
  if random_boundary:
 
  def get_boundary_row():
- this_S = random.randint(
- 0, S
- ) # allow empty sequence
- this_T = random.randint(
- this_S if modified else 1, T
- )
+ this_S = random.randint(0, S) # allow empty sequence
+ this_T = random.randint(this_S if modified else 1, T)
  s_begin = random.randint(0, S - this_S)
  t_begin = random.randint(0, T - this_T)
  s_end = s_begin + this_S
@@ -118,14 +114,10 @@ def get_boundary_row():
  px += 15.0
  if random_py:
  # log of an odds ratio
- py = torch.randn(B, S + 1, T, dtype=dtype).to(
- device
- )
+ py = torch.randn(B, S + 1, T, dtype=dtype).to(device)
  else:
  # log of an odds ratio
- py = torch.zeros(B, S + 1, T, dtype=dtype).to(
- device
- )
+ py = torch.zeros(B, S + 1, T, dtype=dtype).to(device)
  if big_py:
  py += 15.0
 
@@ -206,14 +198,11 @@ def test_mutual_information_deriv(self):
 
  for dtype in self.dtypes:
  for device in self.devices:
-
  if random_boundary:
 
  def get_boundary_row():
  this_S = random.randint(1, S)
- this_T = random.randint(
- this_S if modified else 1, T
- )
+ this_T = random.randint(this_S if modified else 1, T)
  s_begin = random.randint(0, S - this_S)
  t_begin = random.randint(0, T - this_T)
  s_end = s_begin + this_S
@@ -257,14 +246,10 @@ def get_boundary_row():
  px += 15.0
  if random_py:
  # log of an odds ratio
- py = torch.randn(B, S + 1, T, dtype=dtype).to(
- device
- )
+ py = torch.randn(B, S + 1, T, dtype=dtype).to(device)
  else:
  # log of an odds ratio
- py = torch.zeros(B, S + 1, T, dtype=dtype).to(
- device
- )
+ py = torch.zeros(B, S + 1, T, dtype=dtype).to(device)
  if big_py:
  py += 15.0
  else:

fast_rnnt/python/tests/rnnt_loss_test.py

@@ -90,9 +90,7 @@ def test_rnnt_loss_basic(self):
  assert px.shape == (B, S, T + 1)
  assert py.shape == (B, S + 1, T)
  assert symbols.shape == (B, S)
- m = fast_rnnt.mutual_information_recursion(
- px=px, py=py, boundary=None
- )
+ m = fast_rnnt.mutual_information_recursion(px=px, py=py, boundary=None)
 
  if device == torch.device("cpu"):
  expected = -m
@@ -225,9 +223,7 @@ def test_rnnt_loss_random(self):
  rnnt_type=rnnt_type,
  )
  assert (
- px.shape == (B, S, T)
- if rnnt_type != "regular"
- else (B, S, T + 1)
+ px.shape == (B, S, T) if rnnt_type != "regular" else (B, S, T + 1)
  )
  assert py.shape == (B, S + 1, T)
  assert symbols.shape == (B, S)
@@ -343,7 +339,6 @@ def test_rnnt_loss_gradient(self):
  boundary_[:, 3] = frames
 
  for device in self.devices:
-
  # lm: [B][S+1][C]
  lm = lm_.to(device)
  # am: [B][T][C]
@@ -374,13 +369,9 @@ def test_rnnt_loss_gradient(self):
  torch_grad = torch.autograd.grad(torch_loss, logits2)
  torch_grad = torch_grad[0]
 
- assert torch.allclose(
- fast_loss, torch_loss, atol=1e-2, rtol=1e-2
- )
+ assert torch.allclose(fast_loss, torch_loss, atol=1e-2, rtol=1e-2)
 
- assert torch.allclose(
- fast_grad, torch_grad, atol=1e-2, rtol=1e-2
- )
+ assert torch.allclose(fast_grad, torch_grad, atol=1e-2, rtol=1e-2)
 
  def test_rnnt_loss_smoothed(self):
  B = 1