Invokes . in a pure functional fashion. The invocation will not depend on external inputs or have any side effects. The results only depend on the given inputs. All outputs are reflected in the return value. Args: module: The Module to invoke. prng_key:
(
module: Module,
prng_key: Optional[Tensor],
state: NestedTensor,
inputs: Union[Sequence[Any], dict[str, Any]],
*,
method: str = "forward",
is_training: bool,
drop_output_collections: Sequence[str] = ("module_outputs",),
copy_args_tree: bool = True,
)
| 1201 | raise_for_cycles(dict(context=self.context, args=args, kwargs=kwargs)) |
| 1202 | context = self.context |
| 1203 | if self.copy_args_tree: |
| 1204 | context, args, kwargs = jax.tree.map(lambda x: x, (self.context, args, kwargs)) |
| 1205 | |
| 1206 | with set_current_context(context, require_parent=self.require_parent): |
| 1207 | # pylint: disable-next=not-an-iterable,not-a-mapping |
| 1208 | method_outputs = self.method_fn(*args, **kwargs) |
| 1209 | return method_outputs, context.output_collection |
| 1210 | |
| 1211 | |
| 1212 | def functional( |
| 1213 | module: Module, |
| 1214 | prng_key: Optional[Tensor], |
| 1215 | state: NestedTensor, |
| 1216 | inputs: Union[Sequence[Any], dict[str, Any]], |
| 1217 | *, |
| 1218 | method: str = "forward", |
| 1219 | is_training: bool, |
| 1220 | drop_output_collections: Sequence[str] = ("module_outputs",), |
| 1221 | copy_args_tree: bool = True, |
| 1222 | ) -> tuple[Any, OutputCollection]: |
| 1223 | """Invokes <module>.<method> in a pure functional fashion. |
| 1224 | |
| 1225 | The invocation will not depend on external inputs or have any side effects. The results only |
| 1226 | depend on the given inputs. All outputs are reflected in the return value. |
| 1227 | |
| 1228 | Args: |
| 1229 | module: The Module to invoke. |
| 1230 | prng_key: The pseudo-random number generator key (can be None if the computation does not |
| 1231 | require random numbers). |
| 1232 | state: The input state of the module, including model parameters if the module contains a |
| 1233 | model. |
| 1234 | inputs: The inputs for the method. If it's a sequence, it represents the positional args. |
| 1235 | If it's a dict, it represents keyword args. |
| 1236 | method: The Module method to invoke. |
| 1237 | is_training: Whether the invocation should run in the training mode. |
| 1238 | drop_output_collections: The output collection types to drop. |
| 1239 | Defaults to dropping all module outputs. |
| 1240 | copy_args_tree: Whether to copy the `inputs` pytree to prevent method_fn from mutating the |
| 1241 | original. Defaults to True. |
| 1242 | |
| 1243 | Returns: |
| 1244 | (method_outputs, output_collection), where |
| 1245 | - method_outputs are the return value of the method. |
| 1246 | - output_collection is an OutputCollection containing summaries and state updates. |
| 1247 | |
| 1248 | Raises: |
| 1249 | ValueError: If there are circular references in args, kwargs, or context. |
| 1250 | """ |
| 1251 | context = InvocationContext( |
| 1252 | name="root", |
| 1253 | parent=None, |
| 1254 | module=module, |
| 1255 | state=state, |
| 1256 | output_collection=new_output_collection(), |
| 1257 | is_training=is_training, |
| 1258 | prng_key=prng_key, |
| 1259 | ) |
| 1260 |