dctype

Contents

dctype#

class brainstate.environ.dctype(*, env=None)[source]#

Get the default complex data type.

This function returns the appropriate complex type based on the current precision setting.

Parameters:

env (EnvironmentState | None) – The environment state to query. If None, uses the global environment.

Returns:

Default complex data type.

Return type:

str | type[Any] | dtype | SupportsDType

Examples

>>> import brainstate.environ as env
>>> import jax.numpy as jnp
>>>
>>> # With 32-bit precision
>>> env.set(precision=32)
>>> z = jnp.array([1+2j, 3+4j], dtype=env.dctype())
>>> print(z.dtype)  # complex64
>>>
>>> # With 64-bit precision
>>> with env.context(precision=64):
...     w = jnp.array([5+6j], dtype=env.dctype())
...     print(w.dtype)  # complex128

Using custom environment:

>>> import brainstate.environ as env
>>>
>>> custom_env = env.EnvironmentState()
>>> env.set(precision=64, env=custom_env)
>>> print(env.dctype(env=custom_env))  # complex128

Notes

Complex128 is only available with 64-bit precision. All other precisions use complex64.