| 896 | |
| 897 | |
| 898 | class FlaxLLaMABlockCollection(nn.Module): |
| 899 | config: LLaMAConfig |
| 900 | dtype: jnp.dtype = jnp.float32 |
| 901 | param_dtype: jnp.dtype=jnp.float32 |
| 902 | precision: Optional[Union[jax.lax.Precision, str]]=None |
| 903 | |
| 904 | @nn.compact |
| 905 | def __call__( |
| 906 | self, |
| 907 | hidden_states, |
| 908 | attention_mask=None, |
| 909 | segment_ids=None, |
| 910 | position_ids=None, |
| 911 | deterministic: bool = True, |
| 912 | init_cache: bool = False, |
| 913 | output_attentions: bool = False, |
| 914 | output_hidden_states: bool = False, |
| 915 | return_dict: bool = True, |
| 916 | ): |
| 917 | all_attentions = () if output_attentions else None |
| 918 | all_hidden_states = () if output_hidden_states else None |
| 919 | |
| 920 | block = FlaxLLaMABlock |
| 921 | if self.config.scan_layers: |
| 922 | initializing = self.is_mutable_collection('params') |
| 923 | params_spec = ( |
| 924 | self.config.param_scan_axis if initializing else |
| 925 | nn_partitioning.ScanIn(self.config.param_scan_axis)) |
| 926 | cache_spec = 0 |
| 927 | hidden_states, _ = nn.scan( |
| 928 | block, |
| 929 | variable_axes={ |
| 930 | 'params': params_spec, |
| 931 | 'cache': cache_spec, |
| 932 | 'intermediates': 0 |
| 933 | }, |
| 934 | split_rngs={ |
| 935 | 'params': True, |
| 936 | 'dropout': True |
| 937 | }, |
| 938 | in_axes=(nn.broadcast, nn.broadcast, nn.broadcast, nn.broadcast, nn.broadcast, nn.broadcast), |
| 939 | length=self.config.num_hidden_layers, |
| 940 | metadata_params={nn.PARTITION_NAME: 'scan_decoder_layer'}, |
| 941 | )(self.config, name='scan_decoder', dtype=self.dtype, param_dtype=self.param_dtype,)( |
| 942 | hidden_states, |
| 943 | attention_mask, |
| 944 | segment_ids, |
| 945 | position_ids, |
| 946 | deterministic, |
| 947 | init_cache, |
| 948 | output_attentions, |
| 949 | ) |
| 950 | else: |
| 951 | blocks = [ |
| 952 | block( |
| 953 | self.config, |
| 954 | name=str(i), |
| 955 | dtype=self.dtype, |