ditype

Contents

ditype#

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

Get the default integer data type.

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

Parameters:

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

Returns:

Default integer 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)
>>> indices = jnp.arange(10, dtype=env.ditype())
>>> print(indices.dtype)  # int32
>>>
>>> # With 64-bit precision
>>> with env.context(precision=64):
...     big_indices = jnp.arange(1000, dtype=env.ditype())
...     print(big_indices.dtype)  # int64

Using custom environment:

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

See also

dftype

Default floating-point type

dutype

Default unsigned integer type