Internal only change for handling pytype errors.

PiperOrigin-RevId: 609621790

chirp/models/separation_model.py

@@ -50,6 +50,8 @@ def time_reduce_logits(self, reduction: str = 'AVG') -> 'SeparatorOutput':
     elif reduction == 'MAX':
       reduce_fn = lambda x: jnp.max(x, axis=1)
     elif reduction == 'MIDPOINT':
+      if self.label is None:
+        raise ValueError('SeperatorOutput.label is None')
       midpt = self.label.shape[1] // 2
       reduce_fn = lambda x: x[:, midpt, :]
     else:
@@ -169,6 +171,8 @@ def bottleneck_classifier(self, bottleneck, train: bool):
     )(classify_hiddens)
     classify_hiddens = nn.swish(classify_hiddens)
     classify_outputs = {}
+    if self.num_classes is None:
+      raise ValueError('SeparatorModel.num_classes is None')
     for k, n in self.num_classes.items():
       classify_outputs[k] = nn.Conv(n, (1,), (1,), 'SAME')(classify_hiddens)
     classify_outputs['embedding'] = classify_hiddens

chirp/train/hubert.py

@@ -757,6 +757,8 @@ def train(
         "reload_quantizer being True."
     )
 
+  if train_dataset is None:
+    raise ValueError("train_dataset is None.")
   train_iterator = train_dataset.as_numpy_iterator()
   taxonomy_keys = ["label"]
   taxonomy_loss_weight = model_bundle.model.taxonomy_loss_weight

chirp/train/separator.py

@@ -263,6 +263,8 @@ def train(
     ] = optax.sigmoid_binary_cross_entropy,
 ) -> None:
   """Train a model."""
+  if train_dataset is None:
+    raise ValueError('train_dataset is None')
   train_iterator = train_dataset.as_numpy_iterator()
   train_metrics_collection = train_utils.NestedCollection.create(
       **TRAIN_METRICS
@@ -318,6 +320,8 @@ def update_step(params, model_state):
         train_state.params, train_state.model_state
     )
     grads = jax.lax.pmean(grads, axis_name='batch')
+    if model_bundle.optimizer is None:
+      raise ValueError('model_bundle.optimizer is None')
     updates, opt_state = model_bundle.optimizer.update(
         grads, train_state.opt_state
     )