(self)
| 664 | def _create_layer_parameter_specs(self) -> dict[str, ParameterSpec]: |
| 665 | cfg = self.config |
| 666 | return { |
| 667 | "scale": ParameterSpec(shape=[cfg.input_dim], mesh_axes=(None,)), |
| 668 | "bias": ParameterSpec(shape=[cfg.input_dim], mesh_axes=(None,)), |
| 669 | "moving_mean": ParameterSpec( |
| 670 | shape=[cfg.input_dim], |
| 671 | dtype=jnp.float32, |
| 672 | mesh_axes=(None,), |
| 673 | initializer=constant_initializer(0.0), |
| 674 | weight_decay_scale=0, |
| 675 | ), |
| 676 | "moving_variance": ParameterSpec( |
| 677 | shape=[cfg.input_dim], |
| 678 | dtype=jnp.float32, |
| 679 | mesh_axes=(None,), |
| 680 | initializer=constant_initializer(1.0), |
| 681 | weight_decay_scale=0, |
| 682 | ), |
| 683 | } |
| 684 | |
| 685 | def forward(self, x: Tensor, *, segment_ids: Optional[Tensor] = None) -> Tensor: |
| 686 | cfg = self.config |
| 687 | x_dtype = x.dtype |
| 688 | if cfg.forward_dtype is not None: |
nothing calls this directly
no test coverage detected