set_platform

Contents

set_platform#

class brainstate.environ.set_platform(platform)[source]#

Set the computing platform.

Parameters:

platform (str) – Platform to use (‘cpu’, ‘gpu’, or ‘tpu’).

Raises:

ValueError – If platform is not supported.

Return type:

None

Examples

>>> import brainstate.environ as env
>>>
>>> # Set to GPU
>>> env.set_platform('gpu')
>>>
>>> # Verify platform
>>> print(env.get_platform())  # 'gpu'

Notes

Platform changes only take effect at program start. Changing platform after JAX initialization may not have the expected effect.

Contents