Internal

PiperOrigin-RevId: 868681655

Author: init2winit Team

init2winit/model_lib/deepspeech.py

@@ -343,7 +343,7 @@ def __call__(self, inputs):
     var = jnp.mean(jnp.square(inputs - mean), axis=-1, keepdims=True)
 
     normed_inputs = (inputs - mean) * jax.lax.rsqrt(var + self.epsilon)
-    normed_inputs *= (1 + self.scale)
+    normed_inputs *= 1 + self.scale
     normed_inputs += self.bias
 
     return normed_inputs
@@ -963,6 +963,17 @@ def evaluate_batch(self, params, batch_stats, batch):
         targets=labels,
         target_paddings=label_paddings)
 
+  def apply_on_batch(self, params, batch_stats, batch, **apply_kwargs):
+    """Wrapper around flax_module.apply."""
+    return self.flax_module.apply(
+        {
+            'params': params,
+            'batch_stats': batch_stats
+        },
+        batch['inputs'],
+        batch['input_paddings'],
+        **apply_kwargs)
+
   def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
     """Return CTC loss."""