Creates MethodRunner for the specified method and arguments. Args: method: the method name of self.model to invoke. The method should take an `input_batch` arg and return a NestedTensor. Both `input_batch` and the returned value are NestedTensors
(
self,
*,
method: str,
prng_key: Optional[Tensor] = None,
drop_module_outputs: Callable[[str], bool] = lambda _: True,
**kwargs,
)
| 323 | } |
| 324 | |
| 325 | def create_method_runner( |
| 326 | self, |
| 327 | *, |
| 328 | method: str, |
| 329 | prng_key: Optional[Tensor] = None, |
| 330 | drop_module_outputs: Callable[[str], bool] = lambda _: True, |
| 331 | **kwargs, |
| 332 | ) -> MethodRunner: |
| 333 | """Creates MethodRunner for the specified method and arguments. |
| 334 | |
| 335 | Args: |
| 336 | method: the method name of self.model to invoke. The method should take an |
| 337 | `input_batch` arg and return a NestedTensor. Both `input_batch` and the |
| 338 | returned value are NestedTensors containing Tensors with a leading dimension of |
| 339 | `batch_size` and will be partitioned with `input_batch_partition_spec` and |
| 340 | `output_batch_partition_spec` respectively. |
| 341 | prng_key: the random key used for inference. Use restored key if None. |
| 342 | drop_module_outputs: A callable that takes a path and outputs a decision of whether to |
| 343 | drop the module output at the given path, where True means we drop. By default, the |
| 344 | callable always returns True, meaning all module outputs are dropped. |
| 345 | Warning: Returned module outputs are fully replicated. |
| 346 | kwargs: Keyword arguments to pass to the method. |
| 347 | |
| 348 | Returns: |
| 349 | MethodRunner for computing output results. |
| 350 | |
| 351 | Raises: |
| 352 | AttributeError: if method is not found at self.model. |
| 353 | """ |
| 354 | |
| 355 | cfg: InferenceRunner.Config = self.config |
| 356 | available_methods = {m for m in dir(self.model) if not m.startswith("_")} |
| 357 | if method not in available_methods: |
| 358 | raise AttributeError( |
| 359 | f"{self.path()}.model does not have method {method}. " |
| 360 | f"Available methods are {available_methods}." |
| 361 | ) |
| 362 | |
| 363 | with self.mesh(): |
| 364 | input_partition_spec = utils.data_partition_type_to_spec(cfg.input_batch_partition_spec) |
| 365 | output_partition_spec = utils.data_partition_type_to_spec( |
| 366 | cfg.output_batch_partition_spec |
| 367 | ) |
| 368 | |
| 369 | def reshard_fn(x: Tensor, partition_spec: PartitionSpec) -> Tensor: |
| 370 | if partition_spec in [PartitionSpec(), PartitionSpec(None)]: |
| 371 | return jax.lax.with_sharding_constraint(x, PartitionSpec(None)) |
| 372 | elif x.ndim == 0: |
| 373 | return jax.lax.with_sharding_constraint(x, PartitionSpec(None)) |
| 374 | elif x.ndim == 1: |
| 375 | return jax.lax.with_sharding_constraint(x, PartitionSpec(*partition_spec[:1])) |
| 376 | else: |
| 377 | return jax.lax.with_sharding_constraint(x, partition_spec) |
| 378 | |
| 379 | def inference_iter(model_params, prng_key, input_batch): |
| 380 | model_params = jax.lax.with_sharding_constraint( |
| 381 | model_params, self._inference_runner_state_partition_specs.model |
| 382 | ) |