pmap to single GPU

[eadadi]

Hi,
I know that pmap is used to parallelize a function over multiple instances of GPU.
Is it possible to do so with single instance too?
I mean - to consider the single instance of the GPU as multiple instances (similar to what xla_force_host_platform_device_count does)?
Assuming the memory footprint of that function is small and multiple instances of it can fit the GPU - the performance is going to be better?

EDIT:
I am reading about the Jax sharding api (which is amazing), to try figure out myself an answer.
Meanwhile, I want to add another question which is related to my original one-
Is sharding capabilities relevant for me in any way if I am using only one gpu? (again, if we can consider one gpu as multiple instances, maybe it is?)

[yashk2810]

pmap is in maintenance mode. We recommend using jax.jit and shard_map combined with the JAX sharding APIs.

Is sharding capabilities relevant for me in any way if I am using only one gpu?

Nope. sharding only comes into play for >=2 devices :)