Merge branch 'main' of github.com:Borklet-Labs/axlearn into jax_py3.12_nightly

Author: camiloCienet

.circleci/config.yml

@@ -111,6 +111,7 @@ disable=abstract-method,
         no-member, # Config objects use dynamic members generated by attrs.
         no-name-in-module,
         no-self-use,
+        not-callable,
         nonzero-method,
         oct-method,
         old-division,
@@ -163,6 +164,8 @@ disable=abstract-method,
         xrange-builtin,
         zip-builtin-not-iterating,
 
+# Increase default from 5 to 9
+max-positional-arguments=9
 
 [REPORTS]
 

.pylintrc

@@ -35,7 +35,7 @@
 from axlearn.common.layers import Embedding, Linear
 from axlearn.common.logit_modifiers import LogitsToLogitsFn
 from axlearn.common.loss import cross_entropy
-from axlearn.common.metrics import WeightedScalar
+from axlearn.common.metrics import WeightedSummary
 from axlearn.common.module import Module, child_context
 from axlearn.common.rnn import BaseRNNCell, LSTMCell
 from axlearn.common.transducer import Transducer, log_probs_from_blank_and_tokens
@@ -238,7 +238,7 @@ def beam_search_decode(
 
     def _input_stats_summaries(
         self, input_batch: Nested[Tensor], *, target_paddings: Tensor, is_valid_example: Tensor
-    ) -> dict[str, Union[WeightedScalar, Tensor]]:
+    ) -> dict[str, Union[WeightedSummary, Tensor]]:
         """Computes input lengths stats.
 
         Args:
@@ -258,13 +258,13 @@ def _input_stats_summaries(
         total_num_examples = jnp.maximum(is_valid_example.sum(), 1.0)
         total_num_frames = jnp.maximum(jnp.size(input_batch["paddings"]), 1)
         input_stats = {
-            "input_stats/average_target_length": WeightedScalar(
+            "input_stats/average_target_length": WeightedSummary(
                 total_target_lengths / total_num_examples, total_num_examples
             ),
-            "input_stats/average_source_length": WeightedScalar(
+            "input_stats/average_source_length": WeightedSummary(
                 total_source_lengths / total_num_examples, total_num_examples
             ),
-            "input_stats/frame_packing_efficiency": WeightedScalar(
+            "input_stats/frame_packing_efficiency": WeightedSummary(
                 total_source_lengths / total_num_frames, total_num_frames
             ),
         }
@@ -324,7 +324,7 @@ def _loss_summaries(
         per_example_weight: Tensor,
         paddings: Tensor,
         target_paddings: Tensor,
-    ) -> dict[str, Union[WeightedScalar, Tensor]]:
+    ) -> dict[str, Union[WeightedSummary, Tensor]]:
         valid_frame_mask = (1.0 - paddings) * per_example_weight[:, None]
         valid_label_mask = (1.0 - target_paddings) * per_example_weight[:, None]
 
@@ -336,17 +336,17 @@ def _loss_summaries(
 
         ret_dict = {}
         # 1. loss/example_weight
-        ret_dict["loss/example_weight"] = WeightedScalar(jnp.mean(per_example_weight), batch_size)
+        ret_dict["loss/example_weight"] = WeightedSummary(jnp.mean(per_example_weight), batch_size)
         # 2. loss/ctc_loss
-        ret_dict["loss/ctc_loss"] = WeightedScalar(
+        ret_dict["loss/ctc_loss"] = WeightedSummary(
             total_ctc_loss / jnp.maximum(per_example_weight.sum(), 1),
             jnp.maximum(per_example_weight.sum(), 1),
         )
         # 3. loss/invalid_seq_percent, per_frame_ctc_loss, per_label_ctc_loss
         invalid_example_percent = 1.0 - jnp.sum(per_example_weight) / batch_size
         ret_dict["loss/invalid_seq_percent"] = invalid_example_percent
-        ret_dict["loss/per_frame_ctc_loss"] = WeightedScalar(per_frame_loss, num_valid_frames)
-        ret_dict["loss/per_label_ctc_loss"] = WeightedScalar(per_label_loss, num_valid_labels)
+        ret_dict["loss/per_frame_ctc_loss"] = WeightedSummary(per_frame_loss, num_valid_frames)
+        ret_dict["loss/per_label_ctc_loss"] = WeightedSummary(per_label_loss, num_valid_labels)
 
         return ret_dict
 
@@ -896,11 +896,11 @@ def forward(self, input_batch: Nested[Tensor]) -> tuple[Tensor, Nested[Tensor]]:
 
         self.add_summary(
             "loss/example_weight",
-            WeightedScalar(jnp.mean(per_example_weight), per_example_weight.shape[0]),
+            WeightedSummary(jnp.mean(per_example_weight), per_example_weight.shape[0]),
         )
         self.add_summary(
             "loss/rnnt_loss",
-            WeightedScalar(loss, jnp.maximum(1, per_example_weight.sum())),
+            WeightedSummary(loss, jnp.maximum(1, per_example_weight.sum())),
         )
         return loss, aux_outputs
 
@@ -1013,7 +1013,7 @@ def tokens_to_scores(
 
         return tokens_to_scores
 
-    def beam_search_decode(
+    def beam_search_decode(  # pytype: disable=signature-mismatch
         self,
         input_batch: Nested[Tensor],
         num_decodes: int,
@@ -1198,9 +1198,9 @@ def forward(self, input_batch: Nested[Tensor]) -> tuple[Tensor, Nested[Tensor]]:
         # Number of valid tokens per example.
         per_example_weight = jnp.sum(live_targets, axis=-1)
         num_targets = per_example_weight.sum()
-        self.add_summary("loss", WeightedScalar(loss, num_targets))
-        self.add_summary("perplexity", WeightedScalar(jnp.exp(loss), num_targets))
-        self.add_summary("token_accuracy", WeightedScalar(loss_dict["accuracy"], num_targets))
+        self.add_summary("loss", WeightedSummary(loss, num_targets))
+        self.add_summary("perplexity", WeightedSummary(jnp.exp(loss), num_targets))
+        self.add_summary("token_accuracy", WeightedSummary(loss_dict["accuracy"], num_targets))
         return loss, dict(per_example_loss=per_example_loss, per_example_weight=per_example_weight)
 
     def _compute_attention_logit_biases(
@@ -1225,7 +1225,7 @@ def _validate_decode_batch(self, input_batch: Nested[Tensor]):
         if "prefix" not in input_batch:
             raise ValueError("Input batch is expected to contain `prefix`.")
 
-    def beam_search_decode(
+    def beam_search_decode(  # pytype: disable=signature-mismatch
         self,
         input_batch: Nested[Tensor],
         num_decodes: int,

axlearn/audio/decoder_asr.py

@@ -30,7 +30,7 @@
 from axlearn.common.decoder import _scores_from_logits
 from axlearn.common.decoding import NEG_INF
 from axlearn.common.logit_modifiers import top_k_logits
-from axlearn.common.metrics import WeightedScalar
+from axlearn.common.metrics import WeightedSummary
 from axlearn.common.module import Module
 from axlearn.common.module import functional as F
 from axlearn.common.param_converter import as_torch_tensor
@@ -522,7 +522,7 @@ def jit_forward(input_batch):
         assert_allclose(np.sum(ref_per_example_loss) / np.sum(per_example_weight), loss)
 
     def _check_summary(
-        self, summary_collection: dict[str, Any], name: str, value: Union[Tensor, WeightedScalar]
+        self, summary_collection: dict[str, Any], name: str, value: Union[Tensor, WeightedSummary]
     ):
         self.assertIn(name, summary_collection)
         msg = f"mismatch in {name}: {summary_collection[name]} vs {value}"
@@ -570,8 +570,8 @@ def test_forward_summary(self):
         )
         summaries = output_collections.summaries
         # 6 out of 8 examples are valid, therefore the average example weight is 0.75
-        self._check_summary(summaries, "loss/example_weight", WeightedScalar(0.75, 8))
-        self._check_summary(summaries, "loss/ctc_loss", WeightedScalar(6972.135, 6))
+        self._check_summary(summaries, "loss/example_weight", WeightedSummary(0.75, 8))
+        self._check_summary(summaries, "loss/ctc_loss", WeightedSummary(6972.135, 6))
         self._check_summary(summaries, "loss/invalid_seq_percent", 0.25)
         total_ctc_loss = summaries["loss/ctc_loss"].weight * summaries["loss/ctc_loss"].mean
         num_valid_frames = jnp.sum(safe_not(paddings) * per_example_weight[:, None])
@@ -580,28 +580,28 @@ def test_forward_summary(self):
         self._check_summary(
             summaries,
             "loss/per_frame_ctc_loss",
-            WeightedScalar(total_ctc_loss / num_valid_frames, num_valid_frames),
+            WeightedSummary(total_ctc_loss / num_valid_frames, num_valid_frames),
         )
         self._check_summary(
             summaries,
             "loss/per_label_ctc_loss",
-            WeightedScalar(total_ctc_loss / num_valid_labels, num_valid_labels),
+            WeightedSummary(total_ctc_loss / num_valid_labels, num_valid_labels),
         )
 
         self._check_summary(
             summaries,
             "input_stats/average_target_length",
-            WeightedScalar(num_valid_labels / num_valid_examples, num_valid_examples),
+            WeightedSummary(num_valid_labels / num_valid_examples, num_valid_examples),
         )
         self._check_summary(
             summaries,
             "input_stats/average_source_length",
-            WeightedScalar(num_valid_frames / num_valid_examples, num_valid_examples),
+            WeightedSummary(num_valid_frames / num_valid_examples, num_valid_examples),
         )
         self._check_summary(
             summaries,
             "input_stats/frame_packing_efficiency",
-            WeightedScalar(num_valid_frames / paddings.size, paddings.size),
+            WeightedSummary(num_valid_frames / paddings.size, paddings.size),
         )
 
     def _check_paddings(self, outputs: DecodeOutputs, *, blank_id: int):

axlearn/audio/decoder_asr_test.py

@@ -27,7 +27,7 @@
     config_for_function,
 )
 from axlearn.common.evaler import ModelSummaryAccumulator
-from axlearn.common.metrics import WeightedScalar
+from axlearn.common.metrics import WeightedSummary
 from axlearn.common.module import Module
 from axlearn.common.utils import Nested, Tensor, replicate_to_local_data
 
@@ -225,17 +225,17 @@ def _compute_word_error_rate(
             # TODO(zhiyunlu): investigate whether we should keep empty references.
             self._metric_accumulator.update(
                 {
-                    "word_errors/wer": WeightedScalar(metrics.num_total / denom, num_words),
-                    "word_errors/deletions": WeightedScalar(
+                    "word_errors/wer": WeightedSummary(metrics.num_total / denom, num_words),
+                    "word_errors/deletions": WeightedSummary(
                         metrics.num_deletions / denom, num_words
                     ),
-                    "word_errors/insertions": WeightedScalar(
+                    "word_errors/insertions": WeightedSummary(
                         metrics.num_insertions / denom, num_words
                     ),
-                    "word_errors/substitutions": WeightedScalar(
+                    "word_errors/substitutions": WeightedSummary(
                         metrics.num_substitutions / denom, num_words
                     ),
-                    "word_errors/sentence_accuracy": WeightedScalar(
+                    "word_errors/sentence_accuracy": WeightedSummary(
                         int(metrics.num_total == 0), int(num_words > 0)
                     ),
                 }